Support 2d rho arrays in xarray time-varying-array.

docs/beginner_guides.rst

@@ -14,4 +14,3 @@ installation instructions, and a basic primer for running and plotting.
    overview
    installation
    quickstart
-   faq

docs/configuration.rst

@@ -93,58 +93,40 @@ To extend configuration parameters where time-dependence is not enabled, to have
 
 Time-varying arrays
 -------------------
-Time-varying arrays can be defined using either primitives, an
-``xarray.DataArray`` or a ``tuple`` of ``Array``.
+Parameters marked as **time-varying-array** are interpolated on a grid (time, :math:`\hat{\rho}`).
 
-Specifying interpolation methods
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-By default piecewise linear interpolation is used to interpolate values in time.
-To specify a different interpolation method, use the following syntax of a tuple
-with two elements. The first element in the tuple is the usual value for the
-time-varying-array (as defined below), the second value is a dict with keys
-``time_interpolation_mode`` and ``rho_interpolation_mode`` and values the
-desired interpolation modes.
-
-.. code-block:: python
-
-  (time_varying_array_value, {'time_interpolation_mode': 'STEP', 'rho_interpolation_mode': 'PIECEWISE_LINEAR'})
-
-Currently two interpolation modes are supported:
-
-* ``'STEP'``
-* ``'PIECEWISE_LINEAR'``
-
-Using primitives
-^^^^^^^^^^^^^^^^
+**time-varying-array** parameters can be defined using either a nested dictionary,
+or in the form of arrays (represented as a ``xarray.DataArray`` object or a ``tuple`` of arrays).
 
-For fields labelled with **time-varying-array** time-dependence is set by assigning a dict of dicts to the parameter.
+**Note**: :math:`\hat{\rho}` is normalized and will take values between 0 and 1.
 
-The outer dict defines a time-series with ``{time: value}`` pairs.
-The ``value`` itself is interpreted as a radial profile, being made up of {rho: value} pairs.
-It behaves similarly to the **time-varying-scalar** but any interpolation will happen along the
-:math:`\hat{\rho}` axis and can take any of the formats defined for a **time-varying-scalar** above.
-
-Note: :math:`\hat{\rho}` is normalized and will take values between 0 and 1.
-
-None of the keys need to be sorted in order of time. Ordering is carried out internally.
 In the case of non-evolving parameters for each evaluation of the TORAX solver (PDE solver), time-dependent variables
 are interpolated first along the :math:`\hat{\rho}` axis at the cell grid centers and then linearly interpolated in time
 at both time :math:`t` and time :math:`t+dt`..
 
 For :math:`t` greater than or less than the largest or smallest defined time then the interpolation scheme
 will be applied from the closest time value.
 
+Using a nested dictionary
+^^^^^^^^^^^^^^^^
+This is of the form:
+.. code-block:: python
+
+  {time_1: {rho_11: value_11, rho_12: value_12, ...}, time_2: ...}
+
+At each ``time_i``, we have a radial profile composed of  ``{rho: value}`` pairs.
+The ordering of the dict does not matter.
+
+
 Shortcuts:
 
-Passing a single float value is interpreted as defining a constant profile for all times.
+* Passing a single float value is interpreted as defining a constant profile for all times.
 For example ``T_i: 6.0`` would be equivalent to passing in ``T_i: {0.0: {0.0: 6.0}}``.
-
-Passing a single dict (instead of dict of dicts) is a shortcut for defining the rho profile
-for :math:`t=0.0`. For example ``T_i: {0.0: 18.0, 0.95: 5.0, 1.0: 0.2}`` is a shortcut for
+* Passing a single dict (instead of dict of dicts) is a shortcut for defining the rho profile
+for :math:`t=0.0`. For example, ``T_i: {0.0: 18.0, 0.95: 5.0, 1.0: 0.2}`` is a shortcut for
 ``T_i: {0.0: {0: 18.0, 0.95: 5.0, 1.0: 0.2}}`` where :math:`t=0.0` is arbitrary
 (due to constant extrapolation for any input :math:`t=0.0`).
 
-
 Examples:
 
 1. Define an initial profile (at :math:`t=0.0`) for :math:`T_{i}` with a pedestal.
@@ -153,7 +135,7 @@ Examples:
 
   T_i = {0.0: {0.0: 15.0, 0.95: 3.0, 1.0: 1.0}}
 
-Note: due to constant extrapolation the t=0.0 here is an arbitrary number and could be anything.
+Note: due to constant extrapolation, the ``t=0.0`` here is an arbitrary number and could be anything.
 
 2. Define a time-dependent :math:`T_{i}` profile initialised with a pedestal and, if the ion equation is not being
 evolved by the PDE, to have a prescribed time evolution which decays to a
@@ -163,30 +145,47 @@ constant :math:`T_{i}=1` by :math:`t=80.0`.
 
   T_i = {0.0: {0.0: 15.0, 0.95: 3.0, 1.0: 1.0}, 80: 1.0}
 
-Using ``xarray.DataArray``
+
+Using arrays
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
-If a ``xarray.DataArray`` is specified then it is expected to have a
-``time`` and ``rho_norm`` coordinate. The values of the data array are the values
-at each time and rho_norm.
-
-Using ``tuple`` of ``Array``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-If a ``tuple`` of ``Array`` is used, the tuple must have structure of,
-``(time_array, rho_norm_array, values_array)`` or ``(rho_norm_array, values_array)``.
-The latter is a useful shortcut for defining an initial condition or a constant profile.
-
-In the case of ``(time_array, rho_norm_array, values_array)``:
-``time_array`` and ``rho_norm_array`` are expected to map to 1D array values and
-represent the time and rho_norm coordinates.
-``values_array`` is expected to map to a 2D array with shape
-``(len(time_array), len(rho_norm_array))`` and represent the values at the given
-time and rho_norm.
-
-In the case of ``(rho_norm_array, values_array)``:
-``rho_norm_array`` is expected to map to a 1D array values and represent the
-rho_norm coordinates.
-``values_array`` is expected to map to a 1D array with shape
-``len(rho_norm_array)`` and represent the values at the given rho_norm.
+This can be a tuple of arrays ``(time_array, rho_norm_array, values_array)``, or
+equivalently an ``xarray.DataArray`` object of the form:
+.. code-block:: python
+
+  xarray.DataArray(
+    data=values_array,
+    coords={'time': time_array, 'rho_norm': rho_norm_array}
+  )
+
+All arrays can be represented as NumPy arrays or lists. The shapes:
+* ``time_array`` is a 1D array of times.
+* ``values_array`` is a 2D array of shape ``(len(time_array), num_values)``.
+* ``rho_norm_array`` either a 1D array of shape ``(num_values,)``, or a 2D array
+  of shape ``(len(time_array), num_values)``.
+
+Shortcuts:
+* ``(rho_norm_array, values_array)``: constant in time profile, useful for
+  defining an initial condition or a constant profile. Note that both arrays
+  are now 1D arrays.
+
+
+Specifying interpolation methods
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+By default, piecewise linear interpolation is used to interpolate values both in
+time and in :math:`\hat{\rho}`. To specify a different interpolation method:
+
+.. code-block:: python
+
+  (time_varying_array_value, {'time_interpolation_mode': 'STEP', 'rho_interpolation_mode': 'PIECEWISE_LINEAR'})
+
+where ``time_varying_array_value`` is any of the above inputs
+(nested dictionary, arrays, etc.).
+
+Currently two interpolation modes are supported:
+
+* ``'PIECEWISE_LINEAR'``: linear interpolation of the input time-series (default).
+* ``'STEP'``: stepwise change in values following each traversal above a time value in the time-series.
+
 
 .. _config_details:
 

docs/index.rst

@@ -65,7 +65,6 @@ design workflows. TORAX is written in Python using the JAX_ library.
    overview
    installation
    quickstart
-   faq
 
 .. toctree::
    :hidden:

docs/running_programmatically.rst

@@ -38,4 +38,4 @@ We can then run the simulation:
   # to access the data. Example below shows how to access the fusion gain
   # at time=2 seconds.
   data_tree = results.simulation_output_to_xr()
-  Q_fusion_t2 = data_tree.post_processed_outputs.Q_fusion.sel(time=2, method='nearest')
+  Q_fusion_t2 = data_tree.scalars.Q_fusion.sel(time=2, method='nearest')

torax/torax_pydantic/interpolated_param_2d.py

