BOED (Bayesian Optimal Experiment Design)

inspeqtor.boed

AuxEntry

The auxillary entry returned by loss function

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

class AuxEntry(typing.NamedTuple):
    """The auxillary entry returned by loss function"""

    terms: jnp.ndarray | None
    eig: jnp.ndarray

safe_shape

safe_shape(a: Any) -> tuple[int, ...] | str

Safely get the shape of the object

Parameters:

Name	Type	Description	Default
a	Any	Expect the object to be jnp.ndarray	required

Returns:

Type	Description
tuple[int, ...] | str	tuple[int, ...] | str: Either return the shape of a
tuple[int, ...] | str	or string representation of the type

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

def safe_shape(a: typing.Any) -> tuple[int, ...] | str:
    """Safely get the shape of the object

    Args:
        a (typing.Any): Expect the object to be jnp.ndarray

    Returns:
        tuple[int, ...] | str: Either return the shape of `a`
        or string representation of the type
    """
    try:
        assert isinstance(a, jnp.ndarray)
        return a.shape
    except AttributeError:
        return str(type(a))

report_shape

report_shape(a: PyTree) -> PyTree

Report the shape of pytree

Parameters:

Name	Type	Description	Default
a	PyTree	The pytree to be report.	required

Returns:

Type	Description
PyTree	jaxtyping.PyTree: The shape of pytree.

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

def report_shape(a: jaxtyping.PyTree) -> jaxtyping.PyTree:
    """Report the shape of pytree

    Args:
        a (jaxtyping.PyTree): The pytree to be report.

    Returns:
        jaxtyping.PyTree: The shape of pytree.
    """
    return jax.tree.map(safe_shape, a)

lexpand

lexpand(a: ndarray, *dimensions: int) -> ndarray

Expand tensor, adding new dimensions on left.

Parameters:

Name	Type	Description	Default
a	ndarray	expand the dimension on the left with given dimension arguments.	required

Returns:

Type	Description
ndarray	jnp.ndarray: New array with shape (*dimension + a.shape)

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

def lexpand(a: jnp.ndarray, *dimensions: int) -> jnp.ndarray:
    """Expand tensor, adding new dimensions on left.

    Args:
        a (jnp.ndarray): expand the dimension on the left with given dimension arguments.

    Returns:
        jnp.ndarray: New array with shape (*dimension + a.shape)
    """
    return jnp.broadcast_to(a, dimensions + a.shape)

random_split_index

random_split_index(
    rng_key: ndarray, num_samples: int, test_size: int
) -> tuple[ndarray, ndarray]

Create the randomly spilt of indice to two set, with one of test_size and another as the rest.

Parameters:

Name	Type	Description	Default
rng_key	ndarray	The random key	required
num_samples	int	The size of total sample size	required
test_size	int	The size of test set	required

Returns:

Type	Description
tuple[ndarray, ndarray]	tuple[jnp.ndarray, jnp.ndarray]: Array of train indice and array of test indice.

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

