tfep.nn.transformers.moebius.symmetrized_moebius_transformer

tfep.nn.transformers.moebius.symmetrized_moebius_transformer(x: Tensor, w: Tensor, max_radius: float = 0.99) → tuple[Tensor][source]

Symmetrized Moebius transformer.

This implements a generalization of the symmetrized Moebius transformation proposed in [1] to non-unit spheres. The transformer will expand/contract the distribution on the sphere of radius \(r\), where \(r\) is the norm of the input vector.

The transformation has the form

\(y = ||f(x; w)|| \frac{f(x; w) + f(x; -w)}{||f(x; w) + f(x; -w)||}\)

where \(f\) is the Moebius transform (see :class:.MoebiusTransformer), and \(y, x, w\) are all dimension-dimensional vectors with \(||w|| < ||x||\). The function automatically rescales the w argument following the same strategy as in [2] to satisfy the condition on the norm. Consequently, ``w``s of any norm can be passed.

Parameters:

  • x (torch.Tensor) – Shape (batch_size, n_vectors, dimension). The input coordinates.

  • w (torch.Tensor) – Shape (batch_size, n_vectors, dimension). The transformation parameters. These parameter vectors are automatically rescaled so that |w| < |x|.

  • max_radius (float) – Must be strictly less than 1. Rescaling of the w vectors will be performed so that its maximum norm will be max_radius * |x|.

Returns:

  • y (torch.Tensor) – Shape (batch_size, n_vectors, dimension). The transformed vectors.

  • log_det_J (torch.Tensor) – Shape (batch_size,). The logarithm of the absolute value of the Jacobian determinant dy / dx.

References

[1] Köhler J, Invernizzi M, De Haan P, Noé F. Rigid body flows for sampling

molecular crystal structures. arXiv preprint arXiv:2301.11355. 2023 Jan 26.

[2] Rezende DJ, Papamakarios G, Racanière S, Albergo MS, Kanwar G,

Shanahan PE, Cranmer K. Normalizing Flows on Tori and Spheres. arXiv preprint arXiv:2002.02428. 2020 Feb 6.