Atm flux (#863)

This is a way to turn a netcdf file in array: data-array format into an
atmosphere.

There are also some other changes to make the process smoother.

cmake/modules/ArtsTestcases.cmake

@@ -30,7 +30,7 @@ macro (ARTS_TEST_RUN_PYFILE TESTNAME PYFILE)
     )
   set_tests_properties(
     ${TESTNAME_LONG} PROPERTIES
-    ENVIRONMENT "PYTHONPATH=${ARTS_BINARY_DIR}/python/src;ARTS_HEADLESS=1;ARTS_XML_DATA_DIR=${CMAKE_CURRENT_SOURCE_DIR}/${CFILESUBDIR}${DELIM}${ARTS_XML_DATA_DIR};ARTS_CAT_DATA_DIR=${ARTS_CAT_DATA_DIR}"
+    ENVIRONMENT "PYTHONPATH=${ARTS_BINARY_DIR}/python/src;ARTS_HEADLESS=1;ARTS_XML_DATA_DIR=${ARTS_XML_DATA_DIR};ARTS_CAT_DATA_DIR=${ARTS_CAT_DATA_DIR};ARTS_DATA_PATH=${CMAKE_CURRENT_SOURCE_DIR}/${CFILESUBDIR}"
     DEPENDS python_tests
   )
 endmacro()

examples/CMakeLists.txt

@@ -3,5 +3,6 @@
 ### Python Examples ###
 
 collect_test_subdir(arts-catalogue-data)
+collect_test_subdir(atmosphere)
 collect_test_subdir(getting-started)
 collect_test_subdir(recipes)

examples/arts-catalogue-data/cia/cia.py

@@ -93,7 +93,7 @@
 ws.frequency_grid = pyarts.arts.convert.wavelen2freq(np.linspace(6900e-9, 5900e-9, 1001))
 ws.atmospheric_point.temperature = 295  # At room temperature
 ws.atmospheric_point.pressure = 1e5  # At 1 bar
-ws.atmospheric_point[ws.absorption_species[0]] = 0.21  # At 21% atmospheric Oxygen
+ws.atmospheric_point["O2"] = 0.21  # At 21% atmospheric Oxygen
 ws.ray_path_point # No particular POSLOS
 
 # Call the agenda with inputs above

examples/arts-catalogue-data/lines/lines.py

@@ -142,7 +142,7 @@
 ws.atmospheric_pointInit()
 ws.atmospheric_point.temperature = 295  # At room temperature
 ws.atmospheric_point.pressure = 1e5  # At 1 bar
