Skip to content

Commit

Permalink
move entire matched filter flag to so3g
Browse files Browse the repository at this point in the history
  • Loading branch information
@skhrg
skhrg committed Jun 21, 2024
1 parent e86c073 commit e0642d3
Show file tree
Hide file tree
Showing 2 changed files with 160 additions and 97 deletions.
250 changes: 156 additions & 94 deletions src/array_ops.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -556,6 +556,34 @@ void get_gap_fill_poly(const bp::object ranges,
free(a);
}

template <typename T>
void _moment(T* data, T* output, int moment, bool central, int start, int stop)
{
int bsize = stop - start;
// Could replace the loops with boost accumulators?
T center = 0.0;
if(central == 1 | moment == 1) {
for(int si = start; si < stop; si++) {
center = center + data[si];
}
center = center / bsize;
}
T val = 0;
if(moment == 1) {
val = center;
}
else {
for(int si = start; si < stop; si++) {
val = val + pow(data[si] - center, moment);
}
val = val / bsize;
}
for(int si = start; si < stop; si++) {
output[si] = val;
}

}

template <typename T>
void _block_moment(T* tod_data, T* output, int bsize, int moment, bool central, int ndet, int nsamp, int shift)
{
Expand All @@ -564,35 +592,15 @@ void _block_moment(T* tod_data, T* output, int bsize, int moment, bool central,
for(int di = 0; di < ndet; di++)
{
int ioff = di*nsamp;
// do the the pre-shift portion
if(shift > 0){
_moment(tod_data, output, moment, central, ioff, ioff+shift);
}

for(int bi = 0; bi < nblock; bi++) {
int start = (bi * bsize) + shift;
int stop = start + bsize;
int _bsize = bsize;
if(bi == nblock - 1) {
stop = nsamp;
_bsize = stop - start;
}
// Could replace the loops with boost accumulators?
T center = 0.0;
if(central == 1 | moment == 1) {
for(int si = start; si < stop; si++) {
center = center + tod_data[ioff+si];
}
center = center / _bsize;
}
T val = 0;
if(moment == 1) {
val = center;
}
else {
for(int si = start; si < stop; si++) {
val = val + pow(tod_data[ioff+si] - center, moment);
}
val = val / _bsize;
}
for(int si = start; si < stop; si++) {
output[ioff+si] = val;
}
int stop = min(start + bsize, nsamp);
_moment(tod_data, output, moment, central, ioff+start, ioff+stop);
}
}
}
Expand All @@ -613,62 +621,149 @@ void block_moment(const bp::object & tod, const bp::object & out, int bsize, int
_block_moment(tod_data, output, bsize, moment, central, ndet, nsamp, shift);
}


void _clean_flag(int* flag_data, int width, int ndet, int nsamp)
{
int* buffer = new int[ndet * nsamp];
#pragma omp parallel for
for(int di = 0; di < ndet; di++) {
int ioff = di*nsamp;
int val = 0;
for(int si = 0; si < nsamp; si++) {
int i = ioff + si;
buffer[i] = flag_data[i];
int comp = 0;
if(i>=width) {
comp = buffer[i-width];
}
val = val + buffer[i] - comp;
if(val >= width) {
for(int j = max(1+i-width, ioff); j <= i; j++) {
flag_data[j] = 1;
}
}
else {
flag_data[i] = 0;
}
}
}
delete buffer;
}

void clean_flag(const bp::object & flag, int width)
{
BufferWrapper<int> flag_buf ("flag", flag, false, std::vector<int>{-1, -1});
int ndet = flag_buf->shape[0];
int nsamp = flag_buf->shape[1];
if (flag_buf->strides[1] != flag_buf->itemsize || flag_buf->strides[0] != flag_buf->itemsize*nsamp)
throw buffer_exception("flag must be C-contiguous along last axis");
int* flag_data = (int*)flag_buf->buf;
_clean_flag(flag_data, width, ndet, nsamp);

}


template <typename T>
void matched_jumps(const bp::object & tod, const bp::object & out, int bsize)
void _jumps_matched_filter(T* tod_data, T* output, int bsize, int shift, int ndet, int nsamp)
{
// Get the matched filter, this is basically convolving with a step
_block_moment(tod_data, output, bsize, 1, 0, ndet, nsamp, shift);
#pragma omp parallel for
for(int di = 0; di < ndet; di++) {
int ioff = di*nsamp;
T val = 0;
for(int si = 0; si < nsamp; si++) {
int i = ioff + si;
val = val + tod_data[i] - output[i];
output[i] = val;
}
}
}

template <typename T>
void matched_jumps(const bp::object & tod, const bp::object & out, const bp::object & min_size, int bsize)
{
BufferWrapper<T> tod_buf ("tod", tod, false, std::vector<int>{-1, -1});
int ndet = tod_buf->shape[0];
int nsamp = tod_buf->shape[1];
if (tod_buf->strides[1] != tod_buf->itemsize || tod_buf->strides[0] != tod_buf->itemsize*nsamp)
throw buffer_exception("tod must be C-contiguous along last axis");
T* tod_data = (T*)tod_buf->buf;
BufferWrapper<T> out_buf ("out", out, false, std::vector<int>{ndet, nsamp});
BufferWrapper<int> out_buf ("out", out, false, std::vector<int>{ndet, nsamp});
if (out_buf->strides[1] != out_buf->itemsize || out_buf->strides[0] != out_buf->itemsize*nsamp)
throw buffer_exception("out must be C-contiguous along last axis");
T* output = (T*)out_buf->buf;

// Get the matched filters, this is basically convolving with a step
_block_moment(tod_data, output, bsize, 1, 0, ndet, nsamp, 0);
#pragma omp parallel for
for(int di = 0; di < ndet; di++) {
int* output = (int*)out_buf->buf;
BufferWrapper<T> size_buf ("min_size", min_size, false, std::vector<int>{ndet});
if (size_buf->strides[0] != size_buf->itemsize)
throw buffer_exception("min_size must be C-contiguous along last axis");
T* size = (T*)size_buf->buf;
T* buffer = new T[ndet * nsamp];

int half_win = bsize / 2;
int quarter_win = bsize / 4;
float size_fac = 3.0*bsize/32.0;
float size_fac_adj = 2.0 - (4.0/bsize);

// Get the first matched filter, this is basically convolving with a step
_jumps_matched_filter(tod_data, buffer, bsize, 0, ndet, nsamp);
// Because of the shift the closest to a window edge we can be in win_size/4
// In this case the slope of the shorter segment is ~3*height/4
// So the peaks should be at least (3*win_size*min_size)/16
// We want to include peaks of that size so we use a denominator of 32
// Note that after this filtering we are left with at least win_size/4 width
# pragma omp parallel for
for(int di = 0; di < ndet; di++){
int ioff = di*nsamp;
T val = 0;
for(int si = 0; si < nsamp; si++) {
int i = ioff + si;
val = val + tod_data[i] - output[i];
output[i] = val;
T thresh = size_fac*size[di];
for(int si = 0; si < nsamp; si++){
output[ioff+si] = abs(buffer[ioff+si]) > thresh;
}
}
// Clean spurs
_clean_flag(output, quarter_win, ndet, nsamp);
// Recall that we set _min_size to be half the actual peak min above
// We allow for .5 samples worth of uncertainty here
# pragma omp parallel for
for(int di = 0; di < ndet; di++){
int ioff = di*nsamp;
T thresh = size_fac_adj*size_fac*size[di];
for(int si = 0; si < nsamp; si++){
output[ioff+si] = (abs(buffer[ioff+si]) > thresh) && (output[ioff+si] == 1);
}
}

// Now do the shifted filter
int half_win = bsize / 2;
T* buffer = new T[ndet * nsamp];
_block_moment(tod_data, buffer, bsize, 1, 0, ndet, nsamp, half_win);
#pragma omp parallel for
for(int di = 0; di < ndet; di++) {
_jumps_matched_filter(tod_data, buffer, bsize, half_win, ndet, nsamp);
int* shift_flag = new int[ndet * nsamp];
# pragma omp parallel for
for(int di = 0; di < ndet; di++){
int ioff = di*nsamp;
for(int si = 0; si < half_win; si++) {
int i = ioff + si;
buffer[i] = 0.;
T thresh = size_fac*size[di];
for(int si = 0; si < nsamp; si++){
shift_flag[ioff+si] = abs(buffer[ioff+si]) > thresh;
}
T val = 0;
for(int si = half_win; si < nsamp; si++) {
int i = ioff + si;
val = val + tod_data[i] - buffer[i];
buffer[i] = val;
}
_clean_flag(shift_flag, quarter_win, ndet, nsamp);
# pragma omp parallel for
for(int di = 0; di < ndet; di++){
int ioff = di*nsamp;
T thresh = size_fac_adj*size_fac*size[di];
for(int si = 0; si < nsamp; si++){
shift_flag[ioff+si] = (abs(buffer[ioff+si]) > thresh) && (shift_flag[ioff+si] == 1);
}
}
delete buffer;

// Now we combine
#pragma omp parallel for
for(int di = 0; di < ndet; di++) {
int ioff = di*nsamp;
for(int si = 0; si < nsamp; si++) {
int i = ioff + si;
output[i] = max(abs(output[i]), abs(buffer[i]));
output[i] = output[i] || shift_flag[i];
}
}

delete shift_flag;
}

template <typename T>
Expand Down Expand Up @@ -715,40 +810,6 @@ void scale_jumps(const bp::object & tod, const bp::object & out, const bp::objec
}
}

void clean_flag(const bp::object & flag, int width)
{
BufferWrapper<int> flag_buf ("flag", flag, false, std::vector<int>{-1, -1});
int ndet = flag_buf->shape[0];
int nsamp = flag_buf->shape[1];
if (flag_buf->strides[1] != flag_buf->itemsize || flag_buf->strides[0] != flag_buf->itemsize*nsamp)
throw buffer_exception("flag must be C-contiguous along last axis");
int* flag_data = (int*)flag_buf->buf;

int* buffer = new int[ndet * nsamp];
#pragma omp parallel for
for(int di = 0; di < ndet; di++) {
int ioff = di*nsamp;
int val = 0;
for(int si = 0; si < nsamp; si++) {
int i = ioff + si;
buffer[i] = flag_data[i];
int comp = 0;
if(i>=width) {
comp = buffer[i-width];
}
val = val + buffer[i] - comp;
if(val >= width) {
for(int j = max(1+i-width, ioff); j <= i; j++) {
flag_data[j] = 1;
}
}
else {
flag_data[i] = 0;
}
}
}
}


PYBINDINGS("so3g")
{
Expand Down Expand Up @@ -793,18 +854,19 @@ PYBINDINGS("so3g")
"\n"
"See details in docstring for block_moment.\n");
bp::def("matched_jumps", matched_jumps<float>,
"matched_jumps(tod, out, bsize)\n"
"matched_jumps(tod, out, min_size, bsize)\n"
"\n"
"Compute the matched filter for a unit jump in a float32 array.\n"
"Flag jumps with the matched filter for a unit jump in a float32 array.\n"
"\n"
"Args:\n"
" tod: data array (float32) with shape (ndet, nsamp)\n"
" out: output array (float32) with shape (ndet, nsamp)\n"
" out: output array (int32) with shape (ndet, nsamp)\n"
" min_size: minimum jump size for each det, shape (ndet,)\n"
" bsize: number of samples in each block\n");
bp::def("matched_jumps64", matched_jumps<double>,
"matched_jumps64(tod, out, bsize)\n"
"matched_jumps64(tod, out, min_size, bsize)\n"
"\n"
"Compute the matched filter for a unit jump in a float64 array.\n"
"Flag jumps with the matched filter for a unit jump in a float64 array.\n"
"\n"
"See details in docstring for matched_jumps.\n");
bp::def("scale_jumps", scale_jumps<float>,
Expand Down
7 changes: 4 additions & 3 deletions test/test_array_ops.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,10 +92,11 @@ def test_00_clean_flag(self):
def test_01_matched_jump(self):
a = np.zeros((1, 100), dtype="float32", order="C")
a[0, 50:] += 1
b = np.zeros((1, 100), dtype="float32", order="C")
so3g.matched_jumps(a, b, 10)
b = np.zeros((1, 100), dtype="int32", order="C")
min_size = np.ones(1, dtype="float32", order="C")*.5
so3g.matched_jumps(a, b, min_size, 10)
flagged = np.where(b)[1]
np.testing.assert_array_equal(flagged, np.arange(45, 54))
np.testing.assert_array_equal(flagged, np.arange(46, 52))

def test_02_scale_jump(self):
a = np.zeros((1, 100), dtype="float32", order="C")
Expand Down

0 comments on commit e0642d3

Please sign in to comment.