core.posterior_models.cflow_base.ContinuousFlowPosteriorModel

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel(**kwargs)

Class for posterior models based on continuous normalizing flows (CNFs).

CNFs are parameterized by a vector field v(theta_t, t), that transports a simple base distribution (typically a gaussian N(0,1) with same dimension as theta) at time t=0 to the target distribution at time t=1. This vector field defines the flow via the ODE

            d/dt f(theta, t) = v(f(theta, t), t).

The vector field v is parameterized with a neural network. It is impractical to train this neural network (and thereby the CNF) directly with log-likelihood maximization, as solving the full ODE for each training iteration, requires thousands of vector field evaluations.

Several alternative methods have been developed to make training CNFs more efficient. These directly regress on the vector field v (or a scaled version of v, such as the score). It has been shown that this can be done on a per-sample basis by adding noise to the parameters at various scales t. Specifically, a parameter sample theta is transformed as follows.

t               ~ U[0, 1-eps)                               noise level
theta_0         ~ N(0, 1)                                   sampled noise
theta_1         = theta                                     pure sample
theta_t         = c1(t) * theta_1 + c0(t) * theta_0         noisy sample

Within that framework, one can employ different methods to learn the vector field v, such as flow matching or score matching. These have slightly different coefficients c1(t), c2(t) and training objectives.

This class is intended to construct and hold a neural network for estimating the posterior density, as well as saving / loading, and training. It also has functionality for sampling and density evaluation.

Attributes

Name	Description
eps
integration_range	Integration range for ODE. We integrate in the range [0, 1-self.eps]. For score
theta_dim
time_prior_exponent

Methods

Name	Description
evaluate_vector_field	Evaluate the vector field v(t, theta_t, context_data) that generates the flow
initialize_network
log_prob	Evaluate the log posterior density,
rhs_of_joint_ode	Returns the right hand side of the neural ODE that is used to evaluate the
sample	Sample parameters theta from the posterior model,
sample_and_log_prob	Sample parameters theta from the posterior model,
sample_t
sample_theta_0	Sample theta_0 from the gaussian prior.

evaluate_vector_field

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.evaluate_vector_field(
    t,
    theta_t,
    *context_data,
)

Evaluate the vector field v(t, theta_t, context_data) that generates the flow via the ODE

d/dt f(theta_t, t, context) = v(f(theta_t, t, context), t, context).

Parameters

Name	Type	Description	Default
t		time (noise level)	required
theta_t		noisy parameters, perturbed with noise level t	required
*context_data		list with context data (GW data)	()

initialize_network

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.initialize_network(
)

log_prob

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.log_prob(
    theta,
    *context,
    hutchinson=False,
)

Evaluate the log posterior density,

log p(theta | context)

For this we solve an ODE backwards in time until we reach the initial pure noise distribution.

There are two contributions, the log_prob of theta_0 (which is uniquely determined by theta) under the base distribution, and the integrated divergence of the vector field.

Parameters

Name	Type	Description	Default
theta	torch.Tensor	Parameter values at which to evaluate the density. Should have a batch dimension (even if size B = 1).	required
context	torch.Tensor	Context information (typically observed data). Must have context.shape[0] = B.	()
hutchinson			False

Returns

Name	Type	Description
log_prob	torch.Tensor	Shape (B,)

rhs_of_joint_ode

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.rhs_of_joint_ode(
    t,
    theta_and_div_t,
    *context_data,
    hutchinson=False,
)

Returns the right hand side of the neural ODE that is used to evaluate the log_prob of theta samples. This is a joint ODE over the vector field and the divergence. By integrating this ODE, one can simultaneously trace the parameter sample theta_t and integrate the divergence contribution to the log_prob, see e.g., https://arxiv.org/abs/1806.07366 or Appendix C in https://arxiv.org/abs/2210.02747.

Parameters

Name	Type	Description	Default
t		time (noise level)	required
theta_and_div_t		concatenated tensor of (theta_t, div). theta_t: noisy parameters, perturbed with noise level t	required
*context_data		list with context data (GW data)	()

Returns

Name	Type	Description
	torch.Tensor	vector field that generates the flow and its divergence (required for likelihood evaluation).

sample

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.sample(
    *context,
    num_samples=None,
)

Sample parameters theta from the posterior model,

theta ~ p(theta | context)

by solving an ODE forward in time.

Parameters

Name	Type	Description	Default
context	torch.Tensor	Context information (typically observed data). Should have a batch dimension (even if size B = 1).	()
num_samples	int	Number of samples to generate.	None

Returns

Name	Type	Description
samples	torch.Tensor	Shape (B, num_samples, dim(theta))

sample_and_log_prob

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.sample_and_log_prob(
    *context,
    num_samples=None,
)

Sample parameters theta from the posterior model,

theta ~ p(theta | context)

and also return the log_prob. This is more efficient than calling sample_batch and log_prob_batch separately.

If d/dt [phi(t), f(t)] = rhs joint with initial conditions [theta_0, log p(theta_0)], where theta_0 ~ p_0(theta_0), then [phi(1), f(1)] = [theta_1, log p(theta_0) + log p_1(theta_1) - log p(theta_0)] = [theta_1, log p_1(theta_1)].

Parameters

Name	Type	Description	Default
context	torch.Tensor	Context information (typically observed data). Should have a batch dimension (even if size B = 1).	()
num_samples	int	Number of samples to generate.	None

Returns

Name	Type	Description
	samples, log_prob: torch.Tensor, torch.Tensor	Shapes (B, num_samples, dim(theta)), (B, num_samples)

sample_t

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.sample_t(
    batch_size,
)

sample_theta_0

core.posterior_models.cflow_base.ContinuousFlowPosteriorModel.sample_theta_0(
    batch_size,
)

Sample theta_0 from the gaussian prior.