aeolus.const.const — aeolus 0.4.9+9.g5a71939 documentation

# Source code for aeolus.const.const

    # -*- coding: utf-8 -*-
    """Main interface to the physical constants store."""
    import json
    from dataclasses import make_dataclass
    from pathlib import Path

    import iris

    import numpy as np

    from ..exceptions import ArgumentError, LoadError, _warn

    __all__ = ("add_planet_conf_to_cubes", "get_planet_radius", "init_const")

    CONST_DIR = Path(__file__).parent / "store"

    DERIVED_CONST = {
        "dry_air_gas_constant": (
            lambda slf: slf.molar_gas_constant / slf.dry_air_molecular_weight,
            "J kg-1 K-1",
        ),
        "molecular_weight_ratio": (
            lambda slf: slf.condensible_molecular_weight / slf.dry_air_molecular_weight,
            "1",
        ),
        "poisson_exponent": (
            lambda slf: slf.dry_air_gas_constant / slf.dry_air_spec_heat_press,
            "1",
        ),
        "planet_rotation_rate": (lambda slf: (slf.day / (2 * np.pi)) ** (-1), "s-1"),
    }

    class ScalarCube(iris.cube.Cube):
        """Cube without coordinates."""

        def __init__(self, *args, **kw):
            """Initialise aeolus.const.const.ScalarCube."""
            _warn(
                "ScalarCube is deprecated and will be removed in the next release. "
                "Use iris.cube.Cube instead.",
            )
            super(ScalarCube, self).__init__(*args, **kw)

        def __repr__(self):
            """Repr of this class."""
            return f"<ScalarCube of {self.long_name} [{self.units}]>"

        def __deepcopy__(self, memo):
            """Deep copy of a scalar cube."""
            return self.from_cube(self._deepcopy(memo))

        @property
        def asc(self):
            """Convert cube to AuxCoord for math ops."""
            return iris.coords.AuxCoord(
                np.asarray(self.data), units=self.units, long_name=self.long_name
            )

        @classmethod
        def from_cube(cls, cube):
            """Convert iris cube to ScalarCube."""
            return cls(**{k: getattr(cube, k) for k in ["data", "units", "long_name"]})

    class ConstContainer:
        """Base class for creating dataclasses and storing planetary constants."""

        def __repr__(self):
            """Create custom repr."""
            cubes_str = ", ".join(
                [
                    f"{getattr(self, _field).long_name} [{getattr(self, _field).units}]"
                    for _field in self.__dataclass_fields__
                ]
            )
            return f"{self.__class__.__name__}({cubes_str})"

        def __post_init__(self):
            """Do things automatically after __init__()."""
            self._convert_to_iris_cubes()
            self._derive_const()

        def _convert_to_iris_cubes(self):
            """Loop through fields and convert each of them to `iris.cube.Cube`."""
            for name in self.__dataclass_fields__:
                _field = getattr(self, name)
                cube = iris.cube.Cube(
                    data=_field.get("value"), units=_field.get("units", 1), long_name=name
                )
                object.__setattr__(self, name, cube)

        def _derive_const(self):
            """Not fully implemented yet."""
            for name, recipe in DERIVED_CONST.items():
                func, units = recipe
                try:
                    cube = func(self)
                    cube.convert_units(units)
                    cube.rename(name)
                    object.__setattr__(self, name, cube)
                except AttributeError:
                    pass

    def _read_const_file(name, directory=CONST_DIR):
        """Read constants from the JSON file."""
        if not isinstance(directory, Path):
            raise ArgumentError("directory must be a pathlib.Path object")
        try:
            with (directory / name).with_suffix(".json").open("r") as fp:
                list_of_dicts = json.load(fp)
            # transform the list of dictionaries into a dictionary
            const_dict = {}
            for vardict in list_of_dicts:
                const_dict[vardict["name"]] = {k: v for k, v in vardict.items() if k != "name"}
            return const_dict
        except FileNotFoundError:
            raise LoadError(
                f"JSON file for {name} configuration not found, check the directory: {directory}"
            )

    def init_const(name="general", directory=None):
        """
        Create a dataclass with a given set of constants.

        Parameters
        ----------
        name: str, optional
            Name of the constants set.
            Should be identical to the JSON file name (without the .json extension).
            If not given, only general physical constants are returned.
        directory: pathlib.Path, optional
            Path to a folder with JSON files containing constants for a specific planet.

        Returns
        -------
        Dataclass with constants as iris cubes.

        Examples
        --------
        >>> c = init_const('earth')
        >>> c
        EarthConstants(gravity [m s-2], radius [m], day [s], solar_constant [W m-2], ...)
        >>> c.gravity
        <iris 'Cube' of gravity / (m s-2) (scalar cube)>
        """
        cls_name = f"{name.capitalize()}Constants"
        if directory is None:
            # use default directory
            kw = {}
        else:
            kw = {"directory": directory}
        # transform the list of dictionaries into a dictionary
        const_dict = _read_const_file("general")  # TODO: make this more flexible?
        if name != "general":
            const_dict.update(_read_const_file(name, **kw))
        kls = make_dataclass(
            cls_name,
            fields=[*const_dict.keys()],
            bases=(ConstContainer,),
            frozen=True,
            repr=False,
        )
        return kls(**const_dict)

    def get_planet_radius(cube, default=iris.fileformats.pp.EARTH_RADIUS):
        """Get planet radius in metres from cube attributes or coordinate system."""
        cs = cube.coord_system("CoordSystem")
        if cs is not None:
            r = cs.semi_major_axis
        else:
            try:
                r = cube.attributes["planet_conf"].radius.copy()
                r.convert_units("m")
                r = float(r.data)
            except (KeyError, LoadError):
                _warn("Using default radius")
                r = default
        return r

    def add_planet_conf_to_cubes(cubelist, const):
        """
        Add constants container to the cube attributes and replace its coordinate system.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            List of cubes containing a cube of zonal velocity (u).
        const: aeolus.const.const.ConstContainer, optional
            Constainer with the relevant planetary constants.
        """
        const.radius.convert_units("m")
        _coord_system = iris.coord_systems.GeogCS(semi_major_axis=const.radius.data)
        for cube in cubelist:
            # add constants to cube attributes
            cube.attributes["planet_conf"] = const
            for coord in cube.coords():
                if coord.coord_system:
                    # Replace coordinate system with the planet radius given in `self.const`
                    coord.coord_system = _coord_system