Positional Encodings

class spherical_inr.positional_encoding.FourierPE(num_atoms: int, input_dim: int, bias: bool = True, seed: int | None = None, omega0: float = 1.0)[source]

Bases: _PositionalEncoding

Fourier Positional Encoding.

Computes the positional encoding \(\psi(x)\) by applying a learnable linear transformation followed by a sinusoidal activation. For an input \(x\), the encoding is given by

\[z = \Omega(x), \quad \psi(x) = \sin\bigl(\omega_0\,z\bigr),\]

where \(\Omega\) is a linear mapping from \(\mathbb{R}^{\text{input_dim}}\) to \(\mathbb{R}^{\text{num_atoms}}\) and \(\omega_0\) is a frequency factor.

Parameters:

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

  • input_dim (int) – Dimensionality of the input.

  • bias (bool, optional) – If True, the linear mapping includes a bias term (default: True).

  • seed (Optional[int], optional) – Seed for reproducibility.

  • omega0 (float, optional) – Frequency factor applied to the sinusoidal activation (default: 1.0).

omega0

Buffer holding the frequency factor.

Type:

torch.Tensor

Omega

Linear layer mapping \(\mathbb{R}^{\text{input_dim}}\) to \(\mathbb{R}^{\text{num_atoms}}\).

Type:

nn.Linear

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class spherical_inr.positional_encoding.HerglotzPE(num_atoms: int | None = None, input_dim: int = 3, *, bias: bool = True, L: int | None = None, seed: int | None = None, rref: float = 1.0, init_exponents: bool = True, normalize: bool = True)[source]

Bases: _PositionalEncoding

Herglotz Positional Encoding.

This module implements the Herglotz map with optional quaternion‐based rotation. Each atom is a complex n dim vector whose real and imaginary parts are unit‐length and orthogonal, initialized either by specifying:

  • num_atoms: an explicit total count, or

  • L: a “stacking depth” so that num_atoms = (L+1)**2.

A radial reference rref scales the projections so that for ‖x‖<rref the encoding remains bounded by 1.

The forward mapping is given by

\[\psi(x) \;=\;\sin\bigl(w_{R}\,a_{\Im} + b_{I}\bigr)\;\exp\!\bigl(w_{R}\,(a_{\Im} - a_{norm}) + b_{R}\bigr)\]
where :

  • \(a_{\Re} = \frac{\Re \{x^\top A\}}{rref}\).

  • \(a_{\Im} = \frac{\Im\{x^\top A\}}{rref}\).

  • \(a_{norm} = \texttt{norm\_const}\).

The model consider input x = (x_1, …, x_n) in R^n. However, there are not theoretical guarantees that the model will perform well in n != 3. If your input is in spherical coordinates, you must convert them in cartesian before feeding them to the model.

Parameters:

  • num_atoms (Optional[int]) – Explicit number of atoms; if set, overrides L.

  • input_dim (int, optional) – Dimensionality of the input. Default: 3.

  • *

  • bias (bool, optional) – If True, includes bias terms for the sine and exponential components. Default: True.

  • L (Optional[int], keyword-only) – Initialize a number of atoms equals to the number of spherical harmonics up to order L. Used when num_atoms is None.

  • seed (Optional[int], keyword-only) – RNG seed for reproducible atom initialization.

  • rref (float, keyword-only) – Radial reference scale; for ‖x‖<rref, outputs ≤ 1. Default: 1.0.

  • init_exponents (bool, keyword-only) – If True, initialize w_R such that it activates the first L spherical harmonic orders (moments included). Recommended to use with L specified and normalize set to True.

  • normalize (bool, keyword-only) – If False, uses 1/√2 as the internal normalization constant; else 0. Bounding up the atoms by <= exp(b_R) if r <=rref. Default: True.

  • rotation (bool, keyword-only) – If True, applies per-atom quaternion rotation. Default: True.

A

Complex atoms of shape (num_atoms, 3), dtype torch.complex64.

Type:

Tensor, buffer

rref

Radial reference.

Type:

Tensor, buffer

w_R

Frequency weights for both sine and exponential terms.

Type:

Parameter

b_I

Bias for the sine term.

Type:

Parameter

b_R

Bias for the exponential term.

Type:

Parameter

norm_const_buf

Normalization constant (0.0 or 1/√2).

Type:

Tensor, buffer

quaternion_rotation

Applies each atom’s quaternion in forward().

Type:

Module or callable

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class spherical_inr.positional_encoding.IrregularHerglotzPE(num_atoms: int | None = None, *, bias: bool = True, L: int | None = None, seed: int | None = None, rref: float = 1.0, init_exponents: bool = False, normalize: bool = True, rotation: bool = True)[source]

Bases: RegularHerglotzPE

Irregular Herglotz Map Positional Encoding.

This variant of the Herglotz map applies a 1/‖x‖ decay so that features smoothly vanish at large radius. All initialization options (num_atoms vs. L, rref, init_exponents, normalize, rotation) carry over from the regular version.

After (optional) quaternion rotation, we compute

\[\psi(x) \;=\; \frac{1}{r}\,\sin\bigl(w_{R}\,a_{\Im} + b_{I}\bigr)\;\exp\!\bigl(w_{R}\,(a_{\Re} - a_{norm}) + b_{R}\bigr)\]
where :

  • \(a_{\Re} = \frac{\Re \{x^\top A\}}{r} \frac{r_{ref}}{r}\).

  • \(a_{\Im} = \frac{\Im\{x^\top A\}}{r} \frac{r_{ref}}{r}\).

  • \(a_{norm} = \texttt{norm\_const}\)

We only consider input \(x = (x, y, z) \in \mathbb{R}^3\). If your input is in spherical coordinates, you must convert them in cartesian before feeding them to the model.

Here the extra factors of 1/r ensure that as ‖x‖→∞, both the sine and exponential terms decay like 1/‖x‖, yielding a positional encoding that vanishes at infinity.

Parameters:

  • num_atoms (Optional[int]) – Explicit number of atoms; if set, overrides L.

  • *

  • L (Optional[int], keyword-only) – Stacking depth. Used when num_atoms is None.

  • seed (Optional[int], keyword-only) – RNG seed for reproducible atom initialization.

  • rref (float, keyword-only) – Radial reference scale. Default: 1.0.

  • init_exponents (bool, keyword-only) – If True, initialize w_R from harmonic orders up to L.

  • normalize (bool, keyword-only) – If False, uses 1/√2 as the internal normalization constant; else 0. Bounds the atom responses by ≤ exp(b_R) when r ≥ rref. Default: True.

  • rotation (bool, keyword-only) – If True, applies per-atom quaternion rotation. Default: True.

(inherited from `RegularHerglotzPE`)
forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class spherical_inr.positional_encoding.IrregularSolidHarmonicsPE(L: int | None = None, *, num_atoms: int | None = None, seed: int | None = None)[source]

Bases: SphericalHarmonicsPE

Irregular Solid Harmonics Positional Encoding.

Encodes a point \((r,\theta,\phi)\in\mathbb{R}^3\) into the irregular solid harmonics basis

\[I_\ell^m(r,\theta,\phi) = \frac{1}{r^{\ell+1}}\,Y_\ell^m(\theta,\phi),\]

which decays like \(1/r\) for large radius. Useful when you want features that vanish at infinity. Output shape is (…, num_atoms), with num_atoms=(L+1)**2 by default.

Inherits all constructor arguments from RegularSolidHarmonicsPE.

exponents

Float tensor of shape (1, num_atoms) containing :math: ell for each channel, used to compute \(1/r^{\ell+1}\).

Type:

torch.Tensor, buffer

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class spherical_inr.positional_encoding.RegularHerglotzPE(num_atoms: int | None = None, *, bias: bool = True, L: int | None = None, seed: int | None = None, rref: float = 1.0, init_exponents: bool = False, normalize: bool = True, rotation: bool = True)[source]

Bases: _PositionalEncoding

Regular Herglotz Map Positional Encoding.

This module implements the regular Herglotz map with optional quaternion‐based rotation. Each atom is a complex 3D vector whose real and imaginary parts are unit‐length and orthogonal, initialized either by specifying:

  • num_atoms: an explicit total count, or

  • L: a “stacking depth” so that num_atoms = (L+1)**2.

A radial reference rref scales the projections so that for ‖x‖<rref the encoding remains bounded by 1. If rotation=True, each atom is first rotated by a learnable quaternion in ℝ³.

The forward mapping is given by

\[\psi(x) \;=\;\sin\bigl(w_{R}\,a_{\Im} + b_{I}\bigr)\;\exp\!\bigl(w_{R}\,(a_{\Im} - a_{norm}) + b_{R}\bigr)\]
where :

  • \(a_{\Re} = \frac{\Re \{x^\top A\}}{rref}\).

  • \(a_{\Im} = \frac{\Im\{x^\top A\}}{rref}\).

  • \(a_{norm} = \texttt{norm\_const}\).

We only consider input \(x = (x, y, z) \in \mathbb{R}^3\). If your input is in spherical coordinates, you must convert them in cartesian before feeding them to the model.

