furax.obs

class AbstractSEDOperator(AbstractLinearOperator):
    """Abstract base class for Spectral Energy Distribution (SED) operators.

    SED operators model how astrophysical components emit radiation across
    frequencies. They broadcast sky maps to multiple frequency channels.

    Subclasses must implement the ``sed()`` method to define the spectral
    energy distribution.

    Attributes:
        frequencies: Array of observation frequencies.
    """

    frequencies: Float[Array, ' a']

    def __init__(
        self,
        frequencies: Float[Array, '...'],
        *,
        in_structure: PyTree[jax.ShapeDtypeStruct],
    ) -> None:
        input_shape = self._get_input_shape(in_structure)
        n_freq = len(frequencies)
        # The operator broadcasts a map (no frequency axis) to a multi-frequency output: the new
        # frequency axis lands just before the trailing spatial axis, with any leading batch axes
        # (e.g. the leading Stokes axis of a single-array Stokes map) broadcasting through.
        n_batch = len(input_shape) - 1
        frequencies = frequencies.reshape((1,) * n_batch + (n_freq, 1))
        object.__setattr__(self, 'frequencies', frequencies)
        object.__setattr__(self, 'in_structure', in_structure)
        # sanity-check shapes at construction time
        _ = jax.eval_shape(self.mv, in_structure)

    def mv(self, x: PyTree[Inexact[Array, '...']]) -> PyTree[Inexact[Array, '...']]:
        # Insert the frequency axis before the trailing spatial axis, unless already present
        # (e.g. a genuine multi-frequency map passed in directly).
        def func(leaf: Inexact[Array, '...']) -> Inexact[Array, '...']:
            if leaf.ndim < self.frequencies.ndim:
                leaf = jnp.expand_dims(leaf, axis=-2)
            return self.sed() * leaf

        return jax.tree.map(func, x)

    @staticmethod
    def _get_input_shape(in_structure: PyTree[jax.ShapeDtypeStruct]) -> tuple[int, ...]:
        """Determine the shape of the input leaves in the PyTree.

        Args:
            in_structure (PyTree): The PyTree structure.

        Returns:
            tuple[int, ...]: The common shape of the leaves.

        Raises:
            ValueError: If the shapes of the leaves are not consistent.
        """
        input_shapes = set(leaf.shape for leaf in jax.tree.leaves(in_structure))
        if len(input_shapes) != 1:
            raise ValueError(f'the leaves of the input do not have the same shape: {in_structure}')
        return input_shapes.pop()  # type: ignore[no-any-return]

    def _broadcast_over_maps(self, x: Any) -> Float[Array, '...']:
        """Reshape a per-frequency (or scalar) array to broadcast like `frequencies`.

        The frequency axis is placed just before the trailing spatial axis, with batch axes
        broadcasting.
        """
        arr = jnp.asarray(x)
        if arr.ndim == 0:
            return arr.reshape((1,) * self.frequencies.ndim)
        return arr.reshape(self.frequencies.shape)

    @abstractmethod
    def sed(self) -> Float[Array, '...']:
        """Define the spectral energy distribution transformation.

        Returns:
            Float[Array, '...']: The transformed SED.
        """
        ...

    @staticmethod
    def _get_at(
        values: Float[Array, '...'], indices: Int[Array, '...'] | None
    ) -> Float[Array, '...']:
        """Retrieve values at specified indices, or return all values if indices are None.

        Args:
            values (Array): Input array.
            indices (Array | None): Indices to retrieve values from.

        Returns:
            Array: Subset of values or the entire array.
        """
        if indices is None:
            return values
        return values[..., indices]

