Probabilistic

inspeqtor.models.probabilistic

LearningModel

The learning model.

Source code in src/inspeqtor/v1/probabilistic.py

class LearningModel(StrEnum):
    """The learning model."""

    TruncatedNormal = auto()
    BernoulliProbs = auto()

make_probabilistic_model

make_probabilistic_model(
    predictive_model: Callable[..., ndarray],
    shots: int = 1,
    block_graybox: bool = False,
    separate_observables: bool = False,
    log_expectation_values: bool = False,
)

Make probabilistic model from the Statistical model with priors

Parameters:

Name	Type	Description	Default
base_model	Module	The statistical based model, currently only support flax.linen module	required
model_prediction_to_expvals_fn	Callable[..., ndarray]	Function to convert output from model to expectation values array	required
bnn_prior	dict[str, Distribution] | Distribution	The priors of BNN. Defaults to dist.Normal(0.0, 1.0).	required
shots	int	The number of shots forcing PGM to sample. Defaults to 1.	1
block_graybox	bool	If true, the latent variables in Graybox model will be hidden, i.e. not traced by numpyro. Defaults to False.	False
enable_bnn	bool	If true, the statistical model will be convert to probabilistic model. Defaults to True.	required
separate_observables	bool	If true, the observable will be separate into dict form. Defaults to False.	False

Returns:

Type	Description
	typing.Callable: Probabilistic Graybox Model

Source code in src/inspeqtor/v1/probabilistic.py

def make_probabilistic_model(
    predictive_model: typing.Callable[..., jnp.ndarray],
    shots: int = 1,
    block_graybox: bool = False,
    separate_observables: bool = False,
    log_expectation_values: bool = False,
):
    """Make probabilistic model from the Statistical model with priors

    Args:
        base_model (nn.Module): The statistical based model, currently only support flax.linen module
        model_prediction_to_expvals_fn (typing.Callable[..., jnp.ndarray]): Function to convert output from model to expectation values array
        bnn_prior (dict[str, dist.Distribution] | dist.Distribution, optional): The priors of BNN. Defaults to dist.Normal(0.0, 1.0).
        shots (int, optional): The number of shots forcing PGM to sample. Defaults to 1.
        block_graybox (bool, optional): If true, the latent variables in Graybox model will be hidden, i.e. not traced by `numpyro`. Defaults to False.
        enable_bnn (bool, optional): If true, the statistical model will be convert to probabilistic model. Defaults to True.
        separate_observables (bool, optional): If true, the observable will be separate into dict form. Defaults to False.

    Returns:
        typing.Callable: Probabilistic Graybox Model
    """

    def block_graybox_fn(
        control_parameters: jnp.ndarray,
        unitaries: jnp.ndarray,
    ):
        key = numpyro.prng_key()
        with handlers.block(), handlers.seed(rng_seed=key):
            expvals = predictive_model(control_parameters, unitaries)

        return expvals

    graybox_fn = block_graybox_fn if block_graybox else predictive_model

    def bernoulli_model(
        control_parameters: jnp.ndarray,
        unitaries: jnp.ndarray,
        observables: jnp.ndarray | None = None,
    ):
        expvals = graybox_fn(control_parameters, unitaries)

        if log_expectation_values:
            numpyro.deterministic("expectation_values", expvals)

        if observables is None:
            sizes = control_parameters.shape[:-1] + (18,)
            if shots > 1:
                sizes = (shots,) + sizes
        else:
            sizes = observables.shape

        # The plate is for the shots prediction to work properly
        with numpyro.util.optional(
            shots > 1, numpyro.plate_stack("plate", sizes=list(sizes)[:-1])
        ):
            if separate_observables:
                expvals_samples = {}

                for idx, exp in enumerate(default_expectation_values_order):
                    s = numpyro.sample(
                        f"obs/{exp.initial_state}/{exp.observable}",
                        dist.BernoulliProbs(
                            probs=expectation_value_to_prob_minus(
                                jnp.expand_dims(expvals[..., idx], axis=-1)
                            )
                        ).to_event(1),  # type: ignore
                        obs=(
                            observables[..., idx] if observables is not None else None
                        ),
                    )

                    expvals_samples[f"obs/{exp.initial_state}/{exp.observable}"] = s

            else:
                expvals_samples = numpyro.sample(
                    "obs",
                    dist.BernoulliProbs(
                        probs=expectation_value_to_prob_minus(expvals)
                    ).to_event(1),  # type: ignore
                    obs=observables,
                    infer={"enumerate": "parallel"},
                )

        return expvals_samples

    return bernoulli_model

get_args_of_distribution

get_args_of_distribution(x)

Get the arguments used to construct Distribution, if the provided parameter is not Distribution, return it. So that the function can be used with jax.tree.map.

Parameters:

Name	Type	Description	Default
x	Any	Maybe Distribution	required

Returns:

Type	Description
	typing.Any: Argument of Distribution if Distribution is provided.

Source code in src/inspeqtor/v1/probabilistic.py

def get_args_of_distribution(x):
    """Get the arguments used to construct Distribution, if the provided parameter is not Distribution, return it.
    So that the function can be used with `jax.tree.map`.

    Args:
        x (typing.Any): Maybe Distribution

    Returns:
        typing.Any: Argument of Distribution if Distribution is provided.
    """
    warnings.warn(
        "This function is deprecated",
        category=DeprecationWarning,
        stacklevel=2,
    )

    if isinstance(x, dist.Distribution):
        return x.get_args()
    else:
        return x

construct_normal_priors

construct_normal_priors(posterior)

Construct a dict of Normal Distributions with posterior

Parameters:

Name	Type	Description	Default
posterior	Any	Dict of Normal distribution arguments	required

Returns:

Type	Description
	typing.Any: dict of Normal distributions

Source code in src/inspeqtor/v1/probabilistic.py

def construct_normal_priors(posterior):
    """Construct a dict of Normal Distributions with posterior

    Args:
        posterior (typing.Any): Dict of Normal distribution arguments

    Returns:
        typing.Any: dict of Normal distributions
    """

    warnings.warn(
        "This function is deprecated",
        category=DeprecationWarning,
        stacklevel=2,
    )

    posterior_distributions = {}
    assert isinstance(posterior, dict)
    for name, value in posterior.items():
        assert isinstance(name, str)
        assert isinstance(value, dict)
        posterior_distributions[name] = dist.Normal(value["loc"], value["scale"])  # type: ignore
    return posterior_distributions

