Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

reproject #457

Merged
merged 12 commits into from
Feb 8, 2024
Next Next commit
initial commit for reprojecting WorkUnit
maxwest-uw committed Jan 31, 2024
commit 289d381b641fa2ca8cac8456aa92b26c94e09465
92 changes: 92 additions & 0 deletions src/kbmod/reprojection.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
from astropy.wcs import WCS
from astropy.nddata import CCDData
import reproject
import numpy as np

from kbmod.work_unit import WorkUnit
from kbmod.search import RawImage, LayeredImage, ImageStack

def reproject_raw_image(image, original_wcs, common_wcs, obs_time):
"""Given an ndarray representing image data (either science or variance,
when used with `reproject_work_unit`), as well as a common wcs, return the reprojected
RawImage.

Attributes
----------
image : `kbmod.search.RawImage`
The image data to be reprojected.
original_wcs : `astropy.wcs.WCS`
The WCS of the original image.
common_wcs : `astropy.wcs.WCS`
The WCS to reproject all the images into.
obs_time : float
The MJD of the observation.
Returns
----------
maxwest-uw marked this conversation as resolved.
Show resolved Hide resolved
A `kbmod.search.RawImage` reprojected with a common `astropy.wcs.WCS`.
"""
image_data = CCDData(image.image, unit="adu")
image_data.wcs = original_wcs

new_image, _ = reproject.reproject_interp(
image_data, common_wcs, shape_out=common_wcs.array_shape, order="bicubic"
)

return RawImage(img=new_image.astype("float32"), obs_time=obs_time)

def reproject_work_unit(work_unit, common_wcs):
"""Given a WorkUnit and a WCS, reproject all of the images in the ImageStack
into a common WCS.

Attributes
----------
work_unit : `kbmod.WorkUnit`
The WorkUnit to be reprojected.
common_wcs : `astropy.wcs.WCS`
The WCS to reproject all the images into.

Returns
----------
A `kbmod.WorkUnit` reprojected with a common `astropy.wcs.WCS`.
"""
height, width = common_wcs.array_shape
images = work_unit.im_stack.get_images()

if len(work_unit.per_image_wcs) != len(images):
raise ValueError("no per_image_wcs provided for WorkUnit")

image_list = []

for index, image in enumerate(images):
science = image.get_science()
variance = image.get_variance()
obs_time = image.get_obstime()
original_wcs = work_unit.per_image_wcs[index]

reprojected_science = reproject_raw_image(
science, original_wcs, common_wcs, obs_time
)

reprojected_variance = reproject_raw_image(
variance, original_wcs, common_wcs, obs_time
)

mask = image.get_mask()
psf = image.get_psf()

new_layered_image = LayeredImage(
reprojected_science,
reprojected_science,
mask,
psf
)

image_list.append(new_layered_image)

stack = ImageStack(image_list)
new_wunit = WorkUnit(im_stack=stack, config=work_unit.config)
new_wunit.wcs = common_wcs

return new_wunit


4 changes: 2 additions & 2 deletions src/kbmod/work_unit.py
Original file line number Diff line number Diff line change
@@ -24,8 +24,8 @@ class WorkUnit:
needed for a full run of KBMOD, including the: the search parameters,
data files, and the data provenance metadata.

Atributes
---------
Attributes
----------
im_stack : `kbmod.search.ImageStack`
The image data for the KBMOD run.
config : `kbmod.configuration.SearchConfiguration`