Skip to content

Commit

Permalink
Move obstime from RawImage to LayeredImage
Browse files Browse the repository at this point in the history
  • Loading branch information
jeremykubica committed Sep 20, 2024
1 parent 36e739c commit b6ffa34
Show file tree
Hide file tree
Showing 15 changed files with 60 additions and 103 deletions.
4 changes: 0 additions & 4 deletions src/kbmod/file_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,10 +75,6 @@ class FileUtils:
Examples
--------
* Load an external file of visit ID to timestamp mappings.
``time_dict = FileUtils.load_time_dictionary("kbmod/data/demo_times.dat")``
* Load the results of a KBMOD run as trajectory objects.
``FileUtils.load_results_file_as_trajectories("results_DEMO.txt")``
Expand Down
15 changes: 11 additions & 4 deletions src/kbmod/search/layered_image.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@

namespace search {

LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, const PSF& psf)
: psf(psf) {
LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, const PSF& psf,
double obs_time)
: obstime(obs_time), psf(psf) {
// Get the dimensions of the science layer and check for consistency with
// the other two layers.
width = sci.get_width();
Expand All @@ -20,9 +21,10 @@ LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawIm
}

LayeredImage::LayeredImage(Image& sci, Image& var, Image& msk, PSF& psf, double obs_time)
: science(sci, obs_time), variance(var, obs_time), mask(msk, obs_time), psf(psf) {
: science(sci), variance(var), mask(msk), psf(psf) {
width = science.get_width();
height = science.get_height();
obstime = obs_time;
assert_sizes_equal(variance.get_width(), width, "variance layer width");
assert_sizes_equal(variance.get_height(), height, "variance layer height");
assert_sizes_equal(mask.get_width(), width, "mask layer width");
Expand All @@ -33,6 +35,7 @@ LayeredImage::LayeredImage(Image& sci, Image& var, Image& msk, PSF& psf, double
LayeredImage::LayeredImage(const LayeredImage& source) noexcept {
width = source.width;
height = source.height;
obstime = source.obstime;
science = source.science;
mask = source.mask;
variance = source.variance;
Expand All @@ -43,6 +46,7 @@ LayeredImage::LayeredImage(const LayeredImage& source) noexcept {
LayeredImage::LayeredImage(LayeredImage&& source) noexcept
: width(source.width),
height(source.height),
obstime(source.obstime),
science(std::move(source.science)),
mask(std::move(source.mask)),
variance(std::move(source.variance)),
Expand All @@ -52,6 +56,7 @@ LayeredImage::LayeredImage(LayeredImage&& source) noexcept
LayeredImage& LayeredImage::operator=(const LayeredImage& source) noexcept {
width = source.width;
height = source.height;
obstime = source.obstime;
science = source.science;
mask = source.mask;
variance = source.variance;
Expand All @@ -64,6 +69,7 @@ LayeredImage& LayeredImage::operator=(LayeredImage&& source) noexcept {
if (this != &source) {
width = source.width;
height = source.height;
obstime = source.obstime;
science = std::move(source.science);
mask = std::move(source.mask);
variance = std::move(source.variance);
Expand Down Expand Up @@ -252,7 +258,8 @@ static void layered_image_bindings(py::module& m) {
using pf = search::PSF;

py::class_<li>(m, "LayeredImage", pydocs::DOC_LayeredImage)
.def(py::init<const ri&, const ri&, const ri&, pf&>())
.def(py::init<const ri&, const ri&, const ri&, const pf&, double>(), py::arg("sci"),
py::arg("var"), py::arg("msk"), py::arg("psf"), py::arg("obs_time") = -1.0)
.def(py::init<search::Image&, search::Image&, search::Image&, pf&, double>(),
py::arg("sci").noconvert(true), py::arg("var").noconvert(true),
py::arg("msk").noconvert(true), py::arg("psf"), py::arg("obs_time"))
Expand Down
8 changes: 5 additions & 3 deletions src/kbmod/search/layered_image.h
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@
namespace search {
class LayeredImage {
public:
explicit LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, const PSF& psf);
explicit LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, const PSF& psf,
double obs_time = -1.0);

// Build a layered image from the underlying matrices, taking ownership of the image data.
explicit LayeredImage(Image& sci, Image& var, Image& msk, PSF& psf, double obs_time);
Expand All @@ -33,8 +34,8 @@ class LayeredImage {
unsigned get_width() const { return width; }
unsigned get_height() const { return height; }
uint64_t get_npixels() const { return width * height; }
double get_obstime() const { return science.get_obstime(); }
void set_obstime(double obstime) { science.set_obstime(obstime); }
double get_obstime() const { return obstime; }
void set_obstime(double new_obstime) { obstime = new_obstime; }

// Getter functions for the data in the individual layers.
RawImage& get_science() { return science; }
Expand Down Expand Up @@ -92,6 +93,7 @@ class LayeredImage {
private:
unsigned width;
unsigned height;
double obstime;

PSF psf;
RawImage science;
Expand Down
2 changes: 2 additions & 0 deletions src/kbmod/search/pydocs/layered_image_docs.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ static const auto DOC_LayeredImage = R"doc(
The `RawImage` for the mask layer.
psf : `PSF`
The PSF for the image.
obstime : `float`
The time of the image.
Raises
------
Expand Down
14 changes: 4 additions & 10 deletions src/kbmod/search/pydocs/raw_image_docs.h
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,19 @@ namespace pydocs {
static const auto DOC_RawImage = R"doc(
Image and the time it was observed at.
Instantiated from a pair of obstime and image, or from image dimensions and
obstime. When instantiating from image dimensions and obstime, it's possible
to provide a default value used to fill the array with. Otherwise the array is
filled with zeros.
Instantiated from an image or from image dimensions and and a value (which
defaults to zero).
Parameters
----------
image : `numpy.array`, optional
Image, row-major a 2D array. The array *must* be of dtype `numpy.single`.
obstime : `float`, optional
MJD stamp, time the image was observed at.
width : `int`, optional
Width of the image.
height : `int`, optional
Height of the image.
value : `float`, optional
When instantiated from dimensions and obstime, value that fills the array.
When instantiated from dimensions, value that fills the array.
Default is 0.
Attributes
Expand All @@ -32,8 +28,6 @@ static const auto DOC_RawImage = R"doc(
Image width, in pixels.
npixels : `int`
Number of pixels in the image, equivalent to ``width*height``.
obstime : `float`
MJD time of observation.
image : `np.array[np,single]`
Image array.
Expand Down Expand Up @@ -61,7 +55,7 @@ static const auto DOC_RawImage = R"doc(
>>> from kbmod.search import RawImage
>>> import numpy as np
>>> ri = RawImage()
>>> ri = RawImage(w=2, h=3, value=1, obstime=10)
>>> ri = RawImage(w=2, h=3, value=1)
>>> ri.image
array([[1., 1.],
[1., 1.],
Expand Down
21 changes: 7 additions & 14 deletions src/kbmod/search/raw_image.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,16 @@ namespace search {
using Index = indexing::Index;
using Point = indexing::Point;

RawImage::RawImage() : width(0), height(0), obstime(-1.0), image() {}
RawImage::RawImage() : width(0), height(0), image() {}

RawImage::RawImage(Image& img, double obs_time) {
RawImage::RawImage(Image& img) {
image = std::move(img);
height = image.rows();
width = image.cols();
obstime = obs_time;
}

RawImage::RawImage(unsigned w, unsigned h, float value, double obs_time)
: width(w), height(h), obstime(obs_time) {
RawImage::RawImage(unsigned w, unsigned h, float value)
: width(w), height(h) {
if (value != 0.0f)
image = Image::Constant(height, width, value);
else
Expand All @@ -26,22 +25,19 @@ RawImage::RawImage(const RawImage& old) noexcept {
width = old.get_width();
height = old.get_height();
image = old.get_image();
obstime = old.get_obstime();
}

// Move constructor
RawImage::RawImage(RawImage&& source) noexcept
: width(source.width),
height(source.height),
obstime(source.obstime),
image(std::move(source.image)) {}

// Copy assignment
RawImage& RawImage::operator=(const RawImage& source) noexcept {
width = source.width;
height = source.height;
image = source.image;
obstime = source.obstime;
return *this;
}

Expand All @@ -51,7 +47,6 @@ RawImage& RawImage::operator=(RawImage&& source) noexcept {
width = source.width;
height = source.height;
image = std::move(source.image);
obstime = source.obstime;
}
return *this;
}
Expand Down Expand Up @@ -484,15 +479,13 @@ static void raw_image_bindings(py::module& m) {
py::class_<rie>(m, "RawImage", pydocs::DOC_RawImage)
.def(py::init<>())
.def(py::init<search::RawImage&>())
.def(py::init<search::Image&, double>(), py::arg("img").noconvert(true),
py::arg("obs_time") = -1.0d)
.def(py::init<unsigned, unsigned, float, double>(), py::arg("w"), py::arg("h"),
py::arg("value") = 0.0f, py::arg("obs_time") = -1.0d)
.def(py::init<search::Image&>(), py::arg("img").noconvert(true))
.def(py::init<unsigned, unsigned, float>(), py::arg("w"), py::arg("h"),
py::arg("value") = 0.0f)
// attributes and properties
.def_property_readonly("height", &rie::get_height)
.def_property_readonly("width", &rie::get_width)
.def_property_readonly("npixels", &rie::get_npixels)
.def_property("obstime", &rie::get_obstime, &rie::set_obstime)
.def_property("image", py::overload_cast<>(&rie::get_image, py::const_), &rie::set_image)
.def_property("imref", py::overload_cast<>(&rie::get_image), &rie::set_image)
// pixel accessors and setters
Expand Down
9 changes: 2 additions & 7 deletions src/kbmod/search/raw_image.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@ using ImageIRef = Eigen::Ref<Image>;
class RawImage {
public:
explicit RawImage();
explicit RawImage(Image& img, double obs_time = -1.0);
explicit RawImage(unsigned w, unsigned h, float value = 0.0, double obs_time = -1.0);
explicit RawImage(Image& img);
explicit RawImage(unsigned w, unsigned h, float value = 0.0);

RawImage(const RawImage& old) noexcept; // Copy constructor
RawImage(RawImage&& source) noexcept; // Move constructor
Expand Down Expand Up @@ -70,10 +70,6 @@ class RawImage {

void replace_masked_values(float value = 0.0);

// Functions for locally storing the image time.
double get_obstime() const { return obstime; }
void set_obstime(double new_time) { obstime = new_time; }

// this will be a raw pointer to the underlying array
// we use this to copy to GPU and nowhere else!
float* data() { return image.data(); }
Expand Down Expand Up @@ -129,7 +125,6 @@ class RawImage {
private:
unsigned width;
unsigned height;
double obstime;
Image image;
};

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -412,5 +412,5 @@ def toLayeredImage(self):
for sci, var, mask, psf, t in zip(sciences, variances, masks, psfs, mjds):
# Converts nd array mask from bool to np.float32
mask = mask.astype(np.float32)
imgs.append(LayeredImage(RawImage(sci, t), RawImage(var, t), RawImage(mask, t), psf))
imgs.append(LayeredImage(RawImage(sci), RawImage(var), RawImage(mask), psf, obs_time=t))
return imgs
44 changes: 13 additions & 31 deletions src/kbmod/work_unit.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,22 +389,23 @@ def to_fits(self, filename, overwrite=False):

for i in range(self.im_stack.img_count()):
layered = self.im_stack.get_single_image(i)
obstime = layered.get_obstime()
c_indices = self._per_image_indices[i]
n_indices = len(c_indices)

img_wcs = self.get_wcs(i)
sci_hdu = raw_image_to_hdu(layered.get_science(), img_wcs)
sci_hdu = raw_image_to_hdu(layered.get_science(), obstime, img_wcs)
sci_hdu.name = f"SCI_{i}"
sci_hdu.header["NIND"] = n_indices
for j in range(n_indices):
sci_hdu.header[f"IND_{j}"] = c_indices[j]
hdul.append(sci_hdu)

var_hdu = raw_image_to_hdu(layered.get_variance())
var_hdu = raw_image_to_hdu(layered.get_variance(), obstime)
var_hdu.name = f"VAR_{i}"
hdul.append(var_hdu)

msk_hdu = raw_image_to_hdu(layered.get_mask())
msk_hdu = raw_image_to_hdu(layered.get_mask(), obstime)
msk_hdu.name = f"MSK_{i}"
hdul.append(msk_hdu)

Expand Down Expand Up @@ -463,23 +464,24 @@ def to_sharded_fits(self, filename, directory, overwrite=False):

for i in range(self.im_stack.img_count()):
layered = self.im_stack.get_single_image(i)
obstime = layered.get_obstime()
c_indices = self._per_image_indices[i]
n_indices = len(c_indices)
sub_hdul = fits.HDUList()

img_wcs = self.get_wcs(i)
sci_hdu = raw_image_to_hdu(layered.get_science(), img_wcs)
sci_hdu = raw_image_to_hdu(layered.get_science(), obstime, img_wcs)
sci_hdu.name = f"SCI_{i}"
sci_hdu.header["NIND"] = n_indices
for j in range(n_indices):
sci_hdu.header[f"IND_{j}"] = c_indices[j]
sub_hdul.append(sci_hdu)

var_hdu = raw_image_to_hdu(layered.get_variance())
var_hdu = raw_image_to_hdu(layered.get_variance(), obstime)
var_hdu.name = f"VAR_{i}"
sub_hdul.append(var_hdu)

msk_hdu = raw_image_to_hdu(layered.get_mask())
msk_hdu = raw_image_to_hdu(layered.get_mask(), obstime)
msk_hdu.name = f"MSK_{i}"
sub_hdul.append(msk_hdu)

Expand Down Expand Up @@ -835,13 +837,15 @@ def load_layered_image_from_shard(file_path):
return img


def raw_image_to_hdu(img, wcs=None):
def raw_image_to_hdu(img, obstime, wcs=None):
"""Helper function that creates a HDU out of RawImage.
Parameters
----------
img : `RawImage`
The RawImage to convert.
obstime : `float`
The time of the observation.
wcs : `astropy.wcs.WCS`
An optional WCS to include in the header.
Expand All @@ -857,28 +861,6 @@ def raw_image_to_hdu(img, wcs=None):
append_wcs_to_hdu_header(wcs, hdu.header)

# Set the time stamp.
hdu.header["MJD"] = img.obstime
return hdu


def hdu_to_raw_image(hdu):
"""Helper function that creates a RawImage from a HDU.
Parameters
----------
hdu : `astropy.io.fits.hdu.image.ImageHDU`
The image extension.
hdu.header["MJD"] = obstime

Returns
-------
img : `RawImage` or None
The RawImage if there is valid data and None otherwise.
"""
img = None
if isinstance(hdu, fits.hdu.image.ImageHDU):
# This will be a copy whenever dtype != np.single including when
# endianness doesn't match the native float.
img = RawImage(hdu.data.astype(np.single))
if "MJD" in hdu.header:
img.obstime = hdu.header["MJD"]
return img
return hdu
3 changes: 0 additions & 3 deletions tests/test_butlerstd.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,9 +223,6 @@ def test_to_layered_image(self):
# times can only be compred approximately, because sometimes we
# calculate the time in the middle of the exposure
self.assertAlmostEqual(expected_mjd, img.get_obstime(), 2)
self.assertAlmostEqual(expected_mjd, img.get_science().obstime, 2)
self.assertAlmostEqual(expected_mjd, img.get_variance().obstime, 2)
self.assertAlmostEqual(expected_mjd, img.get_mask().obstime, 2)


if __name__ == "__main__":
Expand Down
2 changes: 1 addition & 1 deletion tests/test_fake_data_creator.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ def test_create_fake_times(self):
self.assertAlmostEqual(times2[i], expected[i])

def test_add_fake_object(self):
img = RawImage(20, 10, 0.0, 1.0) # All zero image.
img = RawImage(20, 10, 0.0) # All zero image.
p = PSF(np.full((3, 3), 1.0 / 9.0)) # Equal PSF.
add_fake_object(img, 5.5, 3.5, 100.0, p)

Expand Down
Loading

0 comments on commit b6ffa34

Please sign in to comment.