class harmonic.model.FlowModel(ndim_in: int, learning_rate: float = 0.001, momentum: float = 0.9, standardize: bool = False, temperature: float = 0.8)

Normalizing flow model to approximate the log_e posterior by a normalizing flow.

fit(X: Array, batch_size: int = 64, epochs: int = 3, key=Array([0, 1000], dtype=uint32), verbose: bool = False)

Fit the parameters of the model.

Parameters:

  • X (jnp.ndarray (nsamples, ndim)) – Training samples.

  • batch_size (int, optional) – Batch size used when training flow. Default = 64.

  • epochs (int, optional) – Number of epochs flow is trained for. Default = 3.

  • key (Union[jax.Array, jax.random.PRNGKeyArray], optional) – Key used in random number generation process.

  • verbose (bool, optional) – Controls if progress bar and current loss are displayed when training. Default = False.

Raises:

  • ValueError – Raised if the second dimension of X is not the same as ndim.

  • NotImplementedError – If called directly from FlowModel class.

predict(x: Array) → Array

Predict the value of log_e posterior at batched input x.

Parameters:

x (jnp.ndarray (batch_size, ndim)) – Sample for which to predict posterior values.

Returns:

Predicted log_e posterior value.

Return type:

jnp.ndarray (batch_size,)

Raises:

ValueError – If temperature is negative or greater than 1.

sample(n_sample: int, rng_key=Array([0, 0], dtype=uint32)) → Array

Sample from trained flow.

Parameters:

  • nsample (int) – Number of samples generated.

  • rng_key (Union[jax.Array, jax.random.PRNGKeyArray]), optional) – Key used in random number generation process.

Raises:

ValueError – If temperature is negative or greater than 1.

Returns:

Samples from fitted distribution.

Return type:

jnp.array (n_sample, ndim)

class harmonic.model.RQSplineModel(ndim_in: int, n_layers: int = 8, n_bins: int = 8, hidden_size: Sequence[int] = [64, 64], spline_range: Sequence[float] = (-10.0, 10.0), standardize: bool = False, learning_rate: float = 0.001, momentum: float = 0.9, temperature: float = 0.8)

Rational quadratic spline flow model to approximate the log_e posterior by a normalizing flow.

__init__(ndim_in: int, n_layers: int = 8, n_bins: int = 8, hidden_size: Sequence[int] = [64, 64], spline_range: Sequence[float] = (-10.0, 10.0), standardize: bool = False, learning_rate: float = 0.001, momentum: float = 0.9, temperature: float = 0.8)

Constructor setting the hyper-parameters and domains of the model.

Must be implemented by derived class (currently abstract).

Parameters:

  • ndim_in (int) – Dimension of the problem to solve.

  • n_layers (int, optional) – Number of layers in the flow. Defaults to 8.

  • n_bins (int, optional) – Number of bins in the spline. Defaults to 8.

  • hidden_size (Sequence[int], optional) – Size of the hidden layers in the conditioner. Defaults to [64, 64].

  • spline_range (Sequence[float], optional) – Range of the spline. Defaults to (-10.0, 10.0).

  • standardize (bool, optional) – Indicates if mean and variance should be removed from training data when training the flow. Defaults to False.

  • learning_rate (float, optional) – Learning rate for adam optimizer used in the fit method. Defaults to 0.001.

  • momentum (float, optional) – Learning rate for Adam optimizer used in the fit method. Defaults to 0.9.

  • temperature (float, optional) – Scale factor by which the base distribution Gaussian is compressed in the prediction step. Should be positive and <=1. Defaults to 0.8.

Raises:

ValueError – If the ndim_in is not positive.

class harmonic.model.RealNVPModel(ndim_in: int, n_scaled_layers: int = 2, n_unscaled_layers: int = 4, learning_rate: float = 0.001, momentum: float = 0.9, standardize: bool = False, temperature: float = 0.8)

Normalizing flow model to approximate the log_e posterior by a NVP normalizing flow.

__init__(ndim_in: int, n_scaled_layers: int = 2, n_unscaled_layers: int = 4, learning_rate: float = 0.001, momentum: float = 0.9, standardize: bool = False, temperature: float = 0.8)

Constructor setting the hyper-parameters of the model.

Parameters:

  • ndim_in (int) – Dimension of the problem to solve.

  • n_scaled_layers (int, optional) – Number of layers with scaler in RealNVP flow. Default = 2.

  • n_unscaled_layers (int, optional) – Number of layers without scaler in RealNVP flow. Default = 4.

  • learning_rate (float, optional) – Learning rate for adam optimizer used in the fit method. Default = 0.001.

  • momentum (float, optional) – Learning rate for Adam optimizer used in the fit method. Default = 0.9

  • standardize (bool, optional) – Indicates if mean and variance should be removed from training data when training the flow. Default = False

  • temperature (float, optional) – Scale factor by which the base distribution Gaussian is compressed in the prediction step. Should be positive and <=1. Default = 0.8.

Raises:

  • ValueError – If the ndim_in is not positive.

  • ValueError – If n_scaled_layers is not positive.

harmonic.model.make_training_loop(model)

Create a function that trains an NF model.

Parameters:

model – a neural network model with a log_prob method.

Returns:

wrapper function that trains the model.

Return type:

train_flow (Callable)

Note

Adapted from github.com/kazewong/flowMC