Source code for eulerpi.examples.temperature.temperature

import importlib
from typing import Optional

import jax.numpy as jnp
import numpy as np

from eulerpi.core.models import ArtificialModelInterface, BaseModel

# from functools import partial
# from jax import jit


[docs] class Temperature(BaseModel): """The model describes the temperature :math:`y` in degree celsius at a given latitude :math:`q` in degree. .. math:: s: [0^{\\circ}, 90^{\\circ}] \\rightarrow [-30, 30] q \\mapsto y = s(q) := 60 \\cdot \\cos{(q)} - 30 :math:`q=0^{\\circ}` corresponds to the equator and :math:`q=90^{\\circ}` to the north pole. """ param_dim = 1 #: The latitude of the location data_dim = 1 #: The temperature in degree celsius PARAM_LIMITS = np.array( [[0, np.pi / 2]] ) #: The latitude is between :math:`0^{\circ}` and :math:`90^{\circ}` CENTRAL_PARAM = np.array( [np.pi / 4.0] ) #: The central latitude is :math:`\pi/4 = 45^{\circ}` def __init__( self, central_param: np.ndarray = CENTRAL_PARAM, param_limits: np.ndarray = PARAM_LIMITS, name: Optional[str] = None, **kwargs, ) -> None: super().__init__(central_param, param_limits, name=name, **kwargs)
[docs] def forward(self, param): """Evaluates the temperature at the given latitude using the model equation .. math:: y = 60 \\cdot \\cos{(q)} - 30 """ low_T = -30.0 high_T = 30.0 res = jnp.array( [low_T + (high_T - low_T) * jnp.cos(jnp.abs(param[0]))] ) return res
[docs] def jacobian(self, param): """Calculates the analytical jacobian of the model equation at the given latitude .. math:: \\frac{dy}{dq} = -60 \\cdot \\sin{(q)} """ return jnp.array([60.0 * jnp.sin(jnp.abs(param[0]))])
[docs] class TemperatureArtificial(Temperature, ArtificialModelInterface):
[docs] def generate_artificial_params(self, num_data_points: int = -1): param_resource = importlib.resources.files( "eulerpi.examples.temperature" ).joinpath("TemperatureArtificialParams.csv") with importlib.resources.as_file(param_resource) as param_file: true_param_sample = np.loadtxt(param_file, delimiter=",", ndmin=2) return true_param_sample
[docs] class TemperatureWithFixedParams(Temperature): """The model describes the temperature :math:`y` in degree celsius at a given latitude :math:`q` in degree. .. note:: * Additional fixed parameters: The model includes fixed parameters :code:`self.low_T=30.0` and :code:`self.high_T=30.0`. These fixed parameters are passed to the calc_forward function separately. You can create models with different parameters by creating several model objects. The best way to separate the outputs for the parametrized models is to pass a string based on the fixed_params to the attribute :py:attr:`run_name` of the :py:func:`~eulerpi.core.inference` function. * The functions :py:meth:`~eulerpi.examples.temperature.temperature.TemperatureWithFixedParams.calc_forward` is not strictly necessary. However it can help to make it work with jax. """ def __init__( self, low_T: np.double = -30.0, high_T: np.double = 30.0, name: Optional[str] = None, **kwargs, ) -> None: super().__init__(name=name, **kwargs) self.low_T = low_T self.high_T = high_T
[docs] def forward(self, param): return self.calc_forward(param, self.high_T, self.low_T)
[docs] def calc_forward(self, param, high_T, low_T): res = jnp.array( [low_T + (high_T - low_T) * jnp.cos(jnp.abs(param[0]))] ) return res
[docs] def jacobian(self, param): return self.calc_jacobian(param, self.high_T, self.low_T)
[docs] def calc_jacobian(self, param, high_T, low_T): return jnp.array([(high_T - low_T) * jnp.sin(jnp.abs(param[0]))])