# -*- 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"),
}
[docs]class ScalarCube(iris.cube.Cube):
"""Cube without coordinates."""
[docs] 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
)
[docs] @classmethod
def from_cube(cls, cube):
"""Convert iris cube to ScalarCube."""
return cls(**{k: getattr(cube, k) for k in ["data", "units", "long_name"]})
[docs]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}"
)
[docs]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