Skip to content

Commit

Permalink
Part #2 of KBMOSearch refactoring
Browse files Browse the repository at this point in the history 
Breaks the stamp creation and aggregation functions out into RawImage to reduce duplication and improve test coverage.

Rename some of the functions for consistency (which leads to a name change only in analysis_utils.py)
  • Loading branch information
@jeremykubica
jeremykubica committed Jul 22, 2022
1 parent 1c7724c commit 1376d54
Show file tree
Hide file tree
Showing 8 changed files with 236 additions and 176 deletions.
2 changes: 1 addition & 1 deletion analysis/analysis_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -587,7 +587,7 @@ def get_coadd_stamps(self, results, search, keep, radius=10, stamp_type='sum'):
coadd_stamps = [np.array(stamp) for stamp in
search.median_stamps(results, boolean_idx, radius)]
elif stamp_type=='parallel_sum':
coadd_stamps = [np.array(stamp) for stamp in search.summed_stamps(results, radius)]
coadd_stamps = [np.array(stamp) for stamp in search.summed_sci(results, radius)]
else:
# Python stamp generation
coadd_stamps = []
Expand Down
5 changes: 3 additions & 2 deletions search/pybinds/classBindings.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,8 @@ PYBIND11_MODULE(kbmod, m) {
.def("extreme_in_region", &ri::extremeInRegion)
.def("convolve", &ri::convolve)
.def("save_fits", &ri::saveToFile);

m.def("create_median_image", &kbmod::createMedianImage);
m.def("create_summed_image", &kbmod::createSummedImage);
py::class_<li>(m, "layered_image")
.def(py::init<const std::string>())
.def(py::init<std::string, int, int,
Expand Down Expand Up @@ -163,8 +164,8 @@ PYBIND11_MODULE(kbmod, m) {
.def("filter_lh", &ks::filterLH)
.def("stacked_sci", (ri (ks::*)(tj &, int)) &ks::stackedScience, "set")
.def("stacked_sci", (ri (ks::*)(td &, int)) &ks::stackedScience, "set")
.def("summed_sci", (std::vector<ri> (ks::*)(std::vector<tj>, int)) &ks::summedScience)
.def("median_stamps", (std::vector<ri> (ks::*)(std::vector<tj>, std::vector<std::vector<int>>, int)) &ks::medianStamps)
.def("summed_stamps", (std::vector<ri> (ks::*)(std::vector<tj>, int)) &ks::summedStamps)
.def("sci_stamps", (std::vector<ri> (ks::*)(tj &, int)) &ks::scienceStamps, "set")
.def("psi_stamps", (std::vector<ri> (ks::*)(tj &, int)) &ks::psiStamps, "set2")
.def("phi_stamps", (std::vector<ri> (ks::*)(tj &, int)) &ks::phiStamps, "set3")
Expand Down
223 changes: 69 additions & 154 deletions search/src/KBMOSearch.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -276,9 +276,6 @@ void KBMOSearch::gpuSearchFilter(int minObservations)
searchList.data(), results.data(), stack.getTimesDataRef(),
interleavedPsiPhi.data(), width, height,
&percentiles[0], sigmaGCoeff, &centralMomLims[0], minLH);
//filterResultsLH(minLH);
// stamps = createMedianBatch(10, imgs);
//filterStamps(stamps);
}

std::vector<trajRegion> KBMOSearch::resSearch(float xVel, float yVel,
Expand Down Expand Up @@ -607,7 +604,7 @@ float KBMOSearch::findExtremeInRegion(float x, float y,
// Round Upper corner up to align larger pixel size
hx = (hx+sizeToRead-1)/sizeToRead;
hy = (hy+sizeToRead-1)/sizeToRead;
float regionExtreme =
float regionExtreme =
pooledImgs.getImage(depth).extremeInRegion(lx, ly, hx-1, hy-1, poolType);
return regionExtreme;
}
Expand Down Expand Up @@ -659,50 +656,8 @@ trajectory KBMOSearch::convertTraj(trajRegion& t)
return tb;
}


std::vector<RawImage> KBMOSearch::summedStamps(std::vector<trajectory> t_array, int radius)
{
int numResults = t_array.size();
int dim = radius*2+1;

std::vector<RawImage*> imgs;
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());
std::vector<RawImage> stamps;
stamps.resize(numResults);
const std::vector<float>& times = stack.getTimes();

omp_set_num_threads(30);

#pragma omp parallel for
for (int s = 0; s < numResults; ++s)
{
RawImage singleStamp(dim,dim);
trajectory t = t_array[s];
for (int i=0; i<imgs.size(); ++i)
{
for (int x = 0; x < dim; ++x)
{
for (int y = 0; y < dim; ++y)
{
std::vector<float> pixArray;
float pixVal = imgs[i]->getPixel(
t.x + times[i] * t.xVel + static_cast<float>(x-radius),
t.y + times[i] * t.yVel + static_cast<float>(y-radius));
if ((pixVal == NO_DATA) || isnan(pixVal)) pixVal = 0.0;
singleStamp.addToPixel(static_cast<int>(x),static_cast<int>(y), pixVal);

}
}
}
stamps[s] = (singleStamp);
}
omp_set_num_threads(1);

return(stamps);
}

std::vector<RawImage> KBMOSearch::medianStamps(std::vector<trajectory> t_array,
std::vector<std::vector<int>> goodIdx,
std::vector<RawImage> KBMOSearch::medianStamps(const std::vector<trajectory>& t_array,
const std::vector<std::vector<int>>& goodIdx,
int radius)
{
int numResults = t_array.size();
Expand All @@ -711,87 +666,37 @@ std::vector<RawImage> KBMOSearch::medianStamps(std::vector<trajectory> t_array,
std::vector<RawImage*> imgs;
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());
size_t N = imgs.size() / 2;
std::vector<RawImage> stamps;
stamps.resize(numResults);
std::vector<RawImage> results(numResults);
const std::vector<float>& times = stack.getTimes();
omp_set_num_threads(30);

//#pragma omp parallel for
for (int s = 0; s < numResults; ++s)
{
RawImage singleStamp(dim,dim);
// Create stamps around the current trajectory.
std::vector<RawImage> stamps;
trajectory t = t_array[s];
for (int x = 0; x < dim; ++x)
for (int i = 0; i < goodIdx[s].size(); ++i)
{
for (int y = 0; y < dim; ++y)
if (goodIdx[s][i] == 1)
{
std::vector<float> pixArray;
for (int i = 0; i < goodIdx[s].size(); ++i)
{
if (goodIdx[s][i] == 1)
{
float pixVal = imgs[i]->getPixel(
t.x + times[i] * t.xVel + static_cast<float>(x-radius),
t.y + times[i] * t.yVel + static_cast<float>(y-radius));
if ((pixVal == NO_DATA) || (isnan(pixVal))) pixVal = 0.0;
pixArray.push_back(pixVal);
}
}
int M = pixArray.size() / 2;
std::nth_element(
pixArray.begin(), pixArray.begin()+M, pixArray.end());
singleStamp.setPixel(x, y, pixArray[M]);
float x = t.x + times[i] * t.xVel;
float y = t.y + times[i] * t.yVel;
stamps.push_back(imgs[i]->createStamp(x, y, radius, false));
}
}
stamps[s] = (singleStamp);

// Compute the median of those stamps.
results[s] = createMedianImage(stamps);
}
omp_set_num_threads(1);

return(stamps);
}

std::vector<RawImage> KBMOSearch::createMedianBatch(
int radius, std::vector<RawImage*> imgs)
{
omp_set_num_threads(30);
float pixVal = 0;
float median;
int dim = radius*2+1;
size_t N;
int numResults = results.size();
const std::vector<float>& times = stack.getTimes();
trajectory t;
std::vector<RawImage> medianStamps;
medianStamps.reserve(numResults);
std::vector<float> pixArray;
pixArray.resize(imgs.size());
#pragma omp parallel for
for (int s = 0; s < numResults; ++s)
{
medianStamps[s] = RawImage(dim,dim);
t = results[s];
for (int x = 0; x < dim; ++x)
{
for (int y = 0; y < dim; ++y)
{
for (int i = 0; i < imgs.size(); ++i)
{
pixVal = imgs[i]->getPixel(
t.x + times[i] * t.xVel + static_cast<float>(x-radius),
t.y + times[i] * t.yVel + static_cast<float>(y-radius));
if (pixVal == NO_DATA) pixVal = 0.0;
pixArray[i] = pixVal;
}
N = pixArray.size() / 2;
std::nth_element(
pixArray.begin(), pixArray.begin()+N, pixArray.end());
medianStamps[s].setPixel(x, y, pixArray[N]);
}
}
}
return(medianStamps);
return(results);
}

std::vector<RawImage> KBMOSearch::createStamps(trajectory t, int radius, std::vector<RawImage*> imgs)
std::vector<RawImage> KBMOSearch::createStamps(trajectory t, int radius,
const std::vector<RawImage*>& imgs,
bool interpolate)
{
if (radius<0) throw std::runtime_error("stamp radius must be at least 0");
std::vector<RawImage> stamps;
Expand All @@ -800,7 +705,7 @@ std::vector<RawImage> KBMOSearch::createStamps(trajectory t, int radius, std::ve
{
float x_center = t.x + times[i] * t.xVel;
float y_center = t.y + times[i] * t.yVel;
stamps.push_back(imgs[i]->createStamp(x_center, y_center, radius));
stamps.push_back(imgs[i]->createStamp(x_center, y_center, radius, interpolate));
}
return stamps;
}
Expand Down Expand Up @@ -835,87 +740,97 @@ std::vector<float> KBMOSearch::createCurves(trajectory t, std::vector<RawImage*>
return lightcurve;
}

RawImage KBMOSearch::stackedStamps(trajectory t, int radius, std::vector<RawImage*> imgs)
std::vector<RawImage> KBMOSearch::psiStamps(trajRegion& t, int radius)
{
if (radius<0) throw std::runtime_error("stamp radius must be at least 0");
int dim = radius*2+1;
RawImage stamp(dim, dim);
const std::vector<float>& times = stack.getTimes();
for (int i=0; i<imgs.size(); ++i)
{
for (int x=0; x<dim; ++x)
{
for (int y=0; y<dim; ++y)
{
float pixVal = imgs[i]->getPixel(
t.x + times[i] * t.xVel + static_cast<float>(x-radius),
t.y + times[i] * t.yVel + static_cast<float>(y-radius));
if ((pixVal == NO_DATA) || isnan(pixVal)) pixVal = 0.0;
stamp.addToPixel(static_cast<int>(x),static_cast<int>(y), pixVal);
}
}
}
return stamp;
preparePsiPhi();
std::vector<RawImage*> imgs;
for (auto& im : psiImages) imgs.push_back(&im);
return createStamps(convertTraj(t), radius, imgs, true);
}

RawImage KBMOSearch::stackedScience(trajRegion& t, int radius)
std::vector<RawImage> KBMOSearch::phiStamps(trajRegion& t, int radius)
{
preparePsiPhi();
std::vector<RawImage*> imgs;
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());
return stackedStamps(convertTraj(t), radius, imgs);
for (auto& im : phiImages) imgs.push_back(&im);
return createStamps(convertTraj(t), radius, imgs, true);
}

std::vector<RawImage> KBMOSearch::scienceStamps(trajRegion& t, int radius)
std::vector<RawImage> KBMOSearch::scienceStamps(trajectory& t, int radius)
{
std::vector<RawImage*> imgs;
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());
return createStamps(convertTraj(t), radius, imgs);
return createStamps(t, radius, imgs, true);
}

