Update xarray advanced course to mention imas.util.to_xarray

Author: maarten-ic

docs/source/courses/advanced/imas_snippets/tensorized_ids_to_xarray.py

@@ -1,6 +1,7 @@
 import os
 
 import matplotlib
+
 # To avoid possible display issues when Matplotlib uses a non-GUI backend
 if "DISPLAY" not in os.environ:
     matplotlib.use("agg")
@@ -17,6 +18,39 @@
 entry = imas.training.get_training_db_entry()
 cp = entry.get("core_profiles")
 
+#######################################################################################
+# Steps 2, 3 and 4, using imas.util.to_xarray
+# Create an xarray Dataset containing t_i_average and its coordinates
+xrds = imas.util.to_xarray(cp, "profiles_1d/t_i_average")
+# Note that profiles_1d.grid.rho_tor_norm is a 2D coordinate: its values may be
+# different at different times.
+#
+# Since the values at different time slices differ only minutely in this example, we'll
+# rename the `profiles_1d.grid.rho_tor_norm:i` dimension to `rho_tor_norm` and set the
+# values to the values of rho_tor_norm of the first time slice:
+xrds = xrds.rename({"profiles_1d.grid.rho_tor_norm:i": "rho_tor_norm"}).assign_coords(
+    {"rho_tor_norm": xrds["profiles_1d.grid.rho_tor_norm"].isel(time=0).data}
+)
+
+# Extract temperatures as an xarray DataArray
+temperature = xrds["profiles_1d.t_i_average"]
+
+# 5a. Select subset of temperature where 0.4 <= rho_tor_norm < 0.6:
+print(temperature.sel(rho_tor_norm=slice(0.4, 0.6)))
+
+# 5b. Interpolate temperature on a new grid: [0, 0.1, 0.2, ..., 0.9, 1.0]
+print(temperature.interp(rho_tor_norm=numpy.linspace(0, 1, 11)))
+
+# 5c. Interpolate temperature on a new time base: [10, 20]
+print(temperature.interp(time=[10, 20]))
+
+# 5d. Plot
+temperature.plot(x="time", norm=matplotlib.colors.LogNorm())
+plt.show()
+
+#######################################################################################
+# We can also manually build an xarray DataArray, this is shown below:
+
 # 2. Store the temperature of the first time slice
 temperature = cp.profiles_1d[0].t_i_average
 

docs/source/courses/advanced/xarray.rst

@@ -3,9 +3,10 @@ Create ``xarray.DataArray`` from an IDS
 
 .. info::
 
-    In this lesson you will create a ``DataArray`` manually. In a future version of
-    IMAS-Python we plan to include functionality that will automatically do this for you.
-    That should further simplify working with data inside IDSs.
+    This lesson was written before :py:func:`imas.util.to_xarray` was
+    implemented. This lesson is retained for educational purposes, however we
+    recommend to use :py:func:`imas.util.to_xarray` instead of manually creating
+    xarray ``DataArray``\ s.
 
 Let's start with an introduction of Xarray. According to `their website
 <https://docs.xarray.dev/en/stable/getting-started-guide/why-xarray.html>`_ (where you
@@ -61,6 +62,10 @@ Exercise 1: create a ``DataArray`` for ``profiles_1d/temperature``
 
     .. md-tab-item:: Solution
 
+        This exercise was created before the implementation of
+        :py:func:`imas.util.to_xarray`. The original approach is available below
+        for educational purposes.
+
         .. literalinclude:: imas_snippets/ids_to_xarray.py
 
 
@@ -96,4 +101,9 @@ the ``profiles_1d`` array of structures. When the grid is not changing in the ID
 
     .. md-tab-item:: Solution
 
+        This exercise was created before the implementation of
+        :py:func:`imas.util.to_xarray`. Below code sample is updated to provide
+        two alternatives: the first is based on :py:func:`imas.util.to_xarray`,
+        the second is the original, manual approach.
+
         .. literalinclude:: imas_snippets/tensorized_ids_to_xarray.py