rqSpline

MaskedCouplingRQSpline

Bases: NFModel

Rational quadratic spline normalizing flow model using distrax.

Parameters:

Name	Type	Description	Default
n_features	int	Number of features in the data.	required
num_layers	int	Number of layers in the conditioner.	required
hidden_size	Sequence[int]	Hidden size of the conditioner.	required
num_bins	int	Number of bins in the spline.	required
key	PRNGKeyArray	Random key for initialization.	required
spline_range	Sequence[float]	Range of the spline. Defaults to (-10.0, 10.0).	(-10.0, 10.0)

Properties

n_features (int) : Number of features in the data. data_mean (Array) : Mean of the data. data_cov (Array) : Covariance of the data.

Source code in flowMC/nfmodel/rqSpline.py

class MaskedCouplingRQSpline(NFModel):
    r"""Rational quadratic spline normalizing flow model using distrax.

    Args:
        n_features (int):  Number of features in the data.
        num_layers (int): Number of layers in the conditioner.
        hidden_size (Sequence[int]): Hidden size of the conditioner.
        num_bins (int): Number of bins in the spline.
        key (PRNGKeyArray): Random key for initialization.
        spline_range (Sequence[float]): Range of the spline. Defaults to (-10.0, 10.0).

    Properties:
        n_features (int) :  Number of features in the data.
        data_mean (Array) : Mean of the data.
        data_cov (Array) : Covariance of the data.
    """

    base_dist: Distribution
    layers: list[Bijection]
    _n_features: int
    _data_mean: Float[Array, " n_dim"]
    _data_cov: Float[Array, " n_dim n_dim"]

    @property
    def n_features(self):
        return self._n_features

    @property
    def data_mean(self):
        return jax.lax.stop_gradient(self._data_mean)

    @property
    def data_cov(self):
        return jax.lax.stop_gradient(jnp.atleast_2d(self._data_cov))

    def __init__(
        self,
        n_features: int,
        n_layers: int,
        hidden_size: list[int],
        num_bins: int,
        key: PRNGKeyArray,
        spline_range: tuple[float, float] = (-10.0, 10.0),
        **kwargs
    ):
        if kwargs.get("base_dist") is not None:
            dist = kwargs.get("base_dist")
            assert isinstance(dist, Distribution)
            self.base_dist = dist
        else:
            self.base_dist = Gaussian(
                jnp.zeros(n_features), jnp.eye(n_features), learnable=False
            )

        if kwargs.get("data_mean") is not None:
            data_mean = kwargs.get("data_mean")
            assert isinstance(data_mean, Array)
            self._data_mean = data_mean
        else:
            self._data_mean = jnp.zeros(n_features)

        if kwargs.get("data_cov") is not None:
            data_cov = kwargs.get("data_cov")
            assert isinstance(data_cov, Array)
            self._data_cov = data_cov
        else:
            self._data_cov = jnp.eye(n_features)

        self._n_features = n_features

        def make_layer(i: int, key: PRNGKeyArray):
            mlp = MLP(
                [n_features] + hidden_size + [n_features * (num_bins * 3 + 1)],
                key,
                scale=1e-2,
                activation=jax.nn.tanh,
            )
            mask = ((jnp.arange(0, n_features) + i) % 2).astype(bool)
            mask_all = (jnp.zeros(n_features)).astype(bool)
            layer1 = MaskedCouplingLayer(ScalarAffine(0.0, 0.0), mask_all)
            layer2 = MaskedCouplingLayer(
                RQSpline(mlp, spline_range[0], spline_range[1]), mask
            )
            return eqx.nn.Sequential([layer1, layer2])  # type: ignore

        keys = jax.random.split(key, n_layers)
        self.layers = eqx.filter_vmap(make_layer)(jnp.arange(n_layers), keys)

    def __call__(
        self, x: Float[Array, " n_dim"]
    ) -> tuple[Float[Array, " n_dim"], Float]:
        return self.forward(x)

    def forward(
        self, x: Float[Array, " n_dim"]
    ) -> tuple[Float[Array, " n_dim"], Float]:
        log_det = 0.0
        dynamics, statics = eqx.partition(self.layers, eqx.is_array)

        def f(carry, data):
            x, log_det = carry
            layers = eqx.combine(data, statics)
            x, log_det_i = layers[0](x)
            log_det += log_det_i
            x, log_det_i = layers[1](x)
            return (x, log_det + log_det_i), None

        (x, log_det), _ = jax.lax.scan(f, (x, log_det), dynamics)
        return x, log_det

    @partial(jax.vmap, in_axes=(None, 0))
    def inverse(
        self, x: Float[Array, " n_dim"]
    ) -> tuple[Float[Array, " n_dim"], Float]:
        """From latent space to data space"""
        log_det = 0.0
        dynamics, statics = eqx.partition(self.layers, eqx.is_array)

        def f(carry, data):
            x, log_det = carry
            layers = eqx.combine(data, statics)
            x, log_det_i = layers[0].inverse(x)
            log_det += log_det_i
            x, log_det_i = layers[1].inverse(x)
            return (x, log_det + log_det_i), None

        (x, log_det), _ = jax.lax.scan(f, (x, log_det), dynamics, reverse=True)
        return x, log_det

    @eqx.filter_jit
    def sample(
        self, rng_key: PRNGKeyArray, n_samples: int
    ) -> Float[Array, "n_samples n_dim"]:
        samples = self.base_dist.sample(rng_key, n_samples)
        samples = self.inverse(samples)[0]
        samples = samples * jnp.sqrt(jnp.diag(self.data_cov)) + self.data_mean
        return samples

    @eqx.filter_jit
    @partial(jax.vmap, in_axes=(None, 0))
    def log_prob(self, x: Float[Array, "n_sample n_dim"]) -> Float[Array, " n_sample"]:
        """From data space to latent space"""
        x = (x - self.data_mean) / jnp.sqrt(jnp.diag(self.data_cov))
        y, log_det = self.__call__(x)
        log_det = log_det + self.base_dist.log_prob(y)
        return log_det

