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.

Sign up for GitHub

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

Jump to bottom

Move more stamp functions to StampCreator #387

Merged
merged 3 commits into from
Oct 24, 2023
Merged

Move more stamp functions to StampCreator #387

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5294064
Move more functions to StampCreator
jeremykubica Oct 23, 2023
614bad5
Bug fix
jeremykubica Oct 23, 2023
4fb96be
Add pydoc comments
jeremykubica Oct 24, 2023
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
8 changes: 6 additions & 2 deletions src/kbmod/analysis_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,8 +317,12 @@ def apply_stamp_filter(

# Create and filter the results, using the GPU if there is one and enough
# trajectories to make it worthwhile.
stamps_slice = search.get_coadded_stamps(
trj_slice, bool_slice, params, kb.HAS_GPU and len(trj_slice) > 100
stamps_slice = kb.StampCreator.get_coadded_stamps(
search.get_imagestack(),
trj_slice,
bool_slice,
params,
kb.HAS_GPU and len(trj_slice) > 100,
)
for ind, stamp in enumerate(stamps_slice):
if stamp.get_width() > 1:
Expand Down
8 changes: 0 additions & 8 deletions src/kbmod/search/pydocs/stack_search_docs.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,14 +62,6 @@ namespace pydocs {
todo
)doc";

static const auto DOC_StackSearch_get_coadded_stamps = R"doc(
todo
)doc";

static const auto DOC_StackSearch_filter_stamp = R"doc(
todo
)doc";

static const auto DOC_StackSearch_get_trajectory_position = R"doc(
todo
)doc";
Expand Down
8 changes: 8 additions & 0 deletions src/kbmod/search/pydocs/stamp_creator_docs.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,14 @@ namespace pydocs {
The co-added stamp.
)doc";

static const auto DOC_StampCreator_get_coadded_stamps = R"doc(
todo
)doc";

static const auto DOC_StampCreator_filter_stamp = R"doc(
todo
)doc";

} /* pydocs */
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional: Up to you, but it's probably nice to fill these out while we're here

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good call. I added them.


#endif /* STAMP_CREATOR_DOCS */
154 changes: 1 addition & 153 deletions src/kbmod/search/stack_search.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,6 @@ extern "C" void deviceSearchFilter(int num_images, int width, int height, float*
PerImageData img_data, SearchParameters params, int num_trajectories,
Trajectory* trj_to_search, int num_results, Trajectory* best_results);

void deviceGetCoadds(ImageStack& stack, PerImageData image_data, int num_trajectories,
Trajectory* trajectories, StampParameters params,
std::vector<std::vector<bool>>& use_index_vect, float* results);
#endif

StackSearch::StackSearch(ImageStack& imstack) : stack(imstack) {
Expand Down Expand Up @@ -264,148 +261,6 @@ void StackSearch::fill_psi_phi(const std::vector<RawImage>& psi_imgs, const std:
}
}

bool StackSearch::filter_stamp(const RawImage& img, const StampParameters& params) {
// Allocate space for the coadd information and initialize to zero.
const int stamp_width = 2 * params.radius + 1;
const int stamp_ppi = stamp_width * stamp_width;
const std::vector<float>& pixels = img.get_pixels();

// Filter on the peak's position.
PixelPos pos = img.find_peak(true);
if ((abs(pos.x - params.radius) >= params.peak_offset_x) ||
(abs(pos.y - params.radius) >= params.peak_offset_y)) {
return true;
}

// Filter on the percentage of flux in the central pixel.
if (params.center_thresh > 0.0) {
const std::vector<float>& pixels = img.get_pixels();
float center_val = pixels[(int)pos.y * stamp_width + (int)pos.x];
float pixel_sum = 0.0;
for (int p = 0; p < stamp_ppi; ++p) {
pixel_sum += pixels[p];
}

if (center_val / pixel_sum < params.center_thresh) {
return true;
}
}

// Filter on the image moments.
ImageMoments moments = img.find_central_moments();
if ((fabs(moments.m01) >= params.m01_limit) || (fabs(moments.m10) >= params.m10_limit) ||
(fabs(moments.m11) >= params.m11_limit) || (moments.m02 >= params.m02_limit) ||
(moments.m20 >= params.m20_limit)) {
return true;
}

return false;
}

std::vector<RawImage> StackSearch::get_coadded_stamps(std::vector<Trajectory>& t_array,
std::vector<std::vector<bool>>& use_index_vect,
const StampParameters& params, bool use_gpu) {
if (use_gpu) {
#ifdef HAVE_CUDA
return get_coadded_stamps_gpu(t_array, use_index_vect, params);
#else
std::cout << "WARNING: GPU is not enabled. Performing co-adds on the CPU.";
// py::print("WARNING: GPU is not enabled. Performing co-adds on the CPU.");
#endif
}
return get_coadded_stamps_cpu(t_array, use_index_vect, params);
}

std::vector<RawImage> StackSearch::get_coadded_stamps_cpu(std::vector<Trajectory>& t_array,
std::vector<std::vector<bool>>& use_index_vect,
const StampParameters& params) {
const int num_trajectories = t_array.size();
std::vector<RawImage> results(num_trajectories);
std::vector<float> empty_pixels(1, NO_DATA);

for (int i = 0; i < num_trajectories; ++i) {
std::vector<RawImage> stamps =
StampCreator::create_stamps(stack, t_array[i], params.radius, false, true, use_index_vect[i]);

RawImage coadd(1, 1);
switch (params.stamp_type) {
case STAMP_MEDIAN:
coadd = create_median_image(stamps);
break;
case STAMP_MEAN:
coadd = create_mean_image(stamps);
break;
case STAMP_SUM:
coadd = create_summed_image(stamps);
break;
default:
throw std::runtime_error("Invalid stamp coadd type.");
}

// Do the filtering if needed.
if (params.do_filtering && filter_stamp(coadd, params)) {
results[i] = RawImage(1, 1, empty_pixels);
} else {
results[i] = coadd;
}
}

return results;
}

std::vector<RawImage> StackSearch::get_coadded_stamps_gpu(std::vector<Trajectory>& t_array,
std::vector<std::vector<bool>>& use_index_vect,
const StampParameters& params) {
// Right now only limited stamp sizes are allowed.
if (2 * params.radius + 1 > MAX_STAMP_EDGE || params.radius <= 0) {
throw std::runtime_error("Invalid Radius.");
}

const int num_images = stack.img_count();
const int width = stack.get_width();
const int height = stack.get_height();

// Create a data stucture for the per-image data.
std::vector<float> image_times = stack.build_zeroed_times();
PerImageData img_data;
img_data.num_images = num_images;
img_data.image_times = image_times.data();

// Allocate space for the results.
const int num_trajectories = t_array.size();
const int stamp_width = 2 * params.radius + 1;
const int stamp_ppi = stamp_width * stamp_width;
std::vector<float> stamp_data(stamp_ppi * num_trajectories);

// Do the co-adds.
#ifdef HAVE_CUDA
deviceGetCoadds(stack, img_data, num_trajectories, t_array.data(), params, use_index_vect,
stamp_data.data());
#else
throw std::runtime_error("Non-GPU co-adds is not implemented.");
#endif

// Copy the stamps into RawImages and do the filtering.
std::vector<RawImage> results(num_trajectories);
std::vector<float> current_pixels(stamp_ppi, 0.0);
std::vector<float> empty_pixels(1, NO_DATA);
for (int t = 0; t < num_trajectories; ++t) {
// Copy the data into a single RawImage.
int offset = t * stamp_ppi;
for (unsigned p = 0; p < stamp_ppi; ++p) {
current_pixels[p] = stamp_data[offset + p];
}
RawImage current_image = RawImage(stamp_width, stamp_width, current_pixels);

if (params.do_filtering && filter_stamp(current_image, params)) {
results[t] = RawImage(1, 1, empty_pixels);
} else {
results[t] = RawImage(stamp_width, stamp_width, current_pixels);
}
}
return results;
}

