Skip to content

hssm.distribution_utils

Utility functions for dynamically building pm.Distributions.

hssm.distribution_utils.assemble_callables 

assemble_callables(callable: Op | Callable, missing_data_callable: 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.

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_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.make_distribution 

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

Make a pymc.Distribution.

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

Parameters:

  • rv (str | Type[RandomVariable]) –

    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[Distribution], list_params: list[str], link: str | dict[str, Link], parent: str = 'v', likelihood_name: str = 'SSM Likelihood', family_name='SSM Family') -> 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: 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) -> 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: Op | Callable | PathLike | str, backend: Literal['pytensor', 'jax', 'other'] | None = 'jax', params_is_reg: list[bool] | None = None, params_only: bool | None = None) -> 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_ssm_rv 

make_ssm_rv(model_name: str, list_params: list[str], lapse: Prior | None = None) -> Type[RandomVariable]

Build a RandomVariable Op according to the list of parameters.

Parameters:

  • model_name (str) –

    The name of the model. If the model_name 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.

  • list_params (list[str]) –

    A list of str of all parameters for this RandomVariable.

  • lapse (optional, default: None ) –

    A bmb.Prior object representing the lapse distribution.

Returns:

  • Type[RandomVariable] –

    A class of RandomVariable that are to be used in a pm.Distribution.