# -*- coding: utf-8 -*-
"""Main interface to the physical constants store."""
import json
from dataclasses import make_dataclass
from pathlib import Path
from warnings import warn
import iris
import numpy as np
from ..exceptions import AeolusWarning, ArgumentError, LoadError
__all__ = ("init_const", "get_planet_radius")
CONST_DIR = Path(__file__).parent / "store"
[docs]class ScalarCube(iris.cube.Cube):
"""Cube without coordinates."""
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 = ScalarCube(
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."""
derivatives = {
"dry_air_gas_constant": lambda slf: slf.molar_gas_constant
/ slf.dry_air_molecular_weight
}
for name, func in derivatives.items():
try:
cube = ScalarCube.from_cube(func(self))
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 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
except (KeyError, LoadError):
warn("Using default Earth radius", AeolusWarning)
r = default
return r