ENH: Add Output enumeration to make it easier to access variables

CHANGELOG.md

@@ -10,11 +10,15 @@ All notable changes to this project will be documented in this file.
     to the use of many global/save variables. There is a lock around the
     extension modules so that only one thread will be calling the routines
     at a time, so the Python library is safe to use in a multi-threaded context.
+- **ADDED** Variable enumeration for easier indexing into output arrays.
+  - This can be used as `msis.Variable.O2` for getting the `O2` species index.
+    For example, `output_array[..., msis.Variable.O2]`.
 - **MAINTENANCE** Default `-O1` optimization level for all builds.
   - Previously, this
     was only done on Windows machines. Users can change this by updating
-    environment variables before building with `FFLAGS=-Ofast`, but note that
-    some machines produce invalid results when higher optimizations are used.
+    environment variables before building with `FFLAGS=-Ofast pip install .`,
+    but note that some machines produce invalid results when higher
+    optimizations are used.
 - **PERFORMANCE** Cache options state between subsequent runs.
   - Avoid calling initialization switches unless they have changed between runs
 - **PERFORMANCE** Speed up numpy function calls.

docs/source/reference/index.rst

@@ -14,6 +14,17 @@ To run the code, use the :mod:`pymsis.msis` module.
     msis.create_options
     msis.create_input
 
+The output of the model is stored in basic numpy arrays with the final
+dimension being the variable/species. To get the output in a more human-readable
+format, use the :class:`~.Variable` enumeration that provides
+easier indexing into the output arrays.
+
+.. autosummary::
+    :toctree: generated/
+    :nosignatures:
+
+    msis.Variable
+
 To get input data for historical events, use the :mod:`pymsis.utils` module.
 
 .. autosummary::

examples/plot_altitude_profiles.py

@@ -33,27 +33,13 @@
 output_midnight = np.squeeze(output_midnight)
 output_noon = np.squeeze(output_noon)
 
-variables = [
-    "Total mass density",
-    "N2",
-    "O2",
-    "O",
-    "He",
-    "H",
-    "Ar",
-    "N",
-    "Anomalous O",
-    "NO",
-    "Temperature",
-]
-
 _, ax = plt.subplots()
-for i, label in enumerate(variables):
-    if label in ("NO", "Total mass density", "Temperature"):
-        # There is currently no NO data, also ignore non-number densities
+for variable in msis.Variable:
+    if variable.name in ("Total mass density", "Temperature"):
+        # Ignore non-number densities
         continue
