pycontrails 0.47.3__cp311-cp311-macosx_11_0_arm64.whl → 0.48.1__cp311-cp311-macosx_11_0_arm64.whl

pycontrails/models/accf.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass
 from typing import Any, overload
 
@@ -9,7 +10,7 @@ import xarray as xr
 
 import pycontrails
 from pycontrails.core.flight import Flight
-from pycontrails.core.met import MetDataArray, MetDataset, standardize_variables
+from pycontrails.core.met import MetDataset, standardize_variables
 from pycontrails.core.met_var import (
     AirTemperature,
     EastwardWind,
@@ -23,28 +24,46 @@ from pycontrails.core.vector import GeoVectorDataset
 from pycontrails.datalib import ecmwf
 from pycontrails.utils import dependencies
 
-WideBodyJets = {
-    "A332",
-    "A333",
-    "A338",
-    "A339",
-    "A342",
-    "A343",
-    "A345",
-    "A356",
-    "A359",
-    "A388",
-    "B762",
-    "B763",
-    "B764",
-    "B772",
-    "B773",
-    "B778",
-    "B779",
-    "B788",
-    "B789",
-}
-RegionalJets = {"CRJ1", "CRJ2", "CRJ7", "CRJ8", "CRJ9", "CRJX", "E135", "E145", "E170", "E190"}
+
+def wide_body_jets() -> set[str]:
+    """Return a set of wide body jets."""
+    return {
+        "A332",
+        "A333",
+        "A338",
+        "A339",
+        "A342",
+        "A343",
+        "A345",
+        "A356",
+        "A359",
+        "A388",
+        "B762",
+        "B763",
+        "B764",
+        "B772",
+        "B773",
+        "B778",
+        "B779",
+        "B788",
+        "B789",
+    }
+
+
+def regional_jets() -> set[str]:
+    """Return a set of regional jets."""
+    return {
+        "CRJ1",
+        "CRJ2",
+        "CRJ7",
+        "CRJ8",
+        "CRJ9",
+        "CRJX",
+        "E135",
+        "E145",
+        "E170",
+        "E190",
+    }
 
 
 @dataclass
@@ -102,13 +121,16 @@ class ACCF(Model):
     """Compute Algorithmic Climate Change Functions (ACCF).
 
     This class is a wrapper over the DLR / UMadrid library
-    `climaccf <https://github.com/dlr-pa/climaccf>`__,
-    `DOI: 10.5281/zenodo.6977272 <https://doi.org/10.5281/zenodo.6977272>`__
+    `climaccf <https://github.com/dlr-pa/climaccf>`_,
+    `DOI: 10.5281/zenodo.6977272 <https://doi.org/10.5281/zenodo.6977272>`_
 
     Parameters
     ----------
     met : MetDataset
         Dataset containing "air_temperature" and "specific_humidity" variables
+    surface : MetDataset, optional
+        Dataset containing "surface_solar_downward_radiation" and
+        "top_net_thermal_radiation" variables
 
     References
     ----------
@@ -132,47 +154,32 @@ class ACCF(Model):
     sur_variables = (ecmwf.SurfaceSolarDownwardRadiation, ecmwf.TopNetThermalRadiation)
     default_params = ACCFParams
 
-    short_vars = [v.short_name for v in met_variables + sur_variables]
+    short_vars = {v.short_name for v in (*met_variables, *sur_variables)}
 
-    ds_met: xr.Dataset | None
-    ds_sur: xr.Dataset | None
+    # This variable won't get used since we are not writing the output
+    # anywhere, but the library will complain if it's not defined
+    path_lib = "./"
 
     def __init__(
         self,
         met: MetDataset,
         surface: MetDataset | None = None,
-        params: dict[str, Any] = {},
+        params: dict[str, Any] | None = None,
         **params_kwargs: Any,
     ) -> None:
         # Normalize ECMWF variables
         met = standardize_variables(met, self.met_variables)
 
-        if surface:
-            surface = standardize_variables(surface, self.sur_variables)
+        # Ignore humidity scaling warning
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", module="pycontrails.core.models")
+            super().__init__(met, params=params, **params_kwargs)
 
-        # Surpress warning about humdity scaling because that should be eet
-        # using ACCF config variables for this model
-        try:
-            del met.attrs["history"]
-        except KeyError:
-            pass
-
-        source = met.attrs["met_source"]
-        met.attrs["met_source"] = "not_ecmwf"
-        super().__init__(met, params=params, **params_kwargs)
-        if self.met:
-            self.met.attrs["met_source"] = source
-
-        self._update_accf_config()
         if surface:
-            self.surface = surface.copy()
-
-        # This variable won't get used since we are not writing the output
-        # anywhere, but the library will complain if it's not defined
-        self.path_lib = "./"
-
-        self.ds_met = None
-        self.ds_sur = None
+            surface = surface.copy()
+            surface = standardize_variables(surface, self.sur_variables)
+            surface.data = _rad_instantaneous_to_accumulated(surface.data)
+            self.surface = surface
 
     @overload
     def eval(self, source: Flight, **params: Any) -> Flight: ...
@@ -181,11 +188,11 @@ class ACCF(Model):
     def eval(self, source: GeoVectorDataset, **params: Any) -> GeoVectorDataset: ...
 
     @overload
-    def eval(self, source: MetDataset | None = ..., **params: Any) -> MetDataArray: ...
+    def eval(self, source: MetDataset | None = ..., **params: Any) -> MetDataset: ...
 
     def eval(
         self, source: GeoVectorDataset | Flight | MetDataset | None = None, **params: Any
-    ) -> GeoVectorDataset | Flight | MetDataArray:
+    ) -> GeoVectorDataset | Flight | MetDataset:
         """Evaluate accfs along flight trajectory or on meteorology grid.
 
         Parameters
@@ -225,32 +232,38 @@ class ACCF(Model):
                 self.surface = self.source.downselect_met(self.surface)
 
         if isinstance(self.source, MetDataset):
-            if self.source["longitude"].size > 1:
-                hres = abs(self.source["longitude"].data[1] - self.source["longitude"].data[0])
-                self.params["horizontal_resolution"] = float(hres)
-            elif self.source["latitude"].size > 1:
-                hres = abs(self.source["latitude"].data[1] - self.source["latitude"].data[0])
+            # Overwrite horizontal resolution to match met
+            longitude = self.source.data["longitude"].values
+            if longitude.size > 1:
+                hres = abs(longitude[1] - longitude[0])
                 self.params["horizontal_resolution"] = float(hres)
 
+            else:
+                latitude = self.source.data["latitude"].values
+                if latitude.size > 1:
+                    hres = abs(latitude[1] - latitude[0])
+                    self.params["horizontal_resolution"] = float(hres)
+
+        p_settings = _get_accf_config(self.params)
+
         self.set_source_met()
-        self._update_accf_config()
-        self._generate_weather_store()
+        self._generate_weather_store(p_settings)
 
         # check aircraft type and set in config if needed
         if self.params["nox_ei"] != "TTV":
             if isinstance(self.source, Flight):
                 ac = self.source.attrs["aircraft_type"]
-                if ac in WideBodyJets:
-                    self.p_settings["ac_type"] = "wide-body"
-                elif ac in RegionalJets:
-                    self.p_settings["ac_type"] = "regional"
+                if ac in wide_body_jets():
+                    p_settings["ac_type"] = "wide-body"
+                elif ac in regional_jets():
+                    p_settings["ac_type"] = "regional"
                 else:
-                    self.p_settings["ac_type"] = "single-aisle"
+                    p_settings["ac_type"] = "single-aisle"
             else:
-                self.p_settings["ac_type"] = "wide-body"
+                p_settings["ac_type"] = "wide-body"
 
         clim_imp = GeTaCCFs(self)
-        clim_imp.get_accfs(**self.p_settings)
+        clim_imp.get_accfs(**p_settings)
         aCCFs, _ = clim_imp.get_xarray()
 
         # assign ACCF outputs to source
@@ -259,66 +272,53 @@ class ACCF(Model):
             # skip met variables
             if key in self.short_vars:
                 continue
-            if not isinstance(key, str):
-                continue
 
+            assert isinstance(key, str)
             if isinstance(self.source, GeoVectorDataset):
                 self.source[key] = self.source.intersect_met(maCCFs[key])
             else:
-                self.source[key] = (maCCFs.dim_order, arr.data)
-
-            # Tag output with additional attrs when source is MetDataset
-            if isinstance(self.source, MetDataset):
-                attrs: dict[str, Any] = {
-                    "description": self.long_name,
-                    "pycontrails_version": pycontrails.__version__,
-                }
-                if self.met is not None:
-                    attrs["met_source"] = self.met.attrs.get("met_source", "unknown")
+                self.source[key] = arr
 
-                self.source[key].data.attrs.update(attrs)
+        self.transfer_met_source_attrs()
+        self.source.attrs["pycontrails_version"] = pycontrails.__version__
 
-        return self.source  # type: ignore[return-value]
+        return self.source
 
-    def _generate_weather_store(self) -> None:
+    def _generate_weather_store(self, p_settings: dict[str, Any]) -> None:
         from climaccf.weather_store import WeatherStore
 
         # The library does not call the coordinates by name, it just slices the
         # underlying data array, so we need to put them in the expected order.
         # It also needs variables to have the ECMWF short name
         if isinstance(self.met, MetDataset):
-            mt = self.met.data.transpose("time", "level", "latitude", "longitude")
-            if mt is None or isinstance(mt, xr.Dataset):
-                self.ds_met = mt
-
-        if self.ds_met:
-            for var in self.ds_met.data_vars:
-                matching_variable = [v for v in self.met_variables if var == v.standard_name]
-                if matching_variable:
-                    self.ds_met = self.ds_met.rename({var: matching_variable[0].short_name})
+            ds_met = self.met.data.transpose("time", "level", "latitude", "longitude")
+            name_dict = {v.standard_name: v.short_name for v in self.met_variables}
+            ds_met = ds_met.rename(name_dict)
+        else:
+            ds_met = None
 
         if hasattr(self, "surface"):
-            self.ds_sur = self.surface.data.squeeze().transpose("time", "latitude", "longitude")
-            for var in self.ds_sur.data_vars:
-                matching_variable = [v for v in self.sur_variables if var == v.standard_name]
-                if matching_variable:
-                    self.ds_sur = self.ds_sur.rename({var: matching_variable[0].short_name})
+            ds_sur = self.surface.data.squeeze().transpose("time", "latitude", "longitude")
+            name_dict = {v.standard_name: v.short_name for v in self.sur_variables}
+            ds_sur = ds_sur.rename(name_dict)
         else:
-            self.ds_sur = None
+            ds_sur = None
 
         ws = WeatherStore(
-            self.ds_met,
-            self.ds_sur,
-            ll_resolution=self.p_settings["horizontal_resolution"],
-            forecast_step=self.p_settings["forecast_step"],
+            ds_met,
+            ds_sur,
+            ll_resolution=p_settings["horizontal_resolution"],
+            forecast_step=p_settings["forecast_step"],
         )
-        if self.p_settings["lat_bound"] and self.p_settings["lon_bound"]:
+
+        if p_settings["lat_bound"] and p_settings["lon_bound"]:
             ws.reduce_domain(
                 {
-                    "latitude": self.p_settings["lat_bound"],
-                    "longitude": self.p_settings["lon_bound"],
+                    "latitude": p_settings["lat_bound"],
+                    "longitude": p_settings["lon_bound"],
                 }
             )
+
         self.ds = ws.get_xarray()
         self.variable_names = ws.variable_names
         self.pre_variable_names = ws.pre_variable_names
@@ -329,51 +329,78 @@ class ACCF(Model):
         self.axes = ws.axes
         self.var_xr = ws.var_xr
 
-    def _update_accf_config(self) -> None:
-        # a good portion of these will get ignored since we are not producing an
-        # output file, but the library will complain if they aren't defined
-        self.p_settings = {
-            "lat_bound": self.params["lat_bound"],
-            "lon_bound": self.params["lon_bound"],
-            "time_bound": None,
-            "horizontal_resolution": self.params["horizontal_resolution"],
-            "forecast_step": self.params["forecast_step"],
-            "NOx_aCCF": True,
-            "NOx&inverse_EIs": self.params["nox_ei"],
-            "output_format": "netCDF",
-            "mean": False,
-            "std": False,
-            "merged": self.params["merged"],
-            "aCCF-V": self.params["accf_v"],
-            "efficacy": self.params["efficacy"],
-            "efficacy-option": self.params["efficacy_option"],
-            "emission_scenario": self.params["emission_scenario"],
-            "climate_indicator": self.params["climate_indicator"],
-            "TimeHorizon": self.params["time_horizon"],
-            "ac_type": "wide-body",
-            "sep_ri_rw": self.params["sep_ri_rw"],
-            "PMO": self.params["PMO"],
-            "aCCF-scalingF": {
-                "CH4": self.params["ch4_scaling"],
-                "CO2": self.params["co2_scaling"],
-                "Cont.": self.params["cont_scaling"],
-                "H2O": self.params["h2o_scaling"],
-                "O3": self.params["o3_scaling"],
-            },
-            "PCFA": self.params["pfca"],
-            "PCFA-ISSR": {
-                "rhi_threshold": self.params["issr_rhi_threshold"],
-                "temp_threshold": self.params["issr_temp_threshold"],
-            },
-            "PCFA-SAC": {
-                "EI_H2O": self.params["sac_ei_h2o"],
-                "Q": self.params["sac_q"],
-                "eta": self.params["sac_eta"],
-            },
-            "Chotspots": False,
-            "hotspots_binary": True,
-            "color": "Reds",
-            "geojson": False,
-            "save_path": "./",
-            "save_format": "netCDF",
-        }
+
+def _get_accf_config(params: dict[str, Any]) -> dict[str, Any]:
+    # a good portion of these will get ignored since we are not producing an
+    # output file, but the library will complain if they aren't defined
+    return {
+        "lat_bound": params["lat_bound"],
+        "lon_bound": params["lon_bound"],
+        "time_bound": None,
+        "horizontal_resolution": params["horizontal_resolution"],
+        "forecast_step": params["forecast_step"],
+        "NOx_aCCF": True,
+        "NOx&inverse_EIs": params["nox_ei"],
+        "output_format": "netCDF",
+        "mean": False,
+        "std": False,
+        "merged": params["merged"],
+        "aCCF-V": params["accf_v"],
+        "efficacy": params["efficacy"],
+        "efficacy-option": params["efficacy_option"],
+        "emission_scenario": params["emission_scenario"],
+        "climate_indicator": params["climate_indicator"],
+        "TimeHorizon": params["time_horizon"],
+        "ac_type": "wide-body",
+        "sep_ri_rw": params["sep_ri_rw"],
+        "PMO": params["PMO"],
+        "aCCF-scalingF": {
+            "CH4": params["ch4_scaling"],
+            "CO2": params["co2_scaling"],
+            "Cont.": params["cont_scaling"],
+            "H2O": params["h2o_scaling"],
+            "O3": params["o3_scaling"],
+        },
+        "PCFA": params["pfca"],
+        "PCFA-ISSR": {
+            "rhi_threshold": params["issr_rhi_threshold"],
+            "temp_threshold": params["issr_temp_threshold"],
+        },
+        "PCFA-SAC": {
+            "EI_H2O": params["sac_ei_h2o"],
+            "Q": params["sac_q"],
+            "eta": params["sac_eta"],
+        },
+        "Chotspots": False,
+        "hotspots_binary": True,
+        "color": "Reds",
+        "geojson": False,
+        "save_path": "./",
+        "save_format": "netCDF",
+    }
+
+
+def _rad_instantaneous_to_accumulated(ds: xr.Dataset) -> xr.Dataset:
+    """Convert instantaneous radiation to accumulated radiation."""
+
+    for name, da in ds.items():
+        try:
+            unit = da.attrs["units"]
+        except KeyError as e:
+            msg = (
+                f"Radiation data contains '{name}' variable "
+                "but units are not specified. Provide units in the "
+                f"rad['{name}'].attrs passed into ACCF."
+            )
+            raise KeyError(msg) from e
+
+        if unit == "J m**-2":
+            continue
+        if unit != "W m**-2":
+            msg = f"Unexpected units '{unit}' for '{name}'. Expected 'J m**-2' or 'W m**-2'."
+            raise ValueError(msg)
+
+        # Convert from W m**-2 to J m**-2
+        ds[name] = da.assign_attrs(units="J m**-2") * 3600.0
+
+    return ds