piecewise_quadratic module

Implementation of the piecewise quadratic coupling cell This means that the variable transform is piecewise-quadratic.

class ElementWisePWQuadraticTransform(*args, **kwargs)[source]

Bases: zunis.models.flows.coupling_cells.transforms.InvertibleTransform

Invertible piecewise-quadratic transformations over the unit hypercube

Implements a batched bijective transformation h from the d-dimensional unit hypercube to itself, in an element-wise fashion (each coordinate transformed independently)

In each direction, the bijection is a piecewise-quadratic transform with b bins where the forward transform has bins with adjustable width. The transformation in each bin is then a quadratic spline. The network predicts the bin width w_tilde and the vertex height v_tilde of the derivative of the transform for each direction and each point of the batch. They are normalized such that: 1. h(0) = 0 2. h(1) = 1 3. h is monotonous 4. h is continuous

Conditions 1. to 3. ensure the transformation is a bijection and therefore invertible. The inverse is also an element-wise, piece-wise quadratic transformation.

Constructor for InvertibleTransform

No arguments, no returns, this is just encapsulation for a function and its inverse

static backward(x, wv_tilde, compute_jacobian=True)[source]

Apply an element-wise piecewise-quadratic transformation to some variables

Parameters

  • x (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable spans the k-dimensional unit hypercube

  • wv_tilde (torch.Tensor) – is a tensor with shape (N,k,2b+1) where b is the number of bins. This contains the un-normalized widths and heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using a modified softmax.

  • compute_jacobian (bool, optional) – determines whether the jacobian should be compute or None is returned

Returns

pair (y,j). - y is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor

static forward(y, wv_tilde, compute_jacobian=True)[source]

Apply the inverse of an element-wise piecewise-linear transformation to some variables

Parameters

  • y (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable span the k-dimensional unit hypercube

  • wv_tilde (torch.Tensor) – is a tensor with shape (N,k,2b+1) where b is the number of bins. This contains the un-normalized widths and heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using a modified softmax.

compute_jacobianbool, optional

determines whether the jacobian should be compute or None is returned

Returns

pair (x,j). - x is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor

class GeneralPWQuadraticCoupling(*args, **kwargs)[source]

Bases: zunis.models.flows.coupling_cells.general_coupling.InvertibleCouplingCell

Abstract class implementing a coupling cell based on PW quadratic transformations

A specific way to predict the parameters of the transform must be implemented in child classes.

Generator for the abstract class GeneralPWQuadraticCoupling

Parameters

  • d (int) – dimension of the space

  • mask (list of bool) – variable mask which variables are transformed (False) or used as parameters of the transform (True)

training: bool[source]
class PWQuadraticCoupling(*args, **kwargs)[source]

Bases: zunis.models.flows.coupling_cells.piecewise_coupling.piecewise_quadratic.GeneralPWQuadraticCoupling

Piece-wise Quadratic coupling

Coupling cell using an element-wise piece-wise quadratic transformation as a change of variables. The transverse neural network is a rectangular dense neural network.

Notes

Transformation used: zunis.models.flows.coupling_cells.piecewise_coupling.piecewise_quadratic.ElementWisePWQuadraticTransform Neural network used: zunis.models.layers.trainable.ArbitraryShapeRectangularDNN

Generator for PWQuadraticCoupling

Parameters

  • d (int) –

  • mask (list of bool) – variable mask: which dimension are transformed (False) and which are not (True)

  • n_bins (int) – number of bins in each dimensions

  • d_hidden (int) – dimension of the hidden layers of the DNN

  • n_hidden (int) – number of hidden layers in the DNN

  • input_activation (optional) – pytorch activation function before feeding into the DNN. must be a callable generator without arguments (i.e. a classname or a function)

  • hidden_activation (optional) – pytorch activation function between hidden layers of the DNN. must be a callable generator without arguments (i.e. a classname or a function)

  • output_activation (optional) – pytorch activation function at the output of the DNN. must be a callable generator without arguments (i.e. a classname or a function)

  • use_batch_norm (bool) – whether batch normalization should be used in the DNN.

training: bool[source]
modified_softmax(v, w)[source]
piecewise_quadratic_inverse_transform(y, wv_tilde, compute_jacobian=True)[source]

Apply the inverse of an element-wise piecewise-linear transformation to some variables

Parameters

  • y (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable span the k-dimensional unit hypercube

  • wv_tilde (torch.Tensor) – is a tensor with shape (N,k,2b+1) where b is the number of bins. This contains the un-normalized widths and heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using a modified softmax.

compute_jacobianbool, optional

determines whether the jacobian should be compute or None is returned

Returns

pair (x,j). - x is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor

piecewise_quadratic_transform(x, wv_tilde, compute_jacobian=True)[source]

Apply an element-wise piecewise-quadratic transformation to some variables

Parameters

  • x (torch.Tensor) – a tensor with shape (N,k) where N is the batch dimension while k is the dimension of the variable space. This variable spans the k-dimensional unit hypercube

  • wv_tilde (torch.Tensor) – is a tensor with shape (N,k,2b+1) where b is the number of bins. This contains the un-normalized widths and heights of the bins of the piecewise-constant PDF for dimension k, i.e. q_tilde lives in all of R and we don’t impose a constraint on their sum yet. Normalization is imposed in this function using a modified softmax.

  • compute_jacobian (bool, optional) – determines whether the jacobian should be compute or None is returned

Returns

pair (y,j). - y is a tensor with shape (N,k) living in the k-dimensional unit hypercube - j is the jacobian of the transformation with shape (N,) if compute_jacobian==True, else None.

Return type

tuple of torch.Tensor