@symmetric
class BeamOperator(AbstractLinearOperator):
    """Beam operator with a single transfer function shared across all Stokes components.

    Input PyTree leaves are expected to have shape ``(nfreq, npix)``.  They are
    round-tripped through a spherical harmonic transform and the beam is applied
    in alm space as a per-multipole, per-frequency scaling.

    The operator is symmetric (self-adjoint): applying it twice is equivalent to
    squaring the transfer function.  Use [`inverse`][] to obtain the deconvolution
    operator (``beam_fl`` replaced by its reciprocal).

    Algebraic rule: ``BeamOperator @ BeamOperator`` reduces to a single
    [`BeamOperator`][] whose ``beam_fl`` is the element-wise product of both
    transfer functions (see `BeamRule`).

    Attributes:
        lmax: Maximum spherical harmonic degree.
        beam_fl: Beam transfer function, shape ``(nfreq, lmax+1)`` for a
            per-frequency beam, or ``(lmax+1,)`` for a beam shared across all
            frequencies (promoted to 2-D via `jnp.atleast_2d` before
            broadcasting over frequencies).

    Examples:
        >>> import jax.numpy as jnp
        >>> from furax.obs.stokes import StokesIQU
        >>> from furax.obs.operators import BeamOperator
        >>> nside, lmax, nfreq = 16, 31, 2
        >>> npix = 12 * nside ** 2
        >>> structure = StokesIQU.structure_for((nfreq, npix), jnp.float64)
        >>> beam_fl = jnp.ones((nfreq, lmax + 1))
        >>> op = BeamOperator(lmax=lmax, beam_fl=beam_fl, in_structure=structure)
    """

    lmax: int
    beam_fl: Float[Array, '...']

    def mv(self, x: PyTree[Inexact[Array, 'nfreq npix']]) -> PyTree[Inexact[Array, 'nfreq npix']]:
        """Apply the beam to input sky maps.

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

        Returns:
            Beam-smoothed PyTree with the same structure and leaf shapes as *x*.
        """
        first_leaf = jax.tree_util.tree_leaves(x)[0]
        input_is_1d = first_leaf.ndim == 1
        first_leaf = jnp.atleast_2d(first_leaf)  # (nfreq, npix)
        nside = jhp.npix2nside(first_leaf.shape[-1])

        map2alm = Map2Alm(lmax=self.lmax, nside=nside, in_structure=self.in_structure)
        alm = map2alm.mv(x)

        # Apply the per-multipole transfer function, broadcasting it over every leading axis of the
        # alm leaf (frequency and, for a Stokes-backed map, the leading Stokes axis).
        def apply_beam(alm_leaf: jax.Array) -> jax.Array:
            fl = jnp.broadcast_to(self.beam_fl, alm_leaf.shape[:-2] + (self.lmax + 1,))
            return jnp.asarray(jhp.almxfl(alm_leaf, fl, healpy_ordering=False))

        alm_beam = jax.tree.map(apply_beam, alm)

        out = Alm2Map(lmax=self.lmax, nside=nside, in_structure=map2alm.out_structure).mv(alm_beam)
        if input_is_1d:
            out = jax.tree.map(lambda leaf: jnp.squeeze(leaf, axis=0), out)
        return out

    def inverse(self) -> 'BeamOperator':
        """Return a [`BeamOperator`][] with the reciprocal transfer function.

        Returns:
            A new [`BeamOperator`][] whose ``beam_fl`` equals ``1 / self.beam_fl``.
        """
        return BeamOperator(
            lmax=self.lmax, beam_fl=1.0 / self.beam_fl, in_structure=self.in_structure
        )

@symmetric
class BeamOperatorIQU(AbstractLinearOperator):
    """Beam operator with an independent transfer function per Stokes component.

    Input PyTree leaves are expected to have shape ``(nfreq, npix)``.  Each leaf
    is smoothed with its own transfer function taken from the corresponding leaf
    of ``beam_fl``, enabling different beams for I, Q, and U (or any other
    Stokes combination).

    The operator is symmetric (self-adjoint). Use [`inverse`][] to obtain the
    deconvolution operator (each ``beam_fl`` leaf replaced by its reciprocal).

    Algebraic rule: ``BeamOperatorIQU @ BeamOperatorIQU`` reduces to a single
    operator whose per-leaf transfer functions are the element-wise products of
    both (see `BeamIQURule`).

    Attributes:
        lmax: Maximum spherical harmonic degree.
        beam_fl: PyTree with the same structure as the input maps; each leaf has
            shape ``(nfreq, lmax+1)`` and carries the per-frequency beam transfer
            function for that Stokes component.

    Examples:
        >>> import jax.numpy as jnp
        >>> from furax.obs.stokes import StokesIQU
        >>> from furax.obs.operators import BeamOperatorIQU
        >>> nside, lmax, nfreq = 16, 31, 2
        >>> npix = 12 * nside ** 2
        >>> structure = StokesIQU.structure_for((nfreq, npix), jnp.float64)
        >>> fl = jnp.ones((nfreq, lmax + 1))
        >>> beam_fl = StokesIQU(i=fl, q=fl, u=fl)
        >>> op = BeamOperatorIQU(lmax=lmax, beam_fl=beam_fl, in_structure=structure)
    """

    lmax: int
    beam_fl: PyTree[Float[Array, 'nfreq lp1']]

    def mv(self, x: PyTree[Inexact[Array, 'nfreq npix']]) -> PyTree[Inexact[Array, 'nfreq npix']]:
        """Apply per-Stokes beams to input sky maps.

        Args:
            x: Input PyTree whose leaves have shape ``(nfreq, npix)`` or ``(npix,)``
                for a single-frequency input. The PyTree structure must match that
                of ``self.beam_fl``.

        Returns:
            Beam-smoothed PyTree with the same structure and leaf shapes as *x*.
        """
        first_leaf = jax.tree_util.tree_leaves(x)[0]
        input_is_1d = first_leaf.ndim == 1
        first_leaf = jnp.atleast_2d(first_leaf)  # (nfreq, npix)
        nside = jhp.npix2nside(first_leaf.shape[-1])

        map2alm = Map2Alm(lmax=self.lmax, nside=nside, in_structure=self.in_structure)
        alm = map2alm.mv(x)

        alm_beam = jax.tree.map(
            lambda alm_leaf, fl: jhp.almxfl(alm_leaf, fl, healpy_ordering=False),
            alm,
            self.beam_fl,
        )
        out = Alm2Map(lmax=self.lmax, nside=nside, in_structure=map2alm.out_structure).mv(alm_beam)
        if input_is_1d:
            out = jax.tree.map(lambda leaf: jnp.squeeze(leaf, axis=0), out)
        return out

    def inverse(self) -> 'BeamOperatorIQU':
        """Return a [`BeamOperatorIQU`][] with per-Stokes reciprocal transfer functions.

        Returns:
            A new [`BeamOperatorIQU`][] whose ``beam_fl`` leaves equal
            ``1 / leaf`` for each leaf in ``self.beam_fl``.
        """
        inv_beam = jax.tree.map(lambda fl: 1.0 / fl, self.beam_fl)
        return BeamOperatorIQU(lmax=self.lmax, beam_fl=inv_beam, in_structure=self.in_structure)

class CMBOperator(AbstractSEDOperator):
    """Operator for Cosmic Microwave Background (CMB) spectral energy distribution.

    The CMB has a blackbody spectrum at T_CMB ~ 2.725 K. In K_CMB units, the
    SED is constant (unity) across frequencies. In K_RJ units, a frequency-dependent
    conversion factor is applied.

    Attributes:
        frequencies: Observation frequencies [GHz].
        units: Output units ('K_CMB' or 'K_RJ').
        factor: Unit conversion factor.

    Examples:
        >>> nu = jnp.array([30, 40, 100])  # GHz
        >>> cmb_op = CMBOperator(frequencies=nu, in_structure=landscape.structure)
        >>> tod = cmb_op(sky_map)  # Broadcasts CMB map to all frequencies
    """

    factor: Float[Array, '...'] | float
    units: str = field(metadata={'static': True})

    def __init__(
        self,
        frequencies: Float[Array, '...'],
        *,
        in_structure: PyTree[jax.ShapeDtypeStruct],
        units: str = 'K_CMB',
    ) -> None:
        factor: Float[Array, ...] | float
        if units == 'K_CMB':
            factor = 1.0
        elif units == 'K_RJ':
            factor = K_RK_2_K_CMB(frequencies)
        else:
            raise ValueError(f"Unknown units: {units}. Expected 'K_CMB' or 'K_RJ'.")
        object.__setattr__(self, 'factor', factor)
        object.__setattr__(self, 'units', units)
        super().__init__(frequencies, in_structure=in_structure)

    def sed(self) -> Float[Array, '...']:
        """Compute the spectral energy distribution for the CMB.

        Returns:
            Float[Array, '...']: The SED for the CMB.
        """
        return jnp.ones_like(self.frequencies) / self._broadcast_over_maps(self.factor)

