Skip to content

furax.math

Modules:

  • bspline –
  • quaternion –

    Functions for quaternion operations using JAX.

  • sht –

    Operators for spherical harmonic transforms backed by jax-healpy.

furax.math.bspline

Functions:

  • cubic_bspline –

    Cubic cardinal B-spline, support [0,4).

  • spline_basis –

    Dense cubic B-spline basis on K = n_knots + 2 uniform knots.

  • spline_window –

    Banded form of spline_basis: the 4 nonzero knots under each sample.

cubic_bspline(u)

Cubic cardinal B-spline, support [0,4).

Source code in src/furax/math/bspline.py 
 6
 7
 8
 9
10
11
12
13
def cubic_bspline(u: Float[Array, ' samp']) -> Float[Array, ' samp']:
    """Cubic cardinal B-spline, support [0,4)."""
    out = jnp.zeros_like(u)
    out = jnp.where((u >= 0) & (u < 1), (1 / 6) * u**3, out)
    out = jnp.where((u >= 1) & (u < 2), (1 / 6) * (-3 * u**3 + 12 * u**2 - 12 * u + 4), out)
    out = jnp.where((u >= 2) & (u < 3), (1 / 6) * (3 * u**3 - 24 * u**2 + 60 * u - 44), out)
    out = jnp.where((u >= 3) & (u < 4), (1 / 6) * (4 - u) ** 3, out)
    return out

spline_basis(times, n_knots)

Dense cubic B-spline basis on K = n_knots + 2 uniform knots.

Returns:

  • B ( Float[Array, 'k samp'] ) –

    (K, N) basis matrix, B[j] = cubic_bspline(position - j + 1).

Source code in src/furax/math/bspline.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def spline_basis(
    times: Float[Array, ' samp'],
    n_knots: int,
) -> Float[Array, 'k samp']:
    """Dense cubic B-spline basis on ``K = n_knots + 2`` uniform knots.

    Returns:
        B: (K, N) basis matrix, ``B[j] = cubic_bspline(position - j + 1)``.
    """
    K = _n_grid_knots(n_knots)
    p = _spline_position(times, K)
    # knot j peaks at p = j + 1 (where u = 2); evaluate all knots at once against every sample.
    u = p[None, :] - jnp.arange(K)[:, None] + 1.0
    return cubic_bspline(u)

spline_window(times, n_knots)

Banded form of spline_basis: the 4 nonzero knots under each sample.

A cubic B-spline reaches only 4 consecutive knots, so each sample's column of spline_basis has just 4 nonzero entries. This returns those directly, avoiding the mostly-zero `(K, N) matrix (seefurax.mapmaking.templates.WindowedBasis`).

Returns:

  • offset ( Int[Array, ' samp'] ) –

    (N,) index of the first of the 4 knots, clamped to [0, K - 4].

  • weights ( Float[Array, 'samp 4'] ) –

    (N, 4) the B-spline values at knots offset + 0..3; weights of knots whose support misses the sample (at the time-range edges) are exactly zero.

Source code in src/furax/math/bspline.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def spline_window(
    times: Float[Array, ' samp'],
    n_knots: int,
) -> tuple[Int[Array, ' samp'], Float[Array, 'samp 4']]:
    """Banded form of `spline_basis`: the 4 nonzero knots under each sample.

    A cubic B-spline reaches only 4 consecutive knots, so each sample's column of
    `spline_basis` has just 4 nonzero entries. This returns those directly, avoiding the
    mostly-zero ``(K, N) matrix (see `furax.mapmaking.templates.WindowedBasis`).

    Returns:
        offset: (N,) index of the first of the 4 knots, clamped to ``[0, K - 4]``.
        weights: (N, 4) the B-spline values at knots ``offset + 0..3``; weights of knots
            whose support misses the sample (at the time-range edges) are exactly zero.
    """
    K = _n_grid_knots(n_knots)
    p = _spline_position(times, K)
    offset = jnp.clip(jnp.floor(p).astype(jnp.int32) - 2, 0, K - 4)
    # knot offset+o contributes cubic_bspline(p - (offset+o) + 1); out-of-support -> 0.
    u = p[:, None] - (offset[:, None] + jnp.arange(4)) + 1.0
    return offset, cubic_bspline(u)

furax.math.quaternion

Functions for quaternion operations using JAX.

This module provides vectorized functions for multiplying quaternions, and rotating 3D vectors by quaternions.

We use scalar-vector storage, i.e. (1,i,j,k) with the scalar part first.

Functions:

  • euler –

    The quaternion representing an Euler rotation.

  • qmul –

    Compute quaternion multiplication.

  • qrot –

    Rotate vector by quaternion.

  • qrot_zaxis –

    Rotate the Z axis [0,0,1] by a given quaternion.

  • qrot_xaxis –

    Rotate the X axis [1,0,0] by a given quaternion.

  • to_polarization_angle_cos_sin –

    Compute cos and sin of the polarization angle from the rotation quaternion.

  • to_polarization_angle –

    Compute the polarization angle from the rotation quaternion using the COSMO convention.

euler(axis, angle)

The quaternion representing an Euler rotation.

For example, if axis=2 the computed quaternion(s) will have components:

\[ q = (\cos(\theta/2), 0, 0, \sin(\theta/2)) \]

Parameters:

  • axis (int) –

    The index of the cartesian axis of the rotation (x, y, z). Must be 0, 1, or 2.

  • angle (Float[Array, ...]) –

    Angle of rotation, in radians.

Returns:

  • Quat –

    Quaternion array of shape (..., 4).

Source code in src/furax/math/quaternion.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@partial(jit, static_argnums=(0,))
def euler(axis: int, angle: Float[Array, '...']) -> Quat:
    r"""The quaternion representing an Euler rotation.

    For example, if axis=2 the computed quaternion(s) will have components:

    $$
    q = (\cos(\theta/2), 0, 0, \sin(\theta/2))
    $$

    Args:
        axis: The index of the cartesian axis of the rotation (x, y, z).
            Must be 0, 1, or 2.
        angle: Angle of rotation, in radians.

    Returns:
        Quaternion array of shape (..., 4).
    """
    angle = jnp.asarray(angle)
    c = jnp.cos(angle / 2)
    s = jnp.sin(angle / 2)
    zeros = jnp.zeros_like(angle)
    components = [c, zeros, zeros, zeros]
    components[axis + 1] = s
    return jnp.stack(components, axis=-1)

qmul(q1, q2)

Compute quaternion multiplication.

Source code in src/furax/math/quaternion.py 
61
62
63
64
65
66
67
68
69
70
71
72
@jit
@partial(jnp.vectorize, signature='(4),(4)->(4)')
def qmul(q1: Quat, q2: Quat) -> Quat:
    """Compute quaternion multiplication."""
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    # https://en.wikipedia.org/wiki/Quaternion#Hamilton_product
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
    return jnp.array([w, x, y, z])

qrot(q, vec)

Rotate vector by quaternion.

Source code in src/furax/math/quaternion.py 
83
84
85
86
87
88
@jit
@partial(jnp.vectorize, signature='(4),(3)->(3)')
def qrot(q: Quat, vec: Vec3) -> Vec3:
    """Rotate vector by quaternion."""
    rot = to_rotation_matrix(q)
    return rot @ vec  # type: ignore[no-any-return]

qrot_zaxis(q)

Rotate the Z axis [0,0,1] by a given quaternion.

Source code in src/furax/math/quaternion.py 
118
119
120
121
122
123
124
125
@jit
@partial(jnp.vectorize, signature='(4)->(3)')
def qrot_zaxis(q: Quat) -> Vec3:
    """Rotate the Z axis [0,0,1] by a given quaternion."""
    w, x, y, z = q
    n = jnp.sum(q**2)
    s = jnp.where(n > 1e-8, 2 / n, 0)
    return jnp.array([0, 0, 1]) + s * jnp.array([x * z + w * y, y * z - w * x, -x * x - y * y])

qrot_xaxis(q)

Rotate the X axis [1,0,0] by a given quaternion.

Source code in src/furax/math/quaternion.py 
128
129
130
131
132
133
134
135
@jit
@partial(jnp.vectorize, signature='(4)->(3)')
def qrot_xaxis(q: Quat) -> Vec3:
    """Rotate the X axis [1,0,0] by a given quaternion."""
    w, x, y, z = q
    n = jnp.sum(q**2)
    s = jnp.where(n > 1e-8, 2 / n, 0)
    return jnp.array([1, 0, 0]) + s * jnp.array([-y * y - z * z, x * y + w * z, x * z - w * y])

to_polarization_angle_cos_sin(q)

Compute cos and sin of the polarization angle from the rotation quaternion.

Equivalent to (cos(pa), sin(pa)) where pa = to_polarization_angle(q), but avoids transcendental functions by using quaternion algebra directly.

See to_polarization_angle for the definition and convention.

Source code in src/furax/math/quaternion.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
@jit
@partial(jnp.vectorize, signature='(4)->(),()')
def to_polarization_angle_cos_sin(q: Quat) -> tuple[Ang, Ang]:
    """Compute cos and sin of the polarization angle from the rotation quaternion.

    Equivalent to ``(cos(pa), sin(pa))`` where ``pa = to_polarization_angle(q)``,
    but avoids transcendental functions by using quaternion algebra directly.

    See [`to_polarization_angle`][] for the definition and convention.
    """
    a, b, c, d = q
    cos_theta = a**2 - b**2 - c**2 + d**2
    # clip to avoid numerical issues giving cos_theta**2 > 1
    half_sin_theta = 0.5 * jnp.sqrt(jnp.clip(1 - cos_theta**2, 0.0, None))
    at_pole = half_sin_theta == 0
    safe = jnp.where(at_pole, 1.0, half_sin_theta)
    # angle undefined at the pole, use pa = 0
    # this matches the result of to_polarization_angle (atan2(0, 0) = 0)
    cos_pa = jnp.where(at_pole, 1.0, (a * c - b * d) / safe)
    sin_pa = jnp.where(at_pole, 0.0, (a * b + c * d) / safe)
    return cos_pa, sin_pa

to_polarization_angle(q)

Compute the polarization angle from the rotation quaternion using the COSMO convention.

The polarization angle is measured from the South through the East.

The rotation quaternion q transforms detector coordinates to celestial (equatorial) coordinates. In detector coordinates: - The detector points in the z direction. - The detector is sensitive to electric fields in the x direction.

After applying the rotation: - The vector v identifies the point on the equatorial sphere where the detector is pointing. - The vector u defines the polarization-sensitive direction, tangent to the unit sphere at v.

The unit vector toward the South in the tangent plane is: - w = -z - (-z · v) v

The polarization angle pa between w and u is computed as: - cos(pa) = w · u = -u_z (since u · v = 0) - sin(pa) = (w × u) · v = (u × v) · w = (v × u) · z = v_x u_y - v_y u_x - Therefore, pa = atan2(v_x u_y - v_y u_x, -u_z)

Parameters:

  • q (Quat) –

    Rotation quaternion array of shape (*dims, 4).

Returns:

  • Ang –

    Polarization angle array of shape (*dims).

Source code in src/furax/math/quaternion.py 
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
@jit
@partial(jnp.vectorize, signature='(4)->()')
def to_polarization_angle(q: Quat) -> Ang:
    """Compute the polarization angle from the rotation quaternion using the COSMO convention.

    The polarization angle is measured from the South through the East.

    The rotation quaternion `q` transforms detector coordinates to celestial (equatorial) coordinates.
    In detector coordinates:
    - The detector points in the z direction.
    - The detector is sensitive to electric fields in the x direction.

    After applying the rotation:
    - The vector `v` identifies the point on the equatorial sphere where the detector is pointing.
    - The vector `u` defines the polarization-sensitive direction, tangent to the unit sphere at `v`.

    The unit vector toward the South in the tangent plane is:
    - `w = -z - (-z · v) v`

    The polarization angle `pa` between `w` and `u` is computed as:
    - `cos(pa) = w · u = -u_z` (since `u · v = 0`)
    - `sin(pa) = (w × u) · v = (u × v) · w = (v × u) · z = v_x u_y - v_y u_x`
    - Therefore, `pa = atan2(v_x u_y - v_y u_x, -u_z)`

    Args:
        q: Rotation quaternion array of shape (*dims, 4).

    Returns:
        Polarization angle array of shape (*dims).
    """
    v = qrot_zaxis(q)
    u = qrot_xaxis(q)
    return jnp.arctan2(v[0] * u[1] - v[1] * u[0], -u[2])