construct_normal_prior_from_samples

construct_normal_prior_from_samples(
    posterior_samples: dict[str, ndarray],
) -> dict[str, Distribution]

Construct a dict of Normal Distributions with posterior sample

Parameters:

Name	Type	Description	Default
posterior_samples	dict[str, ndarray]	Posterior sample	required

Returns:

Type	Description
dict[str, Distribution]	dict[str, dist.Distribution]: dict of Normal Distributions

Source code in src/inspeqtor/v1/probabilistic.py

def construct_normal_prior_from_samples(
    posterior_samples: dict[str, jnp.ndarray],
) -> dict[str, dist.Distribution]:
    """Construct a dict of Normal Distributions with posterior sample

    Args:
        posterior_samples (dict[str, jnp.ndarray]): Posterior sample

    Returns:
        dict[str, dist.Distribution]: dict of Normal Distributions
    """

    posterior_mean = jax.tree.map(lambda x: jnp.mean(x, axis=0), posterior_samples)
    posterior_std = jax.tree.map(lambda x: jnp.std(x, axis=0), posterior_samples)

    prior = {}
    for name, mean in posterior_mean.items():
        prior[name] = dist.Normal(mean, posterior_std[name])

    return prior

make_normal_posterior_dist_fn_from_svi_result

make_normal_posterior_dist_fn_from_svi_result(
    key: ndarray,
    guide: Callable,
    params: dict[str, ndarray],
    num_samples: int,
    prefix: str,
) -> Callable[[str, tuple[int, ...]], Distribution]

This function create a get posterior function to be used with numpyro.contrib.module.

Parameters:

Name	Type	Description	Default
key	ndarray	The random key	required
guide	Callable	The guide (variational distribution)	required
params	dict[str, ndarray]	The variational parameters	required
num_samples	int	The number of sample for approxiatation the posterior distributions	required

Examples:

prefix = "graybox"
prior_fn = make_normal_posterior_dist_fn_from_svi_result(
    jax.random.key(0), guide, result.params, 10_000, prefix
)
graybox_model = sq.probabilistic.make_flax_probabilistic_graybox_model(
    name=prefix,
    base_model=base_model,
    adapter_fn=sq.probabilistic.observable_to_expvals,
    prior=prior_fn,
)
posterior_model = sq.probabilistic.make_probabilistic_model(
    predictive_model=graybox_model, log_expectation_values=True
)

Returns:

Type	Description
Callable[[str, tuple[int, ...]], Distribution]	typing.Callable[[str, tuple[int, ...]], dist.Distribution]: The function that return posterior distribution

Source code in src/inspeqtor/v1/probabilistic.py

def make_normal_posterior_dist_fn_from_svi_result(
    key: jnp.ndarray,
    guide: typing.Callable,
    params: dict[str, jnp.ndarray],
    num_samples: int,
    prefix: str,
) -> typing.Callable[[str, tuple[int, ...]], dist.Distribution]:
    """This function create a get posterior function to be used with `numpyro.contrib.module`.

    Args:
        key (jnp.ndarray): The random key
        guide (typing.Callable): The guide (variational distribution)
        params (dict[str, jnp.ndarray]): The variational parameters
        num_samples (int): The number of sample for approxiatation the posterior distributions

    Examples:
        ```python
        prefix = "graybox"
        prior_fn = make_normal_posterior_dist_fn_from_svi_result(
            jax.random.key(0), guide, result.params, 10_000, prefix
        )
        graybox_model = sq.probabilistic.make_flax_probabilistic_graybox_model(
            name=prefix,
            base_model=base_model,
            adapter_fn=sq.probabilistic.observable_to_expvals,
            prior=prior_fn,
        )
        posterior_model = sq.probabilistic.make_probabilistic_model(
            predictive_model=graybox_model, log_expectation_values=True
        )
        ```

    Returns:
        typing.Callable[[str, tuple[int, ...]], dist.Distribution]: The function that return posterior distribution
    """
    posterior_samples = Predictive(model=guide, params=params, num_samples=num_samples)(
        key
    )

    def posterior_dist_fn(name: str, shape: tuple[int, ...]) -> dist.Distribution:
        site_name = prefix + "/" + name
        return construct_normal_prior_from_samples(posterior_samples)[site_name]

    return posterior_dist_fn

make_predictive_fn

make_predictive_fn(
    posterior_model, learning_model: LearningModel
)

Construct predictive model from the probabilsitic model. This function does not relied on guide and the variational parameters

Examples:

characterized_result = sq.probabilistic.SVIResult.from_file(
    PGM_model_path / "model.json"
)

base_model = sq.models.library.linen.WoModel(
    shared_layers=characterized_result.config["model_config"]["hidden_sizes"][0],
    pauli_layers=characterized_result.config["model_config"]["hidden_sizes"][1],
)
graybox_model = sq.probabilistic.make_flax_probabilistic_graybox_model(
    name="graybox",
    base_model=base_model,
    adapter_fn=sq.probabilistic.observable_to_expvals,
    prior=dist.Normal(0, 1),
)
model = sq.probabilistic.make_probabilistic_model(
    graybox_probabilistic_model=graybox_model,
)
# initialize guide
guide = sq.probabilistic.auto_diagonal_normal_guide(
    model,
    ml.custom_feature_map(loaded_data.control_parameters),
    loaded_data.unitaries,
    jnp.zeros(shape=(shots, loaded_data.control_parameters.shape[0], 18)),
)
priors = {
    k.strip("graybox/"): v
    for k, v in make_prior_from_params(guide, characterized_result.params).items()
}
graybox_model = sq.probabilistic.make_flax_probabilistic_graybox_model(
    name="graybox",
    base_model=base_model,
    adapter_fn=sq.probabilistic.observable_to_expvals,
    prior=priors,
)
posterior_model = sq.probabilistic.make_probabilistic_model(
    graybox_probabilistic_model=graybox_model,
    shots=shots,
    block_graybox=True,
    log_expectation_values=True,
)
predicetive_fn = sq.probabilistic.make_predictive_fn(
    posterior_model, sq.probabilistic.LearningModel.BernoulliProbs
)

