Cleanups and documentation updates

PiperOrigin-RevId: 659960040

Author: jcitrin

README.md

@@ -2,7 +2,7 @@
 
 # What is TORAX?
 
-TORAX is a differentiable tokamak core transport simulator aimed for fast and accurate forward modelling, pulse-design, trajectory optimization, and controller design workflows. TORAX is written in Python-JAX, with the following motivations:
+TORAX is a differentiable tokamak core transport simulator aimed for fast and accurate forward modelling, pulse-design, trajectory optimization, and controller design workflows. TORAX is written in Python using JAX, with the following motivations:
 
 - Open-source and extensible, aiding with flexible workflow coupling
 - JAX provides auto-differentiation capabilities and code compilation for fast runtimes. Differentiability allows for gradient-based nonlinear PDE solvers for fast and accurate modelling, and for sensitivity analysis of simulation results to arbitrary parameter inputs, enabling applications such as trajectory optimization and data-driven parameter identification for semi-empirical models. Auto-differentiability allows for these applications to be easily extended with the addition of new physics models, or new parameter inputs, by avoiding the need to hand-derive Jacobians
@@ -14,14 +14,15 @@ TORAX now has the following physics feature set:
 
 - Coupled PDEs of ion and electron heat transport, electron particle transport, and current diffusion
     - Finite-volume-method
