hssm.distribution_utils

The hssm.distribution_utils contains useful functions for building pm.Distribution classes. Other than the download_hf function that downloads ONNX models shared on our huggingface model repository, you will generally not have to use these functions. For advanced users who want to build their own PyMC models, they can use these functions to create pm.Distribution and RandomVariable classes that they desire.

hssm.distribution_utils

Utility functions for dynamically building pm.Distributions.

Functions:

• download_hf –

  Download a file from a HuggingFace repository.

• make_distribution –

  Make a pymc.Distribution.

• make_family –

  Build a family in bambi.

• make_likelihood_callable –

  Make a callable for the likelihood function.

• make_missing_data_callable –

  Make a secondary network for the likelihood function.

• make_blackbox_op –

  Wrap an arbitrary function in a pytensor Op.

• assemble_callables –

  Assemble the likelihood callables into a single callable.

hssm.distribution_utils.download_hf

download_hf(path: str)

Download a file from a HuggingFace repository.

Parameters:

• path (str) –

  The path of the file to download in the repository.

Returns:

• str –

  The local path where the file is downloaded.

Notes

The repository is specified by the REPO_ID constant, which should be a valid HuggingFace.co repository ID. The file is downloaded using the HuggingFace Hub's hf_hub_download function.

hssm.distribution_utils.make_distribution

make_distribution(
    rv: str | Type[RandomVariable] | RandomVariable | Callable,
    loglik: LogLikeFunc | pytensor.graph.Op,
    list_params: list[str],
    bounds: dict | None = None,
    lapse: bmb.Prior | None = None,
    extra_fields: list[np.ndarray] | None = None,
) -> Type[pm.Distribution]

Make a pymc.Distribution.

Constructs a pymc.Distribution from a log-likelihood function and a RandomVariable op.

Parameters:

• model_name –

  The name of the model.

• choices –

  A list of integers indicating the choices.

• rv (str | Type[RandomVariable] | RandomVariable | Callable) –

  A RandomVariable Op (a class, not an instance) or a string indicating the model. If a string, a RandomVariable class will be created automatically with its rng_fn class method generated using the simulator identified with this string from the ssm_simulators package. If the string is not one of the supported models in the ssm_simulators package, a warning will be raised, and any attempt to sample from the RandomVariable will result in an error.

• loglik (LogLikeFunc | Op) –

  A loglikelihood function. It can be any Callable in Python.

• list_params (list[str]) –

  A list of parameters that the log-likelihood accepts. The order of the parameters in the list will determine the order in which the parameters are passed to the log-likelihood function.

• bounds (optional, default: None ) –

  A dictionary with parameters as keys (a string) and its boundaries as values. Example: {"parameter": (lower_boundary, upper_boundary)}.

• lapse (optional, default: None ) –

  A bmb.Prior object representing the lapse distribution.

• extra_fields (optional, default: None ) –

  An optional list of arrays that are stored in the class created and will be used in likelihood calculation. Defaults to None.

Returns:

• Type[Distribution] –

  A pymc.Distribution that uses the log-likelihood function.

hssm.distribution_utils.make_family

make_family(
    dist: Type[pm.Distribution],
    list_params: list[str],
    link: str | dict[str, bmb.families.Link],
    parent: str = "v",
    likelihood_name: str = "SSM Likelihood",
    family_name="SSM Family",
) -> bmb.Family

Build a family in bambi.

Parameters:

• dist (Type[Distribution]) –

  A pm.Distribution class (not an instance).

• list_params (list[str]) –

  A list of parameters for the likelihood function.

• link (str | dict[str, Link]) –

  A link function or a dictionary of parameter: link functions.

• parent (str, default: 'v' ) –

  The parent parameter of the likelihood function. Defaults to v.

• likelihood_name (str, default: 'SSM Likelihood' ) –

  the name of the likelihood function. Defaults to "SSM Likelihood".

• family_name –

  the name of the family. Defaults to "SSM Family".

Returns:

• Family –

  An instance of a bambi family.

hssm.distribution_utils.make_likelihood_callable

make_likelihood_callable(
    loglik: pytensor.graph.Op | Callable | PathLike | str,
    loglik_kind: Literal["analytical", "approx_differentiable", "blackbox"],
    backend: Literal["pytensor", "jax", "other"] | None,
    params_is_reg: list[bool] | None = None,
    params_only: bool | None = None,
) -> pytensor.graph.Op | Callable

Make a callable for the likelihood function.

This function is intended to be general to support different kinds of likelihood functions.

Parameters:

• loglik (Op | Callable | PathLike | str) –

  The log-likelihood function. It can be a string, a path to an ONNX model, a pytensor Op, or a python function.

• loglik_kind (Literal['analytical', 'approx_differentiable', 'blackbox']) –

  The kind of the log-likelihood for the model. This parameter controls how the likelihood function is wrapped.

• backend (Optional) –

  The backend to use for the log-likelihood function.

• params_is_reg (Optional, default: None ) –

  A list of boolean values indicating whether the parameters are regression parameters. Defaults to None.

• params_only (Optional, default: None ) –

  Whether the missing data likelihood is takes its first argument as the data. Defaults to None.

hssm.distribution_utils.make_missing_data_callable

make_missing_data_callable(
    loglik: pytensor.graph.Op | Callable | PathLike | str,
    backend: Literal["pytensor", "jax", "other"] | None = "jax",
    params_is_reg: list[bool] | None = None,
    params_only: bool | None = None,
) -> pytensor.graph.Op | Callable

Make a secondary network for the likelihood function.

Please refer to the documentation of make_likelihood_callable for more.

hssm.distribution_utils.make_blackbox_op

make_blackbox_op(logp: Callable) -> Op

Wrap an arbitrary function in a pytensor Op.

Parameters:

• logp (Callable) –

  A python function that represents the log-likelihood function. The function needs to have signature of logp(data, *dist_params) where data is a two-column numpy array and dist_paramsrepresents all parameters passed to the function.

Returns:

• Op –

  An pytensor op that wraps the log-likelihood function.

hssm.distribution_utils.assemble_callables

assemble_callables(
    callable: pytensor.graph.Op | Callable,
    missing_data_callable: pytensor.graph.Op | Callable,
    params_only: bool,
    has_deadline: bool,
) -> Callable

Assemble the likelihood callables into a single callable.

Parameters:

• callable (Op | Callable) –

  The callable for the likelihood function.

• missing_data_callable (Op | Callable) –

  The callable for the secondary network for the likelihood function.

• params_only (bool) –

  Whether the missing data likelihood is takes its first argument as the data.

• has_deadline (bool) –

  Whether the model has a deadline.