Positional Encodings

Positional encodings for spherical data via harmonics, Fourier features, and Herglotz kernels.

class spherical_inr.positional_encoding.FourierPE(num_atoms: int, input_dim: int = 3, bias: bool = True, omega0: float = 1.0)

Bases: Module

Learned Fourier positional encoding.

This module implements a learnable sinusoidal feature map of the form

\[\psi^{\mathrm{F}}(x) = \sin\bigl(\omega_0 (x \Omega^\top + b)\bigr),\]

where: - \(W \in \mathbb{R}^{N \times d}\) is a learnable weight matrix, - \(b \in \mathbb{R}^N\) is an optional learnable bias, - \(\omega_0 > 0\) is a fixed frequency scaling factor.

This corresponds to a standard Fourier-feature embedding with learned frequencies.

Parameters:

• num_atoms (int) – Number of output features.

• input_dim (int) – Dimension \(d\) of the input space.

• bias (bool, optional) – Whether to include a learnable bias term \(b\). Default = True

• omega0 (float, optional) – Frequency scaling factor \(\omega_0\). Default = 1.0

forward(x: Tensor) → Tensor

Forward pass of the learned Fourier positional encoding.

This method applies a learned linear projection of the input coordinates, followed by sinusoidal activation with fixed frequency scaling.

Parameters:

x (torch.Tensor) – Input tensor of shape (..., input_dim).

Returns:

Learned Fourier features with shape (..., num_atoms).

Return type:

torch.Tensor

class spherical_inr.positional_encoding.HerglotzPE(num_atoms: int, L_init: int, rot: bool = False)

Bases: Module

Herglotz positional encoding with learnable phase and magnitude.

This module implements a real-valued Herglotz-type feature map defined on Cartesian coordinates \(x \in \mathbb{R}^3\).

Each atom \(k\) is defined by two orthonormal vectors \(a_{k, \Re}, a_{k, \Im} \in \mathbb{R}^3\), forming an implicit complex direction \(a_k = a_{k, \Re} + i\,a_{k, \Im}\).

For an input point \(x\), we compute the projections

\[u_k = \langle x, a_{k, \Re} \rangle, \qquad v_k = \langle x, a_{k, \Im} \rangle.\]

Each atom is parameterized by learnable parameter \(\sigma_k\) with \(\sigma_k \sim \mathcal{U}(0, L_{\text{init}})\).

The Herglotz feature associated with atom \(k\) is defined in closed form as

\[\psi^{\mathrm{H}}_k(x) = \frac{1}{1 + 2L_{\text{init}}} e^{\rho_k (u_k - 1)} \Bigl[ (1 + 2\rho_k u_k)\cos(\rho_k v_k) - (2\rho_k v_k)\sin(\rho_k v_k) \Bigr].\]

Optionally, a learnable quaternion rotation may be applied to all atoms before evaluation, allowing the encoding to learn a global orientation.

Parameters:

• num_atoms (int) – Number of Herglotz atoms (output features).

• L_init (int) – Upper bound used to initialize the magnitude parameters \(\sigma_k \sim \mathcal{U}(0, L_{\mathrm{init}})\).

• rot (bool, optional) – If True, applies a learnable quaternion rotation to all atoms. Default = False

Notes

This module is Cartesian-only. If your data is given in spherical coordinates \((\theta,\phi)\), use a wrapper to convert inputs to Cartesian coordinates before applying this encoding.

forward(x: Tensor) → Tensor

Forward pass of the Herglotz positional encoding.

This method computes the closed-form Herglotz feature map by projecting Cartesian input points onto the learned atom directions, applying the learned phase and magnitude parameters.

Parameters:

x (torch.Tensor) – Tensor of shape (..., 3) containing Cartesian coordinates.

Returns:

Herglotz features evaluated at the input points, with shape (..., num_atoms).

Return type:

torch.Tensor

Raises:

ValueError – If x.shape[-1] != 3.

class spherical_inr.positional_encoding.SphericalHarmonicsPE(num_atoms: int)

Bases: Module

Real spherical harmonics positional encoding.

This module maps spherical angles \(x = (\theta, \phi) \in [0,\pi] \times [-\pi,\pi]\) to a vector of real spherical harmonics

\[\psi^{\mathrm{SH}}(x) = \bigl( Y_{\ell_1}^{m_1}(\theta,\phi), \dots, Y_{\ell_N}^{m_N}(\theta,\phi) \bigr),\]

where the index pairs \((\ell_k, m_k)\) follow the standard ordering

\[(0,0), (1,-1),(1,0),(1,1),(2,-2),\dots\]

and only the first num_atoms = N basis functions are retained.

The real spherical harmonics are defined as

\[\begin{split}Y_\ell^m(\theta,\phi) = N_{\ell m}\,P_\ell^{|m|}(\cos\theta) \begin{cases} \cos(m\phi), & m \ge 0, \\ \sin(|m|\phi), & m < 0, \end{cases}\end{split}\]

where \(P_\ell^m\) are the associated Legendre polynomials and \(N_{\ell m}\) is a normalization constant.

Parameters:

num_atoms (int) – Number of spherical harmonic basis functions returned.

forward(x: Tensor) → Tensor

Forward pass of the spherical harmonics encoding.

This method evaluates the preselected real spherical harmonic basis functions at the input angular coordinates and returns the first num_atoms coefficients in standard ordering.

Parameters:

x (torch.Tensor) – Tensor of shape (..., 2) containing spherical angles (theta, phi) in radians.

Returns:

Real spherical harmonics evaluated at the input angles, with shape (..., num_atoms).

Return type:

torch.Tensor

Raises:

ValueError – If x.shape[-1] != 2.