furax.math.sht

Operators for spherical harmonic transforms backed by jax-healpy.

The two core operators are:

  • Map2Alm: analysis transform (pixel map → spherical harmonic coefficients).
  • Alm2Map: synthesis transform (spherical harmonic coefficients → pixel map).

Both operators accept PyTree inputs whose leaves are 2-D arrays of shape (nfreq, npix) or (nfreq, lmax+1, 2*lmax+1). A 1-D leaf is promoted to 2-D via jnp.atleast_2d before processing so that single-frequency inputs work without special-casing.

SHTRule registers an algebraic simplification so that the composition Map2Alm @ Alm2Map (synthesis followed by analysis) is reduced to an IdentityOperator at operator-construction time.

Classes:

  • Map2Alm –

    Analysis spherical harmonic transform: pixel map → alm coefficients.

  • Alm2Map –

    Synthesis spherical harmonic transform: alm coefficients → pixel map.

  • SHTRule –

    Algebraic rule reducing Map2Alm @ Alm2Map to an identity.

Map2Alm dataclass

Bases: AbstractLinearOperator

Analysis spherical harmonic transform: pixel map → alm coefficients.

Each leaf of the input PyTree must be a 2-D array of shape (nfreq, npix). A 1-D leaf is silently promoted to (1, npix) via jnp.atleast_2d. The transform are applied to all frequency rows, via jhp.map2alm producing an output leaf of shape (nfreq, lmax+1, 2*lmax+1).

Construction is guarded by a hasattr check against jax_healpy: if jhp.map2alm is missing (e.g. on an older jax-healpy release), instantiation raises ImportError.

Attributes:

  • lmax (int) –

    Maximum spherical harmonic degree.

  • nside (int) –

    HEALPix resolution parameter, carried so that inverse can construct Alm2Map with the correct resolution.

  • iter (int) –

    Number of iterations for the map2alm solver; more iterations, which are more expensive, yield more accurate alm coefficients. (default is 3)

  • pol (bool) –

    Reserved for the polarization-aware SHT. Routed to jhp.map2alm but only False is supported today; passing True raises NotImplementedError at construction.

Examples:

>>> import jax.numpy as jnp
>>> from furax.math.sht import Map2Alm
>>> from furax.obs.stokes import StokesIQU
>>> nside, lmax, nfreq = 32, 63, 3
>>> npix = 12 * nside ** 2
>>> structure = StokesIQU.structure_for((nfreq, npix), jnp.float64)
>>> op = Map2Alm(lmax=lmax, nside=nside, in_structure=structure)

Methods:

Source code in src/furax/math/sht.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class Map2Alm(AbstractLinearOperator):
    """Analysis spherical harmonic transform: pixel map → alm coefficients.

    Each leaf of the input PyTree must be a 2-D array of shape
    ``(nfreq, npix)``.  A 1-D leaf is silently promoted to ``(1, npix)`` via
    ``jnp.atleast_2d``.  The transform are applied to all frequency rows,
    via ``jhp.map2alm`` producing an output leaf of shape
    ``(nfreq, lmax+1, 2*lmax+1)``.

    Construction is guarded by a ``hasattr`` check against ``jax_healpy``:
    if ``jhp.map2alm`` is missing (e.g. on an older ``jax-healpy`` release),
    instantiation raises [`ImportError`][].

    Attributes:
        lmax: Maximum spherical harmonic degree.
        nside: HEALPix resolution parameter, carried so that
            [`inverse`][] can construct
            [`Alm2Map`][] with the correct resolution.
        iter: Number of iterations for the map2alm solver; more iterations,
            which are more expensive, yield more accurate alm coefficients.
            (default is 3)
        pol: Reserved for the polarization-aware SHT.  Routed to
            ``jhp.map2alm`` but only ``False`` is supported today; passing
            ``True`` raises [`NotImplementedError`][] at construction.

    Examples:
        >>> import jax.numpy as jnp
        >>> from furax.math.sht import Map2Alm
        >>> from furax.obs.stokes import StokesIQU
        >>> nside, lmax, nfreq = 32, 63, 3
        >>> npix = 12 * nside ** 2
        >>> structure = StokesIQU.structure_for((nfreq, npix), jnp.float64)
        >>> op = Map2Alm(lmax=lmax, nside=nside, in_structure=structure)
    """

    lmax: int = field(metadata={'static': True})
    nside: int = field(metadata={'static': True})
    iter: int = field(default=3, metadata={'static': True})
    pol: bool = field(default=False, metadata={'static': True})

    def __post_init__(self) -> None:
        super().__post_init__()
        if self.pol:
            raise NotImplementedError(
                'Polarization-aware SHT (pol=True) is not yet supported; pass pol=False.'
            )
        if not hasattr(jhp, 'map2alm'):
            raise ImportError(
                'jax_healpy.map2alm is unavailable; upgrade jax-healpy to a version '
                'that ships map2alm.'
            )

    def mv(self, x: PyTree[Inexact[Array, ' _a']]) -> PyTree[Inexact[Array, ' _b']]:
        """Apply the analysis SHT to every leaf of *x*.

        Args:
            x: Input PyTree whose leaves have shape ``(nfreq, npix)`` or
                ``(npix,)`` for a single frequency.

        Returns:
            PyTree with the same structure as *x*; each leaf has shape
            ``(nfreq, lmax+1, 2*lmax+1)`` and complex dtype.
        """

        def func(value: jax.Array) -> jax.Array:
            value = jnp.atleast_2d(value)  # (*batch, npix)
            return _batch_flatten_apply(
                value,
                1,
                lambda flat: jhp.map2alm(flat, iter=self.iter, lmax=self.lmax, pol=self.pol),
            )

        return jax.tree.map(func, x)

    def inverse(self) -> 'Alm2Map':
        """Return the pseudo-inverse operator [`Alm2Map`][].

        Returns:
            An [`Alm2Map`][] whose ``in_structure`` matches the output
            structure of this operator (i.e. alm space).
        """
        return Alm2Map(
            lmax=self.lmax,
            nside=self.nside,
            iter=self.iter,
            in_structure=self.out_structure,
            pol=self.pol,
        )

lmax = field(metadata={'static': True}) class-attribute instance-attribute

nside = field(metadata={'static': True}) class-attribute instance-attribute

iter = field(default=3, metadata={'static': True}) class-attribute instance-attribute

pol = field(default=False, metadata={'static': True}) class-attribute instance-attribute

__post_init__()

Source code in src/furax/math/sht.py 
84
85
86
87
88
89
90
91
92
93
94
def __post_init__(self) -> None:
    super().__post_init__()
    if self.pol:
        raise NotImplementedError(
            'Polarization-aware SHT (pol=True) is not yet supported; pass pol=False.'
        )
    if not hasattr(jhp, 'map2alm'):
        raise ImportError(
            'jax_healpy.map2alm is unavailable; upgrade jax-healpy to a version '
            'that ships map2alm.'
        )

mv(x)

Apply the analysis SHT to every leaf of x.

Parameters:

  • x (PyTree[Inexact[Array, ' _a']]) –

    Input PyTree whose leaves have shape (nfreq, npix) or (npix,) for a single frequency.

Returns:

  • PyTree[Inexact[Array, ' _b']] –

    PyTree with the same structure as x; each leaf has shape

  • PyTree[Inexact[Array, ' _b']] –

    (nfreq, lmax+1, 2*lmax+1) and complex dtype.

Source code in src/furax/math/sht.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def mv(self, x: PyTree[Inexact[Array, ' _a']]) -> PyTree[Inexact[Array, ' _b']]:
    """Apply the analysis SHT to every leaf of *x*.

    Args:
        x: Input PyTree whose leaves have shape ``(nfreq, npix)`` or
            ``(npix,)`` for a single frequency.

    Returns:
        PyTree with the same structure as *x*; each leaf has shape
        ``(nfreq, lmax+1, 2*lmax+1)`` and complex dtype.
    """

    def func(value: jax.Array) -> jax.Array:
        value = jnp.atleast_2d(value)  # (*batch, npix)
        return _batch_flatten_apply(
            value,
            1,
            lambda flat: jhp.map2alm(flat, iter=self.iter, lmax=self.lmax, pol=self.pol),
        )

    return jax.tree.map(func, x)