std::vector<float> StackSearch::create_curves(Trajectory t, const std::vector<RawImage>& imgs) {
/*Create a lightcurve from an image along a trajectory
*
Expand Down Expand Up @@ -513,14 +368,7 @@ static void stack_search_bindings(py::module& m) {
.def("get_image_npixels", &ks::get_image_npixels, pydocs::DOC_StackSearch_get_image_npixels)
.def("get_imagestack", &ks::get_imagestack, py::return_value_policy::reference_internal,
pydocs::DOC_StackSearch_get_imagestack)
// Science Stamp Functions
.def("get_coadded_stamps", // wth is happening here
(std::vector<ri>(ks::*)(std::vector<tj>&, std::vector<std::vector<bool>>&,
const search::StampParameters&, bool)) &
ks::get_coadded_stamps,
pydocs::DOC_StackSearch_get_coadded_stamps)
// For testing
.def("filter_stamp", &ks::filter_stamp, pydocs::DOC_StackSearch_filter_stamp)
// For testings
.def("get_trajectory_position", &ks::get_trajectory_position,
pydocs::DOC_StackSearch_get_trajectory_position)
.def("get_psi_curves", (std::vector<float>(ks::*)(tj&)) & ks::get_psi_curves,
Expand Down
28 changes: 0 additions & 28 deletions src/kbmod/search/stack_search.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,27 +53,6 @@ class StackSearch {
void filter_results(int min_observations);
void filter_results_lh(float min_lh);

// Functions for creating science stamps for filtering, visualization, etc. User can specify
// the radius of the stamp, whether to interpolate among pixels, whether to keep NO_DATA values
// or replace them with zero, and what indices to use.
// The indices to use are indicated by use_index: a vector<bool> indicating whether to use
// each time step. An empty (size=0) vector will use all time steps.
std::vector<RawImage> get_stamps(const Trajectory& t, int radius);
RawImage get_median_stamp(const Trajectory& trj, int radius, const std::vector<bool>& use_index);
RawImage get_mean_stamp(const Trajectory& trj, int radius, const std::vector<bool>& use_index);
RawImage get_summed_stamp(const Trajectory& trj, int radius, const std::vector<bool>& use_index);

// Compute a mean or summed stamp for each trajectory on the GPU or CPU.
// The GPU implementation is slower for small numbers of trajectories (< 500), but performs
// relatively better as the number of trajectories increases. If filtering is applied then
// the code will return a 1x1 image with NO_DATA to represent each filtered image.
std::vector<RawImage> get_coadded_stamps(std::vector<Trajectory>& t_array,
std::vector<std::vector<bool> >& use_index_vect,
const StampParameters& params, bool use_cpu);

// Function to do the actual stamp filtering.
bool filter_stamp(const RawImage& img, const StampParameters& params);

// Getters for the Psi and Phi data.
std::vector<RawImage>& get_psi_images();
std::vector<RawImage>& get_phi_images();
Expand Down Expand Up @@ -108,13 +87,6 @@ class StackSearch {
void create_search_list(int angle_steps, int velocity_steps, float min_ang, float max_ang, float min_vel,
float max_vel);

std::vector<RawImage> get_coadded_stamps_gpu(std::vector<Trajectory>& t_array,
std::vector<std::vector<bool> >& use_index_vect,
const StampParameters& params);

std::vector<RawImage> get_coadded_stamps_cpu(std::vector<Trajectory>& t_array,
std::vector<std::vector<bool> >& use_index_vect,
const StampParameters& params);

bool psi_phi_generated;
bool debug_info;
Expand Down
Loading