Use reproject for image registration #317
Conversation
|
Let me know if you are interested in trying to optimize performance here - there's been a lot of dask-related work in the last couple of years in reproject, and it's now easy to do the reprojection with multi-threading for example. |
nabobalis
force-pushed
the
reproject-aiaprep
branch
from
f91380f to
5ab4ead
Compare
nabobalis
force-pushed
the
reproject-aiaprep
branch
from
5ab4ead to
8c21593
Compare
What is the best way to make it "faster", we have a 4k by 4k image and at least the way the code is written, we do one file at a time. Is providing |
Providing |
I got as the fastest on my local PC. Since its a sunpy map, its a FITS file read from astropy.io.fits but map does support dask arrays. This function is specifically for 4k by 4k or 5k by 5k 2D images only if that makes much of a difference? |
|
Worried about optimizing for my own computer but yolo. |
|
Could you show me a simple example of reprojection of this data I could run locally? (just in case I have suggestions for speeding up further). I'm always interested in different real-life cases. |
|
If you are ok with running the aiapy from this branch: import sunpy.map as sm
import aiapy.calibrate
m_l1 = sm.Map("~/Downloads/aia.lev1_euv_12s.2026-01-19T235959Z.211.image_lev1.fits")
m_l15 = aiapy.calibrate.register(m_l1)Is how I testing this out. The WCS its being reprojected too is basically one which is slightly shifted, to account for alignment issues with the telescope. It is not too dissimilar to the idea in this example from the sunpy gallery: https://docs.sunpy.org/en/stable/generated/gallery/map_transformations/reprojection_align_aia_hmi.html |
nabobalis
force-pushed
the
reproject-aiaprep
branch
3 times, most recently
from
d7e1220 to
e41759f
Compare
nabobalis
force-pushed
the
reproject-aiaprep
branch
2 times, most recently
from
039a00a to
1e357a5
Compare
nabobalis
force-pushed
the
reproject-aiaprep
branch
from
1e357a5 to
b4e24af
Compare
| newmap.meta["lvl_num"] = 1.5 | ||
| newmap.meta["bitpix"] = -64 | ||
| return newmap | ||
| wcs_l15_full_disk = astropy.wcs.WCS(header_l15_full_disk, preserve_units=True) |
There was a problem hiding this comment.
So before the PR was doing:
# Find the bottom left corner of the map in the full-frame WCS
blc_full_disk = pixel_to_pixel(smap.wcs, wcs_l15_full_disk, 0, 0) * u.pix
# Calculate distance between full-frame center and bottom left corner
# This is the location of disk center in the aligned WCS of this map
ref_pixel = ref_pixel_full_disk - blc_full_disk
# Construct the L1.5 WCS for this map and reproject
wcs_l15 = astropy.wcs.WCS(
make_fitswcs_header(
smap.data.shape,
ref_coord,
reference_pixel=ref_pixel,
scale=scale,
rotation_matrix=np.eye(2),
)
)which I have removed but as this code stands, I don't think its doing anything useful without this.
There was a problem hiding this comment.
The reason I originally implemented it this way is so this works on submaps. As it stands, your implementation will reproject a submap into a 4K by 4K image always. There are two options:
- We either need to check if the input map is a submap and raise an exception if it is
- restore this codepath so that it works on submaps too.
One of the main reasons I started this PR was to avoid option 1 so I'd suggest 2. If we go with that, we should double check that is the most optimal solution.
There was a problem hiding this comment.
If you need help figuring out submaps feel free to to show me a minimal example
There was a problem hiding this comment.
I have restored it back to how it was but none of the unit tests now pass.
|
@nabobalis - thanks for the MWE! With this example I get: Just to check, what kind of times are you aiming for ideally? |
|
There is no time in mind, it's just that we want to avoid it being noticeable slower than the original version which uses affine_transfrom from scipy. The reason for this is that someone might want to call this code on say 1000 images to construct a new cube of data (used for a DEM analysis for example). At least in my local testing on my PC, it's only 2x slower which for me is good enough. I know that someone will complain to me about it when we release it but yolo. For comparison, |
wtbarnes
left a comment
wtbarnes
left a comment
There was a problem hiding this comment.
I can't request changes on my own PR, but consider this review blocking. The biggest issue I see is that as now implemented, this does not work on submaps. I think we also need to be sure we aren't making this dramatically slower.
| __all__ = ["correct_degradation", "degradation", "register"] | ||
|
|
||
|
|
||
| @add_common_docstring(rotation_function_names=_rotation_function_names) |
There was a problem hiding this comment.
This is no longer relevant
| newmap.meta["lvl_num"] = 1.5 | ||
| newmap.meta["bitpix"] = -64 | ||
| return newmap | ||
| wcs_l15_full_disk = astropy.wcs.WCS(header_l15_full_disk, preserve_units=True) |
There was a problem hiding this comment.
The reason I originally implemented it this way is so this works on submaps. As it stands, your implementation will reproject a submap into a 4K by 4K image always. There are two options:
- We either need to check if the input map is a submap and raise an exception if it is
- restore this codepath so that it works on submaps too.
One of the main reasons I started this PR was to avoid option 1 so I'd suggest 2. If we go with that, we should double check that is the most optimal solution.
nabobalis
force-pushed
the
reproject-aiaprep
branch
from
53213dc to
42fe50a
Compare
nabobalis
force-pushed
the
reproject-aiaprep
branch
from
42fe50a to
21b8d38
Compare
nabobalis
force-pushed
the
reproject-aiaprep
branch
from
b71c1c9 to
d957473
Compare
Reviewer's GuideReimplements AIA image registration to use WCS-based reprojection via sunpy/reproject, enabling registration of submaps, centralizing detector dimensions, and cleaning up PSF/spike utilities, tests, and documentation to match the new behavior and dependencies. Sequence diagram for the new WCS-based register image registrationsequenceDiagram
actor User
participant ClientCode
participant AIA_Map as AIAMap_or_HMIMap
participant Register as register
participant Utils as detector_dimensions
participant WCS as Astropy_WCS
participant Reproject as reproject_to
User->>ClientCode: call register(smap, missing, algorithm, kwargs)
ClientCode->>Register: register(smap, missing, algorithm, kwargs)
Register->>Register: check smap.processing_level
Register-->>Register: maybe warn AIApyUserWarning
Register->>Utils: detector_dimensions()
Utils-->>Register: shape_full_disk
Register->>WCS: make_fitswcs_header(shape_full_disk, ref_coord=Sun_center, reference_pixel=center_full_frame, scale=0.6arcsec, rotation_matrix=I)
WCS-->>Register: WCS_L15_full_disk
Register->>WCS: pixel_to_pixel(smap.wcs, WCS_L15_full_disk, 0, 0)
WCS-->>Register: blc_full_disk
Register->>WCS: make_fitswcs_header(smap.shape, ref_coord=Sun_center, reference_pixel=ref_pixel, scale=0.6arcsec, rotation_matrix=I)
WCS-->>Register: WCS_L15
Register->>Reproject: smap.reproject_to(WCS_L15, algorithm, return_footprint=False, parallel=True, block_size=(1024,1024), kwargs)
Reproject-->>Register: smap_L15
Register->>Register: fill NaN in smap_L15.data with missing
Register->>Register: new_meta = deepcopy(smap.meta)
Register->>Register: remove CROTA2 from new_meta
Register->>Register: new_meta.update(smap_L15.meta)
Register->>Register: set new_meta[lvl_num] = 1.5
Register->>AIA_Map: smap._new_instance(data, new_meta, plot_settings=smap.plot_settings)
AIA_Map-->>Register: registered_map
Register-->>ClientCode: registered_map
ClientCode-->>User: return registered_map
Class diagram for shared detector_dimensions and registration-related utilitiesclassDiagram
class aiapy_utils {
}
class detector_dimensions {
+Quantity detector_dimensions()
}
class register {
+register(smap, missing, algorithm, kwargs)
}
class fetch_spikes {
+fetch_spikes(smap, as_coords)
}
class _psf {
+_psf(meshinfo, angles, diffraction_orders, focal_plane)
}
class AIAMap_or_HMIMap {
+wcs
+meta
+data
+coordinate_frame
+plot_settings
+processing_level
+reproject_to(target_wcs, algorithm, return_footprint, parallel, block_size, kwargs)
+_new_instance(data, meta, plot_settings)
}
aiapy_utils <|-- detector_dimensions
aiapy_utils <|-- register
aiapy_utils <|-- fetch_spikes
aiapy_utils <|-- _psf
register --> detector_dimensions : uses
fetch_spikes --> detector_dimensions : uses
_psf --> detector_dimensions : uses
register --> AIAMap_or_HMIMap : reprojects
fetch_spikes --> AIAMap_or_HMIMap : analyzes
File-Level Changes
Assessment against linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- The new
registerimplementation no longer setsr_sunandbitpixin the output metadata (unlike the old version), but the tests now assume these keys are present; consider explicitly restoring them alongsidelvl_numto preserve the previous behavior. - Resetting the
reproject.commonlogger level toWARNINGinsideregisterintroduces a global side effect on logging; it would be safer to either respect any existing log level or make this behavior optional/configurable instead of hard-coding it. - The previous
registerimplementation enforced AIAMap/HMIMap types and full-disk inputs, while the new one accepts arbitrary maps and submaps; if this widened scope is intentional, it might be worth double-checking that all assumptions in the WCS construction hold for non-SDO maps, or reintroducing some lightweight validation.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The new `register` implementation no longer sets `r_sun` and `bitpix` in the output metadata (unlike the old version), but the tests now assume these keys are present; consider explicitly restoring them alongside `lvl_num` to preserve the previous behavior.
- Resetting the `reproject.common` logger level to `WARNING` inside `register` introduces a global side effect on logging; it would be safer to either respect any existing log level or make this behavior optional/configurable instead of hard-coding it.
- The previous `register` implementation enforced AIAMap/HMIMap types and full-disk inputs, while the new one accepts arbitrary maps and submaps; if this widened scope is intentional, it might be worth double-checking that all assumptions in the WCS construction hold for non-SDO maps, or reintroducing some lightweight validation.
## Individual Comments
### Comment 1
<location> `aiapy/calibrate/prep.py:98-99` </location>
<code_context>
- return newmap
+ kwargs["return_footprint"] = kwargs.get("return_footprint", False)
+ # This was selected as the fastest method in local testing
+ kwargs["parallel"] = kwargs.get("parallel", True)
+ kwargs["block_size"] = kwargs.get("block_size", (1024, 1024))
+ # Suppress the reproject INFO logging
+ # TODO: Make this configurable by the user?
</code_context>
<issue_to_address>
**issue (bug_risk):** Unconditionally injecting `parallel` and `block_size` into `reproject_to` kwargs may break algorithms that don’t accept them.
These options are currently added for all algorithms. Some `reproject` backends don’t accept `parallel`/`block_size` and may raise `TypeError`. Consider only passing them for algorithms that support them (e.g. `algorithm == "interpolation"`), or making them opt-in via a dedicated flag so callers can disable them when incompatible.
</issue_to_address>
### Comment 2
<location> `aiapy/calibrate/prep.py:107` </location>
<code_context>
+ smap_l15 = smap.reproject_to(wcs_l15, algorithm=algorithm, **kwargs)
+ # Fill in missing values
+ data = smap_l15.data
+ missing = np.nanmin(smap_l15.data) if missing is None else missing
+ data[np.isnan(data)] = missing
+ # Restore metadata (reproject_to only carries over the WCS keywords)
</code_context>
<issue_to_address>
**issue (bug_risk):** Using `np.nanmin` will fail when the reprojected data are all NaN.
If `smap_l15.data` is entirely NaN (e.g. off-disk cutout, bad WCS, or very tight cropping), `np.nanmin` will raise a `ValueError`. Consider handling the all-NaN case first (e.g. fallback to the original map’s minimum or a configurable default) before calling `nanmin`.
</issue_to_address>
### Comment 3
<location> `aiapy/calibrate/tests/test_prep.py:39-48` </location>
<code_context>
+def test_register_submap(aia_171_submap) -> None:
</code_context>
<issue_to_address>
**suggestion (testing):** Add explicit tests for `missing` handling and NaN filling behaviour in `register`.
The updated `register` now fills NaNs and accepts a `missing` argument, but this isn’t covered by tests. Please add tests that:
* Build a map containing NaNs and call `register` without `missing`; assert (a) the result has no NaNs and (b) the replacement value matches `np.nanmin` of the reprojected data.
* Call `register(..., missing=<sentinel>)` and assert that NaNs in the result are replaced by that sentinel.
This will protect the NaN-handling logic and the `missing` parameter contract from regressions.
</issue_to_address>
### Comment 4
<location> `changelog/317.breaking.1.rst:1` </location>
<code_context>
+Removed the need for cupy from aiapy.
</code_context>
<issue_to_address>
**issue (typo):** Capitalize "CuPy" to match the official project name.
This keeps the changelog aligned with the library’s official naming and branding.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| kwargs["parallel"] = kwargs.get("parallel", True) | ||
| kwargs["block_size"] = kwargs.get("block_size", (1024, 1024)) |
There was a problem hiding this comment.
issue (bug_risk): Unconditionally injecting parallel and block_size into reproject_to kwargs may break algorithms that don’t accept them.
These options are currently added for all algorithms. Some reproject backends don’t accept parallel/block_size and may raise TypeError. Consider only passing them for algorithms that support them (e.g. algorithm == "interpolation"), or making them opt-in via a dedicated flag so callers can disable them when incompatible.
| smap_l15 = smap.reproject_to(wcs_l15, algorithm=algorithm, **kwargs) | ||
| # Fill in missing values | ||
| data = smap_l15.data | ||
| missing = np.nanmin(smap_l15.data) if missing is None else missing |
There was a problem hiding this comment.
issue (bug_risk): Using np.nanmin will fail when the reprojected data are all NaN.
If smap_l15.data is entirely NaN (e.g. off-disk cutout, bad WCS, or very tight cropping), np.nanmin will raise a ValueError. Consider handling the all-NaN case first (e.g. fallback to the original map’s minimum or a configurable default) before calling nanmin.
| def test_register_submap(aia_171_submap) -> None: | ||
| """ | ||
| Test that we can register a submap and that the output has the same shape | ||
| as the input submap and is correctly aligned. | ||
| """ | ||
| lvl_15_submap = register(aia_171_submap) | ||
| np.testing.assert_array_equal(lvl_15_submap.data.shape, aia_171_submap.data.shape) | ||
| assert lvl_15_submap.meta["crpix1"] == lvl_15_submap.data.shape[0] / 2 + 0.5 | ||
| assert lvl_15_submap.meta["crpix2"] == lvl_15_submap.data.shape[1] / 2 + 0.5 | ||
| assert lvl_15_submap.meta["r_sun"] == lvl_15_submap.meta["rsun_obs"] / lvl_15_submap.meta["cdelt1"] |
There was a problem hiding this comment.
suggestion (testing): Add explicit tests for missing handling and NaN filling behaviour in register.
The updated register now fills NaNs and accepts a missing argument, but this isn’t covered by tests. Please add tests that:
- Build a map containing NaNs and call
registerwithoutmissing; assert (a) the result has no NaNs and (b) the replacement value matchesnp.nanminof the reprojected data. - Call
register(..., missing=<sentinel>)and assert that NaNs in the result are replaced by that sentinel.
This will protect the NaN-handling logic and the missing parameter contract from regressions.
|
Ok I fixed up the tests a bit to use derived properties and also
|
Co-authored-by: Will Barnes <will.t.barnes@gmail.com>
3 participants
Fixes #384
This is a first pass at using the
reprojectpackage to perform image registration by constructing a WCS for the aligned, translated, and scaled image and reprojecting the input map into that new WCS. This also now works for submaps. Note that this does not add a dependency asreprojectis an optional dependency of sunpy that we already pick up by requiringsunpy[map].I've marked this as a draft because it still needs quite a bit of cleaning up and I'd like several pairs of eyes on it before we even think about merging (if we do at all).
There is also the issue of performance. This is markedly slower for full-disk images by an OOM.
However, this version does have the advantage that prep on submaps is possible. Additionally,
reprojectoffers several different algorithms for performing the actual reprojection and there are proposed performance improvements.. There is also progress being made on parallelizing themap_coordinatesfunction with Dask which is the underlying scipy function called byreproject_interp. There's also a cupy version ofmap_coordinates.All my comments above are assuming that the performance bottleneck is the
map_coordinatesfunction, i.e. the actual interpolation of the array onto the new coordinates. While this is almost certainly true, it should be confirmed with some profiling.Some TODOs:
registerSummary by Sourcery
Switch AIA image registration to use WCS-based reprojection and make registration compatible with both full-disk maps and submaps while centralizing detector dimension handling.
New Features:
Enhancements:
sunpy.map.GenericMap.reproject_to.detector_dimensionsutility and apply it across calibration and PSF code to standardize detector size usage.Documentation:
Tests: