# Copyright © 2023, UChicago Argonne, LLC
# All Rights Reserved
"""
This module provides functionality to create and test very basic toy models for debugging, performance evaluation, etc.
"""
import numpy as np
from dataclasses import dataclass, field
from functools import partial
from pyoed import utility
from pyoed.models.core import (
SimulationModelConfigs,
TimeIndependentModel,
TimeIndependentModelConfigs,
)
from pyoed.configs import (
validate_key,
aggregate_configurations,
set_configurations,
)
[docs]
@dataclass(kw_only=True, slots=True)
class ToyQuarkConfigs(TimeIndependentModelConfigs):
"""
Configuration class for the ToyQuark model
:param x_grid: discretization of the state space. Default is ``np.linspace(0, 1, 10)``.
Must be a 1D numpy array.
:param Nq: number of quarks. Default is 2. Must be a postiive integer.
:param model_name: name of the model. Default is "Toy-Quark"
"""
x_grid: np.ndarray = field(default_factory=lambda: np.linspace(0, 1, 10))
Nq: int = 2
model_name: str = "Toy-Quark"
def __post_init__(self):
self.validate()
[docs]
def validate(self):
assert isinstance(self.x_grid, np.ndarray), "x_grid must be a numpy array"
assert self.x_grid.ndim == 1, "x_grid must be a 1D numpy array"
assert self.Nq > 0, "Nq must be a positive integer"
[docs]
@set_configurations(ToyQuarkConfigs)
class ToyQuark(TimeIndependentModel):
"""
Implementation of a toy quark model, from the paper [1]_. It is given by the
expression
.. math::
\\begin{align*}
\\sigma(x) &= \\sum_{i=1}^{N_q} q_i(x) \\
q_i(x) &= N_i x^{\\alpha_i} (1-x)^{\\beta_i}
\\end{align*}
where :math:`N_i, \\alpha_i, \\beta_i` are parameters to be estimated. Hence,
this is a toy model with 3*N_q parameters. The model is defined on the interval
:math:`x \\in [0, 1]` and the discretization of the state is left variable (i.e.,
the user can choose the number of discretization points).
.. [1] Hunt-Smith, N. T., Accardi, A., Melnitchouk, W., Sato, N., Thomas, A. W., &
White, M. J. (2022). Determination of uncertainties in parton densities. In
Physical Review D (Vol. 106, Issue 3). American Physical Society (APS).
https://doi.org/10.1103/physrevd.106.036003
"""
[docs]
def __init__(self, configs: ToyQuarkConfigs | dict | None = None):
"""
Create an instance of the toy quark model.
:param configs: configurations of the model. See :py:class:`ToyQuarkConfigs` for
the possible configurations. If None, default configurations are used.
"""
configs = self.configurations_class.data_to_dataclass(configs)
super().__init__(configs)
[docs]
def validate_configurations(
self, configs: dict | ToyQuarkConfigs, raise_for_invalid: bool = True
) -> bool:
"""
Check the passed configuratios and make sure they are conformable with each
other, and with current configurations once combined. This guarantees that any
key-value pair passed in configs can be properly used
.. note::
Here only the locally-defined configurations in :py:class:`ToyQuarkConfigs`
are validated. Finally, super classes validators are called.
:param configs: full or partial (subset) configurations to be validated
:param raise_for_invalid: if `True` raise :py:class:`TypeError` for invalid
configrations type/key. Default `True`
:returns: flag indicating whether passed configurations dictionary is valid or
not
:raises AttributeError: if any (or a group) of the configurations does not exist
in the model configurations :py:class:`ToyQuarkConfigs`.
:raises PyOEDConfigsValidationError: if the configurations are invalid and
`raise_for_invalid` is set to True.
"""
# Fuse configs into current/default configurations
aggregated_configs = aggregate_configurations(
obj=self,
configs=configs,
configs_class=self.configurations_class,
)
if validate_key(
aggregated_configs,
configs,
"x_grid",
lambda x: isinstance(x, np.ndarray) and (x.ndim == 1),
message="x_grid must be a 1D numpy array!",
):
return False
if validate_key(
aggregated_configs,
configs,
"Nq",
lambda x: (isinstance(x, int) and (x == int(x))) and (x > 0),
message="Nq must be a positive integer!",
):
return False
return super().validate_configurations(configs, raise_for_invalid)
[docs]
def parameter_vector(self, init_val=0):
"""
Create an instance of model parameter vector. This is a vector of size
:math:`3N_q`, and corresponds to parameters in the order
:math:`[N_1, \\alpha_1, \\beta_1, \\dots, N_{N_q}, \\alpha_{N_q}, \\beta_{N_q}]`.
:param float init_val: (optional) value assigned to entries of
the state vector upon initialization
"""
assert utility.isnumber(
init_val
), "Initial values of entries `init_val` must be a number"
parameter = np.empty(3 * self.Nq)
parameter[:] = init_val
return parameter
[docs]
def is_parameter_vector(self, parameter):
"""Test whether the passed parameter vector is valid or not"""
valid = False
if isinstance(parameter, np.ndarray):
if parameter.ndim == 1 and parameter.size == 3 * self.Nq:
valid = True
return valid
[docs]
def state_vector(self, init_val=0):
"""
Create an instance of model state vector
:param float init_val: (optional) value assigned to entries of
the state vector upon initialization
"""
assert utility.isnumber(
init_val
), "Initial values of entries `init_val` must be a number"
state = np.empty_like(self.x_grid)
state[:] = init_val
return state
[docs]
def is_state_vector(self, state):
"""Test whether the passed state vector is valid or not"""
valid = False
if isinstance(state, np.ndarray):
if state.ndim == 1 and state.size == self.x_grid.size:
valid = True
return valid
[docs]
def q_i(self, parameter, i):
"""
Helper function to retrieve the i-th quark distribution function from the
parameter vector. The parameter is assumed to be a vector of size :math:`3N_q`,
and corresponds to parameters in the order :math:`[N_1, \\alpha_1, \\beta_1,
\\dots, N_{N_q}, \\alpha_{N_q}, \\beta_{N_q}]`.
:param parameter: parameter from which the i-th quark distribution function is
retrieved.
:param int i: index of the quark distribution function to retrieve. Note, this
is 1-based indexing.
:return: :math:`q_i(x; p)` where :math:`p` is the parameter vector. This is of
size ``len(self.x_grid)``.
"""
p = np.split(parameter, self.Nq)[i - 1]
return p[0] * self.x_grid ** p[1] * (1 - self.x_grid) ** p[2]
[docs]
def solve_forward(self, parameter, verbose=False):
"""
Apply (solve the forward model) to the given parameter, and return the resulting
state. The parameter is assumed to be a vector of size :math:`3N_q`, and
corresponds to parameters in the order :math:`[N_1, \\alpha_1, \\beta_1, \\dots,
N_{N_q}, \\alpha_{N_q}, \\beta_{N_q}]`.
:param parameter: parameter to which the forward model is applied.
:returns: :math:`\\sigma(x; p)` where :math:`p` is the parameter vector. This is
of size ``len(self.x_grid)``.
"""
assert self.is_parameter_vector(
parameter
), "passed parameter is not a valid parameter vector!"
q = lambda p: p[0] * self.x_grid ** p[1] * (1 - self.x_grid) ** p[2]
return sum([q(p) for p in np.split(parameter, self.Nq)])
[docs]
def Jacobian_T_matvec(self, state, eval_at=None, verbose=False):
"""
Multiply the Jacobian of the model (derivative of model equation w.r.t
parameter) by the passed state
:param state: state to which the forward sensitivities are applied to.
:returns: result resulting from applying the forward sensitivities to
model to the passed `state`. This is of size :math:`3N_q`.
"""
assert self.is_state_vector(state), "passed state is not valid!"
assert eval_at is not None, "eval_at must be provided!"
x = self.x_grid
JqT_v = lambda p, v: np.array( # Jacobian of q_i w.r.t parameters
[
np.dot(x ** p[1] * (1 - x) ** p[2], v),
np.dot(np.log(x) * p[0] * x ** p[1] * (1 - x) ** p[2], v),
np.dot(np.log(1 - x) * p[0] * x ** p[1] * (1 - x) ** p[2], v),
]
)
return np.concatenate([JqT_v(p, state) for p in np.split(eval_at, self.Nq)])
[docs]
def get_model_grid(self):
"""retrieve the model grid"""
return self.x_grid.copy()
@property
def x_grid(self):
"""retrieve the model grid"""
return self.configurations.x_grid
@property
def Nq(self):
"""retrieve the number of quarks"""
return self.configurations.Nq
# ================================================================================ #
# Example Functions to help inistantiate instances of the models implemented here #
# ================================================================================ #
[docs]
def create_ToyQuark_model(
Nq=2,
x_grid=np.linspace(0.1, 0.9, 10),
):
return ToyQuark(ToyQuarkConfigs(Nq=Nq, x_grid=x_grid))