std::vector<RawImage> KBMOSearch::psiStamps(trajRegion& t, int radius)
std::vector<RawImage> KBMOSearch::scienceStamps(trajRegion& t, int radius)
{
preparePsiPhi();
std::vector<RawImage*> imgs;
for (auto& im : psiImages) imgs.push_back(&im);
return createStamps(convertTraj(t), radius, imgs);
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());
return createStamps(convertTraj(t), radius, imgs, true);
}

std::vector<RawImage> KBMOSearch::phiStamps(trajRegion& t, int radius)
RawImage KBMOSearch::stackedScience(trajectory& t, int radius)
{
preparePsiPhi();
std::vector<RawImage*> imgs;
for (auto& im : phiImages) imgs.push_back(&im);
return createStamps(convertTraj(t), radius, imgs);
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());

std::vector<RawImage> stamps = createStamps(t, radius, imgs, false);
RawImage summedStamp = createSummedImage(stamps);
return summedStamp;
}

RawImage KBMOSearch::stackedScience(trajectory& t, int radius)
RawImage KBMOSearch::stackedScience(trajRegion& t, int radius)
{
std::vector<RawImage*> imgs;
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());
return stackedStamps(t, radius, imgs);

std::vector<RawImage> stamps = createStamps(convertTraj(t), radius, imgs, false);
RawImage summedStamp = createSummedImage(stamps);
return summedStamp;
}

