aeolus.core — aeolus 0.4.9+9.g5a71939 documentation

# Source code for aeolus.core

    """Core submodule of aeolus package."""
    from cached_property import cached_property

    import iris
    from iris.exceptions import ConstraintMismatchError as ConMisErr

    from .calc import diag
    from .calc.meta import copy_doc
    from .const import add_planet_conf_to_cubes, init_const
    from .coord import CoordContainer
    from .exceptions import _warn
    from .io import load_data, save_cubelist
    from .model import um
    from .region import Region
    from .subset import DimConstr

    __all__ = (
        "AtmoSim",
        "AtmoSimBase",
        "Run",
    )

    class AtmoSimBase:
        """
        Base class for creating atmospheric model simulation classes in aeolus.

        Used to store and calculate atmospheric fields from gridded model output.
        Derived quantities are stored as cached properties to save computational time.

        Assumes the data are in spherical coordinates on a regular longitude-latitude grid.

        Attributes
        ----------
        name: str
            The run's name.
        description: str
            A description of the run.
        const: aeolus.const.ConstContainer
            Physical constants used in calculations for this run.
        model: aeolus.model.Model, optional
            Model class with relevant coordinate and variable names.
        """

        def __init__(
            self,
            cubes=None,
            name="",
            description="",
            planet="",
            const_dir=None,
            model=um,
            model_type=None,
            timestep=None,
            vert_coord=None,
        ):
            """
            Instantiate an `AtmoSimBase` object.

            Parameters
            ----------
            cubes: iris.cube.CubeList
                Atmospheric fields.
            name: str, optional
                The name or label of this `AtmoSim`.
            description: str, optional
                This is not used internally; it is solely for the user's information.
            planet: str, optional
                Planet configuration. This is used to get appropriate physical constants.
                If not given, Earth physical constants are initialised.
            const_dir: pathlib.Path, optional
                Path to a folder with files containing constants for a specific planet.
            model: aeolus.model.Model, optional
                Model class with relevant coordinate and variable names.
            model_type: str, optional
                Type of the model run, global or LAM.
            timestep: int, optional
                Model time step in s.
            vert_coord: str, optional
                Character identificator for the type of the model's vertical coordinate.
                "z" - data on "level_height"
                "p" - data on pressure levels

            See also
            --------
            aeolus.const.init_const, aeolus.core.Run
            """
            self._cubes = cubes
            self.name = name
            self.description = description

            # Planetary constants
            self._update_planet(planet=planet, const_dir=const_dir)
            self._add_planet_conf_to_cubes()  # TODO: remove?

            # Model-specific names of variables and coordinates
            self.model = model
            self.dim_constr = DimConstr(model=self.model)

            # If the model is global or LAM (nested) and what its driving model is
            self.model_type = model_type
            self.timestep = timestep

            # Domain
            try:
                cube_yx = self._cubes.extract(self.dim_constr.relax.yx)[0]
                self.domain = Region.from_cube(cube_yx, name=f"{name}_domain", shift_lons=True)
            except IndexError:
                _warn("Initialised without a domain.")

            # Variables as attributes
            self._assign_fields()

            # Common coordinates
            dim_seq = ["tyx", "yx"]
            self.vert_coord = vert_coord
            # TODO: make it more flexible
            if self.vert_coord == "z":
                self.coord = CoordContainer(self._cubes.extract(self.dim_constr.relax.z))
            elif self.vert_coord == "p":
                self.coord = CoordContainer(self._cubes.extract(self.dim_constr.relax.p))
            else:
                self.coord = CoordContainer(self._cubes.extract(self.dim_constr.relax.yx))

            if self.vert_coord is not None:
                dim_seq += [f"t{self.vert_coord}yx", f"{self.vert_coord}yx"]
            for seq in dim_seq:
                try:
                    setattr(
                        self,
                        f"_ref_{seq}",
                        self._cubes.extract(getattr(self.dim_constr.strict, seq))[0],
                    )
                except IndexError:
                    pass

        @classmethod
        def from_parent_class(cls, obj):
            """Dynamically inherit from a similar class."""
            new_obj = cls(
                cubes=obj._cubes,
                name=obj.name,
                description=obj.description,
                planet=obj.planet,
                const_dir=obj.const_dir,
                model=obj.model,
                model_type=obj.model_type,
                timestep=obj.timestep,
                vert_coord=obj.vert_coord,
            )
            return new_obj

        def __getitem__(self, key):
            """Redirect self[key] to self.key."""
            return self.__getattribute__(key)

        def _assign_fields(self):
            """Assign input cubelist items as attributes of this class."""
            self.vars = []
            kwargs = {}
            for key in self.model.__dataclass_fields__:
                try:
                    kwargs[key] = self._cubes.extract_cube(getattr(self.model, key))
                    self.vars.append(key)
                except ConMisErr:
                    pass
            self.__dict__.update(**kwargs)
            del kwargs, key

        def _update_planet(self, planet="", const_dir=None):
            """Add or update planetary constants."""
            self.planet = planet
            self.const_dir = const_dir
            self.const = init_const(self.planet, directory=self.const_dir)
            try:
                self.const.radius.convert_units("m")
                self._coord_system = iris.coord_systems.GeogCS(semi_major_axis=self.const.radius.data)
            except AttributeError:
                self._coord_system = None
                _warn("Run initialised without a coordinate system.")

        def _add_planet_conf_to_cubes(self):
            """Add or update planetary constants container to cube attributes."""
            add_planet_conf_to_cubes(self._cubes, self.const)

        def extract(self, constraints):
            """
            Subset `AtmoSim` using iris constraints.

            Parameters
            ----------
            constraints: iris.Constraint or iterable of constraints
                A single constraint or an iterable.
            """
            new_obj = self.__class__(
                cubes=self._cubes.extract(constraints),
                name=self.name,
                description=self.description,
                planet=self.planet,
                const_dir=self.const_dir,
                model=self.model,
                model_type=self.model_type,
                timestep=self.timestep,
                vert_coord=self.vert_coord,
            )
            return new_obj

    class AtmoSim(AtmoSimBase):
        """
        Main class for dealing with a atmospheric model simulation output in aeolus.

        Used to store and calculate atmospheric fields from gridded model output.
        Derived quantities are stored as cached properties to save computational time.
        """

        @cached_property
        @copy_doc(diag.wind_speed)
        def sigma_p(self):
            # TODO: pass the cube only?
            return diag.sigma_p(self._cubes, const=self.const, model=self.model)

        @cached_property
        @copy_doc(diag.wind_speed)
        def wind_speed(self):
            cmpnts = []
            for cmpnt_key in ["u", "v", "w"]:
                try:
                    cmpnts.append(self[cmpnt_key])
                except AttributeError:
                    pass
            return diag.wind_speed(*cmpnts)

        @cached_property
        @copy_doc(diag.toa_net_energy)
        def toa_net_energy(self):
            return diag.toa_net_energy(self._cubes, model=self.model)

        @cached_property
        @copy_doc(diag.sfc_water_balance)
        def sfc_water_balance(self):
            return diag.sfc_water_balance(self._cubes, const=self.const, model=self.model)

    class Run:
        """
        A single model 'run', i.e. simulation.

        Attributes
        ----------
        name: str
            The run's name.
        const: aeolus.const.ConstContainer
            Physical constants used in calculations for this run.
        model: aeolus.model.Model, optional
            Model class with relevant coordinate names.
        """

        attr_keys = ["name", "planet", "model_type", "timestep"]

        def __init__(
            self,
            files=None,
            name="",
            planet="",
            const_dir=None,
            model=um,
            model_type=None,
            timestep=None,
            processed=False,
        ):
            """
            Instantiate a `Run` object.

            Parameters
            ----------
            files: str or pathlib.Path, optional
                Wildcard for loading files.
            name: str, optional
                The run's name.
            planet: str, optional
                Planet configuration. This is used to get appropriate physical constants.
                If not given, Earth physical constants are initialised.
            const_dir: pathlib.Path, optional
                Path to a folder with JSON files containing constants for a specific planet.
            model: aeolus.model.Model, optional
                Model class with relevant coordinate and variable names.
            model_type: str, optional
                Type of the model run, global or LAM.
            timestep: int, optional
                Model time step in s.
            parent: aeolus.core.Run, optional
                Pointer to this run's driving model if this is a LAM-type simulation.
            children: list, optional
                List of `aeolus.core.Run` objects if this is a driving model.
            processed: bool, optional
                If True, data from `files` is assigned to `proc` attribute.

            See also
            --------
            aeolus.const.init_const
            """
            _warn(
                "Run is deprecated and will be removed in the next release. "
                "Use iris.cube.CubeList instead."
            )
            self.name = name

            # Planetary constants
            self._update_planet(planet=planet, const_dir=const_dir)

            # Model-specific names of variables and coordinates
            self.model = model
            self.dim_constr = DimConstr(model=self.model)

            # If the model is global or LAM (nested) and what its driving model is
            self.model_type = model_type
            self.timestep = timestep
            self.processed = processed

            if files:
                self.load_data(files)
                try:
                    if self.processed:
                        cube_yx = self.proc.extract(self.dim_constr.relax.yx)[0]
                    else:
                        cube_yx = self.raw.extract(self.dim_constr.relax.yx)[0]
                    self.domain = Region.from_cube(cube_yx, name=f"{name}_domain", shift_lons=True)
                except IndexError:
                    _warn("Run initialised without a domain.")
            else:
                _warn("Run initialised without input files")

        def load_data(self, files):
            """Load cubes."""
            if self.processed:
                self.proc = load_data(files)
                self._add_planet_conf_to_cubes()
            else:
                self.raw = load_data(files)

        def _update_planet(self, planet="", const_dir=None):
            """Add or update planetary constants."""
            self.planet = planet
            self.const = init_const(planet, directory=const_dir)
            try:
                self.const.radius.convert_units("m")
                self._coord_system = iris.coord_systems.GeogCS(semi_major_axis=self.const.radius.data)
            except AttributeError:
                self._coord_system = None
                _warn("Run initialised without a coordinate system.")

        def _add_planet_conf_to_cubes(self):
            """Add or update planetary constants container to cube attributes."""
            add_planet_conf_to_cubes(self._cubes, self.const)

        def proc_data(self, func=None, **func_args):
            """
            Post-process data for easier analysis and store it in `self.proc` attribute.

            Parameters
            ----------
            func: callable
                Function that takes `iris.cube.CubeList` as its first argument.
            **func_args: dict-like, optional
                Keyword arguments passed to `func`.
            """
            if self.processed:
                _warn("Run's data is already processed. Skipping.")
            else:
                self.proc = iris.cube.CubeList()
                if callable(func):
                    self.proc = func(self.raw, **func_args)
                self._add_planet_conf_to_cubes()
                self.processed = True

        def add_data(self, func=None, **func_args):
            """
            Calculate additional diagnostics (of type `iris.cube.Cube`) and add them to `self.proc`.

            Parameters
            ----------
            func: callable
                Function that takes `iris.cube.CubeList` (`self.proc`) as its first argument
                and appends new cubes to it (and does not return anything).
            **func_args: dict-like, optional
                Keyword arguments passed to `func`.
            """
            if callable(func):
                func(self.proc, **func_args)

        def to_file(self, path):
            """
            Save `proc` cubelist to a file with appropriate metadata.

            Parameters
            ----------
            path: str or pathlib.Path
                File path.
            """
            run_attrs = {}
            for key in self.attr_keys:
                if getattr(self, key):
                    run_attrs[key] = str(getattr(self, key))
            save_cubelist(self.proc, path, **run_attrs)