Parameters:

Name	Type	Description	Default
posterior_model	Any	probabilsitic model	required
learning_model	LearningModel	description	required

Source code in src/inspeqtor/v1/probabilistic.py

def make_predictive_fn(
    posterior_model,
    learning_model: LearningModel,
):
    """Construct predictive model from the probabilsitic model.
    This function does not relied on guide and the variational parameters

    Examples:
        ```python
        characterized_result = sq.probabilistic.SVIResult.from_file(
            PGM_model_path / "model.json"
        )

        base_model = sq.models.library.linen.WoModel(
            shared_layers=characterized_result.config["model_config"]["hidden_sizes"][0],
            pauli_layers=characterized_result.config["model_config"]["hidden_sizes"][1],
        )
        graybox_model = sq.probabilistic.make_flax_probabilistic_graybox_model(
            name="graybox",
            base_model=base_model,
            adapter_fn=sq.probabilistic.observable_to_expvals,
            prior=dist.Normal(0, 1),
        )
        model = sq.probabilistic.make_probabilistic_model(
            graybox_probabilistic_model=graybox_model,
        )
        # initialize guide
        guide = sq.probabilistic.auto_diagonal_normal_guide(
            model,
            ml.custom_feature_map(loaded_data.control_parameters),
            loaded_data.unitaries,
            jnp.zeros(shape=(shots, loaded_data.control_parameters.shape[0], 18)),
        )
        priors = {
            k.strip("graybox/"): v
            for k, v in make_prior_from_params(guide, characterized_result.params).items()
        }
        graybox_model = sq.probabilistic.make_flax_probabilistic_graybox_model(
            name="graybox",
            base_model=base_model,
            adapter_fn=sq.probabilistic.observable_to_expvals,
            prior=priors,
        )
        posterior_model = sq.probabilistic.make_probabilistic_model(
            graybox_probabilistic_model=graybox_model,
            shots=shots,
            block_graybox=True,
            log_expectation_values=True,
        )
        predicetive_fn = sq.probabilistic.make_predictive_fn(
            posterior_model, sq.probabilistic.LearningModel.BernoulliProbs
        )
        ```

    Args:
        posterior_model (typing.Any): probabilsitic model
        learning_model (LearningModel): _description_
    """

    def binary_predict_expectation_values(
        key: jnp.ndarray,
        control_params: jnp.ndarray,
        unitary: jnp.ndarray,
    ) -> jnp.ndarray:
        return jnp.mean(
            binary_to_eigenvalue(
                handlers.seed(posterior_model, key)(  # type: ignore
                    control_params, unitary
                )
            ),
            axis=0,
        )

    def normal_predict_expectation_values(
        key: jnp.ndarray,
        control_params: jnp.ndarray,
        unitary: jnp.ndarray,
    ) -> jnp.ndarray:
        return handlers.seed(posterior_model, key)(  # type: ignore
            control_params, unitary
        )

    return (
        binary_predict_expectation_values
        if learning_model == LearningModel.BernoulliProbs
        else normal_predict_expectation_values
    )

make_pdf

make_pdf(sample: ndarray, bins: int, srange=(-1, 1))

Make the numberical PDF from given sample using histogram method

Parameters:

Name	Type	Description	Default
sample	ndarray	Sample to make PDF.	required
bins	int	The number of interval bin.	required
srange	tuple	The range of the pdf. Defaults to (-1, 1).	(-1, 1)

Returns:

Type	Description
	typing.Any: The approximated numerical PDF

Source code in src/inspeqtor/v1/probabilistic.py

def make_pdf(sample: jnp.ndarray, bins: int, srange=(-1, 1)):
    """Make the numberical PDF from given sample using histogram method

    Args:
        sample (jnp.ndarray): Sample to make PDF.
        bins (int): The number of interval bin.
        srange (tuple, optional): The range of the pdf. Defaults to (-1, 1).

    Returns:
        typing.Any: The approximated numerical PDF
    """
    density, bin_edges = jnp.histogram(sample, bins=bins, range=srange, density=True)
    dx = jnp.diff(bin_edges)
    p = density * dx
    return p

safe_kl_divergence

safe_kl_divergence(p: ndarray, q: ndarray)

Calculate the KL divergence where the infinity is converted to zero.

Parameters:

Name	Type	Description	Default
p	ndarray	The left PDF	required
q	ndarray	The right PDF	required

Returns:

Type	Description
	jnp.ndarray: The KL divergence

Source code in src/inspeqtor/v1/probabilistic.py

def safe_kl_divergence(p: jnp.ndarray, q: jnp.ndarray):
    """Calculate the KL divergence where the infinity is converted to zero.

    Args:
        p (jnp.ndarray): The left PDF
        q (jnp.ndarray): The right PDF

    Returns:
        jnp.ndarray: The KL divergence
    """
    return jnp.sum(jnp.nan_to_num(jax.scipy.special.rel_entr(p, q), posinf=0.0))

kl_divergence

kl_divergence(p: ndarray, q: ndarray)

Calculate the KL divergence

Parameters:

Name	Type	Description	Default
p	ndarray	The left PDF	required
q	ndarray	The right PDF	required

Returns:

Type	Description
	jnp.ndarray: The KL divergence

Source code in src/inspeqtor/v1/probabilistic.py

def kl_divergence(p: jnp.ndarray, q: jnp.ndarray):
    """Calculate the KL divergence

    Args:
        p (jnp.ndarray): The left PDF
        q (jnp.ndarray): The right PDF

    Returns:
        jnp.ndarray:  The KL divergence
    """
    return jnp.sum(jax.scipy.special.rel_entr(p, q))

safe_jensenshannon_divergence

safe_jensenshannon_divergence(p: ndarray, q: ndarray)

Calculate Jensen-Shannon Divergnece using KL divergence. Implement following: https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.jensenshannon.html

Parameters:

Name	Type	Description	Default
p	ndarray	The left PDF	required
q	ndarray	The right PDF	required

Returns:

Type	Description
	typing.Any: description

Source code in src/inspeqtor/v1/probabilistic.py