inverse(x)

From latent space to data space

Source code in flowMC/nfmodel/rqSpline.py

@partial(jax.vmap, in_axes=(None, 0))
def inverse(
    self, x: Float[Array, " n_dim"]
) -> tuple[Float[Array, " n_dim"], Float]:
    """From latent space to data space"""
    log_det = 0.0
    dynamics, statics = eqx.partition(self.layers, eqx.is_array)

    def f(carry, data):
        x, log_det = carry
        layers = eqx.combine(data, statics)
        x, log_det_i = layers[0].inverse(x)
        log_det += log_det_i
        x, log_det_i = layers[1].inverse(x)
        return (x, log_det + log_det_i), None

    (x, log_det), _ = jax.lax.scan(f, (x, log_det), dynamics, reverse=True)
    return x, log_det

log_prob(x)

From data space to latent space

Source code in flowMC/nfmodel/rqSpline.py

@eqx.filter_jit
@partial(jax.vmap, in_axes=(None, 0))
def log_prob(self, x: Float[Array, "n_sample n_dim"]) -> Float[Array, " n_sample"]:
    """From data space to latent space"""
    x = (x - self.data_mean) / jnp.sqrt(jnp.diag(self.data_cov))
    y, log_det = self.__call__(x)
    log_det = log_det + self.base_dist.log_prob(y)
    return log_det

RQSpline

Bases: Bijection

Source code in flowMC/nfmodel/rqSpline.py