def random_split_index(
    rng_key: jnp.ndarray, num_samples: int, test_size: int
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Create the randomly spilt of indice to two set, with one of test_size and another as the rest.

    Args:
        rng_key (jnp.ndarray): The random key
        num_samples (int): The size of total sample size
        test_size (int): The size of test set

    Returns:
        tuple[jnp.ndarray, jnp.ndarray]: Array of train indice and array of test indice.
    """
    idx = jax.random.permutation(rng_key, jnp.arange(num_samples))
    return idx[test_size:], idx[:test_size]

marginal_loss

marginal_loss(
    model: Callable,
    marginal_guide: Callable,
    design: ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: int,
    evaluation: bool = False,
) -> Callable[
    [ArrayTree, ndarray], tuple[ndarray, AuxEntry]
]

The marginal loss implemented following https://docs.pyro.ai/en/dev/contrib.oed.html#pyro.contrib.oed.eig.marginal_eig

Parameters:

Name	Type	Description	Default
model	Callable	The probabilistic model	required
marginal_guide	Callable	The custom guide	required
design	ndarray	Possible designs of the experiment	required
observation_labels	list[str]	The list of string of observations	required
target_labels	list[str]	The target latent parameters to be optimized for	required
num_particles	int	The number of independent trials	required
evaluation	bool	True for actual evalution of the EIG. Defaults to False.	False

Returns:

Type	Description
Callable[[ArrayTree, ndarray], tuple[ndarray, AuxEntry]]	typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG

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

def marginal_loss(
    model: typing.Callable,
    marginal_guide: typing.Callable,
    design: jnp.ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: int,
    evaluation: bool = False,
) -> typing.Callable[[chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]]:
    """The marginal loss implemented following
    https://docs.pyro.ai/en/dev/contrib.oed.html#pyro.contrib.oed.eig.marginal_eig

    Args:
        model (typing.Callable): The probabilistic model
        marginal_guide (typing.Callable): The custom guide
        design (jnp.ndarray): Possible designs of the experiment
        observation_labels (list[str]): The list of string of observations
        target_labels (list[str]): The target latent parameters to be optimized for
        num_particles (int): The number of independent trials
        evaluation (bool, optional): True for actual evalution of the EIG. Defaults to False.

    Returns:
        typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG
    """

    # Marginal loss
    def loss_fn(param, key: jnp.ndarray) -> tuple[jnp.ndarray, AuxEntry]:
        expanded_design = lexpand(design, num_particles)
        # vectorized(model, num_particles)
        # Sample from p(y | d)
        key, subkey = jax.random.split(key)
        trace = handlers.trace(handlers.seed(model, subkey)).get_trace(
            expanded_design,
            *args,
        )
        y_dict = {
            observation_label: trace[observation_label]["value"]
            for observation_label in observation_labels
        }

        # Run through q(y | d)
        key, subkey = jax.random.split(key)
        conditioned_marginal_guide = handlers.condition(marginal_guide, data=y_dict)
        cond_trace = handlers.trace(
            handlers.substitute(
                handlers.seed(conditioned_marginal_guide, subkey), data=param
            )
        ).get_trace(
            expanded_design,
            *args,
            observation_labels=observation_labels,
            target_labels=target_labels,
        )
        # Compute the log prob of observing the data
        terms = -1 * jnp.array(
            [
                cond_trace[observation_label]["fn"].log_prob(
                    cond_trace[observation_label]["value"]
                )
                for observation_label in observation_labels
            ]
        ).sum(axis=0)

        if evaluation:
            terms += jnp.array(
                [
                    trace[observation_label]["fn"].log_prob(
                        trace[observation_label]["value"]
                    )
                    for observation_label in observation_labels
                ]
            ).sum(axis=0)

        agg_loss, loss = _safe_mean_terms_v2(terms)
        return agg_loss, AuxEntry(terms=terms, eig=loss)

    return loss_fn

vnmc_eig_loss

vnmc_eig_loss(
    model: Callable,
    marginal_guide: Callable,
    design: ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int],
    evaluation: bool = False,
) -> Callable[
    [ArrayTree, ndarray], tuple[ndarray, AuxEntry]
]

The VNMC loss implemented following https://docs.pyro.ai/en/dev/_modules/pyro/contrib/oed/eig.html#vnmc_eig

Parameters:

Name	Type	Description	Default
model	Callable	The probabilistic model	required
marginal_guide	Callable	The custom guide	required
design	ndarray	Possible designs of the experiment	required
observation_labels	list[str]	The list of string of observations	required
target_labels	list[str]	The target latent parameters to be optimized for	required
num_particles	int	The number of independent trials	required
evaluation	bool	True for actual evalution of the EIG. Defaults to False.	False

Returns:

Type	Description
Callable[[ArrayTree, ndarray], tuple[ndarray, AuxEntry]]	typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG

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

@warn_not_tested_function
def vnmc_eig_loss(
    model: typing.Callable,
    marginal_guide: typing.Callable,
    design: jnp.ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int],
    evaluation: bool = False,
) -> typing.Callable[[chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]]:
    """The VNMC loss implemented following
    https://docs.pyro.ai/en/dev/_modules/pyro/contrib/oed/eig.html#vnmc_eig

    Args:
        model (typing.Callable): The probabilistic model
        marginal_guide (typing.Callable): The custom guide
        design (jnp.ndarray): Possible designs of the experiment
        observation_labels (list[str]): The list of string of observations
        target_labels (list[str]): The target latent parameters to be optimized for
        num_particles (int): The number of independent trials
        evaluation (bool, optional): True for actual evalution of the EIG. Defaults to False.

    Returns:
        typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG
    """

    # Marginal loss
    def loss_fn(param, key: jnp.ndarray) -> tuple[jnp.ndarray, AuxEntry]:
        N, M = num_particles

        expanded_design = lexpand(design, N)

        # Sample from p(y, theta | d)
        key, subkey = jax.random.split(key)
        trace = handlers.trace(handlers.seed(model, subkey)).get_trace(
            expanded_design,
            *args,
        )
        y_dict = {
            observation_label: trace[observation_label]["value"]
            for observation_label in observation_labels
        }

        # Sample M times from q(theta | y, d) for each y
        key, subkey = jax.random.split(key)
        reexpanded_design = lexpand(expanded_design, M)
        conditioned_marginal_guide = handlers.condition(marginal_guide, data=y_dict)
        cond_trace = handlers.trace(
            handlers.substitute(
                handlers.seed(conditioned_marginal_guide, subkey), data=param
            )
        ).get_trace(
            reexpanded_design,
            *args,
            observation_labels=observation_labels,
            target_labels=target_labels,
        )

        theta_y_dict = {
            target_label: cond_trace[target_label]["value"]
            for target_label in target_labels
        }
        theta_y_dict.update(y_dict)

        # Re-run that through the model to compute the joint
        key, subkey = jax.random.split(key)
        conditioned_model = handlers.condition(model, data=theta_y_dict)
        conditioned_model_trace = handlers.trace(
            handlers.seed(conditioned_model, subkey)
        ).get_trace(
            reexpanded_design,
            *args,
        )

        # Compute the log prob of observing the data
        terms = -1 * jnp.array(
            [
                cond_trace[target_label]["fn"].log_prob(
                    cond_trace[target_label]["value"]
                )
                for target_label in target_labels
            ]
        ).sum(axis=0)

        terms += jnp.array(
            [
                conditioned_model_trace[target_label]["fn"].log_prob(
                    conditioned_model_trace[target_label]["value"]
                )
                for target_label in target_labels
            ]
        ).sum(axis=0)

        terms += jnp.array(
            [
                conditioned_model_trace[observation_label]["fn"].log_prob(
                    conditioned_model_trace[observation_label]["value"]
                )
                for observation_label in observation_labels
            ]
        ).sum(axis=0)

        terms = -jax.scipy.special.logsumexp(terms, axis=0) + jnp.log(M)

        if evaluation:
            terms += jnp.array(
                [
                    trace[observation_label]["fn"].log_prob(
                        trace[observation_label]["value"]
                    )
                    for observation_label in observation_labels
                ]
            ).sum(axis=0)

        agg_loss, loss = _safe_mean_terms_v2(terms)
        return agg_loss, AuxEntry(terms=terms, eig=loss)

    return loss_fn

init_params_from_guide

init_params_from_guide(
    marginal_guide: Callable,
    *args,
    key: ndarray,
    design: ndarray,
) -> ArrayTree

Initlalize parameters of marginal guide.

Parameters:

Name	Type	Description	Default
marginal_guide	Callable	Marginal guide to be used with marginal eig	required
key	ndarray	Random Key	required
design	ndarray	Example of the designs of the experiment	required

Returns:

Type	Description
ArrayTree	chex.ArrayTree: Random parameters for marginal guide to be optimized.

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

def init_params_from_guide(
    marginal_guide: typing.Callable,
    *args,
    key: jnp.ndarray,
    design: jnp.ndarray,
) -> chex.ArrayTree:
    """Initlalize parameters of marginal guide.

    Args:
        marginal_guide (typing.Callable): Marginal guide to be used with marginal eig
        key (jnp.ndarray): Random Key
        design (jnp.ndarray): Example of the designs of the experiment

    Returns:
        chex.ArrayTree: Random parameters for marginal guide to be optimized.
    """
    key, subkey = jax.random.split(key)
    # expanded_design = lexpand(design, num_particles)
    marginal_guide_trace = handlers.trace(
        handlers.seed(marginal_guide, subkey)
    ).get_trace(design, *args, observation_labels=[], target_labels=[])

    # Get only nodes that are parameters
    params = {
        name: node["value"]
        for name, node in marginal_guide_trace.items()
        if node["type"] == "param"
    }

    return params

opt_eig_ape_loss

opt_eig_ape_loss(
    loss_fn: Callable[
        [ArrayTree, ndarray], tuple[ndarray, AuxEntry]
    ],
    params: ArrayTree,
    num_steps: int,
    optim: GradientTransformation,
    key: ndarray,
    callbacks: list = [],
) -> ArrayTree

Optimize the EIG loss function.

Parameters:

Name	Type	Description	Default
loss_fn	Callable[[ArrayTree, ndarray], tuple[ndarray, AuxEntry]]	Loss function	required
params	ArrayTree	Initial parameter	required
num_steps	int	Number of optimization step	required
optim	GradientTransformation	Optax Optimizer	required
key	ndarray	Random key	required

Returns:

Type	Description
ArrayTree	tuple[chex.ArrayTree, list[typing.Any]]: Optimized parameters, and optimization history.

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

def opt_eig_ape_loss(
    loss_fn: typing.Callable[
        [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]
    ],
    params: chex.ArrayTree,
    num_steps: int,
    optim: optax.GradientTransformation,
    key: jnp.ndarray,
    callbacks: list = [],
) -> chex.ArrayTree:
    """Optimize the EIG loss function.

    Args:
        loss_fn (typing.Callable[[chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]]): Loss function
        params (chex.ArrayTree): Initial parameter
        num_steps (int): Number of optimization step
        optim (optax.GradientTransformation): Optax Optimizer
        key (jnp.ndarray): Random key

    Returns:
        tuple[chex.ArrayTree, list[typing.Any]]: Optimized parameters, and optimization history.
    """
    # Initialize the optimizer
    opt_state = optim.init(params)
    # jit the loss function
    loss_fn = jax.jit(loss_fn)

    for step in range(num_steps):
        key, subkey = jax.random.split(key)
        # Compute the loss and its gradient
        (loss, aux), grad = jax.value_and_grad(loss_fn, has_aux=True)(params, subkey)
        # Update the optimizer and params
        updates, opt_state = optim.update(grad, opt_state, params)
        params = optax.apply_updates(params, updates)

        # entry = (step, loss, aux)
        entry = HistoryEntry(step=step, loss=loss, aux=aux)

        for callback in callbacks:
            callback(entry)

    return params

estimate_eig

estimate_eig(
    key: ndarray,
    model: Callable,
    marginal_guide: Callable,
    design: ndarray,
    *args,
    optimizer: GradientTransformation,
    num_optimization_steps: int,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int] | int,
    final_num_particles: tuple[int, int]
    | int
    | None = None,
    loss_fn: Callable = marginal_loss,
    callbacks: list = [],
) -> tuple[ndarray, dict[str, Any]]

Optimize for marginal EIG

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required
model	Callable	Probabilistic model of the experiment	required
marginal_guide	Callable	The marginal guide of the experiment	required
design	ndarray	Possible designs of the experiment	required
optimizer	GradientTransformation	Optax optimizer	required
num_optimization_steps	int	Number of the optimization step	required
observation_labels	list[str]	The list of string of observations	required
target_labels	list[str]	The target latent parameters to be optimized for	required
num_particles	int	The number of independent trials	required
final_num_particles	int | None	Final independent trials to calculate marginal EIG. Defaults to None.	None

Returns:

Type	Description
tuple[ndarray, dict[str, Any]]	tuple[jnp.ndarray, dict[str, typing.Any]]: EIG, and tuple of optimized parameters and optimization history.

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

def estimate_eig(
    key: jnp.ndarray,
    model: typing.Callable,
    marginal_guide: typing.Callable,
    design: jnp.ndarray,
    *args,
    optimizer: optax.GradientTransformation,
    num_optimization_steps: int,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int] | int,
    final_num_particles: tuple[int, int] | int | None = None,
    loss_fn: typing.Callable = marginal_loss,
    callbacks: list = [],
) -> tuple[jnp.ndarray, dict[str, typing.Any]]:
    """Optimize for marginal EIG

    Args:
        key (jnp.ndarray): Random key
        model (typing.Callable): Probabilistic model of the experiment
        marginal_guide (typing.Callable): The marginal guide of the experiment
        design (jnp.ndarray): Possible designs of the experiment
        optimizer (optax.GradientTransformation): Optax optimizer
        num_optimization_steps (int): Number of the optimization step
        observation_labels (list[str]): The list of string of observations
        target_labels (list[str]): The target latent parameters to be optimized for
        num_particles (int): The number of independent trials
        final_num_particles (int | None, optional): Final independent trials to calculate marginal EIG. Defaults to None.

    Returns:
        tuple[jnp.ndarray, dict[str, typing.Any]]: EIG, and tuple of optimized parameters and optimization history.
    """
    # NOTE: In final evalution, if final_num_particles != num_particles,
    # the code will error because we train params with num_particles
    # the shape will mismatch
    # final_num_particles = final_num_particles or num_particles

    # Initialize the parameters by using trace from the marginal_guide
    key, subkey = jax.random.split(key)
    params = init_params_from_guide(
        marginal_guide,
        *args,
        key=subkey,
        design=design,
    )

    # Optimize the loss function first to get the optimal parameters
    # for marginal guide
    params = opt_eig_ape_loss(
        loss_fn=loss_fn(
            model,
            marginal_guide,
            design,
            *args,
            observation_labels=observation_labels,
            target_labels=target_labels,
            num_particles=num_particles,
            evaluation=False,
        ),
        params=params,
        num_steps=num_optimization_steps,
        optim=optimizer,
        key=subkey,
        callbacks=callbacks,
    )

    key, subkey = jax.random.split(key)
    # Evaluate the loss
    _, aux = loss_fn(
        model,
        marginal_guide,
        design,
        *args,
        observation_labels=observation_labels,
        target_labels=target_labels,
        num_particles=final_num_particles,
        evaluation=True,
    )(params, subkey)

    return aux.eig, {
        "params": params,
    }

vectorized_for_eig

vectorized_for_eig(model)

Vectorization function for the EIG function

Parameters:

Name	Type	Description	Default
model	Any	Probabilistic model.	required

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

def vectorized_for_eig(model):
    """Vectorization function for the EIG function

    Args:
        model (typing.Any): Probabilistic model.
    """

    def wrapper(
        design: jnp.ndarray,
        *args,
        **kwargs,
    ):
        # This wrapper has the same call signature as the probabilistic graybox model
        # Expect the design to has shape == (extra, design, feature)
        with plate_stack(prefix="vectorized_plate", sizes=[*design.shape[:2]]):
            return model(design, *args, **kwargs)

    return wrapper

inspeqtor.boed.safe_shape

safe_shape(a: Any) -> tuple[int, ...] | str

Safely get the shape of the object

Parameters:

Name	Type	Description	Default
a	Any	Expect the object to be jnp.ndarray	required

Returns:

Type	Description
tuple[int, ...] | str	tuple[int, ...] | str: Either return the shape of a
tuple[int, ...] | str	or string representation of the type

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

def safe_shape(a: typing.Any) -> tuple[int, ...] | str:
    """Safely get the shape of the object

    Args:
        a (typing.Any): Expect the object to be jnp.ndarray

    Returns:
        tuple[int, ...] | str: Either return the shape of `a`
        or string representation of the type
    """
    try:
        assert isinstance(a, jnp.ndarray)
        return a.shape
    except AttributeError:
        return str(type(a))

inspeqtor.boed.report_shape

report_shape(a: PyTree) -> PyTree

Report the shape of pytree

Parameters:

Name	Type	Description	Default
a	PyTree	The pytree to be report.	required

Returns:

Type	Description
PyTree	jaxtyping.PyTree: The shape of pytree.

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

def report_shape(a: jaxtyping.PyTree) -> jaxtyping.PyTree:
    """Report the shape of pytree

    Args:
        a (jaxtyping.PyTree): The pytree to be report.

    Returns:
        jaxtyping.PyTree: The shape of pytree.
    """
    return jax.tree.map(safe_shape, a)

inspeqtor.boed.lexpand

lexpand(a: ndarray, *dimensions: int) -> ndarray

Expand tensor, adding new dimensions on left.

Parameters:

Name	Type	Description	Default
a	ndarray	expand the dimension on the left with given dimension arguments.	required

Returns:

Type	Description
ndarray	jnp.ndarray: New array with shape (*dimension + a.shape)

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

def lexpand(a: jnp.ndarray, *dimensions: int) -> jnp.ndarray:
    """Expand tensor, adding new dimensions on left.

    Args:
        a (jnp.ndarray): expand the dimension on the left with given dimension arguments.

    Returns:
        jnp.ndarray: New array with shape (*dimension + a.shape)
    """
    return jnp.broadcast_to(a, dimensions + a.shape)

inspeqtor.boed.random_split_index

random_split_index(
    rng_key: ndarray, num_samples: int, test_size: int
) -> tuple[ndarray, ndarray]

Create the randomly spilt of indice to two set, with one of test_size and another as the rest.

Parameters:

Name	Type	Description	Default
rng_key	ndarray	The random key	required
num_samples	int	The size of total sample size	required
test_size	int	The size of test set	required

Returns:

Type	Description
tuple[ndarray, ndarray]	tuple[jnp.ndarray, jnp.ndarray]: Array of train indice and array of test indice.

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

def random_split_index(
    rng_key: jnp.ndarray, num_samples: int, test_size: int
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Create the randomly spilt of indice to two set, with one of test_size and another as the rest.

    Args:
        rng_key (jnp.ndarray): The random key
        num_samples (int): The size of total sample size
        test_size (int): The size of test set

    Returns:
        tuple[jnp.ndarray, jnp.ndarray]: Array of train indice and array of test indice.
    """
    idx = jax.random.permutation(rng_key, jnp.arange(num_samples))
    return idx[test_size:], idx[:test_size]

inspeqtor.boed.AuxEntry

The auxillary entry returned by loss function

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

class AuxEntry(typing.NamedTuple):
    """The auxillary entry returned by loss function"""

    terms: jnp.ndarray | None
    eig: jnp.ndarray

inspeqtor.boed.vectorized

vectorized(fn, *shape, name='vectorization_plate')

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

def vectorized(fn, *shape, name="vectorization_plate"):
    def wrapper_fn(*args, **kwargs):
        with plate_stack(name, [*shape]):
            return fn(*args, **kwargs)

    return wrapper_fn

inspeqtor.boed.marginal_loss

marginal_loss(
    model: Callable,
    marginal_guide: Callable,
    design: ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: int,
    evaluation: bool = False,
) -> Callable[
    [ArrayTree, ndarray], tuple[ndarray, AuxEntry]
]

The marginal loss implemented following https://docs.pyro.ai/en/dev/contrib.oed.html#pyro.contrib.oed.eig.marginal_eig

Parameters:

Name	Type	Description	Default
model	Callable	The probabilistic model	required
marginal_guide	Callable	The custom guide	required
design	ndarray	Possible designs of the experiment	required
observation_labels	list[str]	The list of string of observations	required
target_labels	list[str]	The target latent parameters to be optimized for	required
num_particles	int	The number of independent trials	required
evaluation	bool	True for actual evalution of the EIG. Defaults to False.	False

Returns:

Type	Description
Callable[[ArrayTree, ndarray], tuple[ndarray, AuxEntry]]	typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG

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

def marginal_loss(
    model: typing.Callable,
    marginal_guide: typing.Callable,
    design: jnp.ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: int,
    evaluation: bool = False,
) -> typing.Callable[[chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]]:
    """The marginal loss implemented following
    https://docs.pyro.ai/en/dev/contrib.oed.html#pyro.contrib.oed.eig.marginal_eig

    Args:
        model (typing.Callable): The probabilistic model
        marginal_guide (typing.Callable): The custom guide
        design (jnp.ndarray): Possible designs of the experiment
        observation_labels (list[str]): The list of string of observations
        target_labels (list[str]): The target latent parameters to be optimized for
        num_particles (int): The number of independent trials
        evaluation (bool, optional): True for actual evalution of the EIG. Defaults to False.

    Returns:
        typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG
    """

    # Marginal loss
    def loss_fn(param, key: jnp.ndarray) -> tuple[jnp.ndarray, AuxEntry]:
        expanded_design = lexpand(design, num_particles)
        # vectorized(model, num_particles)
        # Sample from p(y | d)
        key, subkey = jax.random.split(key)
        trace = handlers.trace(handlers.seed(model, subkey)).get_trace(
            expanded_design,
            *args,
        )
        y_dict = {
            observation_label: trace[observation_label]["value"]
            for observation_label in observation_labels
        }

        # Run through q(y | d)
        key, subkey = jax.random.split(key)
        conditioned_marginal_guide = handlers.condition(marginal_guide, data=y_dict)
        cond_trace = handlers.trace(
            handlers.substitute(
                handlers.seed(conditioned_marginal_guide, subkey), data=param
            )
        ).get_trace(
            expanded_design,
            *args,
            observation_labels=observation_labels,
            target_labels=target_labels,
        )
        # Compute the log prob of observing the data
        terms = -1 * jnp.array(
            [
                cond_trace[observation_label]["fn"].log_prob(
                    cond_trace[observation_label]["value"]
                )
                for observation_label in observation_labels
            ]
        ).sum(axis=0)

        if evaluation:
            terms += jnp.array(
                [
                    trace[observation_label]["fn"].log_prob(
                        trace[observation_label]["value"]
                    )
                    for observation_label in observation_labels
                ]
            ).sum(axis=0)

        agg_loss, loss = _safe_mean_terms_v2(terms)
        return agg_loss, AuxEntry(terms=terms, eig=loss)

    return loss_fn

inspeqtor.boed.vnmc_eig_loss

vnmc_eig_loss(
    model: Callable,
    marginal_guide: Callable,
    design: ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int],
    evaluation: bool = False,
) -> Callable[
    [ArrayTree, ndarray], tuple[ndarray, AuxEntry]
]