def safe_jensenshannon_divergence(p: jnp.ndarray, q: jnp.ndarray):
    """Calculate Jensen-Shannon Divergnece using KL divergence.
    Implement following: https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.jensenshannon.html

    Args:
        p (jnp.ndarray): The left PDF
        q (jnp.ndarray): The right PDF

    Returns:
        typing.Any: _description_
    """
    # Compute pointwise mean of p and q
    m = (p + q) / 2
    return (safe_kl_divergence(p, m) + safe_kl_divergence(q, m)) / 2

jensenshannon_divergence_from_pmf

jensenshannon_divergence_from_pmf(p: ndarray, q: ndarray)

Calculate the Jensen-Shannon Divergence from PMF

Example

key = jax.random.key(0)
key_1, key_2 = jax.random.split(key)
sample_1 = jax.random.normal(key_1, shape=(10000, ))
sample_2 = jax.random.normal(key_2, shape=(10000, ))

# Determine srange from sample
merged_sample = jnp.concat([sample_1, sample_2])
srange = jnp.min(merged_sample), jnp.max(merged_sample)

# https://stats.stackexchange.com/questions/510699/discrete-kl-divergence-with-decreasing-bin-width
# Recommend this book: https://catalog.lib.uchicago.edu/vufind/Record/6093380/TOC
bins = int(2 * (sample_2.shape[0]) ** (1/3))
# bins = 10
dis_1 = sq.probabilistic.make_pdf(sample_1, bins=bins, srange=srange)
dis_2 = sq.probabilistic.make_pdf(sample_2, bins=bins, srange=srange)

jsd = sq.probabilistic.jensenshannon_divergence_from_pdf(dis_1, dis_2)

Parameters:

Name	Type	Description	Default
p	ndarray	The 1st probability mass function	required
q	ndarray	The 1st probability mass function	required

Returns:

Type	Description
	jnp.ndarray: The Jensen-Shannon Divergence of p and q

Source code in src/inspeqtor/v1/probabilistic.py

def jensenshannon_divergence_from_pmf(p: jnp.ndarray, q: jnp.ndarray):
    """Calculate the Jensen-Shannon Divergence from PMF

    Example
    ```python
    key = jax.random.key(0)
    key_1, key_2 = jax.random.split(key)
    sample_1 = jax.random.normal(key_1, shape=(10000, ))
    sample_2 = jax.random.normal(key_2, shape=(10000, ))

    # Determine srange from sample
    merged_sample = jnp.concat([sample_1, sample_2])
    srange = jnp.min(merged_sample), jnp.max(merged_sample)

    # https://stats.stackexchange.com/questions/510699/discrete-kl-divergence-with-decreasing-bin-width
    # Recommend this book: https://catalog.lib.uchicago.edu/vufind/Record/6093380/TOC
    bins = int(2 * (sample_2.shape[0]) ** (1/3))
    # bins = 10
    dis_1 = sq.probabilistic.make_pdf(sample_1, bins=bins, srange=srange)
    dis_2 = sq.probabilistic.make_pdf(sample_2, bins=bins, srange=srange)

    jsd = sq.probabilistic.jensenshannon_divergence_from_pdf(dis_1, dis_2)

    ```

    Args:
        p (jnp.ndarray): The 1st probability mass function
        q (jnp.ndarray): The 1st probability mass function

    Returns:
        jnp.ndarray: The Jensen-Shannon Divergence of p and q
    """
    # Note for JSD: https://medium.com/data-science/how-to-understand-and-use-jensen-shannon-divergence-b10e11b03fd6
    # Implement following: https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.jensenshannon.html
    # Compute pointwise mean of p and q
    m = (p + q) / 2
    return (kl_divergence(p, m) + kl_divergence(q, m)) / 2

jensenshannon_divergence_from_sample

jensenshannon_divergence_from_sample(
    sample_1: ndarray, sample_2: ndarray
) -> ndarray

Calculate the Jensen-Shannon Divergence from sample

Parameters:

Name	Type	Description	Default
sample_1	ndarray	The left PDF	required
sample_2	ndarray	The right PDF	required

Returns:

Type	Description
ndarray	jnp.ndarray: The Jensen-Shannon Divergence of p and q

Source code in src/inspeqtor/v1/probabilistic.py

def jensenshannon_divergence_from_sample(
    sample_1: jnp.ndarray, sample_2: jnp.ndarray
) -> jnp.ndarray:
    """Calculate the Jensen-Shannon Divergence from sample

    Args:
        sample_1 (jnp.ndarray): The left PDF
        sample_2 (jnp.ndarray): The right PDF

    Returns:
        jnp.ndarray: The Jensen-Shannon Divergence of p and q
    """
    merged_sample = jnp.concat([sample_1, sample_2])
    bins = int(2 * (sample_2.shape[0]) ** (1 / 3))
    srange = jnp.min(merged_sample), jnp.max(merged_sample)

    dis_1 = make_pdf(sample_1, bins=bins, srange=srange)
    dis_2 = make_pdf(sample_2, bins=bins, srange=srange)

    return jensenshannon_divergence_from_pmf(dis_1, dis_2)

batched_matmul

batched_matmul(x, w, b)

A specialized batched matrix multiplication of weight and input x, then add the bias. This function is intended to be used in dense_layer

Parameters:

Name	Type	Description	Default
x	ndarray	The input x	required
w	ndarray	The weight to multiply to x	required
b	ndarray	The bias	required

Returns:

Type	Description
	jnp.ndarray: Output of the operations.

Source code in src/inspeqtor/v1/probabilistic.py

def batched_matmul(x, w, b):
    """A specialized batched matrix multiplication of weight and input x, then add the bias.
    This function is intended to be used in `dense_layer`

    Args:
        x (jnp.ndarray): The input x
        w (jnp.ndarray): The weight to multiply to x
        b (jnp.ndarray): The bias

    Returns:
        jnp.ndarray: Output of the operations.
    """
    return jnp.einsum(x, (..., 0), w, (..., 0, 1), (..., 1)) + b

get_trace

get_trace(fn, key=key(0))

Convinent function to get a trace of the probabilistic model in numpyro.

Parameters:

Name	Type	Description	Default
fn	function	The probabilistic model in numpyro.	required
key	ndarray	The random key. Defaults to jax.random.key(0).	key(0)

Source code in src/inspeqtor/v1/probabilistic.py

def get_trace(fn, key=jax.random.key(0)):
    """Convinent function to get a trace of the probabilistic model in numpyro.

    Args:
        fn (function): The probabilistic model in numpyro.
        key (jnp.ndarray, optional): The random key. Defaults to jax.random.key(0).
    """

    def inner(*args, **kwargs):
        return handlers.trace(handlers.seed(fn, key)).get_trace(*args, **kwargs)

    return inner

default_priors_fn

default_priors_fn(
    name: str, shape: tuple[int, ...]
) -> Distribution

This is a default prior function for the dense_layer

Parameters:

Name	Type	Description	Default
name	str	The site name of the parameters, if end with sigma will return Log Normal distribution, otherwise, return Normal distribution	required

Returns:

Type	Description
Distribution	typing.Any: description

Source code in src/inspeqtor/v1/probabilistic.py

def default_priors_fn(name: str, shape: tuple[int, ...]) -> dist.Distribution:
    """This is a default prior function for the `dense_layer`

    Args:
        name (str): The site name of the parameters, if end with `sigma` will return Log Normal distribution,
                          otherwise, return Normal distribution

    Returns:
        typing.Any: _description_
    """
    if name.endswith("bias"):
        return dist.LogNormal(0, 1).expand(shape)

    return dist.Normal(0, 1).expand(shape)

dense_layer

dense_layer(
    x: ndarray,
    name: str,
    in_features: int,
    out_features: int,
    priors_fn: Callable[
        [str, tuple[int, ...]], Distribution
    ] = default_priors_fn,
)

A custom probabilistic dense layer for neural network model. This function intended to be used with numpyro

Parameters:

Name	Type	Description	Default
x	ndarray	The input x	required
name	str	Site name of the layer	required
in_features	int	The size of the feature.	required
out_features	int	The desired size of the output feature.	required
priors_fn	Callable[[str], Distribution]	The prior function to be used for initializing prior distribution. Defaults to default_priors_fn.	default_priors_fn

Returns:

Type	Description
	typing.Any: Output of the layer given x.

Source code in src/inspeqtor/v1/probabilistic.py

def dense_layer(
    x: jnp.ndarray,
    name: str,
    in_features: int,
    out_features: int,
    priors_fn: typing.Callable[
        [str, tuple[int, ...]], dist.Distribution
    ] = default_priors_fn,
):
    """A custom probabilistic dense layer for neural network model.
    This function intended to be used with `numpyro`

    Args:
        x (jnp.ndarray): The input x
        name (str): Site name of the layer
        in_features (int): The size of the feature.
        out_features (int): The desired size of the output feature.
        priors_fn (typing.Callable[[str], dist.Distribution], optional): The prior function to be used for initializing prior distribution. Defaults to default_priors_fn.

    Returns:
        typing.Any: Output of the layer given x.
    """
    w_name = f"{name}.kernel"
    w = numpyro.sample(
        w_name,
        priors_fn(w_name, (in_features, out_features)).to_event(2),  # type: ignore
    )
    b_name = f"{name}.bias"
    b = numpyro.sample(
        b_name,
        priors_fn(b_name, (out_features,)).to_event(1),  # type: ignore
    )
    return batched_matmul(x, w, b)  # type: ignore

init_default

init_default(params_name: str)

The initialization function for deterministic dense layer

Parameters:

Name	Type	Description	Default
params_name	str	The site name	required

Raises:

Type	Description
ValueError	Unsupport site name

Returns:

Type	Description
	typing.Any: The function to be used for parameters init given site name.

Source code in src/inspeqtor/v1/probabilistic.py

def init_default(params_name: str):
    """The initialization function for deterministic dense layer

    Args:
        params_name (str): The site name

    Raises:
        ValueError: Unsupport site name

    Returns:
        typing.Any: The function to be used for parameters init given site name.
    """
    if params_name.endswith("kernel"):
        return jnp.ones
    elif params_name.endswith("bias"):
        return lambda x: 0.1 * jnp.ones(x)
    else:
        raise ValueError("Unsupport param name")

dense_deterministic_layer

dense_deterministic_layer(
    x,
    name: str,
    in_features: int,
    out_features: int,
    batch_shape: tuple[int, ...] = (),
    init_fn=init_default,
)

The deterministic dense layer, to be used with SVI optimizer.

Parameters:

Name	Type	Description	Default
x	Any	The input feature	required
name	str	The site name	required
in_features	int	The size of the input features	required
out_features	int	The desired size of the output features.	required
batch_shape	tuple[int, ...]	The batch size of the x. Defaults to ().	()
init_fn	Any	Initilization function of the model parameters. Defaults to init_default.	init_default

Returns:

Type	Description
	typing.Any: The output of the layer given x.

Source code in src/inspeqtor/v1/probabilistic.py

def dense_deterministic_layer(
    x,
    name: str,
    in_features: int,
    out_features: int,
    batch_shape: tuple[int, ...] = (),
    init_fn=init_default,
):
    """The deterministic dense layer, to be used with SVI optimizer.

    Args:
        x (typing.Any): The input feature
        name (str): The site name
        in_features (int): The size of the input features
        out_features (int): The desired size of the output features.
        batch_shape (tuple[int, ...], optional): The batch size of the x. Defaults to ().
        init_fn (typing.Any, optional): Initilization function of the model parameters. Defaults to init_default.

    Returns:
        typing.Any: The output of the layer given x.
    """
    # Sample weights - shape (in_features, out_features)
    weight_shape = batch_shape + (in_features, out_features)
    W_name = f"{name}.kernel"
    W = numpyro.param(
        W_name,
        init_fn(W_name)(shape=weight_shape),  # type: ignore
    )

    # Sample bias - shape (out_features,)
    bias_shape = batch_shape + (out_features,)
    b_name = f"{name}.bias"
    b = numpyro.param(b_name, init_fn(b_name)(shape=bias_shape))  # type: ignore

    return batched_matmul(x, W, b)  # type: ignore

make_posteriors_fn

make_posteriors_fn(
    key: ndarray, guide, params, num_samples=10000
)

Make the posterior distribution function that will return the posterior of parameter of the given name, from guide and parameters.

Parameters:

Name	Type	Description	Default
guide	Any	The guide function	required
params	Any	The parameters of the guide	required
num_samples	int	The sample size. Defaults to 10000.	10000