Parameters:

  • num_atoms (Optional[int]) – Explicit number of atoms; if set, overrides L.

  • *

  • bias (bool, optional) – If True, includes bias terms for the sine and exponential components. Default: True.

  • L (Optional[int], keyword-only) – Initialize a number of atoms equals to the number of spherical harmonics up to order L. Used when num_atoms is None.

  • seed (Optional[int], keyword-only) – RNG seed for reproducible atom initialization.

  • rref (float, keyword-only) – Radial reference scale; for ‖x‖<rref, outputs ≤ 1. Default: 1.0.

  • init_exponents (bool, keyword-only) – If True, initialize w_R such that it activates the first L spherical harmonic orders (moments included). Recommended to use with L specified and normalize set to True.

  • normalize (bool, keyword-only) – If False, uses 1/√2 as the internal normalization constant; else 0. Bounding up the atoms by <= exp(b_R) if r <=rref. Default: True.

  • rotation (bool, keyword-only) – If True, applies per-atom quaternion rotation. Default: True.

A

Complex atoms of shape (num_atoms, 3), dtype torch.complex64.

Type:

Tensor, buffer

rref

Radial reference.

Type:

Tensor, buffer

w_R

Frequency weights for both sine and exponential terms.

Type:

Parameter

b_I

Bias for the sine term.

Type:

Parameter

b_R

Bias for the exponential term.

Type:

Parameter

norm_const_buf

Normalization constant (0.0 or 1/√2).

Type:

Tensor, buffer

quaternion_rotation

Applies each atom’s quaternion in forward().

Type:

Module or callable

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class spherical_inr.positional_encoding.RegularSolidHarmonicsPE(L: int | None = None, *, num_atoms: int | None = None, seed: int | None = None)[source]

Bases: SphericalHarmonicsPE

Regular Solid Harmonics Positional Encoding.

Extends SphericalHarmonicsPE to encode a full 3-D point \((r,\theta,\phi)\in\R^3\) into the regular solid harmonics basis functions

\[R_\ell^m(r,\theta,\phi) = r^\ell\,Y_\ell^m(\theta,\phi),\]

for \(\ell=0,\dots,L\) and \(m=-\ell,\dots,\ell\). Output shape is (…, num_atoms) with num_atoms=(L+1)**2 by default.

Parameters:

  • L (Optional[int]) – Maximum harmonic order; default num_atoms=(L+1)**2.

  • num_atoms (Optional[int], keyword-only) – Override the number of basis functions.

  • seed (Optional[int], keyword-only) – Random seed for reproducible behavior.

exponents

Float tensor of shape (1, num_atoms) containing the exponent ℓ for each channel, used to compute r**ℓ.

Type:

torch.Tensor, buffer

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class spherical_inr.positional_encoding.SphericalHarmonicsPE(L: int | None = None, *, num_atoms: int | None = None, input_dim: int = 2, seed: int | None = None)[source]

Bases: _PositionalEncoding

Real Spherical Harmonics Positional Encoding.

Encodes a direction given by polar angles \((\theta,\phi)\in\mathbb{R}^2\) into a vector of real spherical harmonics up to degree \(L\). By default the number of output channels is \((L+1)^2\), but you may override it via the num_atoms keyword.

The real spherical harmonics are computed as

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

where \(\ell=0,\dots,L\) and \(m=-\ell,\dots,\ell\). Internally this class builds two Python lists ls_list and ms_list of length num_atoms to drive a TorchScript-compatible loop over all (ℓ,m) pairs.

Parameters:

  • L (Optional[int]) – Maximum harmonic order. If provided, the default num_atoms=(L+1)**2. Must be ≥0.

  • num_atoms (Optional[int], keyword-only) – Explicit number of output channels; if given, overrides (L+1)**2.

  • input_dim (int, optional) – Dimensionality of each input x; must be 2 (for (θ,φ)). Default: 2.

  • seed (Optional[int], keyword-only) – Random seed for reproducible behavior (unused here).

num_atoms

Number of output channels (atoms).

Type:

int

input_dim

Expected dimensionality of the input tensor (2).

Type:

int

ls_list

List of ℓ indices for each channel.

Type:

List[int]

ms_list

List of m indices for each channel.

Type:

List[int]

forward(x: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

spherical_inr.positional_encoding.get_positional_encoding(pe: str, **kwargs) → Module[source]

Construct a positional encoding module.

This function returns an instance of a positional encoding module corresponding to the specified

type. The available types are: "herglotz", "solid_herglotz", "irregular_solid_herglotz", "fourier", "spherical_harmonics", "solid_harmonics", and "irregular_solid_harmonics". Additional parameters are forwarded to the constructor of the chosen module.

Parameters:

  • pe (str) – Identifier for the type of positional encoding. Must be one of "herglotz", "solid_herglotz", "irregular_solid_herglotz", "fourier", "spherical_harmonics", "solid_harmonics", and "irregular_solid_harmonics".

  • **kwargs – Additional keyword arguments to configure the positional encoding module. Drop any kwargs not in the constructor.

Returns:

An instance of the specified positional encoding module.

Return type:

nn.Module

Raises:

ValueError – If the specified positional encoding type is not supported.