Source code for anml.parameter.main

from __future__ import annotations

from operator import attrgetter
from typing import List, Optional, Union

import numpy as np
from anml.data.component import Component
from anml.data.validator import NoNans
from anml.parameter.smoothmapping import Identity, SmoothMapping
from anml.prior.main import Prior
from anml.prior.utils import filter_priors, get_prior_type
from anml.variable.main import Variable
from numpy.typing import NDArray
from pandas import DataFrame
from scipy.linalg import block_diag


[docs]class Parameter: """Parameter class contains information from a list of variables that are used to parametrize the distribution parameter. Parameter class also include optional transformation function, offset and list of priors. Parameters ---------- variables A list of variables that parametrized the linear predictor of the parameter. transform A :class:`SmoothMapping` instance that transforms the linear prediction from the variables to the parameter space. Default is `None`, which will be converted into identity mapping. offset A data component contains the offset information for the parameter. Default is `None` which indicates no offset. Offset will be applied onto the linear predictor. priors A list of additional priors that directly apply to the parameter linear predictor. Default is `None`, where no additional priors will be added. """ variables = property(attrgetter("_variables")) """A list of variables that parametrized the linear predictor of the parameter. Raises ------ TypeError Raised if the input variables are not all instances of :class:`Variable`. """ transform = property(attrgetter("_transform")) """A function that transforms the linear prediction to the parameter space. Raises ------ TypeError Raised when the input transform is not an instance of :class:`SmoothMapping`. """ offset = property(attrgetter("_offset")) """Offset for the linear predictor. Raises ------ TypeError Raised when the input offset is not `None`, or a string, or an instance of :class:`DataComponent`. """ priors = property(attrgetter("_priors")) """A list of additional priors that apply to the linear predictor. Raises ------ TypeError Raised when the input priors are not `None` or a list of instances of :class:`Prior`. """ def __init__(self, variables: List[Variable], transform: Optional[SmoothMapping] = None, offset: Optional[Union[str, Component]] = None, priors: Optional[List[Prior]] = None): self.variables = variables self.transform = transform self.offset = offset self.priors = priors self.design_mat = None self.prior_dict = {"direct": {}, "linear": {}} @variables.setter def variables(self, variables: List[Variable]): variables = list(variables) if not all(isinstance(variable, Variable) for variable in variables): raise TypeError("Parameter input variables must be a list of " "instances of Variable.") self._variables = variables @transform.setter def transform(self, transform: Optional[SmoothMapping]): if transform is not None and not isinstance(transform, SmoothMapping): raise TypeError("Parameter input transform must be an instance " "of SmoothMapping or None.") if transform is None: transform = Identity() self._transform = transform @offset.setter def offset(self, offset: Optional[Union[str, Component]]): if offset is not None: if not isinstance(offset, (str, Component)): raise TypeError("Parameter input offset has to be a string or " "an instance of Component.") if isinstance(offset, str): offset = Component(offset, validators=[NoNans()]) self._offset = offset @priors.setter def priors(self, priors: Optional[List[Prior]]): priors = list(priors) if priors is not None else [] if not all(isinstance(prior, Prior) for prior in priors): raise TypeError("Parameter input priors must be a list of " "instances of Prior.") self._priors = priors @property def size(self) -> int: """Size of the parameter coefficients. It is the sum of all sizes for variables. """ return sum([variable.size for variable in self.variables])
[docs] def attach(self, df: DataFrame): """Attach data frame to offset and cache the design matrix and gather the prior information. Parameters ---------- df Given data frame. """ if self.offset is not None: self.offset.attach(df) self.design_mat = np.hstack([ variable.get_design_mat(df) for variable in self.variables ]) for prior_category in ["direct", "linear"]: for prior_type in ["UniformPrior", "GaussianPrior"]: getattr(self, f"_get_{prior_category}_prior")(prior_type)
def _get_direct_prior(self, prior_type: str): """Get the direct prior parameters. The direct prior refers to the priors that do not have a linear map and direct act on the variable. This function will ignore the direct priors provided by the additional priors in the parameter. Please add direct priors on :class:`Variable` instances. Parameters ---------- prior_type Given name of the prior type. """ params = np.hstack([variable.get_direct_prior_params(prior_type) for variable in self.variables]) self.prior_dict["direct"][prior_type] = get_prior_type(prior_type)( params[0], params[1] ) def _get_linear_prior(self, prior_type: str): """Get the linear prior parameters. The linear prior refers to the priors that contain a linear map. This function will combine the linear priors from the list of variables and the ones in the additional priors provided by the parameter. Parameters ---------- prior_type Given name of the prior type. """ params, mat = tuple(zip(*[variable.get_linear_prior_params(prior_type) for variable in self.variables])) params = np.hstack(params) mat = block_diag(*mat) linear_priors = filter_priors(self.priors, prior_type, with_mat=True) if len(linear_priors) == 0: extra_params = np.empty((2, 0)) extra_mat = np.empty((0, self.size)) else: extra_params = np.hstack([prior.params for prior in linear_priors]) extra_mat = np.vstack([prior.mat for prior in linear_priors]) params = np.hstack([params, extra_params]) mat = np.vstack([mat, extra_mat]) self.prior_dict["linear"][prior_type] = get_prior_type(prior_type)( params[0], params[1], mat )
[docs] def get_params(self, x: NDArray, df: Optional[DataFrame] = None, order: int = 0) -> NDArray: """Compute and return the parameter. Denote :math:`x` as the coefficients, :math:`A` as the design matrix, :math:`z` as the offset, :math:`f` as the transformation function, the parameter :math:`p` can be represented as .. math:: p = f(z + Ax) Here we call :math:`Ax` as the linear predictor. Parameters ---------- x Coefficients for the design matrix. df Given data frame used for compute the design matrix. Default is `None`. order Order of the derivative. Default is 0. Returns ------- NDArray When `order=0`, it will return the parameter value. When `order=1`, it will return the Jacobian matrix. And when `order=2`, it will return the second order Jacobian tensor. Raises ------ ValueError Raised when there is not cache of the design matrix and no data frame is provided. """ if df is not None: self.attach(df) if self.design_mat is None: raise ValueError("Must provide a data frame to attach data.") y = self.design_mat.dot(x) if self.offset is not None: y += self.offset.value z = self.transform(y, order=order) if order == 0: return z if order == 1: return z[:, np.newaxis] * self.design_mat return (z[:, np.newaxis, np.newaxis] * (self.design_mat[..., np.newaxis] * self.design_mat[:, np.newaxis, :]))
[docs] def prior_objective(self, x: NDArray) -> float: """Objective function from the prior. Parameters ---------- x Coefficients for the design matrix. Returns ------- float Objective value from the prior. """ value = 0.0 for prior_category in ["direct", "linear"]: prior = self.prior_dict[prior_category]["GaussianPrior"] value += prior.objective(x) return value
[docs] def prior_gradient(self, x: NDArray) -> NDArray: """Gradient function from the prior. Parameters ---------- x Coefficients for the design matrix. Returns ------- NDArray Gradient value from the prior. """ value = np.zeros(x.size, dtype=x.dtype) for prior_category in ["direct", "linear"]: prior = self.prior_dict[prior_category]["GaussianPrior"] value += prior.gradient(x) return value
[docs] def prior_hessian(self, x: NDArray) -> NDArray: """Hessian function from the prior. Parameters ---------- x Coefficients for the design matrix. Returns ------- NDArray Hessian value from the prior. """ value = np.zeros((x.size, x.size), dtype=x.dtype) for prior_category in ["direct", "linear"]: prior = self.prior_dict[prior_category]["GaussianPrior"] value += prior.hessian(x) return value
def __repr__(self) -> str: return (f"{type(self).__name__}(variables={self.variables}, " f"transform={self.transform}, offset={self.offset}, " f"priors={self.priors})")