Returns:

Type	Description
	typing.Any: A function of parameter name that return the sample from the posterior distribution of the parameters.

Source code in src/inspeqtor/v1/probabilistic.py

def make_posteriors_fn(key: jnp.ndarray, guide, params, num_samples=10000):
    """Make the posterior distribution function that will
    return the posterior of parameter of the given name, from guide and parameters.

    Args:
        guide (typing.Any): The guide function
        params (typing.Any): The parameters of the guide
        num_samples (int, optional): The sample size. Defaults to 10000.

    Returns:
        typing.Any: A function of parameter name that return the sample from the posterior distribution of the parameters.
    """

    warnings.warn(
        "This function is deprecated",
        category=DeprecationWarning,
        stacklevel=2,
    )

    posterior_distribution = Predictive(
        model=guide, params=params, num_samples=num_samples
    )(key)

    posterior_dict = construct_normal_prior_from_samples(posterior_distribution)

    def posteriors_fn(param_name: str):
        return posterior_dict[param_name]

    return posteriors_fn

auto_diagonal_normal_guide

auto_diagonal_normal_guide(
    model,
    *args,
    block_sample: bool = False,
    init_loc_fn=zeros,
    key: ndarray = key(0),
)

Automatically generate guide from given model. Expected to be initialized with the example input of the model. The given input should also including the observed site. The blocking capability is intended to be used in the when the guide will be used with its corresponding model in anothe model. This is the avoid site name duplication, while allows for model to use newly sample from the guide.

Parameters:

Name	Type	Description	Default
model	Any	The probabilistic model.	required
block_sample	bool	Flag to block the sample site. Defaults to False.	False
init_loc_fn	Any	Initialization of guide parameters function. Defaults to jnp.zeros.	zeros

Returns:

Type	Description
	typing.Any: description

Source code in src/inspeqtor/v1/probabilistic.py

def auto_diagonal_normal_guide(
    model,
    *args,
    block_sample: bool = False,
    init_loc_fn=jnp.zeros,
    key: jnp.ndarray = jax.random.key(0),
):
    """Automatically generate guide from given model. Expected to be initialized with the example input of the model.
    The given input should also including the observed site.
    The blocking capability is intended to be used in the when the guide will be used with its corresponding model in anothe model.
    This is the avoid site name duplication, while allows for model to use newly sample from the guide.

    Args:
        model (typing.Any): The probabilistic model.
        block_sample (bool, optional): Flag to block the sample site. Defaults to False.
        init_loc_fn (typing.Any, optional): Initialization of guide parameters function. Defaults to jnp.zeros.

    Returns:
        typing.Any: _description_
    """
    model_trace = handlers.trace(handlers.seed(model, key)).get_trace(*args)
    # get the trace of the model
    # Then get only the sample site with observed equal to false
    sample_sites = [v for k, v in model_trace.items() if v["type"] == "sample"]
    non_observed_sites = [v for v in sample_sites if not v["is_observed"]]
    params_sites = [
        {"name": v["name"], "shape": v["value"].shape} for v in non_observed_sites
    ]

    def guide(
        *args,
        **kwargs,
    ):
        params_loc = {
            param["name"]: numpyro.param(
                f"{param['name']}_loc", init_loc_fn(param["shape"])
            )
            for param in params_sites
        }

        params_scale = {
            param["name"]: numpyro.param(
                f"{param['name']}_scale",
                0.1 * jnp.ones(param["shape"]),
                constraint=dist.constraints.softplus_positive,
            )
            for param in params_sites
        }

        samples = {}

        if block_sample:
            with handlers.block():
                # Sample from Normal distribution
                for (k_loc, v_loc), (k_scale, v_scale) in zip(
                    params_loc.items(), params_scale.items(), strict=True
                ):
                    s = numpyro.sample(
                        k_loc,
                        dist.Normal(v_loc, v_scale).to_event(),  # type: ignore
                    )
                    samples[k_loc] = s
        else:
            # Sample from Normal distribution
            for (k_loc, v_loc), (k_scale, v_scale) in zip(
                params_loc.items(), params_scale.items(), strict=True
            ):
                s = numpyro.sample(
                    k_loc,
                    dist.Normal(v_loc, v_scale).to_event(),  # type: ignore
                )
                samples[k_loc] = s

        return samples

    return guide

auto_diagonal_normal_guide_v2

auto_diagonal_normal_guide_v2(
    model,
    *args,
    init_dist_fn=init_normal_dist_fn,
    init_params_fn=init_params_fn,
    block_sample: bool = False,
    key: ndarray = key(0),
)

Automatically generate guide from given model. Expected to be initialized with the example input of the model. The given input should also including the observed site. The blocking capability is intended to be used in the when the guide will be used with its corresponding model in anothe model. This is the avoid site name duplication, while allows for model to use newly sample from the guide.

Parameters:

Name	Type	Description	Default
model	Any	The probabilistic model.	required
block_sample	bool	Flag to block the sample site. Defaults to False.	False
init_loc_fn	Any	Initialization of guide parameters function. Defaults to jnp.zeros.	required

Returns:

Type	Description
	typing.Any: description

Source code in src/inspeqtor/v1/probabilistic.py

def auto_diagonal_normal_guide_v2(
    model,
    *args,
    init_dist_fn=init_normal_dist_fn,
    init_params_fn=init_params_fn,
    block_sample: bool = False,
    key: jnp.ndarray = jax.random.key(0),
):
    """Automatically generate guide from given model. Expected to be initialized with the example input of the model.
    The given input should also including the observed site.
    The blocking capability is intended to be used in the when the guide will be used with its corresponding model in anothe model.
    This is the avoid site name duplication, while allows for model to use newly sample from the guide.

    Args:
        model (typing.Any): The probabilistic model.
        block_sample (bool, optional): Flag to block the sample site. Defaults to False.
        init_loc_fn (typing.Any, optional): Initialization of guide parameters function. Defaults to jnp.zeros.

    Returns:
        typing.Any: _description_
    """
    # get the trace of the model
    model_trace = handlers.trace(handlers.seed(model, key)).get_trace(*args)
    # Then get only the sample site with observed equal to false
    sample_sites = [v for k, v in model_trace.items() if v["type"] == "sample"]
    non_observed_sites = [v for v in sample_sites if not v["is_observed"]]
    params_sites = [
        {"name": v["name"], "shape": v["value"].shape} for v in non_observed_sites
    ]

    def sample_fn(
        params_loc: dict[str, typing.Any], params_scale: dict[str, typing.Any]
    ):
        samples = {}
        # Sample from Normal distribution
        for (k_loc, v_loc), (k_scale, v_scale) in zip(
            params_loc.items(), params_scale.items(), strict=True
        ):
            s = numpyro.sample(
                k_loc,
                init_dist_fn(k_loc)(v_loc, v_scale).to_event(),  # type: ignore
            )
            samples[k_loc] = s

        return samples

    def guide(
        *args,
        **kwargs,
    ):
        params_loc = {
            param["name"]: init_params_fn(f"{param['name']}_loc", param["shape"])
            for param in params_sites
        }

        params_scale = {
            param["name"]: init_params_fn(f"{param['name']}_scale", param["shape"])
            for param in params_sites
        }

        if block_sample:
            with handlers.block():
                samples = sample_fn(params_loc, params_scale)
        else:
            samples = sample_fn(params_loc, params_scale)

        return samples

    return guide

make_predictive_fn_v2

make_predictive_fn_v2(model, guide, params, shots: int)

Make a postirior predictive model function from model, guide, SVI parameters, and the number of shots. This version relied explicitly on the guide and variational parameters. It might be slow than the first version.

Parameters:

Name	Type	Description	Default
model	Any	Probabilistic model.	required
guide	Any	Gudie corresponded to the model	required
params	Any	SVI parameters of the guide	required
shots	int	The number of shots	required

Returns:

Type	Description
	typing.Any: The posterior predictive model.

Source code in src/inspeqtor/v1/probabilistic.py

def make_predictive_fn_v2(
    model,
    guide,
    params,
    shots: int,
):
    """Make a postirior predictive model function from model, guide, SVI parameters, and the number of shots.
    This version relied explicitly on the guide and variational parameters. It might be slow than the first version.

    Args:
        model (typing.Any): Probabilistic model.
        guide (typing.Any): Gudie corresponded to the model
        params (typing.Any): SVI parameters of the guide
        shots (int): The number of shots

    Returns:
        typing.Any: The posterior predictive model.
    """
    predictive = Predictive(
        model, guide=guide, params=params, num_samples=shots, return_sites=["obs"]
    )

    def predictive_fn(*args, **kwargs):
        return predictive(*args, **kwargs)["obs"]

    return predictive_fn

make_predictive_SGM_model

make_predictive_SGM_model(
    model: Module,
    model_params,
    output_to_expectation_values_fn,
    shots: int,
)

Make a predictive model from given SGM model, the model parameters, and number of shots.

Parameters:

Name	Type	Description	Default
model	Module	Flax model	required
model_params	Any	The model parameters.	required
shots	int	The number of shots.	required

Source code in src/inspeqtor/v1/probabilistic.py

def make_predictive_SGM_model(
    model: nn.Module, model_params, output_to_expectation_values_fn, shots: int
):
    """Make a predictive model from given SGM model, the model parameters, and number of shots.

    Args:
        model (nn.Module): Flax model
        model_params (typing.Any): The model parameters.
        shots (int): The number of shots.
    """

    warnings.warn(
        "This function is deprecated",
        category=DeprecationWarning,
        stacklevel=2,
    )

    def predictive_model(
        key: jnp.ndarray, control_param: jnp.ndarray, unitaries: jnp.ndarray
    ):
        output = model.apply(model_params, control_param)
        predicted_expvals = output_to_expectation_values_fn(output, unitaries)

        return binary_to_eigenvalue(
            jax.vmap(jax.random.bernoulli, in_axes=(0, None))(
                jax.random.split(key, shots),
                expectation_value_to_prob_minus(predicted_expvals),
            ).astype(int)
        ).mean(axis=0)

    return predictive_model

make_predictive_MCDGM_model

make_predictive_MCDGM_model(model: Module, model_params)

Make a predictive model from given Monte-Carlo Dropout Graybox model, and the model parameters.

Parameters:

Name	Type	Description	Default
model	Module	Monte-Carlo Dropout Graybox model	required
model_params	Any	The model parameters	required

Source code in src/inspeqtor/v1/probabilistic.py

def make_predictive_MCDGM_model(model: nn.Module, model_params):
    """Make a predictive model from given Monte-Carlo Dropout Graybox model, and the model parameters.

    Args:
        model (nn.Module): Monte-Carlo Dropout Graybox model
        model_params (typing.Any): The model parameters
    """

    def predictive_model(
        key: jnp.ndarray, control_param: jnp.ndarray, unitaries: jnp.ndarray
    ):
        wo_params = model.apply(
            model_params,
            control_param,
            rngs={"dropout": key},
        )

        predicted_expvals = get_predict_expectation_value(
            wo_params,  # type: ignore
            unitaries,
            default_expectation_values_order,
        )

        return predicted_expvals

    return predictive_model

make_predictive_resampling_model

make_predictive_resampling_model(
    predictive_fn: Callable[[ndarray, ndarray], ndarray],
    shots: int,
) -> Callable[[ndarray, ndarray, ndarray], ndarray]

Make a binary predictive model from given SGM model, the model parameters, and number of shots.

Parameters:

Name	Type	Description	Default
predictive_fn	Callable[[ndarray, ndarray], ndarray]	The predictive_fn embeded with the SGM model.	required
shots	int	The number of shots.	required

Returns:

Type	Description
Callable[[ndarray, ndarray, ndarray], ndarray]	typing.Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray], jnp.ndarray]: Binary predictive model.

Source code in src/inspeqtor/v1/probabilistic.py

def make_predictive_resampling_model(
    predictive_fn: typing.Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray], shots: int
) -> typing.Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray], jnp.ndarray]:
    """Make a binary predictive model from given SGM model, the model parameters, and number of shots.

    Args:
        predictive_fn (typing.Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]): The predictive_fn embeded with the SGM model.
        shots (int): The number of shots.

    Returns:
        typing.Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray], jnp.ndarray]: Binary predictive model.
    """

    def predictive_model(
        key: jnp.ndarray, control_parameters: jnp.ndarray, unitaries: jnp.ndarray
    ):
        predicted_expvals = predictive_fn(control_parameters, unitaries)

        return binary_to_eigenvalue(
            jax.vmap(jax.random.bernoulli, in_axes=(0, None))(
                jax.random.split(key, shots),
                expectation_value_to_prob_minus(predicted_expvals),
            ).astype(int)
        ).mean(axis=0)

    return predictive_model