class RQSpline(Bijection):
    _range_min: float
    _range_max: float
    _num_bins: int
    _min_bin_size: float
    _min_knot_slope: float
    conditioner: MLP

    """A rational-quadratic spline bijection.

    This bijection is a piecewise rational-quadratic spline with `num_bins` bins.
    The spline is defined by the bin boundaries on the x and y axes, the slopes
    at the knot points, and the slopes at the boundaries of the spline range.
    The spline is linear outside the spline range.

    Args:
        conditioner (eqx.Module): A conditioner that takes the input and returns
            the parameters of the spline.
        range_min (float): The minimum value of the spline range.
        range_max (float): The maximum value of the spline range.
        num_bins (int): The number of bins in the spline.
        min_bin_size (float): The minimum size of a bin.
        min_knot_slope (float): The minimum slope of the spline at the knot points.

    Attributes:
        range_min (float): The minimum value of the spline range.
        range_max (float): The maximum value of the spline range.
    """

    @property
    def range_min(self):
        return jax.lax.stop_gradient(self._range_min)

    @property
    def range_max(self):
        return jax.lax.stop_gradient(self._range_max)

    @property
    def num_bins(self):
        return jax.lax.stop_gradient(self._num_bins)

    @property
    def min_bin_size(self):
        return jax.lax.stop_gradient(self._min_bin_size)

    @property
    def min_knot_slope(self):
        return jax.lax.stop_gradient(self._min_knot_slope)

    @property
    def dtype(self):
        return self.conditioner.dtype

    def __init__(
        self,
        conditioner: MLP,
        range_min: float,
        range_max: float,
        min_bin_size: float = 1e-4,
        min_knot_slope: float = 1e-4,
    ):
        self._range_min = range_min
        self._range_max = range_max
        self._min_bin_size = min_bin_size
        self._min_knot_slope = min_knot_slope
        self._num_bins = int(conditioner.n_output / conditioner.n_input - 1) // 3

        self.conditioner = conditioner

    def get_params(
        self, x: Float[Array, " n_condition"]
    ) -> tuple[
        Float[Array, " n_param"], Float[Array, " n_param"], Float[Array, " n_param"]
    ]:
        params = self.conditioner(x).reshape(-1, self._num_bins * 3 + 1)
        unnormalized_bin_widths = params[:, : self._num_bins]
        unnormalized_bin_heights = params[:, self._num_bins : 2 * self._num_bins]
        unnormalized_knot_slopes = params[:, 2 * self._num_bins :]
        # Normalize bin sizes and compute bin positions on the x and y axis.
        range_size = self.range_max - self.range_min
        bin_widths = _normalize_bin_sizes(
            unnormalized_bin_widths, range_size, self.min_bin_size
        )
        bin_heights = _normalize_bin_sizes(
            unnormalized_bin_heights, range_size, self.min_bin_size
        )
        x_pos = self.range_min + jnp.cumsum(bin_widths[..., :-1], axis=-1)
        y_pos = self.range_min + jnp.cumsum(bin_heights[..., :-1], axis=-1)
        pad_shape = params.shape[:-1] + (1,)
        pad_below = jnp.full(pad_shape, self.range_min, dtype=self.dtype)
        pad_above = jnp.full(pad_shape, self.range_max, dtype=self.dtype)
        x_pos = jnp.concatenate([pad_below, x_pos, pad_above], axis=-1)
        y_pos = jnp.concatenate([pad_below, y_pos, pad_above], axis=-1)
        # Normalize knot slopes and enforce requested boundary conditions.
        knot_slopes = _normalize_knot_slopes(
            unnormalized_knot_slopes, self.min_knot_slope
        )
        return x_pos, y_pos, knot_slopes

    def __call__(self, x: Array, condition_x: Array) -> tuple[Array, Array]:
        return self.forward(x, condition_x)

    def forward(self, x: Array, condition_x: Array) -> tuple[Array, Array]:
        x_pos, y_pos, knot_slopes = self.get_params(condition_x)
        return _rational_quadratic_spline_fwd(x, x_pos, y_pos, knot_slopes)

    def inverse(self, x: Array, condition_x: Array) -> tuple[Array, Array]:
        x_pos, y_pos, knot_slopes = self.get_params(condition_x)
        return _rational_quadratic_spline_inv(x, x_pos, y_pos, knot_slopes)

conditioner: MLP = conditioner instance-attribute

A rational-quadratic spline bijection.

This bijection is a piecewise rational-quadratic spline with num_bins bins. The spline is defined by the bin boundaries on the x and y axes, the slopes at the knot points, and the slopes at the boundaries of the spline range. The spline is linear outside the spline range.

Parameters:

Name	Type	Description	Default
conditioner	Module	A conditioner that takes the input and returns the parameters of the spline.	required
range_min	float	The minimum value of the spline range.	required
range_max	float	The maximum value of the spline range.	required
num_bins	int	The number of bins in the spline.	required
min_bin_size	float	The minimum size of a bin.	required
min_knot_slope	float	The minimum slope of the spline at the knot points.	required

Attributes:

Name	Type	Description
range_min	float	The minimum value of the spline range.
range_max	float	The maximum value of the spline range.