-    - Multiple solver options: linear with Pereverzev-Corrigan terms, nonlinear with Newton-Raphson, nonlinear with optimization using the jaxopt library
-- Ohmic power, ion-electron heat exchange, fusion power, bootstrap current with the analytical Sauter model
-- Time dependent boundary conditions and sources
+    - Multiple solver options: linear with Pereverzev-Corrigan terms and the predictor-corrector method, nonlinear with Newton-Raphson, nonlinear with optimization using the jaxopt library
+- Ohmic power, ion-electron heat exchange, fusion power, Bremsstrahlung, and bootstrap current with the analytical Sauter model
+- Time dependent boundary conditions, sources, geometry.
 - Coupling to the QLKNN10D [[van de Plassche et al, Phys. Plasmas 2020]](https://doi.org/10.1063/1.5134126) QuaLiKiz-neural-network surrogate for physics-based turbulent transport
-- General geometry, provided via CHEASE equilibrium files
+- General geometry, provided via CHEASE or FBT equilibrium files
     - For testing and demonstration purposes, a single CHEASE equilibrium file is available in the data/geo directory. It corresponds to an ITER hybrid scenario equilibrium based on simulations in [[Citrin et al, Nucl. Fusion 2010]](https://doi.org/10.1088/0029-5515/50/11/115007), and was obtained from [PINT](https://gitlab.com/qualikiz-group/pyntegrated_model). A PINT license file is available in data/geo.
+    - Time dependent geometry is supported by provided a time series of geometry files
 
-Additional heating and current drive sources can be provided by prescribed formulas, or user-provided analytical models.
+Additional heating and current drive sources can be provided by prescribed formulas, user-provided analytical models, or user-provided prescribed data.
 
 Model implementation was verified through direct comparison of simulation outputs to the RAPTOR [[Felici et al, Plasma Phys. Control. Fusion 2012]](https://iopscience.iop.org/article/10.1088/0741-3335/54/2/025002) tokamak transport simulator.
 
@@ -31,10 +32,8 @@ This is not an officially supported Google product.
 
 Short term development plans include:
 
-- Time dependent geometry
-- More flexible initial conditions
+- Extension of and more flexible data structures for prescribed input data
 - Implementation of forward sensitivity calculations w.r.t. control inputs and parameters
-- Implementation of persistent compilation cache for CPU
 - More extensive documentation and tutorials
 - Visualisation improvements
 
@@ -44,7 +43,6 @@ Longer term desired features include:
 - Neoclassical tearing modes (modified Rutherford equation)
 - Radiation sinks
     - Cyclotron radiation
-    - Bremsstrahlung
     - Line radiation
 - Neoclassical transport + multi-ion transport, with a focus on heavy impurities
 - IMAS coupling
@@ -192,6 +190,14 @@ If false, JAX does not compile internal TORAX functions. Used for debugging. Def
 $ export TORAX_COMPILATION_ENABLED=<True/False>
 ```
 
+The following implements the JAX persistent cache and will cause jax to store
+compiled programs to the filesystem, reducing recompilation time in some cases:
+
+```shell
+$ export JAX_COMPILATION_CACHE_DIR=<path of your choice, such as ~/jax_cache>
+$ export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1
+$ export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0.0
+```
 
 ### Set flags
 Output simulation time, dt, and number of stepper iterations (dt backtracking with nonlinear solver) carried out at each timestep.

docs/configuration.rst

@@ -356,7 +356,7 @@ geometry
 --------
 
 ``geometry_type`` (str = 'chease')
-  Geometry model used. There are currently two options:
+  Geometry model used. There are currently three options:
 
 * ``'circular'``
     An ad-hoc circular geometry model. Includes elongation corrections.
@@ -365,7 +365,8 @@ geometry
 * ``'chease'``
     Loads a CHEASE geometry file.
 
-Time dependent geometry is not currently supported, but is on the immediate-term roadmap.
+* ``'fbt'``
+    Loads FBT geometry files.
 
 ``nrho`` (int = 25)
   Number of radial grid points
@@ -374,6 +375,15 @@ Time dependent geometry is not currently supported, but is on the immediate-term
   Only used for ``geometry_type='chease'``. Sets the geometry file loaded.
   The geometry directory is set with the ``TORAX_GEOMETRY_DIR`` environment variable. If none is set, then the default is ``torax/data/third_party/geo``.
 
+``LY_file`` (str = None)
+  Only used for ``geometry_type='fbt'``. Sets the FBT LY geometry file loaded.
+
+``L_file`` (str = None)
+  Only used for ``geometry_type='fbt'``. Sets the FBT L geometry file loaded.
+
+``geometry_dir`` (str = None)
+  Optionally overrides the``TORAX_GEOMETRY_DIR`` environment variable.
+
 ``Ip_from_parameters`` (bool = True)
   Only used for ``geometry_type='chease'``.Toggles whether total plasma current is read from the configuration file,
   or from the geometry file. If True, then the :math:`\psi` calculated from the geometry file is scaled to match the desired :math:`I_p`.
@@ -394,6 +404,44 @@ Time dependent geometry is not currently supported, but is on the immediate-term
   Only used when the initial condition ``psi`` is from plasma current. Sets up a higher resolution mesh
   with ``nrho_hires = nrho * hi_res_fac``, used for ``j`` to ``psi`` conversions.
 
+For setting up time-dependent geometry, a subset of varying geometry parameters
+and input files are defined in a ``geometry_configs`` dict, which is a
+time-series of {time: {configs}} pairs. For example, a time-dependent geometry
+input with 3 time-slices of FBT geometries can be set up as:
+
+.. code-block:: python
+
+  'geometry': {
+      'geometry_type': 'fbt',
+      'Ip_from_parameters': True,
+      'geometry_configs': {
+          20.0: {
+              'LY_file': 'LY_early_rampup.mat',
+              'L_file': 'L_early_rampup.mat',
+              'Rmaj': 6.2,
+              'Rmin': 2.0,
+              'B0': 5,3,
+          },
+          50.0: {
+              'LY_file': 'LY_mid_rampup.mat',
+              'L_file': 'L_mid_rampup.mat',
+              'Rmaj': 6.2,
+              'Rmin': 2.0,
+              'B0': 5,3,
+          },
+          100.0: {
+              'LY_file': 'LY_endof_rampup.mat',
+              'L_file': 'L_endof_rampup.mat',
+              'Rmaj': 6.2,
+              'Rmin': 2.0,
+              'B0': 5,3,
+          },
+      },
+  },
+
+All file loading and geometry processing is done upon simulation initialization.
+The geometry inputs into the TORAX PDE coefficients are then time-interpolated
+on-the-fly onto the TORAX time slices where the PDE calculations are done.
 
 transport
 ---------

docs/overview.rst

@@ -32,7 +32,7 @@ TORAX's physics and solver feature set includes the following:
   * Time dependent boundary conditions and sources
   * Coupling to the QLKNN10D `[van de Plassche et al, Phys. Plasmas 2020] <https://doi.org/10.1063/1.5134126>`_
     QuaLiKiz-neural-network surrogate for physics-based turbulent transport
-  * General geometry, provided via CHEASE equilibrium files
+  * General geometry, provided via CHEASE or FBT equilibrium files
 
     * For testing and demonstration purposes, a single CHEASE equilibrium file is available in the
       data/geo directory. It corresponds to an ITER hybrid scenario equilibrium based on simulations

docs/physics_models.rst

@@ -15,6 +15,9 @@ geometry variables required for the transport equations.
   - **CHEASE:** This model utilizes equilibrium data from the `CHEASE <https://doi.org/10.1016/0010-4655(96)00046-X>`_ fixed boundary
     Grad-Shafranov equilibrium code. Users provide a CHEASE output file, and TORAX extracts the relevant geometric quantities.
 
+  - **FBT:** This model utilizes equilibrium data from the `FBT <https://www.sciencedirect.com/science/article/pii/0010465588900410>`_ free boundary
+    Grad-Shafranov equilibrium code. Users provide FBT output files, and TORAX extracts the relevant geometric quantities.
+
   - **Circular:** For testing and demonstration purposes, TORAX includes a simplified circular geometry model.
     This model assumes a circular plasma cross-section and corrects for elongation to approximate the metric coefficients.
 

docs/roadmap.rst

@@ -5,10 +5,9 @@ Development Roadmap
 
 Short term development plans include:
 
-* Time dependent geometry
-* More flexible functional forms for initial and prescribed conditions
+* Extension of and more flexible datastructures for prescribed input data
 * Implementation of forward sensitivity calculations w.r.t. control inputs and parameters
-* Implementation of persistent compilation cache for CPU
+* More extensive documentation and tutorials
 * Extended visualisation
 
 Longer term desired features include:

torax/core_profile_setters.py

@@ -624,9 +624,9 @@ def initial_core_profiles(
       isinstance(geo, geometry.StandardGeometry)
       and not dynamic_runtime_params_slice.profile_conditions.initial_psi_from_j
   ):
-    # psi is already provided from the CHEASE equilibrium, so no need to first
-    # calculate currents. However, non-inductive currents are still calculated
-    # and used in current diffusion equation.
+    # psi is already provided from a numerical equilibrium, so no need to
+    # first calculate currents. However, non-inductive currents are still
+    # calculated and used in current diffusion equation.
     psi_grad_constraint = _calculate_psi_grad_constraint(
         dynamic_runtime_params_slice,
         geo,

torax/geometry.py

@@ -115,7 +115,7 @@ class GeometryType(enum.Enum):
   """
 
   CIRCULAR = 0
-  CHEASE = 1
+  NUMERICAL = 1
 
 
 # pylint: disable=invalid-name
@@ -891,8 +891,8 @@ def build_standard_geometry(
   g2g3_over_rhon = g2[1:] * g3[1:] / rho_norm_intermediate[1:]
   g2g3_over_rhon = np.concatenate((np.zeros(1), g2g3_over_rhon))
 
-  # make an alternative initial psi, self-consistent with CHEASE Ip profile
-  # needed because CHEASE psi profile has noisy second derivatives
+  # make an alternative initial psi, self-consistent with numerical geometry
+  # Ip profile. Needed since input psi profile may have noisy second derivatives
   dpsidrhon = (
       intermediate.Ip_profile[1:]
       * (16 * constants.CONSTANTS.mu0 * np.pi**3 * intermediate.Phi[-1])
@@ -1012,7 +1012,7 @@ def build_standard_geometry(
   area = rhon_interpolation_func(rho_norm, intermediate.area)
 
   return StandardGeometry(
-      geometry_type=GeometryType.CHEASE.value,
+      geometry_type=GeometryType.NUMERICAL.value,
       drho_norm=np.asarray(drho_norm),
       torax_mesh=mesh,
       Phi=Phi,