The VNMC loss implemented following https://docs.pyro.ai/en/dev/_modules/pyro/contrib/oed/eig.html#vnmc_eig

Parameters:

Name	Type	Description	Default
model	Callable	The probabilistic model	required
marginal_guide	Callable	The custom guide	required
design	ndarray	Possible designs of the experiment	required
observation_labels	list[str]	The list of string of observations	required
target_labels	list[str]	The target latent parameters to be optimized for	required
num_particles	int	The number of independent trials	required
evaluation	bool	True for actual evalution of the EIG. Defaults to False.	False

Returns:

Type	Description
Callable[[ArrayTree, ndarray], tuple[ndarray, AuxEntry]]	typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG

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

@warn_not_tested_function
def vnmc_eig_loss(
    model: typing.Callable,
    marginal_guide: typing.Callable,
    design: jnp.ndarray,
    *args,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int],
    evaluation: bool = False,
) -> typing.Callable[[chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]]:
    """The VNMC loss implemented following
    https://docs.pyro.ai/en/dev/_modules/pyro/contrib/oed/eig.html#vnmc_eig

    Args:
        model (typing.Callable): The probabilistic model
        marginal_guide (typing.Callable): The custom guide
        design (jnp.ndarray): Possible designs of the experiment
        observation_labels (list[str]): The list of string of observations
        target_labels (list[str]): The target latent parameters to be optimized for
        num_particles (int): The number of independent trials
        evaluation (bool, optional): True for actual evalution of the EIG. Defaults to False.

    Returns:
        typing.Callable[ [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry] ]: Loss function that return tuple of (1) Total loss, (2.1) Each terms without the average, (2.2) The EIG
    """

    # Marginal loss
    def loss_fn(param, key: jnp.ndarray) -> tuple[jnp.ndarray, AuxEntry]:
        N, M = num_particles

        expanded_design = lexpand(design, N)

        # Sample from p(y, theta | d)
        key, subkey = jax.random.split(key)
        trace = handlers.trace(handlers.seed(model, subkey)).get_trace(
            expanded_design,
            *args,
        )
        y_dict = {
            observation_label: trace[observation_label]["value"]
            for observation_label in observation_labels
        }

        # Sample M times from q(theta | y, d) for each y
        key, subkey = jax.random.split(key)
        reexpanded_design = lexpand(expanded_design, M)
        conditioned_marginal_guide = handlers.condition(marginal_guide, data=y_dict)
        cond_trace = handlers.trace(
            handlers.substitute(
                handlers.seed(conditioned_marginal_guide, subkey), data=param
            )
        ).get_trace(
            reexpanded_design,
            *args,
            observation_labels=observation_labels,
            target_labels=target_labels,
        )

        theta_y_dict = {
            target_label: cond_trace[target_label]["value"]
            for target_label in target_labels
        }
        theta_y_dict.update(y_dict)

        # Re-run that through the model to compute the joint
        key, subkey = jax.random.split(key)
        conditioned_model = handlers.condition(model, data=theta_y_dict)
        conditioned_model_trace = handlers.trace(
            handlers.seed(conditioned_model, subkey)
        ).get_trace(
            reexpanded_design,
            *args,
        )

        # Compute the log prob of observing the data
        terms = -1 * jnp.array(
            [
                cond_trace[target_label]["fn"].log_prob(
                    cond_trace[target_label]["value"]
                )
                for target_label in target_labels
            ]
        ).sum(axis=0)

        terms += jnp.array(
            [
                conditioned_model_trace[target_label]["fn"].log_prob(
                    conditioned_model_trace[target_label]["value"]
                )
                for target_label in target_labels
            ]
        ).sum(axis=0)

        terms += jnp.array(
            [
                conditioned_model_trace[observation_label]["fn"].log_prob(
                    conditioned_model_trace[observation_label]["value"]
                )
                for observation_label in observation_labels
            ]
        ).sum(axis=0)

        terms = -jax.scipy.special.logsumexp(terms, axis=0) + jnp.log(M)

        if evaluation:
            terms += jnp.array(
                [
                    trace[observation_label]["fn"].log_prob(
                        trace[observation_label]["value"]
                    )
                    for observation_label in observation_labels
                ]
            ).sum(axis=0)

        agg_loss, loss = _safe_mean_terms_v2(terms)
        return agg_loss, AuxEntry(terms=terms, eig=loss)

    return loss_fn