inverse()

Return the pseudo-inverse operator Alm2Map.

Returns:

  • Alm2Map –

    An Alm2Map whose in_structure matches the output

  • Alm2Map –

    structure of this operator (i.e. alm space).

Source code in src/furax/math/sht.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def inverse(self) -> 'Alm2Map':
    """Return the pseudo-inverse operator [`Alm2Map`][].

    Returns:
        An [`Alm2Map`][] whose ``in_structure`` matches the output
        structure of this operator (i.e. alm space).
    """
    return Alm2Map(
        lmax=self.lmax,
        nside=self.nside,
        iter=self.iter,
        in_structure=self.out_structure,
        pol=self.pol,
    )

Alm2Map dataclass

Bases: AbstractLinearOperator

Synthesis spherical harmonic transform: alm coefficients → pixel map.

Each leaf of the input PyTree must be a 2-D array of shape (nfreq, lmax+1, 2*lmax+1). A 2-D leaf (lmax+1, 2*lmax+1) is silently reshaped to (1, lmax+1, 2*lmax+1) via reshape(-1, ...), handling the single-frequency case without branching. The transform are applied to all frequency rows, via jhp.alm2map producing an output leaf of shape (nfreq, npix) where npix = 12 * nside ** 2.

Construction is guarded by a hasattr check against jax_healpy: if jhp.alm2map is missing (e.g. on an older jax-healpy release), instantiation raises ImportError.

Attributes:

  • lmax (int) –

    Maximum spherical harmonic degree.

  • nside (int) –

    HEALPix resolution parameter used by the synthesis step.

  • iter (int) –

    Number of iterations for the map2alm solver used by the inverse. (Default is 3)

  • pol (bool) –

    Reserved for the polarization-aware SHT. Routed to jhp.alm2map but only False is supported today; passing True raises NotImplementedError at construction.

Examples:

>>> import jax.numpy as jnp
>>> from furax.math.sht import Alm2Map
>>> from furax.obs.stokes import StokesIQU
>>> nside, lmax, nfreq = 32, 63, 3
>>> structure = StokesIQU.structure_for((nfreq, lmax+1, 2*lmax+1), jnp.complex128)
>>> op = Alm2Map(lmax=lmax, nside=nside, in_structure=structure)

Methods:

Source code in src/furax/math/sht.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
class Alm2Map(AbstractLinearOperator):
    """Synthesis spherical harmonic transform: alm coefficients → pixel map.

    Each leaf of the input PyTree must be a 2-D array of shape
    ``(nfreq, lmax+1, 2*lmax+1)``.  A 2-D leaf ``(lmax+1, 2*lmax+1)`` is
    silently reshaped to ``(1, lmax+1, 2*lmax+1)`` via ``reshape(-1, ...)``,
    handling the single-frequency case without branching.  The transform are applied
    to all frequency rows, via ``jhp.alm2map`` producing an output leaf of shape
    ``(nfreq, npix)`` where ``npix = 12 * nside ** 2``.

    Construction is guarded by a ``hasattr`` check against ``jax_healpy``:
    if ``jhp.alm2map`` is missing (e.g. on an older ``jax-healpy`` release),
    instantiation raises [`ImportError`][].

    Attributes:
        lmax: Maximum spherical harmonic degree.
        nside: HEALPix resolution parameter used by the synthesis step.
        iter: Number of iterations for the map2alm solver used by
            the inverse. (Default is 3)
        pol: Reserved for the polarization-aware SHT.  Routed to
            ``jhp.alm2map`` but only ``False`` is supported today; passing
            ``True`` raises [`NotImplementedError`][] at construction.

    Examples:
        >>> import jax.numpy as jnp
        >>> from furax.math.sht import Alm2Map
        >>> from furax.obs.stokes import StokesIQU
        >>> nside, lmax, nfreq = 32, 63, 3
        >>> structure = StokesIQU.structure_for((nfreq, lmax+1, 2*lmax+1), jnp.complex128)
        >>> op = Alm2Map(lmax=lmax, nside=nside, in_structure=structure)
    """

    lmax: int = field(metadata={'static': True})
    nside: int = field(metadata={'static': True})
    iter: int = field(default=3, metadata={'static': True})
    pol: bool = field(default=False, metadata={'static': True})

    def __post_init__(self) -> None:
        super().__post_init__()
        if self.pol:
            raise NotImplementedError(
                'Polarization-aware SHT (pol=True) is not yet supported; pass pol=False.'
            )
        if not hasattr(jhp, 'alm2map'):
            raise ImportError(
                'jax_healpy.alm2map is unavailable; upgrade jax-healpy to a version '
                'that ships alm2map.'
            )

    def mv(self, x: PyTree[Inexact[Array, ' _a']]) -> PyTree[Inexact[Array, ' _b']]:
        """Apply the synthesis SHT to every leaf of *x*.

        Args:
            x: Input PyTree whose leaves have shape ``(nfreq, lmax+1, 2*lmax+1)`` or
                ``(lmax+1, 2*lmax+1,)`` for a single frequency.

        Returns:
            PyTree with the same structure as *x*; each leaf has shape
            ``(nfreq, npix)`` and real dtype.
        """

        def func(value: jax.Array) -> jax.Array:
            return _batch_flatten_apply(
                value,
                2,
                lambda flat: jnp.real(
                    jhp.alm2map(flat, nside=self.nside, lmax=self.lmax, pol=self.pol)
                ),
            )

        return jax.tree.map(func, x)

    def inverse(self) -> 'Map2Alm':
        """Return the pseudo-inverse operator [`Map2Alm`][].

        Returns:
            A [`Map2Alm`][] whose ``in_structure`` matches the output
            structure of this operator (i.e. map space).
        """
        return Map2Alm(
            lmax=self.lmax,
            nside=self.nside,
            iter=self.iter,
            in_structure=self.out_structure,
            pol=self.pol,
        )

lmax = field(metadata={'static': True}) class-attribute instance-attribute

nside = field(metadata={'static': True}) class-attribute instance-attribute

iter = field(default=3, metadata={'static': True}) class-attribute instance-attribute

pol = field(default=False, metadata={'static': True}) class-attribute instance-attribute

__post_init__()

Source code in src/furax/math/sht.py 
171
172
173
174
175
176
177
178
179
180
181
def __post_init__(self) -> None:
    super().__post_init__()
    if self.pol:
        raise NotImplementedError(
            'Polarization-aware SHT (pol=True) is not yet supported; pass pol=False.'
        )
    if not hasattr(jhp, 'alm2map'):
        raise ImportError(
            'jax_healpy.alm2map is unavailable; upgrade jax-healpy to a version '
            'that ships alm2map.'
        )

mv(x)

Apply the synthesis SHT to every leaf of x.

Parameters:

  • x (PyTree[Inexact[Array, ' _a']]) –

    Input PyTree whose leaves have shape (nfreq, lmax+1, 2*lmax+1) or (lmax+1, 2*lmax+1,) for a single frequency.

Returns:

  • PyTree[Inexact[Array, ' _b']] –

    PyTree with the same structure as x; each leaf has shape

  • PyTree[Inexact[Array, ' _b']] –

    (nfreq, npix) and real dtype.

Source code in src/furax/math/sht.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def mv(self, x: PyTree[Inexact[Array, ' _a']]) -> PyTree[Inexact[Array, ' _b']]:
    """Apply the synthesis SHT to every leaf of *x*.

    Args:
        x: Input PyTree whose leaves have shape ``(nfreq, lmax+1, 2*lmax+1)`` or
            ``(lmax+1, 2*lmax+1,)`` for a single frequency.

    Returns:
        PyTree with the same structure as *x*; each leaf has shape
        ``(nfreq, npix)`` and real dtype.
    """

    def func(value: jax.Array) -> jax.Array:
        return _batch_flatten_apply(
            value,
            2,
            lambda flat: jnp.real(
                jhp.alm2map(flat, nside=self.nside, lmax=self.lmax, pol=self.pol)
            ),
        )

    return jax.tree.map(func, x)

inverse()

Return the pseudo-inverse operator Map2Alm.

Returns:

  • Map2Alm –

    A Map2Alm whose in_structure matches the output

  • Map2Alm –

    structure of this operator (i.e. map space).

Source code in src/furax/math/sht.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
def inverse(self) -> 'Map2Alm':
    """Return the pseudo-inverse operator [`Map2Alm`][].

    Returns:
        A [`Map2Alm`][] whose ``in_structure`` matches the output
        structure of this operator (i.e. map space).
    """
    return Map2Alm(
        lmax=self.lmax,
        nside=self.nside,
        iter=self.iter,
        in_structure=self.out_structure,
        pol=self.pol,
    )

SHTRule

Bases: AbstractCompositionRule

Algebraic rule reducing Map2Alm @ Alm2Map to an identity.

The composition analysis after synthesis (Map2Alm @ Alm2Map) is exact when the band-limit lmax is consistent with the HEALPix resolution nside, so the rule collapses it to an IdentityOperator on the alm input space.

The reverse composition Alm2Map @ Map2Alm is a band-limited projection, not an identity, and is therefore not reduced.

Methods:

  • apply –

    Reduce left @ right to an identity when the pair is Map2Alm/Alm2Map.

Attributes:

Source code in src/furax/math/sht.py 
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
class SHTRule(AbstractCompositionRule):
    """Algebraic rule reducing ``Map2Alm @ Alm2Map`` to an identity.

    The composition *analysis after synthesis* (``Map2Alm @ Alm2Map``) is
    exact when the band-limit ``lmax`` is consistent with the HEALPix
    resolution ``nside``, so the rule collapses it to an
    [`IdentityOperator`][furax.IdentityOperator] on the alm input space.

    The reverse composition ``Alm2Map @ Map2Alm`` is a band-limited
    projection, **not** an identity, and is therefore not reduced.
    """

    left_operator_class = Map2Alm
    right_operator_class = Alm2Map

    def apply(
        self, left: AbstractLinearOperator, right: AbstractLinearOperator
    ) -> list[AbstractLinearOperator]:
        """Reduce ``left @ right`` to an identity when the pair is Map2Alm/Alm2Map.

        Args:
            left: Left operator in the composition (must be [`Map2Alm`][]).
            right: Right operator in the composition (must be [`Alm2Map`][]).

        Returns:
            A single-element list containing an
            [`IdentityOperator`][furax.IdentityOperator] on the alm input space.

        Raises:
            NoReduction: If the operator pair does not match the expected types or if the
                transforms use different ``lmax`` values.
        """
        if isinstance(left, Map2Alm) and isinstance(right, Alm2Map):
            if left.lmax != right.lmax:
                raise NoReduction
            return [IdentityOperator(in_structure=right.in_structure)]
        raise NoReduction

left_operator_class = Map2Alm class-attribute instance-attribute

right_operator_class = Alm2Map class-attribute instance-attribute

apply(left, right)

Reduce left @ right to an identity when the pair is Map2Alm/Alm2Map.

Parameters:

Returns:

Raises:

  • NoReduction –

    If the operator pair does not match the expected types or if the transforms use different lmax values.

Source code in src/furax/math/sht.py 
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
def apply(
    self, left: AbstractLinearOperator, right: AbstractLinearOperator
) -> list[AbstractLinearOperator]:
    """Reduce ``left @ right`` to an identity when the pair is Map2Alm/Alm2Map.

    Args:
        left: Left operator in the composition (must be [`Map2Alm`][]).
        right: Right operator in the composition (must be [`Alm2Map`][]).

    Returns:
        A single-element list containing an
        [`IdentityOperator`][furax.IdentityOperator] on the alm input space.

    Raises:
        NoReduction: If the operator pair does not match the expected types or if the
            transforms use different ``lmax`` values.
    """
    if isinstance(left, Map2Alm) and isinstance(right, Alm2Map):
        if left.lmax != right.lmax:
            raise NoReduction
        return [IdentityOperator(in_structure=right.in_structure)]
    raise NoReduction