reprojax 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- reprojax-0.2.0/LICENSE +63 -0
- reprojax-0.2.0/MANIFEST.in +8 -0
- reprojax-0.2.0/PKG-INFO +198 -0
- reprojax-0.2.0/README.md +179 -0
- reprojax-0.2.0/pyproject.toml +34 -0
- reprojax-0.2.0/setup.cfg +4 -0
- reprojax-0.2.0/src/reprojax/__init__.py +15 -0
- reprojax-0.2.0/src/reprojax/_astropy_wcs.py +74 -0
- reprojax-0.2.0/src/reprojax/adaptive.py +318 -0
- reprojax-0.2.0/src/reprojax/core.py +55 -0
- reprojax-0.2.0/src/reprojax/interp.py +278 -0
- reprojax-0.2.0/src/reprojax/py.typed +0 -0
- reprojax-0.2.0/src/reprojax/transform.py +231 -0
- reprojax-0.2.0/src/reprojax.egg-info/PKG-INFO +198 -0
- reprojax-0.2.0/src/reprojax.egg-info/SOURCES.txt +16 -0
- reprojax-0.2.0/src/reprojax.egg-info/dependency_links.txt +1 -0
- reprojax-0.2.0/src/reprojax.egg-info/requires.txt +8 -0
- reprojax-0.2.0/src/reprojax.egg-info/top_level.txt +1 -0
reprojax-0.2.0/LICENSE
ADDED
|
@@ -0,0 +1,63 @@
|
|
|
1
|
+
Copyright (c) 2026, Jae-Joon Lee
|
|
2
|
+
All rights reserved.
|
|
3
|
+
|
|
4
|
+
Redistribution and use in source and binary forms, with or without
|
|
5
|
+
modification, are permitted provided that the following conditions are met:
|
|
6
|
+
|
|
7
|
+
* Redistributions of source code must retain the above copyright notice, this
|
|
8
|
+
list of conditions and the following disclaimer.
|
|
9
|
+
|
|
10
|
+
* Redistributions in binary form must reproduce the above copyright notice,
|
|
11
|
+
this list of conditions and the following disclaimer in the documentation
|
|
12
|
+
and/or other materials provided with the distribution.
|
|
13
|
+
|
|
14
|
+
* Neither the name of the reprojax project nor the names of its contributors
|
|
15
|
+
may be used to endorse or promote products derived from this software
|
|
16
|
+
without specific prior written permission.
|
|
17
|
+
|
|
18
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
19
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
20
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
|
21
|
+
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
|
|
22
|
+
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
|
23
|
+
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
|
24
|
+
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
|
25
|
+
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
|
26
|
+
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
|
27
|
+
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
|
28
|
+
POSSIBILITY OF SUCH DAMAGE.
|
|
29
|
+
|
|
30
|
+
---
|
|
31
|
+
|
|
32
|
+
reprojax is a JAX port of astropy's `reproject` package
|
|
33
|
+
(https://github.com/astropy/reproject), which carries the following
|
|
34
|
+
license:
|
|
35
|
+
|
|
36
|
+
Copyright (c) 2011-16, Thomas P. Robitaille
|
|
37
|
+
All rights reserved.
|
|
38
|
+
|
|
39
|
+
Redistribution and use in source and binary forms, with or without
|
|
40
|
+
modification, are permitted provided that the following conditions are met:
|
|
41
|
+
|
|
42
|
+
* Redistributions of source code must retain the above copyright notice, this
|
|
43
|
+
list of conditions and the following disclaimer.
|
|
44
|
+
|
|
45
|
+
* Redistributions in binary form must reproduce the above copyright notice,
|
|
46
|
+
this list of conditions and the following disclaimer in the documentation
|
|
47
|
+
and/or other materials provided with the distribution.
|
|
48
|
+
|
|
49
|
+
* Neither the name of the Astropy project nor the names of its contributors
|
|
50
|
+
may be used to endorse or promote products derived from this software
|
|
51
|
+
without specific prior written permission.
|
|
52
|
+
|
|
53
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
54
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
55
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
|
56
|
+
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
|
|
57
|
+
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
|
58
|
+
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
|
59
|
+
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
|
60
|
+
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
|
61
|
+
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
|
62
|
+
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
|
63
|
+
POSSIBILITY OF SUCH DAMAGE.
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
# The sdist carries the package, not the test suite or the FITS fixtures:
|
|
2
|
+
# tests/ depends on tests/data/ (16 MB of real exposures), which no consumer
|
|
3
|
+
# needs, and shipping tests without their fixtures leaves a suite that cannot
|
|
4
|
+
# even be collected. Tests run from a checkout and from CI instead.
|
|
5
|
+
prune tests
|
|
6
|
+
prune benchmarks
|
|
7
|
+
prune examples
|
|
8
|
+
prune docs
|
reprojax-0.2.0/PKG-INFO
ADDED
|
@@ -0,0 +1,198 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: reprojax
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: JAX-native, differentiable reprojection of astronomical images (a port of astropy's reproject)
|
|
5
|
+
License-Expression: BSD-3-Clause
|
|
6
|
+
Classifier: Programming Language :: Python :: 3
|
|
7
|
+
Classifier: Operating System :: OS Independent
|
|
8
|
+
Requires-Python: >=3.11
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Dist: numpy>=2.4.6
|
|
12
|
+
Requires-Dist: jax>=0.10.2
|
|
13
|
+
Requires-Dist: astropy>=7.2.0
|
|
14
|
+
Requires-Dist: mapprojax<0.4,>=0.3.0
|
|
15
|
+
Provides-Extra: test
|
|
16
|
+
Requires-Dist: pytest; extra == "test"
|
|
17
|
+
Requires-Dist: reproject>=0.21; extra == "test"
|
|
18
|
+
Dynamic: license-file
|
|
19
|
+
|
|
20
|
+
# reprojax
|
|
21
|
+
|
|
22
|
+
A JAX-native, differentiable port of astropy's [`reproject`](https://github.com/astropy/reproject) package.
|
|
23
|
+
|
|
24
|
+
reprojax reprojects astronomical images between WCS frames on top of
|
|
25
|
+
[mapprojax](https://github.com/leejjoon/mapprojax), a JAX port of CDS's
|
|
26
|
+
`mapproj` crate providing fast, differentiable pixel↔sky transforms. The goal
|
|
27
|
+
is GPU-accelerated, batchable, gradient-friendly reprojection, validated
|
|
28
|
+
against the original `reproject` package as the accuracy reference.
|
|
29
|
+
|
|
30
|
+
reprojax is a JAX port of `reproject`'s algorithms and API shape; see
|
|
31
|
+
`LICENSE` for `reproject`'s BSD-3-Clause license, carried alongside
|
|
32
|
+
reprojax's own.
|
|
33
|
+
|
|
34
|
+
## Status
|
|
35
|
+
|
|
36
|
+
0.2.0, first public release. The validated scope is deliberate:
|
|
37
|
+
|
|
38
|
+
- **Projections**: TAN and SIN, with SIP distortion. Anything else raises
|
|
39
|
+
`NotImplementedError`.
|
|
40
|
+
- **Algorithms**: bilinear/nearest interpolation (`reproject_interp`) and the
|
|
41
|
+
DeForest (2004) adaptive anti-aliased algorithm (`reproject_adaptive`).
|
|
42
|
+
- **Reference**: `reproject` 0.21.0 is the golden accuracy reference. It is a
|
|
43
|
+
**test-time** dependency only; nothing in `reprojax` imports it at runtime.
|
|
44
|
+
- **Not included**: cubic interpolation, flux-conserving resampling, mosaic
|
|
45
|
+
coadds, filename/HDU/`NDData` input forms. See the
|
|
46
|
+
[issue tracker](https://github.com/leejjoon/reprojax/issues) for the plan.
|
|
47
|
+
|
|
48
|
+
## Precision: read this before your first call
|
|
49
|
+
|
|
50
|
+
`matrix_precision` selects the transform's coordinate precision, and the
|
|
51
|
+
**default is `"f64"`, which requires JAX x64**. Without it your first call
|
|
52
|
+
raises:
|
|
53
|
+
|
|
54
|
+
```
|
|
55
|
+
ValueError: precision='f64' requires float64. Call
|
|
56
|
+
jax.config.update('jax_enable_x64', True) before constructing it.
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
Two valid choices:
|
|
60
|
+
|
|
61
|
+
| Mode | Requirement | Use it for |
|
|
62
|
+
| --- | --- | --- |
|
|
63
|
+
| `"f64"` (default) | `jax.config.update("jax_enable_x64", True)` before importing/creating anything JAX | Full-precision work, astrometry, gradients with respect to WCS parameters |
|
|
64
|
+
| `"f32"` | none — works with x64 off | GPU throughput; the celestial rotation is still composed in float64 on the host, only the per-pixel path is float32 |
|
|
65
|
+
|
|
66
|
+
`"f64"` is the accurate default. `"f32"` is the fast path and is documented
|
|
67
|
+
in terms of its own measured caps, not as an equivalent to `"f64"` (below).
|
|
68
|
+
|
|
69
|
+
## Quick start
|
|
70
|
+
|
|
71
|
+
No data files needed. This runs exactly as written, from an installed package:
|
|
72
|
+
|
|
73
|
+
```python
|
|
74
|
+
import numpy as np
|
|
75
|
+
import jax
|
|
76
|
+
jax.config.update("jax_enable_x64", True) # required by the default matrix_precision="f64"
|
|
77
|
+
|
|
78
|
+
from astropy.io import fits
|
|
79
|
+
from astropy.wcs import WCS
|
|
80
|
+
from reprojax import reproject_interp
|
|
81
|
+
|
|
82
|
+
ny_in = nx_in = 64
|
|
83
|
+
hdr_in = fits.Header()
|
|
84
|
+
hdr_in["NAXIS"] = 2
|
|
85
|
+
hdr_in["NAXIS1"] = hdr_in["NAXIS2"] = ny_in
|
|
86
|
+
hdr_in["CTYPE1"], hdr_in["CTYPE2"] = "RA---TAN", "DEC--TAN"
|
|
87
|
+
hdr_in["CRPIX1"] = hdr_in["CRPIX2"] = ny_in / 2 + 0.5
|
|
88
|
+
hdr_in["CRVAL1"], hdr_in["CRVAL2"] = 210.0, 54.0
|
|
89
|
+
hdr_in["CDELT1"], hdr_in["CDELT2"] = -1e-3, 1e-3
|
|
90
|
+
|
|
91
|
+
y, x = np.mgrid[:ny_in, :nx_in]
|
|
92
|
+
data = np.exp(-(((x - 32) ** 2 + (y - 32) ** 2) / (2 * 6.0 ** 2))).astype(np.float32)
|
|
93
|
+
|
|
94
|
+
hdr_out = hdr_in.copy()
|
|
95
|
+
hdr_out["CRVAL1"] = 210.02 # shift the output grid
|
|
96
|
+
|
|
97
|
+
array, footprint = reproject_interp((data, WCS(hdr_in)), hdr_out, shape_out=(40, 40))
|
|
98
|
+
|
|
99
|
+
assert array.shape == footprint.shape == (40, 40)
|
|
100
|
+
assert np.isfinite(array).any()
|
|
101
|
+
print(float(np.nanmax(array)))
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
`array` holds the resampled image (NaN where no valid input pixel was
|
|
105
|
+
covered) and `footprint` is 1.0 where a valid input contributed, 0.0
|
|
106
|
+
otherwise. Adding `matrix_precision="f32"` runs the same call without x64.
|
|
107
|
+
|
|
108
|
+
## Architecture
|
|
109
|
+
|
|
110
|
+
- `reprojax.core` — the `PixelTransform` contract: a callable
|
|
111
|
+
`transform(x_out, y_out) -> (x_in, y_in)` mapping output-image pixel
|
|
112
|
+
coordinates to input-image pixel coordinates. This is the seam the rest of
|
|
113
|
+
the package is built around.
|
|
114
|
+
- `reprojax.transform` — builds a `PixelTransform` from a pair of astropy
|
|
115
|
+
WCS objects, using mapprojax for the projection math and astropy for
|
|
116
|
+
celestial frame rotation.
|
|
117
|
+
- Mapprojax-backed transforms also expose an optional
|
|
118
|
+
`transform.pixel_map(x_out, y_out, bounds_shape=(ny, nx))` capability. It
|
|
119
|
+
supplies finite coordinates and geometric validity to interpolation while
|
|
120
|
+
the two-result callable contract remains available for generic providers.
|
|
121
|
+
|
|
122
|
+
## Accuracy
|
|
123
|
+
|
|
124
|
+
Two kinds of numbers appear here, and they are labelled: an **observation**
|
|
125
|
+
is what the two implementations measured on a case, an **enforced tolerance**
|
|
126
|
+
is what the test suite requires.
|
|
127
|
+
|
|
128
|
+
Real SPHEREx pair (2040² ICRS TAN+SIP exposure onto the 1201² Galactic TAN
|
|
129
|
+
template, compared against `reproject` 0.21.0):
|
|
130
|
+
|
|
131
|
+
| Algorithm | Observation | Enforced |
|
|
132
|
+
| --- | --- | --- |
|
|
133
|
+
| `reproject_interp` | 1,345,444 of 1,345,460 jointly-valid pixels agree to `rtol <= 1e-8` | `rtol=1e-6, atol=0.01` over all jointly-valid pixels, plus exact footprint equality |
|
|
134
|
+
| `reproject_adaptive` | median disagreement ~`4e-8`; 1,334,193 of 1,334,200 jointly-valid pixels form the bulk | `rtol=1e-6, atol=0.005`, plus exact footprint equality |
|
|
135
|
+
|
|
136
|
+
In both algorithms the pixels outside the stated bulk are 16 (interp) and 7
|
|
137
|
+
(adaptive) pixels sitting at one real, un-flagged data spike: this fixture
|
|
138
|
+
drops the `FLAGS` extension that would normally mask it, so a sub-micropixel
|
|
139
|
+
position difference next to a ~10⁴-per-pixel local gradient is amplified into
|
|
140
|
+
a value difference of up to ~0.01 (interp) / ~0.003 (adaptive). The fixtures
|
|
141
|
+
are described in [`tests/data/README.md`](tests/data/README.md).
|
|
142
|
+
|
|
143
|
+
Synthetic TAN / TAN-SIP pairs: interpolated images agree within enforced
|
|
144
|
+
`rtol=1e-8, atol=1e-9`; adaptive within enforced `rtol=1e-6, atol=1e-8`.
|
|
145
|
+
|
|
146
|
+
`matrix_precision="f32"` has **enforced coordinate caps** of `6e-4` px
|
|
147
|
+
(forward) and `5e-4` px (reverse) per direction on the real pair. These bound
|
|
148
|
+
the coordinate map only — they are not image-agreement bounds, and they are
|
|
149
|
+
not a claim of float64 equivalence.
|
|
150
|
+
|
|
151
|
+
On a Strix Halo host (ROCm device, JAX 0.11.1 + rocm7 plugins 0.11.1),
|
|
152
|
+
batched and `jax.jit`-compiled reprojax reprojects 8 SPHEREx-scale
|
|
153
|
+
(2040² → 1201²) exposures in 0.026 s. The looped CPU `reproject_interp`
|
|
154
|
+
baseline for the same work measured 4.4–7.9 s across runs on this host, i.e.
|
|
155
|
+
roughly 170–300×. The reprojax figure is stable; the baseline is what varies.
|
|
156
|
+
This is a cross-backend batched-vs-looped comparison, not a like-for-like
|
|
157
|
+
backend comparison, and the CPU skips the batched GPU-only tests.
|
|
158
|
+
|
|
159
|
+
## Development
|
|
160
|
+
|
|
161
|
+
```bash
|
|
162
|
+
python -m pip install -e ".[test]" -c ci/latest.txt
|
|
163
|
+
python -m pytest tests/
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
Using a `ci/` constraints file matters: the golden reference is pinned to
|
|
167
|
+
`reproject==0.21.0` there, and the suite asserts the installed version. Use
|
|
168
|
+
`ci/minimum.txt` for the oldest supported stack (Python 3.11). If a
|
|
169
|
+
constraint conflicts with your environment, install `.[test]` without it and
|
|
170
|
+
expect the version guard to tell you which pin is off.
|
|
171
|
+
|
|
172
|
+
The test suite needs `tests/data/`, which is 16 MB of real SPHEREx
|
|
173
|
+
exposures and is therefore not shipped in the sdist. Install from a checkout
|
|
174
|
+
when you want to run the tests.
|
|
175
|
+
|
|
176
|
+
## License
|
|
177
|
+
|
|
178
|
+
BSD-3-Clause. The SIP and reprojection algorithms are ported from
|
|
179
|
+
[`reproject`](https://github.com/astropy/reproject); its license is carried in
|
|
180
|
+
[`LICENSE`](LICENSE) alongside reprojax's own.
|
|
181
|
+
|
|
182
|
+
## Test data
|
|
183
|
+
|
|
184
|
+
The integration tests use a real SPHEREx exposure and a Galactic template
|
|
185
|
+
header. Only the exposure is mission data: NASA's Science Mission
|
|
186
|
+
Directorate releases mission data into the public domain (SPD-41), while the
|
|
187
|
+
template header is project metadata under reprojax's own licence. If you
|
|
188
|
+
publish work that uses the exposure, include the acknowledgement IRSA
|
|
189
|
+
requests:
|
|
190
|
+
|
|
191
|
+
> This publication makes use of data products from the Spectro-Photometer for
|
|
192
|
+
the History of the Universe, Epoch of Reionization and Ices Explorer
|
|
193
|
+
> (SPHEREx), which is a joint project of the Jet Propulsion Laboratory and the
|
|
194
|
+
> California Institute of Technology, and is funded by the National
|
|
195
|
+
> Aeronautics and Space Administration.
|
|
196
|
+
|
|
197
|
+
Provenance, checksums and the citation DOIs are in
|
|
198
|
+
[`tests/data/README.md`](tests/data/README.md).
|
reprojax-0.2.0/README.md
ADDED
|
@@ -0,0 +1,179 @@
|
|
|
1
|
+
# reprojax
|
|
2
|
+
|
|
3
|
+
A JAX-native, differentiable port of astropy's [`reproject`](https://github.com/astropy/reproject) package.
|
|
4
|
+
|
|
5
|
+
reprojax reprojects astronomical images between WCS frames on top of
|
|
6
|
+
[mapprojax](https://github.com/leejjoon/mapprojax), a JAX port of CDS's
|
|
7
|
+
`mapproj` crate providing fast, differentiable pixel↔sky transforms. The goal
|
|
8
|
+
is GPU-accelerated, batchable, gradient-friendly reprojection, validated
|
|
9
|
+
against the original `reproject` package as the accuracy reference.
|
|
10
|
+
|
|
11
|
+
reprojax is a JAX port of `reproject`'s algorithms and API shape; see
|
|
12
|
+
`LICENSE` for `reproject`'s BSD-3-Clause license, carried alongside
|
|
13
|
+
reprojax's own.
|
|
14
|
+
|
|
15
|
+
## Status
|
|
16
|
+
|
|
17
|
+
0.2.0, first public release. The validated scope is deliberate:
|
|
18
|
+
|
|
19
|
+
- **Projections**: TAN and SIN, with SIP distortion. Anything else raises
|
|
20
|
+
`NotImplementedError`.
|
|
21
|
+
- **Algorithms**: bilinear/nearest interpolation (`reproject_interp`) and the
|
|
22
|
+
DeForest (2004) adaptive anti-aliased algorithm (`reproject_adaptive`).
|
|
23
|
+
- **Reference**: `reproject` 0.21.0 is the golden accuracy reference. It is a
|
|
24
|
+
**test-time** dependency only; nothing in `reprojax` imports it at runtime.
|
|
25
|
+
- **Not included**: cubic interpolation, flux-conserving resampling, mosaic
|
|
26
|
+
coadds, filename/HDU/`NDData` input forms. See the
|
|
27
|
+
[issue tracker](https://github.com/leejjoon/reprojax/issues) for the plan.
|
|
28
|
+
|
|
29
|
+
## Precision: read this before your first call
|
|
30
|
+
|
|
31
|
+
`matrix_precision` selects the transform's coordinate precision, and the
|
|
32
|
+
**default is `"f64"`, which requires JAX x64**. Without it your first call
|
|
33
|
+
raises:
|
|
34
|
+
|
|
35
|
+
```
|
|
36
|
+
ValueError: precision='f64' requires float64. Call
|
|
37
|
+
jax.config.update('jax_enable_x64', True) before constructing it.
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
Two valid choices:
|
|
41
|
+
|
|
42
|
+
| Mode | Requirement | Use it for |
|
|
43
|
+
| --- | --- | --- |
|
|
44
|
+
| `"f64"` (default) | `jax.config.update("jax_enable_x64", True)` before importing/creating anything JAX | Full-precision work, astrometry, gradients with respect to WCS parameters |
|
|
45
|
+
| `"f32"` | none — works with x64 off | GPU throughput; the celestial rotation is still composed in float64 on the host, only the per-pixel path is float32 |
|
|
46
|
+
|
|
47
|
+
`"f64"` is the accurate default. `"f32"` is the fast path and is documented
|
|
48
|
+
in terms of its own measured caps, not as an equivalent to `"f64"` (below).
|
|
49
|
+
|
|
50
|
+
## Quick start
|
|
51
|
+
|
|
52
|
+
No data files needed. This runs exactly as written, from an installed package:
|
|
53
|
+
|
|
54
|
+
```python
|
|
55
|
+
import numpy as np
|
|
56
|
+
import jax
|
|
57
|
+
jax.config.update("jax_enable_x64", True) # required by the default matrix_precision="f64"
|
|
58
|
+
|
|
59
|
+
from astropy.io import fits
|
|
60
|
+
from astropy.wcs import WCS
|
|
61
|
+
from reprojax import reproject_interp
|
|
62
|
+
|
|
63
|
+
ny_in = nx_in = 64
|
|
64
|
+
hdr_in = fits.Header()
|
|
65
|
+
hdr_in["NAXIS"] = 2
|
|
66
|
+
hdr_in["NAXIS1"] = hdr_in["NAXIS2"] = ny_in
|
|
67
|
+
hdr_in["CTYPE1"], hdr_in["CTYPE2"] = "RA---TAN", "DEC--TAN"
|
|
68
|
+
hdr_in["CRPIX1"] = hdr_in["CRPIX2"] = ny_in / 2 + 0.5
|
|
69
|
+
hdr_in["CRVAL1"], hdr_in["CRVAL2"] = 210.0, 54.0
|
|
70
|
+
hdr_in["CDELT1"], hdr_in["CDELT2"] = -1e-3, 1e-3
|
|
71
|
+
|
|
72
|
+
y, x = np.mgrid[:ny_in, :nx_in]
|
|
73
|
+
data = np.exp(-(((x - 32) ** 2 + (y - 32) ** 2) / (2 * 6.0 ** 2))).astype(np.float32)
|
|
74
|
+
|
|
75
|
+
hdr_out = hdr_in.copy()
|
|
76
|
+
hdr_out["CRVAL1"] = 210.02 # shift the output grid
|
|
77
|
+
|
|
78
|
+
array, footprint = reproject_interp((data, WCS(hdr_in)), hdr_out, shape_out=(40, 40))
|
|
79
|
+
|
|
80
|
+
assert array.shape == footprint.shape == (40, 40)
|
|
81
|
+
assert np.isfinite(array).any()
|
|
82
|
+
print(float(np.nanmax(array)))
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
`array` holds the resampled image (NaN where no valid input pixel was
|
|
86
|
+
covered) and `footprint` is 1.0 where a valid input contributed, 0.0
|
|
87
|
+
otherwise. Adding `matrix_precision="f32"` runs the same call without x64.
|
|
88
|
+
|
|
89
|
+
## Architecture
|
|
90
|
+
|
|
91
|
+
- `reprojax.core` — the `PixelTransform` contract: a callable
|
|
92
|
+
`transform(x_out, y_out) -> (x_in, y_in)` mapping output-image pixel
|
|
93
|
+
coordinates to input-image pixel coordinates. This is the seam the rest of
|
|
94
|
+
the package is built around.
|
|
95
|
+
- `reprojax.transform` — builds a `PixelTransform` from a pair of astropy
|
|
96
|
+
WCS objects, using mapprojax for the projection math and astropy for
|
|
97
|
+
celestial frame rotation.
|
|
98
|
+
- Mapprojax-backed transforms also expose an optional
|
|
99
|
+
`transform.pixel_map(x_out, y_out, bounds_shape=(ny, nx))` capability. It
|
|
100
|
+
supplies finite coordinates and geometric validity to interpolation while
|
|
101
|
+
the two-result callable contract remains available for generic providers.
|
|
102
|
+
|
|
103
|
+
## Accuracy
|
|
104
|
+
|
|
105
|
+
Two kinds of numbers appear here, and they are labelled: an **observation**
|
|
106
|
+
is what the two implementations measured on a case, an **enforced tolerance**
|
|
107
|
+
is what the test suite requires.
|
|
108
|
+
|
|
109
|
+
Real SPHEREx pair (2040² ICRS TAN+SIP exposure onto the 1201² Galactic TAN
|
|
110
|
+
template, compared against `reproject` 0.21.0):
|
|
111
|
+
|
|
112
|
+
| Algorithm | Observation | Enforced |
|
|
113
|
+
| --- | --- | --- |
|
|
114
|
+
| `reproject_interp` | 1,345,444 of 1,345,460 jointly-valid pixels agree to `rtol <= 1e-8` | `rtol=1e-6, atol=0.01` over all jointly-valid pixels, plus exact footprint equality |
|
|
115
|
+
| `reproject_adaptive` | median disagreement ~`4e-8`; 1,334,193 of 1,334,200 jointly-valid pixels form the bulk | `rtol=1e-6, atol=0.005`, plus exact footprint equality |
|
|
116
|
+
|
|
117
|
+
In both algorithms the pixels outside the stated bulk are 16 (interp) and 7
|
|
118
|
+
(adaptive) pixels sitting at one real, un-flagged data spike: this fixture
|
|
119
|
+
drops the `FLAGS` extension that would normally mask it, so a sub-micropixel
|
|
120
|
+
position difference next to a ~10⁴-per-pixel local gradient is amplified into
|
|
121
|
+
a value difference of up to ~0.01 (interp) / ~0.003 (adaptive). The fixtures
|
|
122
|
+
are described in [`tests/data/README.md`](tests/data/README.md).
|
|
123
|
+
|
|
124
|
+
Synthetic TAN / TAN-SIP pairs: interpolated images agree within enforced
|
|
125
|
+
`rtol=1e-8, atol=1e-9`; adaptive within enforced `rtol=1e-6, atol=1e-8`.
|
|
126
|
+
|
|
127
|
+
`matrix_precision="f32"` has **enforced coordinate caps** of `6e-4` px
|
|
128
|
+
(forward) and `5e-4` px (reverse) per direction on the real pair. These bound
|
|
129
|
+
the coordinate map only — they are not image-agreement bounds, and they are
|
|
130
|
+
not a claim of float64 equivalence.
|
|
131
|
+
|
|
132
|
+
On a Strix Halo host (ROCm device, JAX 0.11.1 + rocm7 plugins 0.11.1),
|
|
133
|
+
batched and `jax.jit`-compiled reprojax reprojects 8 SPHEREx-scale
|
|
134
|
+
(2040² → 1201²) exposures in 0.026 s. The looped CPU `reproject_interp`
|
|
135
|
+
baseline for the same work measured 4.4–7.9 s across runs on this host, i.e.
|
|
136
|
+
roughly 170–300×. The reprojax figure is stable; the baseline is what varies.
|
|
137
|
+
This is a cross-backend batched-vs-looped comparison, not a like-for-like
|
|
138
|
+
backend comparison, and the CPU skips the batched GPU-only tests.
|
|
139
|
+
|
|
140
|
+
## Development
|
|
141
|
+
|
|
142
|
+
```bash
|
|
143
|
+
python -m pip install -e ".[test]" -c ci/latest.txt
|
|
144
|
+
python -m pytest tests/
|
|
145
|
+
```
|
|
146
|
+
|
|
147
|
+
Using a `ci/` constraints file matters: the golden reference is pinned to
|
|
148
|
+
`reproject==0.21.0` there, and the suite asserts the installed version. Use
|
|
149
|
+
`ci/minimum.txt` for the oldest supported stack (Python 3.11). If a
|
|
150
|
+
constraint conflicts with your environment, install `.[test]` without it and
|
|
151
|
+
expect the version guard to tell you which pin is off.
|
|
152
|
+
|
|
153
|
+
The test suite needs `tests/data/`, which is 16 MB of real SPHEREx
|
|
154
|
+
exposures and is therefore not shipped in the sdist. Install from a checkout
|
|
155
|
+
when you want to run the tests.
|
|
156
|
+
|
|
157
|
+
## License
|
|
158
|
+
|
|
159
|
+
BSD-3-Clause. The SIP and reprojection algorithms are ported from
|
|
160
|
+
[`reproject`](https://github.com/astropy/reproject); its license is carried in
|
|
161
|
+
[`LICENSE`](LICENSE) alongside reprojax's own.
|
|
162
|
+
|
|
163
|
+
## Test data
|
|
164
|
+
|
|
165
|
+
The integration tests use a real SPHEREx exposure and a Galactic template
|
|
166
|
+
header. Only the exposure is mission data: NASA's Science Mission
|
|
167
|
+
Directorate releases mission data into the public domain (SPD-41), while the
|
|
168
|
+
template header is project metadata under reprojax's own licence. If you
|
|
169
|
+
publish work that uses the exposure, include the acknowledgement IRSA
|
|
170
|
+
requests:
|
|
171
|
+
|
|
172
|
+
> This publication makes use of data products from the Spectro-Photometer for
|
|
173
|
+
the History of the Universe, Epoch of Reionization and Ices Explorer
|
|
174
|
+
> (SPHEREx), which is a joint project of the Jet Propulsion Laboratory and the
|
|
175
|
+
> California Institute of Technology, and is funded by the National
|
|
176
|
+
> Aeronautics and Space Administration.
|
|
177
|
+
|
|
178
|
+
Provenance, checksums and the citation DOIs are in
|
|
179
|
+
[`tests/data/README.md`](tests/data/README.md).
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "reprojax"
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "JAX-native, differentiable reprojection of astronomical images (a port of astropy's reproject)"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = "BSD-3-Clause"
|
|
11
|
+
license-files = ["LICENSE"]
|
|
12
|
+
requires-python = ">=3.11"
|
|
13
|
+
dependencies = [
|
|
14
|
+
"numpy>=2.4.6",
|
|
15
|
+
"jax>=0.10.2",
|
|
16
|
+
"astropy>=7.2.0",
|
|
17
|
+
"mapprojax>=0.3.0,<0.4",
|
|
18
|
+
]
|
|
19
|
+
classifiers = [
|
|
20
|
+
"Programming Language :: Python :: 3",
|
|
21
|
+
"Operating System :: OS Independent",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.optional-dependencies]
|
|
25
|
+
test = ["pytest", "reproject>=0.21"]
|
|
26
|
+
|
|
27
|
+
[tool.setuptools.package-data]
|
|
28
|
+
reprojax = ["py.typed"]
|
|
29
|
+
|
|
30
|
+
[tool.setuptools.packages.find]
|
|
31
|
+
where = ["src"]
|
|
32
|
+
|
|
33
|
+
[tool.pytest.ini_options]
|
|
34
|
+
testpaths = ["tests"]
|
reprojax-0.2.0/setup.cfg
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
from .adaptive import reproject_adaptive, reproject_adaptive_core
|
|
2
|
+
from .core import PixelTransform
|
|
3
|
+
from .interp import reproject_interp, reproject_interp_batched, reproject_interp_core
|
|
4
|
+
from .transform import build_transform, compose_transform
|
|
5
|
+
|
|
6
|
+
__all__ = [
|
|
7
|
+
"PixelTransform",
|
|
8
|
+
"build_transform",
|
|
9
|
+
"compose_transform",
|
|
10
|
+
"reproject_adaptive",
|
|
11
|
+
"reproject_adaptive_core",
|
|
12
|
+
"reproject_interp",
|
|
13
|
+
"reproject_interp_batched",
|
|
14
|
+
"reproject_interp_core",
|
|
15
|
+
]
|
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
"""Host-side extraction of mapprojax WCS parameters and frame rotations from astropy.
|
|
2
|
+
|
|
3
|
+
Everything in this module runs eagerly in NumPy/astropy, at transform-build
|
|
4
|
+
time — never inside the jit-compiled :class:`~reprojax.core.PixelTransform`
|
|
5
|
+
itself.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import numpy as np
|
|
9
|
+
from astropy.coordinates import SkyCoord
|
|
10
|
+
from astropy.wcs.utils import wcs_to_celestial_frame
|
|
11
|
+
from mapprojax.sip import Sip
|
|
12
|
+
|
|
13
|
+
#: FITS projection code -> mapprojax provider key (see transform.py).
|
|
14
|
+
SUPPORTED_PROJECTIONS = ("TAN", "SIN")
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def projection_code(wcs):
|
|
18
|
+
"""The 3-letter FITS projection code for an astropy WCS's first axis.
|
|
19
|
+
|
|
20
|
+
Per the WCS papers, CTYPEia's first 8 characters are always
|
|
21
|
+
``"<axis>-<code>"`` with the code in the last 3 characters of that
|
|
22
|
+
8-character field — true whether or not a further distortion suffix
|
|
23
|
+
(``-SIP``, ``-TPV``, ...) follows.
|
|
24
|
+
"""
|
|
25
|
+
ctype = wcs.wcs.ctype[0]
|
|
26
|
+
code = ctype[:8][-3:].strip()
|
|
27
|
+
if code not in SUPPORTED_PROJECTIONS:
|
|
28
|
+
raise NotImplementedError(
|
|
29
|
+
f"Unsupported projection {code!r} (from CTYPE {ctype!r}); "
|
|
30
|
+
f"reprojax currently supports: {SUPPORTED_PROJECTIONS}"
|
|
31
|
+
)
|
|
32
|
+
return code
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def wcs_params_from_astropy(wcs):
|
|
36
|
+
"""Extract ``(crpix, cd, crval, sip)`` from an astropy WCS.
|
|
37
|
+
|
|
38
|
+
``crpix`` is converted to mapprojax's 0-based convention (astropy's
|
|
39
|
+
``wcs.wcs.crpix`` is 1-based, matching the FITS ``CRPIXi`` keywords).
|
|
40
|
+
``cd`` is the CD-matrix-equivalent linear transform regardless of
|
|
41
|
+
whether the header used ``CDi_j`` or ``PCi_j``/``CDELTi`` keywords.
|
|
42
|
+
"""
|
|
43
|
+
crpix = (np.asarray(wcs.wcs.crpix, dtype=float) - 1.0).tolist()
|
|
44
|
+
crval = np.asarray(wcs.wcs.crval, dtype=float).tolist()
|
|
45
|
+
cd = np.asarray(wcs.pixel_scale_matrix, dtype=float).tolist()
|
|
46
|
+
|
|
47
|
+
sip = None
|
|
48
|
+
if wcs.sip is not None:
|
|
49
|
+
sip = Sip(a=wcs.sip.a, b=wcs.sip.b, ap=wcs.sip.ap, bp=wcs.sip.bp)
|
|
50
|
+
|
|
51
|
+
return crpix, cd, crval, sip
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def frame_rotation_matrix(wcs_from, wcs_to):
|
|
55
|
+
"""3x3 matrix rotating Cartesian unit vectors from ``wcs_from``'s celestial
|
|
56
|
+
frame to ``wcs_to``'s.
|
|
57
|
+
|
|
58
|
+
Returns the identity when the two WCSes share the same celestial frame
|
|
59
|
+
(the common case, and the only case that avoids an astropy frame
|
|
60
|
+
transform).
|
|
61
|
+
"""
|
|
62
|
+
frame_from = wcs_to_celestial_frame(wcs_from)
|
|
63
|
+
frame_to = wcs_to_celestial_frame(wcs_to)
|
|
64
|
+
|
|
65
|
+
if frame_from.name == frame_to.name:
|
|
66
|
+
return np.eye(3)
|
|
67
|
+
|
|
68
|
+
basis_lon = [0.0, 90.0, 0.0]
|
|
69
|
+
basis_lat = [0.0, 0.0, 90.0]
|
|
70
|
+
basis = SkyCoord(basis_lon, basis_lat, unit="deg", frame=frame_from)
|
|
71
|
+
cart = basis.transform_to(frame_to).cartesian
|
|
72
|
+
|
|
73
|
+
# Transformed basis vectors as columns: matrix @ v_from == v_to.
|
|
74
|
+
return np.stack([cart.x.value, cart.y.value, cart.z.value], axis=0)
|