inspeqtor.boed.init_params_from_guide

init_params_from_guide(
    marginal_guide: Callable,
    *args,
    key: ndarray,
    design: ndarray,
) -> ArrayTree

Initlalize parameters of marginal guide.

Parameters:

Name	Type	Description	Default
marginal_guide	Callable	Marginal guide to be used with marginal eig	required
key	ndarray	Random Key	required
design	ndarray	Example of the designs of the experiment	required

Returns:

Type	Description
ArrayTree	chex.ArrayTree: Random parameters for marginal guide to be optimized.

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

def init_params_from_guide(
    marginal_guide: typing.Callable,
    *args,
    key: jnp.ndarray,
    design: jnp.ndarray,
) -> chex.ArrayTree:
    """Initlalize parameters of marginal guide.

    Args:
        marginal_guide (typing.Callable): Marginal guide to be used with marginal eig
        key (jnp.ndarray): Random Key
        design (jnp.ndarray): Example of the designs of the experiment

    Returns:
        chex.ArrayTree: Random parameters for marginal guide to be optimized.
    """
    key, subkey = jax.random.split(key)
    # expanded_design = lexpand(design, num_particles)
    marginal_guide_trace = handlers.trace(
        handlers.seed(marginal_guide, subkey)
    ).get_trace(design, *args, observation_labels=[], target_labels=[])

    # Get only nodes that are parameters
    params = {
        name: node["value"]
        for name, node in marginal_guide_trace.items()
        if node["type"] == "param"
    }

    return params

inspeqtor.boed.HistoryEntry

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

class HistoryEntry(typing.NamedTuple):
    step: int
    loss: jnp.ndarray
    aux: AuxEntry

inspeqtor.boed.opt_eig_ape_loss

opt_eig_ape_loss(
    loss_fn: Callable[
        [ArrayTree, ndarray], tuple[ndarray, AuxEntry]
    ],
    params: ArrayTree,
    num_steps: int,
    optim: GradientTransformation,
    key: ndarray,
    callbacks: list = [],
) -> ArrayTree

Optimize the EIG loss function.

Parameters:

Name	Type	Description	Default
loss_fn	Callable[[ArrayTree, ndarray], tuple[ndarray, AuxEntry]]	Loss function	required
params	ArrayTree	Initial parameter	required
num_steps	int	Number of optimization step	required
optim	GradientTransformation	Optax Optimizer	required
key	ndarray	Random key	required

Returns:

Type	Description
ArrayTree	tuple[chex.ArrayTree, list[typing.Any]]: Optimized parameters, and optimization history.

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

def opt_eig_ape_loss(
    loss_fn: typing.Callable[
        [chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]
    ],
    params: chex.ArrayTree,
    num_steps: int,
    optim: optax.GradientTransformation,
    key: jnp.ndarray,
    callbacks: list = [],
) -> chex.ArrayTree:
    """Optimize the EIG loss function.

    Args:
        loss_fn (typing.Callable[[chex.ArrayTree, jnp.ndarray], tuple[jnp.ndarray, AuxEntry]]): Loss function
        params (chex.ArrayTree): Initial parameter
        num_steps (int): Number of optimization step
        optim (optax.GradientTransformation): Optax Optimizer
        key (jnp.ndarray): Random key

    Returns:
        tuple[chex.ArrayTree, list[typing.Any]]: Optimized parameters, and optimization history.
    """
    # Initialize the optimizer
    opt_state = optim.init(params)
    # jit the loss function
    loss_fn = jax.jit(loss_fn)

    for step in range(num_steps):
        key, subkey = jax.random.split(key)
        # Compute the loss and its gradient
        (loss, aux), grad = jax.value_and_grad(loss_fn, has_aux=True)(params, subkey)
        # Update the optimizer and params
        updates, opt_state = optim.update(grad, opt_state, params)
        params = optax.apply_updates(params, updates)

        # entry = (step, loss, aux)
        entry = HistoryEntry(step=step, loss=loss, aux=aux)

        for callback in callbacks:
            callback(entry)

    return params

inspeqtor.boed.estimate_eig

estimate_eig(
    key: ndarray,
    model: Callable,
    marginal_guide: Callable,
    design: ndarray,
    *args,
    optimizer: GradientTransformation,
    num_optimization_steps: int,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int] | int,
    final_num_particles: tuple[int, int]
    | int
    | None = None,
    loss_fn: Callable = marginal_loss,
    callbacks: list = [],
) -> tuple[ndarray, dict[str, Any]]

Optimize for marginal EIG

Parameters:

Name	Type	Description	Default
key	ndarray	Random key	required
model	Callable	Probabilistic model of the experiment	required
marginal_guide	Callable	The marginal guide of the experiment	required
design	ndarray	Possible designs of the experiment	required
optimizer	GradientTransformation	Optax optimizer	required
num_optimization_steps	int	Number of the optimization step	required
observation_labels	list[str]	The list of string of observations	required
target_labels	list[str]	The target latent parameters to be optimized for	required
num_particles	int	The number of independent trials	required
final_num_particles	int | None	Final independent trials to calculate marginal EIG. Defaults to None.	None

Returns:

Type	Description
tuple[ndarray, dict[str, Any]]	tuple[jnp.ndarray, dict[str, typing.Any]]: EIG, and tuple of optimized parameters and optimization history.

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

def estimate_eig(
    key: jnp.ndarray,
    model: typing.Callable,
    marginal_guide: typing.Callable,
    design: jnp.ndarray,
    *args,
    optimizer: optax.GradientTransformation,
    num_optimization_steps: int,
    observation_labels: list[str],
    target_labels: list[str],
    num_particles: tuple[int, int] | int,
    final_num_particles: tuple[int, int] | int | None = None,
    loss_fn: typing.Callable = marginal_loss,
    callbacks: list = [],
) -> tuple[jnp.ndarray, dict[str, typing.Any]]:
    """Optimize for marginal EIG

    Args:
        key (jnp.ndarray): Random key
        model (typing.Callable): Probabilistic model of the experiment
        marginal_guide (typing.Callable): The marginal guide of the experiment
        design (jnp.ndarray): Possible designs of the experiment
        optimizer (optax.GradientTransformation): Optax optimizer
        num_optimization_steps (int): Number of the optimization step
        observation_labels (list[str]): The list of string of observations
        target_labels (list[str]): The target latent parameters to be optimized for
        num_particles (int): The number of independent trials
        final_num_particles (int | None, optional): Final independent trials to calculate marginal EIG. Defaults to None.

    Returns:
        tuple[jnp.ndarray, dict[str, typing.Any]]: EIG, and tuple of optimized parameters and optimization history.
    """
    # NOTE: In final evalution, if final_num_particles != num_particles,
    # the code will error because we train params with num_particles
    # the shape will mismatch
    # final_num_particles = final_num_particles or num_particles

    # Initialize the parameters by using trace from the marginal_guide
    key, subkey = jax.random.split(key)
    params = init_params_from_guide(
        marginal_guide,
        *args,
        key=subkey,
        design=design,
    )

    # Optimize the loss function first to get the optimal parameters
    # for marginal guide
    params = opt_eig_ape_loss(
        loss_fn=loss_fn(
            model,
            marginal_guide,
            design,
            *args,
            observation_labels=observation_labels,
            target_labels=target_labels,
            num_particles=num_particles,
            evaluation=False,
        ),
        params=params,
        num_steps=num_optimization_steps,
        optim=optimizer,
        key=subkey,
        callbacks=callbacks,
    )

    key, subkey = jax.random.split(key)
    # Evaluate the loss
    _, aux = loss_fn(
        model,
        marginal_guide,
        design,
        *args,
        observation_labels=observation_labels,
        target_labels=target_labels,
        num_particles=final_num_particles,
        evaluation=True,
    )(params, subkey)

    return aux.eig, {
        "params": params,
    }

inspeqtor.boed.vectorized_for_eig

vectorized_for_eig(model)

Vectorization function for the EIG function

Parameters:

Name	Type	Description	Default
model	Any	Probabilistic model.	required

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

def vectorized_for_eig(model):
    """Vectorization function for the EIG function

    Args:
        model (typing.Any): Probabilistic model.
    """

    def wrapper(
        design: jnp.ndarray,
        *args,
        **kwargs,
    ):
        # This wrapper has the same call signature as the probabilistic graybox model
        # Expect the design to has shape == (extra, design, feature)
        with plate_stack(prefix="vectorized_plate", sizes=[*design.shape[:2]]):
            return model(design, *args, **kwargs)

    return wrapper

inspeqtor.boed.marginal_init_params_fn

marginal_init_params_fn(
    name: str, shape: tuple[int, ...], num_candiadtes: int
)

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

def marginal_init_params_fn(name: str, shape: tuple[int, ...], num_candiadtes: int):
    if name.endswith("_loc"):
        return numpyro.param(name, jnp.zeros((num_candiadtes,) + shape))
    elif name.endswith("_scale"):
        return numpyro.param(
            name,
            0.1 * jnp.ones((num_candiadtes,) + shape),
            constraint=dist.constraints.softplus_positive,
        )
    else:
        raise ValueError(f"name: {name} is not supported")

inspeqtor.boed.marginal_guide

marginal_guide(
    design, unitaries, observation_labels, target_labels
)

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

def marginal_guide(design, unitaries, observation_labels, target_labels):
    # This shape allows us to learn a different parameter for each candidate design l
    q_logit = numpyro.param("theta", jnp.zeros((design.shape[-2],) + (18,)))
    numpyro.sample("obs", dist.Bernoulli(logits=q_logit).to_event(1))  # type: ignore

inspeqtor.boed.auto_marginal_guide

auto_marginal_guide(model, guide)

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

def auto_marginal_guide(model, guide):
    # The guide should be create for the model with block = True

    def marginal_guide(control_parameters, *args, observation_labels, target_labels):
        # Get the number of candidate designs
        num_candidates = control_parameters.shape[-2]

        with numpyro.plate("candidates", num_candidates):
            params = guide(control_parameters, *args)

            cond_model = numpyro.handlers.condition(model, params)
            expectation_values = cond_model(control_parameters, *args)

        return expectation_values

    return marginal_guide