SHT beam convolution#114
Open
ArtemBasyrov wants to merge 9 commits intomainfrom
Open
Conversation
There was a problem hiding this comment.
Pull request overview
Adds spherical-harmonic-transform (SHT) operators and beam-convolution operators to furax, along with algebraic reduction rules and tests to simplify operator compositions.
Changes:
- Introduce
Map2Alm/Alm2Maplinear operators backed byjax_healpy, plus anSHTRulereduction. - Add
BeamOperator/BeamOperatorIQUfor symmetric-beam convolution in harmonic space, plus beam composition reduction rules. - Add unit tests covering shape/dtype behavior, transpose/inverse expectations, and reduction behavior.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
src/furax/math/sht.py |
New SHT operators and SHTRule for simplifying Map2Alm @ Alm2Map. |
src/furax/obs/operators/_beam_operator.py |
New harmonic-space beam operators and beam composition reduction rules. |
src/furax/obs/operators/__init__.py |
Re-export beam operators from the obs operators package. |
src/furax/obs/__init__.py |
Re-export beam operators from the top-level obs namespace. |
tests/math/test_sht.py |
New tests for SHT operator behavior and SHTRule. |
tests/obs/operators/test_beam_operator.py |
New tests for beam operators and beam reduction rules. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Author
|
Implemented changes, suggested by the copilot review |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 9 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+152
to
+188
| class SHTRule(AbstractBinaryRule): | ||
| """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 | ||
| :class:`~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, op1: AbstractLinearOperator, op2: AbstractLinearOperator | ||
| ) -> list[AbstractLinearOperator]: | ||
| """Reduce ``op1 @ op2`` to an identity when the pair is Map2Alm/Alm2Map. | ||
|
|
||
| Args: | ||
| op1: Left operator in the composition (must be :class:`Map2Alm`). | ||
| op2: Right operator in the composition (must be :class:`Alm2Map`). | ||
|
|
||
| Returns: | ||
| A single-element list containing an | ||
| :class:`~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(op1, Map2Alm) and isinstance(op2, Alm2Map): | ||
| if op1.lmax != op2.lmax: | ||
| raise NoReduction | ||
| return [IdentityOperator(in_structure=op2.in_structure)] | ||
| raise NoReduction |
Author
There was a problem hiding this comment.
The motivation behind the rule is the mathematical expectation, hence I would rather it stays
Comment on lines
+101
to
+113
| first_leaf = jax.tree_util.tree_leaves(x)[0] | ||
| 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) | ||
|
|
||
| nfreq = first_leaf.shape[0] | ||
| fl = jnp.broadcast_to(jnp.atleast_2d(self.beam_fl), (nfreq, self.lmax + 1)) | ||
| alm_beam = jax.tree.map(lambda alm_leaf: _apply_fl_to_alm_leaf(alm_leaf, fl), alm) | ||
|
|
||
| return Alm2Map(lmax=self.lmax, nside=nside, in_structure=map2alm.out_structure).mv(alm_beam) | ||
|
|
Comment on lines
+162
to
+181
| 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)``. 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] | ||
| 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(_apply_fl_to_alm_leaf, alm, self.beam_fl) | ||
| return Alm2Map(lmax=self.lmax, nside=nside, in_structure=map2alm.out_structure).mv(alm_beam) | ||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Added sht operators for jax_healpy map2alm and alm2map. Added two beam operators, they support convolving all Stokes parameters and all frequencies with one b_ell function (harmonic representation of a symmetric beam); all Stokes parameters with different beams for different frequencies (both are BeamOperator); and each Stokes parameter, each frequency with its own beam in BeamOperatorIQU which takes beam_fl as StokesIQU. All new functions come with tests, and custom reduction rules.
In order to run requires: