"""The following is a list of abstract classes that are used to define the interface for all models."""
from abc import abstractmethod
from typing import Generic, TypeVar
from flax import nnx
from jax import Array
from simple_pytree import Pytree
[docs]
class AbstractState(Pytree):
"""Abstract state class to define the interface for all states."""
StateT = TypeVar("StateT", bound=AbstractState)
[docs]
class AbstractRadiationState(AbstractState):
"""Abstract rad state."""
net_rad: Array
"""Net surface rad [W m-2]."""
in_srad: Array
"""Incoming solar rad [W m-2]."""
out_srad: Array
"""Outgoing solar rad [W m-2]."""
in_lrad: Array
"""Incoming longwave rad [W m-2]."""
out_lrad: Array
"""Outgoing longwave rad [W m-2]."""
[docs]
class AbstractLandState(AbstractState):
"""Abstract land state."""
@property
@abstractmethod
def alpha(self) -> Array:
"""surface albedo [-], range 0 to 1."""
raise NotImplementedError
@property
@abstractmethod
def surf_temp(self) -> Array:
"""Surface temperature [K]."""
raise NotImplementedError
@property
@abstractmethod
def rs(self) -> Array:
"""Surface resistance [s m-1]."""
raise NotImplementedError
@property
@abstractmethod
def wg(self) -> Array:
"""No moisture content in the root zone [m3 m-3]."""
raise NotImplementedError
@property
@abstractmethod
def wl(self) -> Array:
"""No water content in the canopy [m]."""
raise NotImplementedError
@property
@abstractmethod
def esat(self) -> Array:
"""Saturation vapor pressure [Pa]."""
raise NotImplementedError
@property
@abstractmethod
def qsat(self) -> Array:
"""Saturation specific humidity [kg/kg]."""
raise NotImplementedError
@property
@abstractmethod
def dqsatdT(self) -> Array:
"""Derivative of saturation specific humidity with respect to temperature [kg/kg/K]."""
raise NotImplementedError
@property
@abstractmethod
def e(self) -> Array:
"""Vapor pressure [Pa]."""
raise NotImplementedError
@property
@abstractmethod
def qsatsurf(self) -> Array:
"""Saturation specific humidity at surface temperature [kg/kg]."""
raise NotImplementedError
@property
@abstractmethod
def wtheta(self) -> Array:
"""Kinematic heat flux [K m/s]."""
raise NotImplementedError
@property
@abstractmethod
def wq(self) -> Array:
"""Kinematic moisture flux [kg/kg m/s]."""
raise NotImplementedError
@property
@abstractmethod
def wCO2(self) -> Array:
"""Kinematic CO2 flux [kg/kg m/s] or [mol m-2 s-1]."""
raise NotImplementedError
[docs]
class AbstractAtmosphereState(AbstractState):
"""Abstract atmos state."""
@property
@abstractmethod
def theta(self) -> Array:
"""Potential temperature [K]."""
raise NotImplementedError
@property
@abstractmethod
def q(self) -> Array:
"""Specific humidity [kg/kg]."""
raise NotImplementedError
@property
@abstractmethod
def co2(self) -> Array:
"""CO2 concentration [ppmv]."""
raise NotImplementedError
@property
@abstractmethod
def surf_pressure(self) -> Array:
"""Surface pressure [Pa]."""
raise NotImplementedError
@property
@abstractmethod
def u(self) -> Array:
"""Zonal wind speed [m/s]."""
raise NotImplementedError
@property
@abstractmethod
def v(self) -> Array:
"""Meridional wind speed [m/s]."""
raise NotImplementedError
@property
@abstractmethod
def ra(self) -> Array:
"""Aerodynamic resistance [s/m]."""
raise NotImplementedError
@property
@abstractmethod
def thetasurf(self) -> Array:
"""Surface potential temperature [K]."""
raise NotImplementedError
@property
@abstractmethod
def h_abl(self) -> Array:
"""Boundary layer height [m]."""
raise NotImplementedError
@property
@abstractmethod
def ustar(self) -> Array:
"""Friction velocity [m/s]."""
raise NotImplementedError
@property
@abstractmethod
def uw(self) -> Array:
"""Zonal momentum flux [m²/s²]."""
raise NotImplementedError
@property
@abstractmethod
def vw(self) -> Array:
"""Meridional momentum flux [m²/s²]."""
raise NotImplementedError
@property
@abstractmethod
def wstar(self) -> Array:
"""Convective velocity scale [m/s]."""
raise NotImplementedError
@property
@abstractmethod
def thetav(self) -> Array:
"""Mixed-layer virtual potential temperature [K]."""
raise NotImplementedError
@property
@abstractmethod
def top_T(self) -> Array:
"""Temperature at top of mixed layer [K]."""
raise NotImplementedError
@property
@abstractmethod
def top_p(self) -> Array:
"""Pressure at top of mixed layer [Pa]."""
raise NotImplementedError
@property
@abstractmethod
def cc_mf(self) -> Array:
"""Cloud core mass flux [kg/kg/s]."""
raise NotImplementedError
@property
@abstractmethod
def cc_qf(self) -> Array:
"""Cloud core moisture flux [kg/kg/s]."""
raise NotImplementedError
@property
@abstractmethod
def wCO2M(self) -> Array:
"""Cloud core CO2 mass flux [mgC/m²/s]."""
raise NotImplementedError
@property
@abstractmethod
def cc_frac(self) -> Array:
"""Cloud core fraction [-]."""
raise NotImplementedError
@property
@abstractmethod
def wthetav(self) -> Array:
"""Virtual potential temperature flux at surface [K m/s]."""
raise NotImplementedError
@property
@abstractmethod
def wqe(self) -> Array:
"""Entrainment moisture flux [kg/kg m/s]."""
raise NotImplementedError
@property
@abstractmethod
def dq(self) -> Array:
"""Specific humidity jump at h [kg/kg]."""
raise NotImplementedError
@property
@abstractmethod
def dz_h(self) -> Array:
"""Transition layer thickness [m]."""
raise NotImplementedError
@property
@abstractmethod
def deltaCO2(self) -> Array:
"""CO2 jump at h [ppm]."""
raise NotImplementedError
@property
@abstractmethod
def wCO2e(self) -> Array:
"""Entrainment CO2 flux [mgC/m²/s]."""
raise NotImplementedError
RadT = TypeVar("RadT", bound=AbstractRadiationState)
LandT = TypeVar("LandT", bound=AbstractLandState)
AtmosT = TypeVar("AtmosT", bound=AbstractAtmosphereState)
[docs]
class AbstractCoupledState(AbstractState, Generic[RadT, LandT, AtmosT]):
"""Abstract coupled state, generic over rad, land and atmos types."""
rad: RadT
land: LandT
atmos: AtmosT
t: Array
total_water_mass: Array
@property
def net_rad(self) -> Array:
"""Net surface rad [W m-2]."""
return self.rad.net_rad
@property
def in_srad(self) -> Array:
"""Incoming shortwave rad [W m-2]."""
return self.rad.in_srad
[docs]
class AbstractModel(nnx.Module, Generic[StateT]):
"""Abstract model class to define the interface for all models."""
[docs]
@abstractmethod
def init_state(self, *args, **kwargs) -> StateT:
raise NotImplementedError
[docs]
class AbstractRadiationModel(AbstractModel, Generic[RadT]):
"""Abstract radiation model class to define the interface for all rad models."""
tstart: Array
"""Start time of the model."""
[docs]
@abstractmethod
def run(
self,
state: AbstractCoupledState[RadT, LandT, AtmosT],
t: Array,
dt: float,
tstart: float,
) -> RadT:
raise NotImplementedError
[docs]
class AbstractLandModel(AbstractModel, Generic[LandT]):
"""Abstract land model class to define the interface for all land models."""
[docs]
@abstractmethod
def run(
self,
state: AbstractCoupledState[RadT, LandT, AtmosT],
) -> LandT:
raise NotImplementedError
[docs]
@abstractmethod
def integrate(self, state: LandT, dt: float) -> LandT:
raise NotImplementedError
[docs]
class AbstractAtmosphereModel(AbstractModel, Generic[AtmosT]):
"""Abstract atmos model class to define the interface for all atmos models."""
[docs]
@abstractmethod
def warmup(
self,
radmodel: AbstractRadiationModel[RadT],
landmodel: AbstractLandModel[LandT],
state: AbstractCoupledState[RadT, LandT, AtmosT],
t: Array,
dt: float,
tstart: float,
) -> AbstractCoupledState[RadT, LandT, AtmosT]:
raise NotImplementedError
[docs]
@abstractmethod
def run(
self,
state: AbstractCoupledState[RadT, LandT, AtmosT],
) -> AtmosT:
raise NotImplementedError
[docs]
@abstractmethod
def statistics(
self,
state: AbstractCoupledState[RadT, LandT, AtmosT],
t: Array,
) -> AtmosT:
raise NotImplementedError
[docs]
@abstractmethod
def integrate(self, state: AtmosT, dt: float) -> AtmosT:
raise NotImplementedError