std::vector<RawImage> KBMOSearch::scienceStamps(trajectory& t, int radius)
std::vector<RawImage> KBMOSearch::summedScience(const std::vector<trajectory>& t_array,
int radius)
{
int numResults = t_array.size();
std::vector<RawImage> results(numResults);

// Extract the science images.
std::vector<RawImage*> imgs;
for (auto& im : stack.getImages()) imgs.push_back(&im.getScience());
return createStamps(t, radius, imgs);

// Build the result for each trajectory.
omp_set_num_threads(30);
#pragma omp parallel for
for (int s = 0; s < numResults; ++s)
{
// Create stamps around the current trajectory.
trajectory t = t_array[s];
std::vector<RawImage> stamps = createStamps(t, radius, imgs, false);

// Compute the summation of those stamps.
results[s] = createSummedImage(stamps);
}
omp_set_num_threads(1);

return(results);
}

std::vector<RawImage> KBMOSearch::psiStamps(trajectory& t, int radius)
{
preparePsiPhi();
std::vector<RawImage*> imgs;
for (auto& im : psiImages) imgs.push_back(&im);
return createStamps(t, radius, imgs);
return createStamps(t, radius, imgs, true);
}

std::vector<RawImage> KBMOSearch::phiStamps(trajectory& t, int radius)
{
preparePsiPhi();
std::vector<RawImage*> imgs;
for (auto& im : phiImages) imgs.push_back(&im);
return createStamps(t, radius, imgs);
return createStamps(t, radius, imgs, true);
}

std::vector<float> KBMOSearch::psiCurves(trajectory& t)
Expand Down
23 changes: 14 additions & 9 deletions search/src/KBMOSearch.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ class KBMOSearch {
int biggestFit(int x, int y, int maxX, int maxY); // inline?
float squareSDF(float scale, float centerX, float centerY,
float pointX, float pointY);
float findExtremeInRegion(float x, float y, int size,
float findExtremeInRegion(float x, float y, int size,
PooledImage& pooledImg, int poolType);

// Converts a trajRegion result into a trajectory result.
Expand All @@ -99,19 +99,18 @@ class KBMOSearch {

// Functions to create and access stamps around proposed trajectories or
// regions. Used to visualize the results.
std::vector<RawImage> createStamps(trajectory t, int radius, std::vector<RawImage*> imgs);
std::vector<RawImage> medianStamps(std::vector<trajectory> t_array,
std::vector<std::vector<int>> goodIdx,
std::vector<RawImage> medianStamps(const std::vector<trajectory>& t_array,
const std::vector<std::vector<int>>& goodIdx,
int radius);
std::vector<RawImage> createMedianBatch(int radius, std::vector<RawImage*> imgs);
std::vector<RawImage> summedStamps(std::vector<trajectory> t_array, int radius);
RawImage stackedStamps(trajectory t, int radius, std::vector<RawImage*> imgs);

// Creates science stamps (or a summed stamp) around a
// trajectory, trajRegion, or vector of trajectories.
std::vector<RawImage> scienceStamps(trajRegion& t, int radius);
std::vector<RawImage> scienceStamps(trajectory& t, int radius);

// Creates stack science stamps around a trajectory or trajRegion.
RawImage stackedScience(trajRegion& t, int radius);
RawImage stackedScience(trajectory& t, int radius);
std::vector<RawImage> summedScience(const std::vector<trajectory>& t_array,
int radius);

// Getters for the Psi and Phi data, including pooled
// and stamped versions.
Expand Down Expand Up @@ -157,6 +156,12 @@ class KBMOSearch {
void sortResults();
std::vector<float> createCurves(trajectory t, std::vector<RawImage*> imgs);

// Functions to create and access stamps around proposed trajectories or
// regions. Used to visualize the results.
std::vector<RawImage> createStamps(trajectory t, int radius,
const std::vector<RawImage*>& imgs,
bool interpolate);

// Creates list of trajectories to search.
void createSearchList(int angleSteps, int veloctiySteps, float minAngle,
float maxAngle, float minVelocity, float maxVelocity);
Expand Down
Loading

0 comments on commit 1376d54

Please sign in to comment.