-ws.atmospheric_point[ws.absorption_species[0]] = (
+ws.atmospheric_point["O2"] = (
     0.21  # At 21% atmospheric Oxygen
 )
 

examples/getting-started/3-disort/3.clearsky-flux-from-dataset.py

@@ -0,0 +1,81 @@
+import pyarts
+import numpy as np
+import xarray as xa
+from dataclasses import dataclass
+import matplotlib.pyplot as plt
+
+if not pyarts.arts.globals.data.is_lgpl:
+    NQuad = 16
+    max_level_step = 1e3
+    atm_latitude = 0.0
+    atm_longitude = 0.0
+    solar_latitude = 0.0
+    solar_longitude = 0.0
+    surface_temperature = 293.0
+    surface_reflectivity = 0.05
+    cutoff = ["ByLine", 750e9]
+    remove_lines_percentile = 70
+    sunfile = "star/Sun/solar_spectrum_QUIET.xml"
+    planet = "Earth"
+
+    xarr = pyarts.data.xarray_open_dataset("atm.nc")
+
+    ws = pyarts.Workspace()
+
+    ws.frequency_grid = pyarts.arts.convert.kaycm2freq(np.linspace(500, 2500, 1001))
+    ws.atmospheric_field = pyarts.data.to_atmospheric_field(xarr)
+
+    v = pyarts.data.to_absorption_species(ws.atmospheric_field)
+
+    ws.absorption_species = v
+    ws.ReadCatalogData(ignore_missing=True)
+    ws.propagation_matrix_agendaAuto(T_extrapolfac=1e9)
+
+    for band in ws.absorption_bands:
+        ws.absorption_bands[band].cutoff = cutoff[0]
+        ws.absorption_bands[band].cutoff_value = cutoff[1]
+
+    ws.absorption_bands.keep_hitran_s(remove_lines_percentile)
+
+    ws.surface_fieldSetPlanetEllipsoid(option=planet)
+
+    sun = pyarts.arts.GriddedField2.fromxml(sunfile)
+
+    ws.surface_field["t"] = surface_temperature
+
+    ws.sunFromGrid(
+        sun_spectrum_raw=sun,
+        latitude=solar_latitude,
+        longitude=solar_longitude,
+    )
+
+    ws.disort_quadrature_dimension = NQuad
+    ws.disort_fourier_mode_dimension = 1
+    ws.disort_legendre_polynomial_dimension = 1
+
+    ws.ray_pathGeometricDownlooking(
+        latitude=atm_latitude,
+        longitude=atm_longitude,
+        max_step=max_level_step,
+    )
+    ws.ray_path_atmospheric_pointFromPath()
+
+    ws.ray_path_frequency_gridFromPath()
+    ws.ray_path_propagation_matrixFromPath()
+    ws.disort_settingsInit()
+    ws.disort_settingsOpticalThicknessFromPath()
+    ws.disort_settingsLayerThermalEmissionLinearInTau()
+    ws.disort_settingsSurfaceEmissionByTemperature(ray_path_point=ws.ray_path[0])
+    ws.disort_settingsCosmicMicrowaveBackgroundRadiation()
+    ws.disort_settingsSurfaceLambertian(value=surface_reflectivity)
+    ws.disort_settingsNoSingleScatteringAlbedo()
+    ws.disort_settingsNoFractionalScattering()
+    ws.disort_settingsNoLegendre()
+    ws.disort_settingsSetSun(ray_path_point=ws.ray_path[-1])
+    ws.disort_spectral_flux_fieldCalc()
+
+    plt.semilogy(pyarts.arts.convert.freq2kaycm(ws.frequency_grid),
+                ws.disort_spectral_flux_field[:, 1])
+
+    f, s = pyarts.plots.AtmField.plot(ws.atmospheric_field,
+                            alts=np.linspace(0, ws.atmospheric_field.top_of_atmosphere))

python/doc/gen_pyarts_arts_autoclass_rst.py

@@ -106,6 +106,8 @@ def is_operator(name):
         "__delitem__",
         "__truediv__",
         "__array__",
+        "__getstate__",
+        "__setstate__",
     ]
     return name in operators
 

python/src/pyarts/data.py

@@ -5,8 +5,10 @@
 import urllib.request
 import zipfile
 from pyarts.arts.globals import parameters
-
+import pyarts
+import numpy as np
 from tqdm import tqdm
+import xarray
 
 
 def download(data=("xml", "cat"), download_dir=None, verbose=False, **kwargs):
@@ -200,3 +202,161 @@ def download_arts_cat_data(download_dir=None, version=None, verbose=False):
             Whether to print info messages. Defaults to False.
     """
     return _download_arts_data("arts-cat-data", download_dir, version, verbose)
+
+
+def to_atmospheric_field(
+    data: xarray.Dataset,
+    remap: None | dict[str, str] = None,
+    ignore: None | list[str] = None,
+    *,
+    atm: None | pyarts.arts.AtmField = None,
+) -> pyarts.arts.AtmField:
+    """
+    Populates a ~pyarts.arts.AtmField from an xarray Dataset-like structure
+
+    Parameters
+    ----------
+    data : xarray.Dataset
+        A dataset.  Coordinates should contain 'alt', 'lat', and 'lon'.
+        All other keys() should be assignable to the atm-field by name.
+    remap : None | dict[str, str], optional
+        All names in this optional dict are renamed when accessing the dataset.
+        For example, if the altitude-grid is called 'Alt' instead of 'alt', the
+        dict should contain {..., 'alt': 'Alt', ...}.
+    ignore : None | list[str], optional
+        Ignore keys listed from assignment into the atmospheric field.
+    atm : None | pyarts.arts.AtmField
+        The default atmospheric field to use.  Defaults to None to use default-
+        constructed atmospheric field object.
+
+    Returns
+    -------
+    atm : pyarts.arts.AtmField
+        An atmospheric field
+
+    """
+    if ignore is None:
+        ignore = []
+
+    if remap is None:
+        remap = {}
+
+    alt = (
+        getattr(data, remap.get("alt", "alt")).data.flatten()
+        if "alt" not in ignore
+        else np.array([0])
+    )
+    lat = (
+        getattr(data, remap.get("lat", "lat")).data.flatten()
+        if "lat" not in ignore
+        else np.array([0])
+    )
+    lon = (
+        getattr(data, remap.get("lon", "lon")).data.flatten()
+        if "lon" not in ignore
+        else np.array([0])
+    )
+
+    GF3 = pyarts.arts.GriddedField3(
+        name="Generic",
+        data=np.zeros((alt.size, lat.size, lon.size)),
+        grid_names=["Altitude", "Latitude", "Longitude"],
+        grids=[alt, lat, lon],
+    )
+
+    if atm is None:
+        atm = pyarts.arts.AtmField()
+        atm.top_of_atmosphere = max(alt)
+
+    for k in data.keys():
+        try:
+            if k in ignore:
+                continue
+
+            k = remap.get(k, k)
+            kstr = str(k)
+
+            GF3.dataname = kstr
+            np.asarray(GF3.data).flat[:] = getattr(data, kstr).data.flat[:]
+
+            atm[k] = GF3
+
+        except Exception as e:
+            raise Exception(f"{e}\n\nFailed to assign key to atm: {k}")
+
+    return atm
+
+
+def to_absorption_species(
+    atm_field: pyarts.arts.AtmField,
+) -> pyarts.arts.ArrayOfArrayOfSpeciesTag:
+    """Scans the ARTS data path for species relevant to the given atmospheric field.
+
+    The scan is done over files in an arts-cat-data like directory structure.
+
+    Args:
+        atm_field (pyarts.arts.AtmField): A relevant atmospheric field.
+
+    Returns:
+        pyarts.arts.ArrayOfArrayOfSpeciesTag: All found species tags.
+        The intent is that this is enough information to use pyarts.workspace.Workspace.ReadCatalogData
+    """
+    species = atm_field.species_keys()
+
+    out = []
+    for spec in species:
+        out.append(f"{spec}")
+
+        if pyarts.arts.file.find_xml(f"xsec/{spec}-XFIT") is not None:
+            out.append(f"{spec}-XFIT")
+
+        for spec2 in species:
+            if pyarts.arts.file.find_xml(f"cia/{spec}-CIA-{spec2}"):
+                out.append(f"{spec}-CIA-{spec2}")
+
+            if pyarts.arts.file.find_xml(f"cia/{spec2}-CIA-{spec}"):
+                out.append(f"{spec2}-CIA-{spec}")
+
+        if pyarts.arts.file.find_xml(f"cia/{spec}-CIA-{spec}") is not None:
+            out.append(f"{spec}-CIA-{spec}")
+
+        if spec == pyarts.arts.SpeciesEnum.Water:
+            out.append("H2O-ForeignContCKDMT400")
+            out.append("H2O-SelfContCKDMT400")
+        elif spec == pyarts.arts.SpeciesEnum.CarbonDioxide:
+            out.append("CO2-CKDMT252")
+        elif spec == pyarts.arts.SpeciesEnum.Oxygen:
+            out.append("O2-SelfContStandardType")
+        elif spec == pyarts.arts.SpeciesEnum.Nitrogen:
+            out.append("N2-SelfContStandardType")
+
+    return pyarts.arts.ArrayOfArrayOfSpeciesTag(np.unique(out))
+
+def xarray_open_dataset(filename_or_obj, *args, **kwargs):
+    """Wraps xarray.open_dataset to search for files in the current and in ARTS' data path.
+
+    All args and kwargs are passed on to xarray.open_dataset directly.  Any FileNotFoundError
+    will be caught and just raised if no path works.
+
+    Args:
+        filename_or_obj (_type_): _description_
+
+    Raises:
+        FileNotFoundError: _description_
+
+    Returns:
+        _type_: _description_
+    """
+
+    try:
+        return xarray.open_dataset(filename_or_obj, *args, **kwargs)
+    except FileNotFoundError:
+        pass
+
+    for p in parameters.datapath:
+        try:
+            return xarray.open_dataset(os.path.join(p, filename_or_obj), *args, **kwargs)
+        except FileNotFoundError:
+            pass
+
+    raise FileNotFoundError(f"File not found in ARTS data path ({parameters.datapath}): {filename_or_obj}")

python/src/pyarts/plots/AtmField.py

@@ -0,0 +1,62 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+def plot(
+    atm_field,
+    *,
+    fig=None,
+    alts=np.linspace(0, 1e5),
+    lats=0,
+    lons=0,
+    ygrid=None,
+    keep_basic=True,
+    keep_specs=True,
+    keep_isots=False,
+    keep_nlte=False,
+    keep_ssprops=True,
+):
+    """Plot the atmospheric field parameters in a default manner.
+
+    Args:
+        atm_field (pyarts.arts.AtmField): An atmospheric field
+        fig (optional): The matplotlib figure to draw on. Defaults to None for new figure.
+        alts (optional): A grid to plot on - must after broadcast with lats and lons be 1D. Defaults to np.linspace(0, 1e5).
+        lats (optional): A grid to plot on - must after broadcast with alts and lons be 1D. Defaults to 0.
+        lons (optional): A grid to plot on - must after broadcast with alts and lats be 1D. Defaults to 0.
+        ygrid (optional): Choice of y-grid for plotting.  Uses broadcasted alts if None. Defaults to None.
+        keep_basic (bool, optional): Forwarded to pyarts.arts.AtmPoint::keys. Defaults to True.
+        keep_specs (bool, optional): Forwarded to pyarts.arts.AtmPoint::keys. Defaults to True.
+        keep_isots (bool, optional): Forwarded to pyarts.arts.AtmPoint::keys. Defaults to False.
+        keep_nlte (bool, optional): Forwarded to pyarts.arts.AtmPoint::keys. Defaults to False.
+        keep_ssprops (bool, optional): Forwarded to pyarts.arts.AtmPoint::keys. Defaults to True.
+
+    Returns:
+        fig: as input or a new figure
+        subs: list of subplots
+    """
+    alts, lats, lons = np.broadcast_arrays(alts, lats, lons)
+    v = atm_field.at(alts, lats, lons)
+
+    keys = v[0].keys(
+        keep_basic=keep_basic,
+        keep_specs=keep_specs,
+        keep_isots=keep_isots,
+        keep_nlte=keep_nlte,
+        keep_ssprops=keep_ssprops,
+    )
+    N = len(keys)
+    n = int(np.ceil(np.sqrt(N))) + 1
+
+    if fig is None:
+        fig = plt.figure(figsize=(5 * n, 5 * n))
+
+    subs = []
+    for i in range(N):
+        subs.append(fig.add_subplot(n, n, i + 1))
+        subs[-1].plot(
+            [x[keys[i]] for x in v],
+            alts if ygrid is None else ygrid,
+            label=keys[i],
+        )
+        subs[-1].legend()
+    return fig, subs

python/src/pyarts/plots/__init__.py

@@ -4,4 +4,6 @@
 from pyarts.plots.ppath import *  # noqa
 from pyarts.plots.ppvar_atm import * # noqa
 
+from . import AtmField
+
 __all__ = [s for s in dir() if not s.startswith("_")]