make_probabilistic_graybox_model

make_probabilistic_graybox_model(model, adapter_fn)

This function make a probabilistic graybox model using custom numpyro BNN model

Parameters:

Name	Type	Description	Default
model	_type_	description	required
adapter_fn	_type_	description	required

Source code in src/inspeqtor/v1/probabilistic.py

def make_probabilistic_graybox_model(model, adapter_fn):
    """This function make a probabilistic graybox model using custom numpyro BNN model

    Args:
        model (_type_): _description_
        adapter_fn (_type_): _description_
    """

    def probabilistic_graybox_model(control_parameters, unitaries):
        samples_shape = control_parameters.shape[:-2]
        unitaries = jnp.broadcast_to(unitaries, samples_shape + unitaries.shape[-3:])

        # Predict from control parameters
        output = model(control_parameters)

        numpyro.deterministic("output", output)

        # With unitary and Wo, calculate expectation values
        expvals = adapter_fn(output, unitaries)

        return expvals

    return probabilistic_graybox_model

auto_diagonal_normal_guide_v3

auto_diagonal_normal_guide_v3(
    model,
    *args,
    init_dist_fn=bnn_init_dist_fn,
    init_params_fn=bnn_init_params_fn,
    dist_transform_fn=default_transform_dist_fn,
    block_sample: bool = False,
    key: ndarray = key(0),
)

Automatically generate guide from given model. Expected to be initialized with the example input of the model. The given input should also including the observed site. The blocking capability is intended to be used in the when the guide will be used with its corresponding model in anothe model. This is the avoid site name duplication, while allows for model to use newly sample from the guide.

Notes

This version enable even more flexible initialization strategy. This function intended to be able to be compatible with auto marginal guide.

Parameters:

Name	Type	Description	Default
model	Any	The probabilistic model.	required
block_sample	bool	Flag to block the sample site. Defaults to False.	False
init_loc_fn	Any	Initialization of guide parameters function. Defaults to jnp.zeros.	required

Returns:

Type	Description
	typing.Any: description

Source code in src/inspeqtor/v1/probabilistic.py

def auto_diagonal_normal_guide_v3(
    model,
    *args,
    init_dist_fn=bnn_init_dist_fn,
    init_params_fn=bnn_init_params_fn,
    dist_transform_fn=default_transform_dist_fn,
    block_sample: bool = False,
    key: jnp.ndarray = jax.random.key(0),
):
    """Automatically generate guide from given model. Expected to be initialized with the example input of the model.
    The given input should also including the observed site.
    The blocking capability is intended to be used in the when the guide will be used with its corresponding model in anothe model.
    This is the avoid site name duplication, while allows for model to use newly sample from the guide.

    Notes:
        This version enable even more flexible initialization strategy.
        This function intended to be able to be compatible with auto marginal guide.

    Args:
        model (typing.Any): The probabilistic model.
        block_sample (bool, optional): Flag to block the sample site. Defaults to False.
        init_loc_fn (typing.Any, optional): Initialization of guide parameters function. Defaults to jnp.zeros.

    Returns:
        typing.Any: _description_
    """
    # get the trace of the model
    model_trace = handlers.trace(handlers.seed(model, key)).get_trace(*args)
    # Then get only the sample site with observed equal to false
    sample_sites = [v for k, v in model_trace.items() if v["type"] == "sample"]
    non_observed_sites = [v for v in sample_sites if not v["is_observed"]]
    params_sites = [
        {"name": v["name"], "shape": v["value"].shape} for v in non_observed_sites
    ]

    def sample_fn(
        params_loc: dict[str, typing.Any], params_scale: dict[str, typing.Any]
    ):
        samples = {}
        # Sample from Normal distribution
        for (k_loc, v_loc), (k_scale, v_scale) in zip(
            params_loc.items(), params_scale.items(), strict=True
        ):
            s = numpyro.sample(
                k_loc,
                dist_transform_fn(k_loc, init_dist_fn(k_loc)(v_loc, v_scale)),  # type: ignore
            )
            samples[k_loc] = s

        return samples

    def guide(
        *args,
        **kwargs,
    ):
        params_loc = {
            param["name"]: init_params_fn(f"{param['name']}_loc", param["shape"])
            for param in params_sites
        }

        params_scale = {
            param["name"]: init_params_fn(f"{param['name']}_scale", param["shape"])
            for param in params_sites
        }

        with numpyro.util.optional(block_sample, handlers.block()):
            samples = sample_fn(params_loc, params_scale)

        return samples

    return guide

make_posterior_fn

make_posterior_fn(
    params, get_dist_fn: Callable[[str], Any]
)

This function create a posterior function to make a posterior predictive model

Examples:

posterior_model = sq.probabilistic.make_probabilistic_model(
    predictive_model=partial(
            probabilistic_graybox_model,
            priors_fn=make_posterior_fn(result.params, init_bnn_dist_fn),
        ),
    log_expectation_values=True,
)

Parameters:

Name	Type	Description	Default
params	_type_	The variational parameters from SVI	required
get_dist_fn	Callable[[str], Distribution]	The function that return function given name	required

Source code in src/inspeqtor/v1/probabilistic.py

def make_posterior_fn(params, get_dist_fn: typing.Callable[[str], typing.Any]):
    """This function create a posterior function to make a posterior predictive model

    Examples:
        ```python
        posterior_model = sq.probabilistic.make_probabilistic_model(
            predictive_model=partial(
                    probabilistic_graybox_model,
                    priors_fn=make_posterior_fn(result.params, init_bnn_dist_fn),
                ),
            log_expectation_values=True,
        )
        ```

    Args:
        params (_type_): The variational parameters from SVI
        get_dist_fn (typing.Callable[[str], dist.Distribution]): The function that return function given name
    """

    def posterior_fn(name: str, shape: tuple[int, ...]):
        return get_dist_fn(name)(params[name + "_loc"], params[name + "_scale"])

    return posterior_fn