-    (line,) = ax.plot(output_midnight[:, i], alts, linestyle="--")
-    ax.plot(output_noon[:, i], alts, c=line.get_color(), label=label)
+    (line,) = ax.plot(output_midnight[:, variable], alts, linestyle="--")
+    ax.plot(output_noon[:, variable], alts, c=line.get_color(), label=variable.name)
 
 ax.legend(
     loc="upper center", bbox_to_anchor=(0.5, 1.15), fancybox=True, shadow=True, ncol=4

examples/plot_annual_variation.py

@@ -35,26 +35,12 @@
 # Lets get the percent variation from the annual mean for each variable
 variation = 100 * (output / output.mean(axis=0) - 1)
 
-variables = [
-    "Total mass density",
-    "N2",
-    "O2",
-    "O",
-    "He",
-    "H",
-    "Ar",
-    "N",
-    "Anomalous O",
-    "NO",
-    "Temperature",
-]
-
 _, ax = plt.subplots()
-for i, label in enumerate(variables):
-    if label == "NO":
+for variable in msis.Variable:
+    if variable.name == "NO":
         # There is currently no NO data
         continue
-    ax.plot(dates, variation[:, i], label=label)
+    ax.plot(dates, variation[:, variable], label=variable.name)
 
 ax.legend(
     loc="upper center", bbox_to_anchor=(0.5, 1.15), fancybox=True, shadow=True, ncol=5

examples/plot_diurnal_variation.py

@@ -36,26 +36,9 @@
 # Lets get the percent variation from the annual mean for each variable
 variation = 100 * (output / output.mean(axis=0) - 1)
 
-variables = [
-    "Total mass density",
-    "N2",
-    "O2",
-    "O",
-    "He",
-    "H",
-    "Ar",
-    "N",
-    "Anomalous O",
-    "NO",
-    "Temperature",
-]
-
 _, ax = plt.subplots()
-for i, label in enumerate(variables):
-    if label == "NO":
-        # There is currently no NO data
-        continue
-    ax.plot(dates, variation[:, i], label=label)
+for variable in msis.Variable:
+    ax.plot(dates, variation[:, variable], label=variable.name)
 
 ax.legend(
     loc="upper center", bbox_to_anchor=(0.5, 1.15), fancybox=True, shadow=True, ncol=5

examples/plot_surface.py

@@ -28,11 +28,9 @@
 # Get rid of the single dimensions
 output = np.squeeze(output)
 
-i = 2  # O2
-
 _, ax = plt.subplots()
 xx, yy = np.meshgrid(lons, lats)
-mesh = ax.pcolormesh(xx, yy, output[:, :, i].T, shading="auto")
+mesh = ax.pcolormesh(xx, yy, output[:, :, msis.Variable.O2].T, shading="auto")
 plt.colorbar(mesh, label="Number density (/m$^3$)")
 
 ax.set_title(f"O$_2$ number density at {alt} km")

examples/plot_surface_animation.py

@@ -35,21 +35,21 @@
 # Get rid of the single dimensions
 output = np.squeeze(output)
 
-i = 7  # N
-
 fig, (ax_time, ax_mesh) = plt.subplots(
     nrows=2, gridspec_kw={"height_ratios": [1, 4]}, constrained_layout=True
 )
 xx, yy = np.meshgrid(lons, lats)
 vmin, vmax = 1e13, 8e13
 norm = matplotlib.colors.Normalize(vmin, vmax)
-mesh = ax_mesh.pcolormesh(xx, yy, output[0, :, :, i].T, shading="auto", norm=norm)
+mesh = ax_mesh.pcolormesh(
+    xx, yy, output[0, :, :, msis.Variable.N].T, shading="auto", norm=norm
+)
 plt.colorbar(
     mesh, label=f"N number density at {alt} km (/m$^3$)", orientation="horizontal"
 )
 
 time_data = output[:, len(lons) // 2, len(lats) // 2, :]
-ax_time.plot(dates, time_data[:, i], c="k")
+ax_time.plot(dates, time_data[:, msis.Variable.N], c="k")
 ax_time.set_xlim(dates[0], dates[-1])
 ax_time.set_ylim(vmin, vmax)
 ax_time.xaxis.set_major_locator(mdates.HourLocator(interval=3))
@@ -74,7 +74,7 @@ def update_surface(t):
     date_string = dates[t].astype("O").strftime("%H:%M")
     title.set_text(f"{date_string}")
     time_line.set_xdata([dates[t]])
-    mesh.set_array(output[t, :, :, i].T)
+    mesh.set_array(output[t, :, :, msis.Variable.N].T)
     sun.set_data([sun_loc[t]], [0])
 
 

examples/plot_version_diff_altitude.py

@@ -39,27 +39,14 @@
 diff_midnight = np.squeeze(diff_midnight)
 diff_noon = np.squeeze(diff_noon)
 
-variables = [
-    "Total mass density",
-    "N2",
-    "O2",
-    "O",
-    "He",
-    "H",
-    "Ar",
-    "N",
-    "Anomalous O",
-    "NO",
-    "Temperature",
-]
-
 _, ax = plt.subplots()
-for i, label in enumerate(variables):
-    if label in ("NO", "Total mass density", "Temperature"):
-        # There is currently no NO data, also ignore non-number densities
+for variable in msis.Variable:
+    if variable.name in ("NO", "Total mass density", "Temperature"):
+        # There is currently no NO data for earlier versions,
+        # also ignore non-number densities
         continue
-    (line,) = ax.plot(diff_midnight[:, i], alts, linestyle="--")
-    ax.plot(diff_noon[:, i], alts, c=line.get_color(), label=label)
+    (line,) = ax.plot(diff_midnight[:, variable], alts, linestyle="--")
+    ax.plot(diff_noon[:, variable], alts, c=line.get_color(), label=variable.name)
 
 ax.legend(
     loc="upper center", bbox_to_anchor=(0.5, 1.15), fancybox=True, shadow=True, ncol=4

examples/plot_version_diff_surface.py

@@ -32,27 +32,18 @@
 # Get rid of the single dimensions
 diff = np.squeeze(diff)
 
-variables = [
-    "Total mass density",
-    "N2",
-    "O2",
-    "O",
-    "He",
-    "H",
-    "Ar",
-    "N",
-    "Anomalous O",
-    "NO",
-    "Temperature",
-]
-
 fig, axarr = plt.subplots(nrows=3, ncols=3, constrained_layout=True, figsize=(8, 6))
 xx, yy = np.meshgrid(lons, lats)
 norm = mpl.colors.Normalize(-50, 50)
 cmap = mpl.colormaps["RdBu_r"]
-for i, ax in enumerate(axarr.flatten()):
-    mesh = ax.pcolormesh(xx, yy, diff[:, :, i].T, shading="auto", norm=norm, cmap=cmap)
-    ax.set_title(f"{variables[i]}")
+for i, variable in enumerate(msis.Variable):
+    if i > 8:
+        break
+    ax = axarr.flatten()[i]
+    mesh = ax.pcolormesh(
+        xx, yy, diff[:, :, variable].T, shading="auto", norm=norm, cmap=cmap
+    )
+    ax.set_title(f"{variable.name}")
 
 plt.colorbar(
     mesh, ax=axarr, label="Change from MSIS-00 to MSIS2 (%)", orientation="horizontal"

pymsis/msis.py

@@ -1,6 +1,7 @@
 """Interface for running and creating input for the MSIS models."""
 
 import threading
+from enum import IntEnum
 from pathlib import Path
 
 import numpy as np
@@ -21,6 +22,54 @@
     lib._lock = threading.Lock()
 
 
+class Variable(IntEnum):
+    """
+    Enumeration of output data indices for the pymsis run calls.
+
+    This can be used to access the data from the output arrays instead of having
+    to remember the order of the output. For example,
+    ``output_array[..., Variable.MASS_DENSITY]``.
+
+    Attributes
+    ----------
+    MASS_DENSITY
+        Index of total mass density (kg/m3).
+    N2
+        Index of N2 number density (m-3).
+    O2
+        Index of O2 number density (m-3).
+    O
+        Index of O number density (m-3).
+    HE
+        Index of He number density (m-3).
+    H
+        Index of H number density (m-3).
+    AR
+        Index of Ar number density (m-3).
+    N
+        Index of N number density (m-3).
+    ANOMALOUS_O
+        Index of anomalous oxygen number density (m-3).
+    NO
+        Index of NO number density (m-3).
+    TEMPERATURE
+        Index of temperature (K).
+
+    """
+
+    MASS_DENSITY = 0
+    N2 = 1
+    O2 = 2
+    O = 3  # noqa: E741 (ambiguous name)
+    HE = 4
+    H = 5
+    AR = 6
+    N = 7
+    ANOMALOUS_O = 8
+    NO = 9
+    TEMPERATURE = 10
+
+
 def run(
     dates: npt.ArrayLike,
     lons: npt.ArrayLike,
@@ -35,13 +84,25 @@ def run(
     **kwargs: dict,
 ) -> npt.NDArray:
     """
-    Call MSIS looping over all possible inputs.
-
-    If ndates is the same as nlons, nlats, and nalts, then a flattened
-    multi-point input array is assumed. Otherwise, the data
-    will be expanded in a grid-like fashion. The possible
-    return shapes are therefore (ndates, 11) and
-    (ndates, nlons, nlats, nalts, 11).
+    Call MSIS to calculate the atmosphere at the provided input points.
+
+    **Satellite Fly-Through Mode:**
+    If ndates is the same length as nlons, nlats, and nalts, then the
+    input arrays are assumed to be aligned and no regridding is done.
+    This is equivalent to a satellite fly-through, producing an output
+    return shape of (ndates, 11).
+
+    **Grid Mode:**
+    If the input arrays have different lengths the data will be expanded
+    in a grid-like fashion broadcasting to a larger shape than the input
+    arrays. This is equivalent to a full atmosphere simulation where you
+    want to calculate the data at every grid point. The output shape will
+    be 5D (ndates, nlons, nlats, nalts, 11), with potentially single element
+    dimensions if you have a single date, lon, lat, or alt.
+
+    The output array can be indexed with the :class:`~.Variable` enum
+    for easier access. ``output_array[..., Variable.MASS_DENSITY]``
+    returns the total mass density.
 
     Parameters
     ----------
@@ -58,7 +119,7 @@ def run(
     f107as : ArrayLike, optional
         F10.7 running 81-day average centered on the given date(s)
     aps : ArrayLike, optional
-        | Ap for the given date(s), (1-6 only used if `geomagnetic_activity=-1`)
+        | Ap for the given date(s), (1-6 only used if ``geomagnetic_activity=-1``)
         | (0) Daily Ap
         | (1) 3 hr ap index for current time
         | (2) 3 hr ap index for 3 hrs before current time
@@ -75,6 +136,8 @@ def run(
         MSIS version number, one of (0, 2.0, 2.1).
     **kwargs : dict
         Single options for the switches can be defined through keyword arguments.
+        For example, run(..., geomagnetic_activity=-1) will set the geomagnetic
+        activity switch to -1 (storm-time ap mode).
 
     Returns
     -------

pyproject.toml

@@ -86,7 +86,7 @@ lint.select = ["B", "D", "E", "F", "I", "N", "S", "W", "PL", "PT", "UP", "RUF",
 lint.ignore = ["B028", "D203", "D212", "PLR0913", "S310"]
 
 [tool.ruff.lint.per-file-ignores]
-"examples/*" = ["ANN", "D"]
+"examples/*" = ["ANN", "D", "PLR2004"]
 "tests/*" = ["ANN", "D", "S"]
 "tools/*" = ["ANN", "S"]
 ".github/*" = ["S"]

tests/test_msis.py

@@ -520,3 +520,23 @@ def run_function(input_items):
     for future in results:
         result, expected_result = future.result()
         assert_allclose(result, expected_result, rtol=1e-5)
+
+
+def test_output_enum(input_data):
+    # Make sure we can access the output enums
+    assert msis.Variable.MASS_DENSITY == 0
+    assert msis.Variable._member_names_ == [
+        "MASS_DENSITY",
+        "N2",
+        "O2",
+        "O",
+        "HE",
+        "H",
+        "AR",
+        "N",
+        "ANOMALOUS_O",
+        "NO",
+        "TEMPERATURE",
+    ]
+    data = msis.run(*input_data)
+    assert data[..., msis.Variable.MASS_DENSITY] == data[..., 0]