class DustOperator(AbstractSEDOperator):
    """Operator for thermal dust spectral energy distribution.

    Models dust emission as a modified blackbody: a power law times a Planck
    function. The SED is: (nu/nu0)^(1+beta) * B(nu,T)/B(nu0,T), where B is
    the Planck function.

    Supports spatially varying spectral parameters via patch indices.

    Attributes:
        frequencies: Observation frequencies [GHz].
        frequency0: Reference frequency [GHz].
        temperature: Dust temperature [K].
        beta: Spectral index (typically ~1.5).
        units: Output units ('K_CMB' or 'K_RJ').

    Examples:
        >>> nu = jnp.array([100, 143, 217, 353])  # GHz
        >>> dust_op = DustOperator(
        ...     frequencies=nu, frequency0=353, beta=1.54, temperature=20.0,
        ...     in_structure=landscape.structure
        ... )
        >>> tod = dust_op(dust_map)
    """

    temperature: Float[Array, '...']
    temperature_patch_indices: Int[Array, '...'] | None
    beta: Float[Array, '...']
    beta_patch_indices: Int[Array, '...'] | None
    factor: Float[Array, '...'] | float
    units: str = field(metadata={'static': True})
    frequency0: float = field(metadata={'static': True})

    def __init__(
        self,
        frequencies: Float[Array, '...'],
        *,
        frequency0: float = 100,
        temperature: float | Float[Array, '...'],
        units: str = 'K_CMB',
        temperature_patch_indices: Int[Array, '...'] | None = None,
        beta: float | Float[Array, '...'],
        beta_patch_indices: Int[Array, '...'] | None = None,
        in_structure: PyTree[jax.ShapeDtypeStruct],
    ) -> None:
        factor: Float[Array, ...] | float
        if units == 'K_CMB':
            factor = K_RK_2_K_CMB(frequencies) / K_RK_2_K_CMB(frequency0)
        elif units == 'K_RJ':
            factor = 1.0
        else:
            raise ValueError(f"Unknown units: {units}. Expected 'K_CMB' or 'K_RJ'.")
        object.__setattr__(self, 'temperature', jnp.asarray(temperature))
        object.__setattr__(self, 'temperature_patch_indices', temperature_patch_indices)
        object.__setattr__(self, 'beta', jnp.asarray(beta))
        object.__setattr__(self, 'beta_patch_indices', beta_patch_indices)
        object.__setattr__(self, 'units', units)
        object.__setattr__(self, 'frequency0', frequency0)
        object.__setattr__(self, 'factor', factor)
        super().__init__(frequencies, in_structure=in_structure)

    def sed(self) -> Float[Array, '...']:
        t = self._get_at(
            jnp.expm1(self.frequency0 / self.temperature * _H_OVER_K_GHZ)
            / jnp.expm1(self.frequencies / self.temperature * _H_OVER_K_GHZ),
            self.temperature_patch_indices,
        )
        b = self._get_at(
            (self.frequencies / self.frequency0) ** (1 + self.beta), self.beta_patch_indices
        )
        sed = (t * b) * self._broadcast_over_maps(self.factor)
        return sed

@diagonal
class HWPOperator(AbstractLinearOperator):
    """Operator for an ideal half-wave plate (HWP).

    A HWP flips the sign of the U (and V) Stokes parameters while leaving
    I and Q unchanged. This models the Mueller matrix diag(1, 1, -1, -1).

    The operator is diagonal and symmetric. For a rotating HWP at angle theta,
    use ``HWPOperator.create(..., angles=theta)`` which computes R(-theta) @ HWP @ R(theta).

    Algebraic rule: R(theta) @ HWP = HWP @ R(-theta).
    """

    @classmethod
    def create(
        cls,
        shape: tuple[int, ...],
        dtype: DTypeLike = np.float64,
        stokes: ValidStokesType = 'IQU',
        *,
        angles: Float[Array, '...'] | None = None,
    ) -> AbstractLinearOperator:
        in_structure = Stokes.class_for(stokes).structure_for(shape, dtype)
        hwp = cls(in_structure=in_structure)
        if angles is None:
            return hwp
        rot = QURotationOperator(angles=angles, in_structure=in_structure)
        rotated_hwp: AbstractLinearOperator = rot.T @ hwp @ rot
        return rotated_hwp

    def mv(self, x: StokesType) -> Stokes:
        if isinstance(x, StokesI):
            return x
        if isinstance(x, StokesQU):
            return StokesQU(x.q, -x.u)
        if isinstance(x, StokesIQU):
            return StokesIQU(x.i, x.q, -x.u)
        if isinstance(x, StokesIQUV):
            return StokesIQUV(x.i, x.q, -x.u, -x.v)
        raise NotImplementedError

class LinearPolarizerOperator(AbstractLinearOperator):
    """Operator for an ideal linear polarizer.

    Extracts the intensity seen by a linear polarizer aligned with the x-axis:
    d = 0.5 * (I + Q). For a polarizer at angle psi, use
    ``LinearPolarizerOperator.create(..., angles=psi)``.

    This implements: d = 0.5 * (I + Q*cos(2*psi) + U*sin(2*psi)).

    Algebraic rule: LinearPolarizer @ HWP = LinearPolarizer.
    """

    @classmethod
    def create(
        cls,
        shape: tuple[int, ...],
        dtype: DTypeLike = np.float64,
        stokes: ValidStokesType = 'IQU',
        *,
        angles: Float[Array, '...'] | None = None,
    ) -> AbstractLinearOperator:
        in_structure = Stokes.class_for(stokes).structure_for(shape, dtype)
        polarizer = cls(in_structure=in_structure)
        if angles is None:
            return polarizer
        rot = QURotationOperator(angles=angles, in_structure=in_structure)
        rotated_polarizer: AbstractLinearOperator = polarizer @ rot
        return rotated_polarizer

    def mv(self, x: StokesType) -> Float[Array, '...']:
        if isinstance(x, StokesI):
            return 0.5 * x.i
        if isinstance(x, StokesQU):
            return 0.5 * x.q
        return 0.5 * (x.i + x.q)