Update documentation to refer to use of non uniform grid.

PiperOrigin-RevId: 861591200

Author: Nush395

docs/configuration.rst

@@ -861,14 +861,42 @@ geometry
 
 Geometry dicts for all geometry types can contain the following additional keys.
 
-``n_rho`` (int [default = 25])
-  Number of radial grid points
+``n_rho`` (int | None [default = 25])
+  Number of radial grid cells. Creates a uniform grid with cell width
+  :math:`d\hat{\rho} = 1/n\_rho`. Must be at least 4.
+  Either ``n_rho`` or ``face_centers`` must be specified.
+
+``face_centers`` (array | None [default = None])
+  Explicit array of face center coordinates in normalized :math:`\hat{\rho}`
+  (ranging from 0 to 1) for defining non-uniform radial grids. For a grid with
+  N cells, there should be N+1 face centers. The array must start at 0.0, end
+  at 1.0, be strictly increasing, and have at least 5 elements (4 cells).
+  This enables finer resolution in regions of interest, such as
+  near the plasma edge.
+
+  Example of a non-uniform grid with finer resolution near the edge:
+
+  .. code-block:: python
+
+    import numpy as np
+
+    # Create non-uniform grid: coarse in core, fine near edge
+    core_faces = np.linspace(0, 0.8, 9)  # 8 cells from rho=0 to rho=0.8
+    edge_faces = np.linspace(0.8, 1.0, 9)[1:]  # 8 cells from rho=0.8 to rho=1.0
+    face_centers = np.concatenate([core_faces, edge_faces])
+
+    'geometry': {
+        'geometry_type': 'chease',
+        'face_centers': face_centers,
+        # ... other geometry parameters
+    }
 
 ``hires_factor`` (int [default = 4])
   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.
 
+
 Geometry dicts for all non-circular geometry types can contain the following
 additional keys.
 

docs/fvm.rst

@@ -26,7 +26,10 @@ arbitrary mesh topologies. `torax.fvm` is much less general ---it
 avoids hard-coded Tokamak physics into `torax.fvm` but has not attempted
 to include any more features than needed for the Tokamak core transport
 problem. `torax.fvm` thus supports only one topology, the 1-D
-grid topology. However, within this restricted domain, `torax.fvm`
+grid topology, though it does support both uniform and non-uniform radial
+grids (see :ref:`configuration` for details on configuring the grid via
+``n_rho`` or ``face_centers``).
+However, within this restricted domain, `torax.fvm`
 offers a variety of solver techniques not available in `fipy`,
 especially solver techniques that are possible only with differentiable
 solvers, in our case enabled by JAX. For example, in `torax.fvm` the

docs/solver_details.rst

@@ -16,10 +16,15 @@ The TORAX JAX 1D FVM library is significantly influenced by
 library API. This section summarizes 1D FVM numerics in general.
 
 The 1D spatial domain, :math:`0 \leq \hat{\rho} \leq 1`, is divided into a
-uniform grid of :math:`N` cells, each with a width of
-:math:`d \hat{\rho} = 1/N`.  The cell centers are denoted by
-:math:`\hat{\rho}_i`, where :math:`0 = 1, 2,..., N-1`, and the :math:`N+1` cell
-faces are located at :math:`\hat{\rho}_{i\pm1/2}`. Both :math:`\hat{\rho}=0` and
+grid of :math:`N` cells. TORAX supports both uniform grids (configured via
+``n_rho``) and non-uniform grids (configured via ``face_centers``). For uniform
+grids, each cell has width :math:`d \hat{\rho} = 1/N`. For non-uniform grids,
+cell widths :math:`d \hat{\rho}_i` vary across the domain, enabling finer
+resolution in regions of interest such as near the plasma edge.
+
+The cell centers are denoted by :math:`\hat{\rho}_i`, where
+:math:`i = 0, 1, 2,..., N-1`, and the :math:`N+1` cell faces are located at
+:math:`\hat{\rho}_{i\pm1/2}`. Both :math:`\hat{\rho}=0` and
 :math:`\hat{\rho}=1` are on the face grid.
 
 For a generic conservation law of the form:
@@ -37,12 +42,14 @@ of finite differences:
 
 .. math::
 
-  \frac{\partial }{\partial t}(x_i) + \frac{1}{d \hat{\rho}}({\Gamma}_{i+1/2}
+  \frac{\partial }{\partial t}(x_i) + \frac{1}{d \hat{\rho}_i}({\Gamma}_{i+1/2}
   - {\Gamma}_{i-1/2})  = S_i
 
 where: :math:`x_i` is the cell-averaged value of :math:`x` in cell :math:`i`,
-:math:`\Gamma_{i+1/2}` is the flux at face :math:`i+1/2`, and :math:`S_i` is the
-cell-averaged source term in cell :math:`i`.
+:math:`\Gamma_{i+1/2}` is the flux at face :math:`i+1/2`, :math:`S_i` is the
+cell-averaged source term in cell :math:`i`, and :math:`d \hat{\rho}_i` is the
+width of cell :math:`i` (which may vary for non-uniform grids).
+
 
 In general, the fluxes in TORAX are decomposed as
 
@@ -55,9 +62,9 @@ convection coefficient, leading to:
 .. math::
 
   \begin{aligned}
-  \Gamma_{i+1/2} &= -D_{i+1/2}\frac{x_{i+1} - x_{i}}{d\hat{\rho}} +
+  \Gamma_{i+1/2} &= -D_{i+1/2}\frac{x_{i+1} - x_{i}}{\hat{\rho}_{i+1} - \hat{\rho}_{i}} +
    V_{i+1/2}x_{i+1/2} \\
-  \Gamma_{i-1/2} &= -D_{i-1/2}\frac{x_{i} - x_{i-1}}{d\hat{\rho}} +
+  \Gamma_{i-1/2} &= -D_{i-1/2}\frac{x_{i} - x_{i-1}}{\hat{\rho}_{i} - \hat{\rho}_{i-1}} +
    V_{i-1/2}x_{i-1/2}
   \end{aligned}
 
@@ -78,9 +85,11 @@ defined as:
 
 .. math::
 
-  Pe = \frac{V d \hat{\rho}}{D}
+  Pe = \frac{V d \hat{\rho}_i}{D}
 
-where :math:`V` is convection and :math:`D` is diffusion. The power-law scheme
+where :math:`V` is convection, :math:`D` is diffusion, and
+:math:`d \hat{\rho}_i` is the local cell width (which varies for non-uniform
+grids). The power-law scheme
 is as follows:
 
 .. math::
@@ -327,9 +336,10 @@ as follows.
 
   .. math::
 
-    \Delta t_{ \mathrm{base}}=\frac{(d\hat{\rho})^2}{2\chi_{\max}}
+    \Delta t_{ \mathrm{base}}=\frac{(d\hat{\rho}_{\min})^2}{2\chi_{\max}}
 
-  where
+  where :math:`d \hat{\rho}_{\min}` is the minimum cell width (for non-uniform
+  grids, this is the smallest cell; for uniform grids, all cells are equal), and
   :math:`\Delta t = c_{ \mathrm{mult}}^{dt} \Delta t_{ \mathrm{base}}`.
 
   :math:`c_{ \mathrm{mult}}^{dt}` is a user-configurable prefactor.

torax/_src/math_utils.py

@@ -68,10 +68,11 @@ def cell_to_face(
   """Convert cell values to face values.
 
   We make three assumptions:
-  1) Inner face values are the average of neighbouring cells.
+  1) Inner face values are the interpolation of neighboring cell values.
   2) The left most face value is linearly extrapolated from the left most cell
-  values.
-  3) The transformation from cell to face is integration preserving.
+    values.
+  3) The transformation from cell to face is integration preserving and the
+    quantity to preserve the integral of is specified by `preserved_quantity`.
 
   Args:
     cell_values: Values defined on the TORAX cell grid.