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

A series of small cleanups for KBMOD #431

Merged
merged 6 commits into from
Jan 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion src/kbmod/search/bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ namespace py = pybind11;
#include "kernel_testing_helpers.cpp"
#include "psi_phi_array.cpp"


PYBIND11_MODULE(search, m) {
m.attr("KB_NO_DATA") = pybind11::float_(search::NO_DATA);
m.attr("HAS_GPU") = pybind11::bool_(search::HAVE_GPU);
Expand Down
34 changes: 16 additions & 18 deletions src/kbmod/search/layered_image.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@ RawImage LayeredImage::generate_psi_image() {
const int num_pixels = get_npixels();
for (int p = 0; p < num_pixels; ++p) {
float var_pix = var_array[p];
if (var_pix != NO_DATA) {
if (var_pix != NO_DATA && var_pix != 0.0 && sci_array[p] != NO_DATA) {
result_arr[p] = sci_array[p] / var_pix;
} else {
result_arr[p] = NO_DATA;
Expand All @@ -251,7 +251,7 @@ RawImage LayeredImage::generate_phi_image() {
const int num_pixels = get_npixels();
for (int p = 0; p < num_pixels; ++p) {
float var_pix = var_array[p];
if (var_pix != NO_DATA) {
if (var_pix != NO_DATA && var_pix != 0.0) {
result_arr[p] = 1.0 / var_pix;
} else {
result_arr[p] = NO_DATA;
Expand Down Expand Up @@ -280,28 +280,26 @@ static void layered_image_bindings(py::module& m) {
.def("contains", &li::contains, pydocs::DOC_LayeredImage_cointains)
.def("get_science_pixel", &li::get_science_pixel, pydocs::DOC_LayeredImage_get_science_pixel)
.def("get_variance_pixel", &li::get_variance_pixel, pydocs::DOC_LayeredImage_get_variance_pixel)
.def(
"contains",
[](li& cls, int i, int j) {
return cls.contains({i, j});
})
.def(
"get_science_pixel",
[](li& cls, int i, int j) {
return cls.get_science_pixel({i, j});
})
.def(
"get_variance_pixel",
[](li& cls, int i, int j) {
return cls.get_variance_pixel({i, j});
})
.def("contains",
[](li& cls, int i, int j) {
return cls.contains({i, j});
})
.def("get_science_pixel",
[](li& cls, int i, int j) {
return cls.get_science_pixel({i, j});
})
.def("get_variance_pixel",
[](li& cls, int i, int j) {
return cls.get_variance_pixel({i, j});
})
.def("set_psf", &li::set_psf, pydocs::DOC_LayeredImage_set_psf)
.def("get_psf", &li::get_psf, py::return_value_policy::reference_internal,
pydocs::DOC_LayeredImage_get_psf)
.def("binarize_mask", &li::binarize_mask, pydocs::DOC_LayeredImage_binarize_mask)
.def("apply_mask", &li::apply_mask, pydocs::DOC_LayeredImage_apply_mask)
.def("union_masks", &li::union_masks, pydocs::DOC_LayeredImage_union_masks)
.def("union_threshold_masking", &li::union_threshold_masking, pydocs::DOC_LayeredImage_union_threshold_masking)
.def("union_threshold_masking", &li::union_threshold_masking,
pydocs::DOC_LayeredImage_union_threshold_masking)
.def("sub_template", &li::subtract_template, pydocs::DOC_LayeredImage_sub_template)
.def("save_layers", &li::save_layers, pydocs::DOC_LayeredImage_save_layers)
.def("get_science", &li::get_science, py::return_value_policy::reference_internal,
Expand Down
8 changes: 4 additions & 4 deletions src/kbmod/search/layered_image.h
Original file line number Diff line number Diff line change
Expand Up @@ -42,13 +42,13 @@ class LayeredImage {
// Getter functions for the pixels of the science and variance layers that check
// the mask layer for any set bits.
inline float get_science_pixel(const Index& idx) const {
return contains(idx) ? (mask.get_pixel(idx) == 0 ? science.get_pixel(idx) : NO_DATA) : NO_DATA;
// The get_pixel() functions perform the bounds checking and will return NO_DATA for out of bounds.
return mask.get_pixel(idx) == 0 ? science.get_pixel(idx) : NO_DATA;
}

inline float get_variance_pixel(const Index& idx) const {
return contains(idx) ?
(mask.get_pixel(idx) == 0 ? variance.get_pixel(idx) : NO_DATA) :
NO_DATA;
// The get_pixel() functions perform the bounds checking and will return NO_DATA for out of bounds.
return mask.get_pixel(idx) == 0 ? variance.get_pixel(idx) : NO_DATA;
}

inline bool contains(const Index& idx) const {
Expand Down
24 changes: 13 additions & 11 deletions src/kbmod/search/psi_phi_array.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -177,11 +177,12 @@ void set_encode_cpu_psi_phi_array(PsiPhiArray& data, const std::vector<RawImage>
throw std::runtime_error("CPU PsiPhi already allocated.");
}
if (debug) {
printf("Allocating CPU memory for encoded PsiPhi array using %lu bytes.\n", data.get_total_array_size());
printf("Allocating CPU memory for encoded PsiPhi array using %lu bytes.\n",
data.get_total_array_size());
}
T* encoded = (T*)malloc(data.get_total_array_size());
if (encoded == nullptr) {
throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
}

// Create a safe maximum that is slightly less than the true max to avoid
Expand Down Expand Up @@ -224,7 +225,7 @@ void set_float_cpu_psi_phi_array(PsiPhiArray& data, const std::vector<RawImage>&
}
float* encoded = (float*)malloc(data.get_total_array_size());
if (encoded == nullptr) {
throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
throw std::runtime_error("Unable to allocate space for CPU PsiPhi array.");
}

int current_index = 0;
Expand Down Expand Up @@ -266,12 +267,10 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect
result_data.set_phi_scaling(phi_params[0], phi_params[1], phi_params[2]);

if (debug) {
printf("Encoding psi to %i bytes min=%f, max=%f, scale=%f\n",
result_data.get_num_bytes(), psi_params[0], psi_params[1],
psi_params[2]);
printf("Encoding phi to %i bytes min=%f, max=%f, scale=%f\n",
result_data.get_num_bytes(), phi_params[0], phi_params[1],
phi_params[2]);
printf("Encoding psi to %i bytes min=%f, max=%f, scale=%f\n", result_data.get_num_bytes(),
psi_params[0], psi_params[1], psi_params[2]);
printf("Encoding phi to %i bytes min=%f, max=%f, scale=%f\n", result_data.get_num_bytes(),
phi_params[0], phi_params[1], phi_params[2]);
}

// Do the local encoding.
Expand All @@ -281,15 +280,18 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect
set_encode_cpu_psi_phi_array<uint16_t>(result_data, psi_imgs, phi_imgs, debug);
}
} else {
if (debug) { printf("Encoding psi and phi as floats.\n"); }
if (debug) {
printf("Encoding psi and phi as floats.\n");
}
// Just interleave psi and phi images.
set_float_cpu_psi_phi_array(result_data, psi_imgs, phi_imgs, debug);
}

#ifdef HAVE_CUDA
// Create a copy of the encoded data in GPU memory.
if (debug) {
printf("Allocating GPU memory for PsiPhi array using %lu bytes.\n", result_data.get_total_array_size());
printf("Allocating GPU memory for PsiPhi array using %lu bytes.\n",
result_data.get_total_array_size());
}
device_allocate_psi_phi_array(&result_data);
if (result_data.get_gpu_array_ptr() == nullptr) {
Expand Down
2 changes: 1 addition & 1 deletion src/kbmod/search/psi_phi_array_utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ namespace search {
std::array<float, 3> compute_scale_params_from_image_vect(const std::vector<RawImage>& imgs, int num_bytes);

void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector<RawImage>& psi_imgs,
const std::vector<RawImage>& phi_imgs, bool debug=false);
const std::vector<RawImage>& phi_imgs, bool debug = false);

} /* namespace search */

Expand Down
22 changes: 20 additions & 2 deletions src/kbmod/search/pydocs/layered_image_docs.h
Original file line number Diff line number Diff line change
Expand Up @@ -238,11 +238,29 @@ static const auto DOC_LayeredImage_get_variance_pixel = R"doc(
)doc";

static const auto DOC_LayeredImage_generate_psi_image = R"doc(
todo
Generates the full psi image where the value of each pixel p in the
resulting image is science[p] / variance[p]. To handle masked bits
apply_mask() must be called before the psi image is generated. Otherwise,
all pixels are used.
Convolves the resulting image with the PSF.

Returns
-------
result : `kbmod.RawImage`
A ``RawImage`` of the same dimensions as the ``LayeredImage``.
)doc";

static const auto DOC_LayeredImage_generate_phi_image = R"doc(
todo
Generates the full phi image where the value of each pixel p in the
resulting image is 1.0 / variance[p]. To handle masked bits
apply_mask() must be called before the phi image is generated. Otherwise,
all pixels are used.
Convolves the resulting image with the PSF.

Returns
-------
result : `kbmod.RawImage`
A ``RawImage`` of the same dimensions as the ``LayeredImage``.
)doc";
} // namespace pydocs

Expand Down
64 changes: 27 additions & 37 deletions src/kbmod/search/raw_image.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -110,8 +110,7 @@ float RawImage::interpolate(const Point& p) const {
return total / sumWeights;
}

RawImage RawImage::create_stamp(const Point& p, const int radius,
const bool keep_no_data) const {
RawImage RawImage::create_stamp(const Point& p, const int radius, const bool keep_no_data) const {
if (radius < 0) throw std::runtime_error("stamp radius must be at least 0");

const int dim = radius * 2 + 1;
Expand All @@ -123,8 +122,7 @@ RawImage RawImage::create_stamp(const Point& p, const int radius,
Image stamp = Image::Constant(dim, dim, NO_DATA);
stamp.block(anchor.i, anchor.j, h, w) = image.block(corner.i, corner.j, h, w);

if (!keep_no_data)
stamp = (stamp.array() == NO_DATA).select(0.0, stamp);
if (!keep_no_data) stamp = (stamp.array() == NO_DATA).select(0.0, stamp);

return RawImage(stamp);
}
Expand Down Expand Up @@ -580,29 +578,25 @@ static void raw_image_bindings(py::module& m) {
.def("set_pixel", &rie::set_pixel, pydocs::DOC_RawImage_set_pixel)
.def("set_all", &rie::set_all, pydocs::DOC_RawImage_set_all)
// python interface adapters (avoids having to construct Index & Point)
.def(
"get_pixel",
[](rie& cls, int i, int j) {
return cls.get_pixel({i, j});
})
.def(
"pixel_has_data",
[](rie& cls, int i, int j) {
return cls.pixel_has_data({i, j});
})
.def(
"set_pixel",
[](rie& cls, int i, int j, double val) {
cls.set_pixel({i, j}, val);
})
.def("get_pixel",
[](rie& cls, int i, int j) {
return cls.get_pixel({i, j});
})
.def("pixel_has_data",
[](rie& cls, int i, int j) {
return cls.pixel_has_data({i, j});
})
.def("set_pixel",
[](rie& cls, int i, int j, double val) {
cls.set_pixel({i, j}, val);
})
// methods
.def("l2_allclose", &rie::l2_allclose, pydocs::DOC_RawImage_l2_allclose)
.def("compute_bounds", &rie::compute_bounds, pydocs::DOC_RawImage_compute_bounds)
.def("find_peak", &rie::find_peak, pydocs::DOC_RawImage_find_peak)
.def("find_central_moments", &rie::find_central_moments,
pydocs::DOC_RawImage_find_central_moments)
.def("center_is_local_max", &rie::center_is_local_max,
pydocs::DOC_RawImage_center_is_local_max)
pydocs::DOC_RawImage_find_central_moments)
.def("center_is_local_max", &rie::center_is_local_max, pydocs::DOC_RawImage_center_is_local_max)
.def("create_stamp", &rie::create_stamp, pydocs::DOC_RawImage_create_stamp)
.def("interpolate", &rie::interpolate, pydocs::DOC_RawImage_interpolate)
.def("interpolated_add", &rie::interpolated_add, pydocs::DOC_RawImage_interpolated_add)
Expand All @@ -620,21 +614,17 @@ static void raw_image_bindings(py::module& m) {
.def("append_fits_extension", &rie::append_to_fits, pydocs::DOC_RawImage_append_to_fits)
.def("load_fits", &rie::from_fits, pydocs::DOC_RawImage_load_fits)
// python interface adapters
.def(
"create_stamp",
[](rie& cls, float x, float y, int radius, bool keep_no_data) {
return cls.create_stamp({x, y}, radius, keep_no_data);
})
.def(
"interpolate",
[](rie& cls, float x, float y) {
return cls.interpolate({x, y});
})
.def(
"interpolated_add",
[](rie& cls, float x, float y, float val) {
cls.interpolated_add({x, y}, val);
});
.def("create_stamp",
[](rie& cls, float x, float y, int radius, bool keep_no_data) {
return cls.create_stamp({x, y}, radius, keep_no_data);
})
.def("interpolate",
[](rie& cls, float x, float y) {
return cls.interpolate({x, y});
})
.def("interpolated_add", [](rie& cls, float x, float y, float val) {
cls.interpolated_add({x, y}, val);
});
}
#endif

Expand Down
34 changes: 32 additions & 2 deletions tests/test_layered_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,29 @@ def test_create(self):
self.assertGreaterEqual(science.get_pixel(y, x), -100.0)
self.assertLessEqual(science.get_pixel(y, x), 100.0)

# Check direct lookup of pixel values matches the RawImage lookup.
self.assertEqual(science.get_pixel(y, x), self.image.get_science_pixel(y, x))
self.assertEqual(variance.get_pixel(y, x), self.image.get_variance_pixel(y, x))

# Check that the LayeredImage pixel lookups work with a masked pixel.
# But the the mask was not applied yet to the images themselves.
mask.set_pixel(5, 6, 1)
self.assertGreater(science.get_pixel(5, 6), KB_NO_DATA)
self.assertGreater(variance.get_pixel(5, 6), KB_NO_DATA)
self.assertEqual(self.image.get_science_pixel(5, 6), KB_NO_DATA)
self.assertEqual(self.image.get_variance_pixel(5, 6), KB_NO_DATA)

# Test that out of bounds pixel lookups are handled correctly.
self.assertEqual(self.image.get_science_pixel(-1, 1), KB_NO_DATA)
self.assertEqual(self.image.get_science_pixel(1, -1), KB_NO_DATA)
self.assertEqual(self.image.get_science_pixel(self.image.get_height() + 1, 1), KB_NO_DATA)
self.assertEqual(self.image.get_science_pixel(1, self.image.get_width() + 1), KB_NO_DATA)

self.assertEqual(self.image.get_variance_pixel(-1, 1), KB_NO_DATA)
self.assertEqual(self.image.get_variance_pixel(1, -1), KB_NO_DATA)
self.assertEqual(self.image.get_variance_pixel(self.image.get_height() + 1, 1), KB_NO_DATA)
self.assertEqual(self.image.get_variance_pixel(1, self.image.get_width() + 1), KB_NO_DATA)

def test_create_from_layers(self):
sci = RawImage(30, 40)
for y in range(40):
Expand Down Expand Up @@ -285,7 +308,11 @@ def test_psi_and_phi_image(self):
for x in range(6):
sci.set_pixel(y, x, float(x))
var.set_pixel(y, x, float(y + 1))

# Mask a single pixel and set another to variance of zero.
sci.set_pixel(3, 1, KB_NO_DATA)
var.set_pixel(3, 1, KB_NO_DATA)
var.set_pixel(3, 2, 0.0)

# Generate and check psi and phi images.
psi = img.generate_psi_image()
Expand All @@ -298,12 +325,15 @@ def test_psi_and_phi_image(self):

for y in range(5):
for x in range(6):
has_data = not (x == 1 and y == 3)
has_data = y != 3 or x == 0 or x > 2
self.assertEqual(psi.pixel_has_data(y, x), has_data)
self.assertEqual(phi.pixel_has_data(y, x), has_data)
if x != 1 or y != 3:
if has_data:
self.assertAlmostEqual(psi.get_pixel(y, x), x / (y + 1))
self.assertAlmostEqual(phi.get_pixel(y, x), 1.0 / (y + 1))
else:
self.assertEqual(psi.get_pixel(y, x), KB_NO_DATA)
self.assertEqual(phi.get_pixel(y, x), KB_NO_DATA)

def test_subtract_template(self):
sci = self.image.get_science()
Expand Down