@@ -205,7 +205,7 @@ def _conform_data(
         data = data[0]
 
     if isinstance(data, xr.DataArray):
-      value = _load_from_xr_array(data)
+      value = _load_from_arrays(data)
     elif isinstance(data, tuple):
       values = []
       for v in data:
@@ -332,47 +332,33 @@ def _load_from_primitives(
   return loaded_values
 
 
-def _load_from_xr_array(
-    xr_array: xr.DataArray,
-) -> Mapping[float, tuple[chex.Array, chex.Array]]:
-  """Loads the data from an xr.DataArray."""
-  if 'time' not in xr_array.coords:
-    raise ValueError('"time" must be a coordinate in given dataset.')
-  if interpolated_param.RHO_NORM not in xr_array.coords:
-    raise ValueError(
-        f'"{interpolated_param.RHO_NORM}" must be a coordinate in given'
-        ' dataset.'
-    )
-  values = {
-      t: (
-          np.asarray(xr_array.rho_norm.data, dtype=jax_utils.get_np_dtype()),
-          np.asarray(
-              xr_array.sel(time=t).values, dtype=jax_utils.get_np_dtype()
-          ),
-      )
-      for t in xr_array.time.data
-  }
-  return values
-
-
 def _load_from_arrays(
-    arrays: tuple[chex.Array, ...],
+    arrays: tuple[chex.Array, ...] | xr.DataArray,
 ) -> Mapping[float, tuple[np.ndarray, np.ndarray]]:
   """Loads the data from numpy arrays.
 
   Args:
-    arrays: A tuple of (times, rho_norm, values) or (rho_norm, values). - In the
-      former case times and rho_norm are assumed to be 1D arrays of equal
-      length, values is a 2D array with shape (len(times), len(rho_norm)). - In
-      the latter case rho_norm and values are assumed to be 1D arrays of equal
-      length (shortcut for initial condition profile).
+    arrays: A tuple of (times, rho_norm, values) or (rho_norm, values), or an
+      xarray.DataArray.
 
   Returns:
     A mapping from time to (rho_norm, values)
   """
+
+  if isinstance(arrays, xr.DataArray):
+    if interpolated_param.RHO_NORM not in arrays.coords:
+      raise ValueError(
+          f'"{interpolated_param.RHO_NORM}" must be a coordinate in given'
+          ' dataset.'
+      )
+    if 'time' in arrays.coords:
+      arrays = (arrays.time.data, arrays.rho_norm.data, arrays.data)
+    else:
+      arrays = (arrays.rho_norm.data, arrays.data)
+
   if len(arrays) == 2:
     # Shortcut for initial condition profile.
-    rho_norm, values = arrays
+    rho_norm, values = arrays  # pytype: disable=bad-unpacking
     return {
         0.0: (
             np.asarray(rho_norm, dtype=jax_utils.get_np_dtype()),
@@ -384,12 +370,16 @@ def _load_from_arrays(
     rho_norm = np.asarray(arrays[1], dtype=jax_utils.get_np_dtype())
     values = np.asarray(arrays[2], dtype=jax_utils.get_np_dtype())
 
-    if values.shape != (len(times), len(rho_norm)):
+    if values.ndim != 2:
       raise ValueError(
-          'values must be of shape (len(times), len(rho_norm)). Given: '
-          f'{values.shape}.'
+          f'The values array must have ndim=2, but got {values.ndim}.'
       )
-    return {t: (rho_norm, values[i, :]) for i, t in enumerate(times)}
+    _, num_vals = values.shape
+
+    if rho_norm.ndim == 1:
+      rho_norm = np.stack([rho_norm] * num_vals)
+
+    return {t: (rho_norm[i], values[i]) for i, t in enumerate(times)}
   else:
     raise ValueError(f'arrays must be length 2 or 3. Given: {len(arrays)}.')
 

torax/torax_pydantic/tests/interpolated_param_2d_test.py

@@ -79,6 +79,20 @@ class InterpolatedParam2dTest(parameterized.TestCase):
           time=0.0,
           expected_output=np.array([1.0, 2.0, 3.0, 4.0]),
       ),
+      dict(
+          testcase_name='3_tuple_input_time_rho_t=0',
+          time_rho_interpolated_input=(
+              np.array([0.0, 1.0]),
+              np.array(
+                  [[0.125, 0.375, 0.625, 0.875], [0.125, 0.375, 0.625, 0.875]]
+              ),
+              np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]),
+          ),
+          nx=4,
+          dx=0.25,
+          time=0.0,
+          expected_output=np.array([1.0, 2.0, 3.0, 4.0]),
+      ),
       dict(
           testcase_name='3_tuple_input_t=1',
           time_rho_interpolated_input=(
@@ -202,6 +216,24 @@ class InterpolatedParam2dTest(parameterized.TestCase):
           time=0.5,
           expected_output=np.array([3.0, 4.0, 5.0, 6.0]),
       ),
+      dict(
+          testcase_name='xarray_input_full_t=0.5',
+          time_rho_interpolated_input=xr.DataArray(
+              data=np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]),
+              coords={
+                  'time': [0.0, 1.0],
+                  'rho_norm': (
+                      ('time', 'value'),
+                      np.array([_RHO_NORM_ARRAY, _RHO_NORM_ARRAY]),
+                  ),
+              },
+              dims=['time', 'value'],
+          ),
+          nx=4,
+          dx=0.25,
+          time=0.5,
+          expected_output=np.array([3.0, 4.0, 5.0, 6.0]),
+      ),
       dict(
           testcase_name='single_dict_t=0',
           time_rho_interpolated_input={
@@ -236,6 +268,17 @@ class InterpolatedParam2dTest(parameterized.TestCase):
           time=0.0,
           expected_output=np.array([18.0, 5.0, 5.0, 5.0]),
       ),
+      dict(
+          testcase_name='nested_dict_t=0.5',
+          time_rho_interpolated_input={
+              0.0: {0.125: 1.0, 0.375: 2.0, 0.625: 3.0, 0.875: 4.0},
+              1.0: {0.125: 5.0, 0.375: 6.0, 0.625: 7.0, 0.875: 8.0},
+          },
+          nx=4,
+          dx=0.25,
+          time=0.5,
+          expected_output=np.array([3.0, 4.0, 5.0, 6.0]),
+      ),
       # Single float represents a constant (in time and rho) profile.
       dict(
           testcase_name='float_t=0.0',