diff --git a/src/array_ops.cxx b/src/array_ops.cxx index 66bc53cd..d4dd17a4 100644 --- a/src/array_ops.cxx +++ b/src/array_ops.cxx @@ -1260,6 +1260,291 @@ void detrend(bp::object & tod, const std::string & method, const int linear_ncou } } +template +int _find_bin_index(const T* bin_edges, T value, int nbins) { + int left = 0; + int right = nbins; + + // Mimic np.clip and assign out-of-bounds points to + // first and last bin + if (value < bin_edges[left]) { + return 0; + } + else if (value >= bin_edges[right]) { + return right - 1; + } + + while (left < right) { + int mid = left + (right - left) / 2; + + if (value >= bin_edges[mid] && value < bin_edges[mid + 1]) { + return mid; + } + else if (value >= bin_edges[mid + 1]) { + left = mid + 1; + } + else { + right = mid; + } + } + + return -1; +} + +template +void _bin_signal(const bp::object & bin_by, const bp::object & signal, + const bp::object & weight, bp::object & binned_sig, + bp::object & binned_sig_sigma, bp::object & bin_counts, + const bp::object & bin_edges, const T lower, const T upper, + const int* flags_data=nullptr, + const std::vector & flags_shape={0,0}, + const int flags_stride=0) +{ + // signal + BufferWrapper signal_buf ("signal", signal, false, std::vector{-1, -1}); + if (signal_buf->strides[1] != signal_buf->itemsize) + throw ValueError_exception("Argument 'signal' must be contiguous in last axis."); + const int ndets = signal_buf->shape[0]; + const int nsamps = signal_buf->shape[1]; + T* signal_data = (T*)signal_buf->buf; + + // Check flags shape + bool is_flags_2d = false; + if (flags_data) { + if (flags_shape.size() == 2) { + is_flags_2d = true; + if (flags_shape[0] != ndets || flags_shape[1] != nsamps) { + throw ValueError_exception("2D 'flags' array has incorrect shape"); + } + } + else if (flags_shape[0] != nsamps) { + throw ValueError_exception("1D 'flags' array has incorrect shape"); + } + } + + // bin_by + BufferWrapper bin_by_buf ("bin_by", bin_by, false, std::vector{nsamps}); + if (bin_by_buf->strides[0] != bin_by_buf->itemsize) + throw ValueError_exception("Argument 'bin_by' must be a C-contiguous 1d array"); + T* bin_by_data = (T*)bin_by_buf->buf; + + // weight + BufferWrapper weight_buf_temp ("weight", weight, true); + if (weight_buf_temp->ndim == 1 && weight_buf_temp->strides[0] != weight_buf_temp->itemsize) + throw ValueError_exception("Argument 'weight' must be a C-contiguous 1d array"); + else if (weight_buf_temp->ndim == 2 && weight_buf_temp->strides[1] != weight_buf_temp->itemsize) + throw ValueError_exception("Argument 'weight' must be contiguous in last axis."); + + // Check weight dimensions + bool is_weight_2d = false; + std::vector weight_dims; + if (weight_buf_temp->ndim == 2) { + weight_dims.push_back(ndets); + is_weight_2d = true; + } + if (weight_buf_temp->ndim >= 1) { + weight_dims.push_back(nsamps); + } + + BufferWrapper weight_buf ("weight", weight, false, weight_dims); + T* weight_data = (T*)weight_buf->buf; + + // bin_edges + BufferWrapper bin_edges_buf ("bin_edges", bin_edges, false, std::vector{-1}); + if (bin_edges_buf->strides[0] != bin_edges_buf->itemsize) + throw ValueError_exception("Argument 'bin_edges' must be a C-contiguous 1d array"); + const int nbins = bin_edges_buf->shape[0] - 1; + T* bin_edges_data = (T*)bin_edges_buf->buf; + + // binned_sig + BufferWrapper binned_sig_buf ("binned_sig", binned_sig, false, std::vector{ndets, nbins}); + if (binned_sig_buf->strides[1] != binned_sig_buf->itemsize) + throw ValueError_exception("Argument 'binned_sig' must be contiguous in last axis."); + T* binned_sig_data = (T*)binned_sig_buf->buf; + + // binned_sig_sigma + BufferWrapper binned_sig_sigma_buf ("binned_sig_sigma", binned_sig_sigma, false, std::vector{ndets, nbins}); + if (binned_sig_sigma_buf->strides[1] != binned_sig_sigma_buf->itemsize) + throw ValueError_exception("Argument 'binned_sig_sigma' must be contiguous in last axis."); + T* binned_sig_sigma_data = (T*)binned_sig_sigma_buf->buf; + + // bin_counts + BufferWrapper bin_counts_buf ("bin_counts", bin_counts, false, std::vector{ndets, nbins}); + if (bin_counts_buf->strides[1] != bin_counts_buf->itemsize) + throw ValueError_exception("Argument 'bin_counts' must be contiguous in last axis."); + int* bin_counts_data = (int*)bin_counts_buf->buf; + + // Strides + int signal_stride = signal_buf->strides[0] / sizeof(T); + int weight_stride = 0; + if (is_weight_2d) { + weight_stride = weight_buf->strides[0] / sizeof(T); + } + int binned_sig_stride = binned_sig_buf->strides[0] / sizeof(T); + int binned_sig_sigma_stride = binned_sig_sigma_buf->strides[0] / sizeof(T); + int bin_counts_stride = bin_counts_buf->strides[0] / sizeof(int); + + // Map from data column to bin index + int* bin_indices = (int*) malloc(nsamps * sizeof(int)); + for (int i = 0; i < nsamps; ++i) { + bin_indices[i] = _find_bin_index(bin_edges_data, bin_by_data[i], nbins); + } + + // Populate bin_counts_data up front if no flag array given since its + // faster + if (!flags_data) { + for (int i = 0; i < nbins; ++i) { + bin_counts_data[i] = 0; + } + for (int i = 0; i < nsamps; ++i) { + if (bin_by[i] < lower || bin_by[i] > upper) + continue; + int bin = bin_indices[i]; + bin_counts_data[bin] += weight_data[i]; + } + + // Set all other detectors to first det bins if no flags + #pragma omp parallel for + for (int i = 1; i < ndets; ++i) { + int binned_ioff = i * bin_counts_stride; + int* bin_counts_row = bin_counts_data + binned_ioff; + for (int j = 0; j < nbins; ++j) { + bin_counts_row[j] = bin_counts_data[j]; + } + } + } + + T* binned_sig_sq_mean = (T*) malloc(nbins * ndets * sizeof(T)); + + #pragma omp parallel for + for (int i = 0; i < ndets; ++i) { + int ioff = i * signal_stride; + int binned_ioff = i * bin_counts_stride; + int weight_ioff = i * weight_stride; + + T* signal_row = signal_data + ioff; + T* binned_sig_row = binned_sig_data + (i * binned_sig_stride); + T* binned_sig_sq_mean_row = binned_sig_sq_mean + (i * nbins); + T* binned_sig_sigma_row = binned_sig_sigma_data + (i * binned_sig_sigma_stride); + T* weight_row = weight_data + weight_stride; + int* bin_counts_row = bin_counts_data + binned_ioff; + + // Zero out binned data + for (int j = 0; j < nbins; ++j) { + binned_sig_row[j] = 0; + binned_sig_sq_mean_row[j] = 0; + + if (flags_data) { + bin_counts_row[j] = 0; + } + } + + // Populate binned data + for (int j = 0; j < nsamps; ++j) { + bool samp_flagged = false; + if (flags_data) { + int flags_ioff = i * flags_stride; + samp_flagged = flags_data[flags_ioff + j]; + } + if (!samp_flagged) { + int bin = bin_indices[j]; + if (flags_data) { + bin_counts_row[bin] += weight_row[j]; + } + + if (bin_counts_row[bin] > 0) { + T ws = weight_row[j] * signal_row[j]; + binned_sig_row[bin] += ws; + binned_sig_sq_mean_row[bin] += ws * ws; + } + } + } + // Normalize + for (int j = 0; j < nbins; ++j) { + if (bin_counts_row[j] > 0) { + binned_sig_row[j] /= bin_counts_row[j]; + binned_sig_sq_mean_row[j] /= bin_counts_row[j]; + binned_sig_sigma_row[j] = + std::sqrt(std::abs(binned_sig_sq_mean_row[j] - + binned_sig_row[j] * binned_sig_row[j])) / + std::sqrt(bin_counts_row[j]); + } + else { + binned_sig_row[j] = std::numeric_limits::quiet_NaN(); + binned_sig_sigma_row[j] = std::numeric_limits::quiet_NaN(); + } + } + } + + free(bin_indices); + free(binned_sig_sq_mean); +} + +void bin_signal(const bp::object & bin_by, const bp::object & signal, + const bp::object & weight, bp::object & binned_sig, + bp::object & binned_sig_sigma, bp::object & bin_counts, + const bp::object & bin_edges, const double lower, + const double upper) +{ + // Get data type + int dtype = get_dtype(signal); + + if (dtype == NPY_FLOAT) { + _bin_signal(bin_by, signal, weight, binned_sig, binned_sig_sigma, + bin_counts, bin_edges, (float)lower, (float)upper); + } + else if (dtype == NPY_DOUBLE) { + _bin_signal(bin_by, signal, weight, binned_sig, binned_sig_sigma, + bin_counts, bin_edges, (double)lower, (double)upper); + } + else { + throw TypeError_exception("Only float32 or float64 arrays are supported."); + } +} + +void bin_flagged_signal(const bp::object & bin_by, const bp::object & signal, + const bp::object & weight, bp::object & binned_sig, + bp::object & binned_sig_sigma, bp::object & bin_counts, + const bp::object & bin_edges, const double lower, + const double upper, const bp::object & flags) +{ + // Get data type + int dtype = get_dtype(signal); + + // flags + BufferWrapper flags_buf ("flags", flags, false); + if (flags_buf->ndim == 1 && flags_buf->strides[0] != flags_buf->itemsize) + throw ValueError_exception("Argument 'flags' must be a C-contiguous 1d array"); + else if (flags_buf->ndim == 2 && flags_buf->strides[1] != flags_buf->itemsize) + throw ValueError_exception("Argument 'flags' must be contiguous in last axis."); + + std::vector flags_shape; + flags_shape.push_back(flags_buf->shape[0]); + + int flags_stride = 0; + if (flags_buf->ndim == 2) { + flags_shape.push_back(flags_buf->shape[1]); + flags_stride = flags_buf->strides[0] / sizeof(int); + } + + int* flags_data = (int*)flags_buf->buf; + + if (dtype == NPY_FLOAT) { + _bin_signal(bin_by, signal, weight, binned_sig, binned_sig_sigma, + bin_counts, bin_edges, (float)lower, (float)upper, + flags_data, flags_shape, flags_stride); + } + else if (dtype == NPY_DOUBLE) { + _bin_signal(bin_by, signal, weight, binned_sig, binned_sig_sigma, + bin_counts, bin_edges, (double)lower, (double)upper, + flags_data, flags_shape, flags_stride); + } + else { + throw TypeError_exception("Only float32 or float64 arrays are supported."); + } +} + + PYBINDINGS("so3g") { bp::def("nmat_detvecs_apply", nmat_detvecs_apply); @@ -1418,4 +1703,44 @@ PYBINDINGS("so3g") " linear_ncount: Number (int) of samples to use on each end, when measuring mean level for 'linear'" " detrend. Must be a positive integer or -1. If -1, nsamps / 2 will be used. Values " " larger than 1 suppress the influence of white noise.\n"); -} \ No newline at end of file + bp::def("bin_signal", bin_signal, + "bin_signal(bin_by, signal, weight, binned_sig, binned_sig_sigma, bin_counts, bin_edges, lower, upper)" + "\n" + "Bin time-ordered data by ``bin_by`` and return the binned signal and its standard deviation.\n" + "This function uses OMP to parallelize over the dets (rows) axis. Supports unequal bin widths.\n" + "Args:\n" + " bin_by: the array (float32/float64) by which signal is binned with shape (nsamp)\n" + " signal: the signal array (float32/float64) to be binned with shape (ndet,nsamp)\n" + " weight: array (float32/float64) of weights for the signal values. May have shapes\n" + " of (nsamps) or (ndets, nsamps)\n" + " binned_sig: binned signal array (float32/float64) with shape (ndet,nbin).\n" + " Modified in place.\n" + " binned_sig_sigma: estimated sigma of binned signal (float32/float64) with shape (ndet,nbin).\n" + " Modified in place." + " bin_counts: counts of binned samples (int32) with shape (ndet,nbin). Modified in place.\n" + " bin_edges: array (float32/float64) of bin edges with length=nbins+1. Must be monotonically increasing but\n" + " but may have different widths.\n" + " lower: lower bin range (float64). Data points falling outside this range will be ignored.\n" + " upper: upper bin range (float64). Data points falling outside this range will be ignored.\n"); + bp::def("bin_flagged_signal", bin_flagged_signal, + "bin_signal(bin_by, signal, weight, binned_sig, binned_sig_sigma, bin_counts, bin_edges, lower, upper, flags)\n" + "\n" + "Bin time-ordered data by ``bin_by`` and return the binned signal and its standard deviation.\n" + "This function uses OMP to parallelize over the dets (rows) axis. Supports unequal bin widths.\n" + "Args:\n" + " bin_by: the array (float32/float64) by which signal is binned with shape (nsamp)\n" + " signal: the signal array (float32/float64) to be binned with shape (ndet,nsamp)\n" + " weight: array (float32/float64) of weights for the signal values. May have shapes\n" + " of (nsamps) or (ndets, nsamps)\n" + " binned_sig: binned signal array (float32/float64) with shape (ndet,nbin).\n" + " Modified in place.\n" + " binned_sig_sigma: estimated sigma of binned signal (float32/float64) with shape (ndet,nbin).\n" + " Modified in place.\n" + " bin_counts: counts of binned samples (int32) with shape (ndet,nbin). Modified in place.\n" + " bin_edges: array (float32/float64) of bin edges with length=nbins+1. Must be monotonically increasing but\n" + " but may have different widths.\n" + " lower: lower bin range (float64). Data points falling outside this range will be ignored.\n" + " upper: upper bin range (float64). Data points falling outside this range will be ignored.\n" + " flags: array (int32) indicating whether to exclude flagged samples when binning the signal.\n" + " Can be of shape (nsamp) or (ndet,nsamp).\n"); +} diff --git a/test/test_array_ops.py b/test/test_array_ops.py index 20c6036a..e38987eb 100644 --- a/test/test_array_ops.py +++ b/test/test_array_ops.py @@ -425,5 +425,188 @@ def test_02_linear_detrending(self): np.testing.assert_allclose(signal_copy, signal, rtol=rtol, atol=atol) +class TestBinning(unittest.TestCase): + """ + Test binning. + """ + # make fake input data + def make_input(self, nsamps, ndets, x_min, x_max, dtype, order="C"): + x = np.linspace(x_min, x_max, nsamps, dtype=dtype) + signal = np.array([(i + 1) * np.sin(2*np.pi*x + i) for i in range(ndets)], + dtype=dtype, order=order) + weight = np.ones((nsamps), dtype=dtype, order=order) + + return x, signal, weight + + # allocate so3g containers + def make_so3g(self, nsamps, ndets, nbins, dtype, order): + binned_signal_so3g = np.zeros((ndets, nbins), dtype=dtype, order=order) + binned_signal_sigma_so3g = np.zeros((ndets, nbins), dtype=dtype, order=order) + bin_counts_so3g = np.zeros((ndets, nbins), dtype=np.int32, order=order) + + return binned_signal_so3g, binned_signal_sigma_so3g, bin_counts_so3g + + def get_bins(self, x, bins, ranges): + bin_edges = np.histogram_bin_edges(x, bins=bins, range=ranges,) + bin_centers = (bin_edges[1] - bin_edges[0])/2. + bin_edges[:-1] # edge to center + nbins = len(bin_centers) + + return bin_edges, nbins + + # equivalent to sotodlib binning function + def numpy_binning(self, x, signal, weight, bin_edges, nbins, ranges, flags=None): + ndets, nsamps = signal.shape + binned_signal = np.full([ndets, nbins], np.nan) + binned_signal_squared_mean = np.full([ndets, nbins], np.nan) + binned_signal_sigma = np.full([ndets, nbins], np.nan) + + # get bin indices + bin_indices = np.digitize(x, bin_edges) - 1 + bin_indices = np.clip(bin_indices, 0, nbins-1) + + if flags is None: + bin_counts, _ = np.histogram(x, bins=nbins, range=ranges, weights = weight) + mcnts = bin_counts > 0 + + for i in range(ndets): + binned_signal[i][mcnts] = np.bincount(bin_indices, weights=signal[i]*weight, minlength=nbins + )[mcnts]/bin_counts[mcnts] + binned_signal_squared_mean[i][mcnts] = np.bincount(bin_indices, weights=(signal[i]*weight)**2, minlength=nbins + )[mcnts]/bin_counts[mcnts] + + binned_signal_sigma[:, mcnts] = np.sqrt(np.abs(binned_signal_squared_mean[:,mcnts] - binned_signal[:,mcnts]**2) + ) / np.sqrt(bin_counts[mcnts]) + bin_counts_dets = np.tile(bin_counts, (ndets, 1)) + + return binned_signal, binned_signal_sigma, bin_counts_dets + + else: + bin_counts_dets = np.full([ndets, nbins], np.nan) + + if flags.shape == (ndets, nsamps): + flag_is_2d = True + m_2d = ~flags.astype(bool) + elif flags.shape == (nsamps, ): + flag_is_2d = False + m = ~flags.astype(bool) + + for i in range(ndets): + if flag_is_2d: + m = m_2d[i] + + if weight.shape == (ndets, nsamps): + weight_det = weight[i] + elif weight.shape == (nsamps, ): + weight_det = weight + + bin_counts_dets[i] = np.bincount(bin_indices[m], weights=weight_det[m], minlength=nbins) + mcnts = bin_counts_dets[i] > 0 + binned_signal[i][mcnts] = np.bincount(bin_indices[m], weights=signal[i][m]*weight_det[m], minlength=nbins + )[mcnts]/bin_counts_dets[i][mcnts] + binned_signal_squared_mean[i][mcnts] = np.bincount(bin_indices[m], weights=(signal[i][m]*weight_det[m])**2, minlength=nbins + )[mcnts]/bin_counts_dets[i][mcnts] + binned_signal_sigma[i][mcnts] = np.sqrt(np.abs(binned_signal_squared_mean[i,mcnts] - binned_signal[i,mcnts]**2) + ) / np.sqrt(bin_counts_dets[i][mcnts]) + + return binned_signal, binned_signal_sigma, bin_counts_dets + + def test_00_binning_no_flags_float64(self): + nsamps = 1000 + ndets = 3 + dtype = "float64" + order = "C" + + x_min = 0.0 + x_max = 1.0 + bins = 10 + ranges = (x_min, x_max) + + x, signal, weight = self.make_input(nsamps, ndets, x_min, x_max, dtype, order="C") + bin_edges, nbins = self.get_bins(x, bins, ranges) + # numpy binning + binned_signal, binned_signal_sigma, bin_counts_dets = self.numpy_binning(x, signal, weight, bin_edges, nbins, ranges) + + # so3g + binned_signal_so3g, binned_signal_sigma_so3g, bin_counts_so3g = self.make_so3g(nsamps, ndets, nbins, dtype, order) + + so3g.bin_signal(x, signal, weight, binned_signal_so3g, binned_signal_sigma_so3g, + bin_counts_so3g, bin_edges, ranges[0], ranges[1]) + + tolerance = 1e-10 + np.testing.assert_allclose(binned_signal_so3g, binned_signal, atol=tolerance) + np.testing.assert_allclose(binned_signal_sigma_so3g, binned_signal_sigma, atol=tolerance) + np.testing.assert_allclose(bin_counts_so3g, bin_counts_dets, atol=tolerance) + + def test_01_binning_flags_float32(self): + nsamps = 1000 + ndets = 3 + dtype = "float32" + order = "C" + + x_min = 0.0 + x_max = 1.0 + bins = 10 + ranges = (x_min, x_max) + + x, signal, weight = self.make_input(nsamps, ndets, x_min, x_max, + dtype, order="C") + + flags = np.zeros((ndets, nsamps), dtype=np.int32, order=order) + flags[0,::2] = 1 + flags[1,::3] = 1 + flags[2,::4] = 1 + + bin_edges, nbins = self.get_bins(x, bins, ranges) + # numpy binning + binned_signal, binned_signal_sigma, bin_counts_dets = self.numpy_binning(x, signal, weight, bin_edges, nbins, ranges, flags) + + # so3g + binned_signal_so3g, binned_signal_sigma_so3g, bin_counts_so3g = self.make_so3g(nsamps, ndets, nbins, dtype, order) + + so3g.bin_flagged_signal(x, signal, weight, binned_signal_so3g, + binned_signal_sigma_so3g, bin_counts_so3g, + bin_edges, ranges[0], ranges[1], flags) + + tolerance = 1e-4 + np.testing.assert_allclose(binned_signal_so3g, binned_signal, atol=tolerance) + np.testing.assert_allclose(binned_signal_sigma_so3g, binned_signal_sigma, atol=tolerance) + np.testing.assert_allclose(bin_counts_so3g, bin_counts_dets, atol=tolerance) + + def test_02_binning_flags_float64(self): + nsamps = 1000 + ndets = 3 + dtype = "float64" + order = "C" + + x_min = 0.0 + x_max = 1.0 + bins = 10 + ranges = (x_min, x_max) + + x, signal, weight = self.make_input(nsamps, ndets, x_min, x_max, + dtype, order="C") + + flags = np.zeros((ndets, nsamps), dtype=np.int32, order=order) + flags[0,::2] = 1 + flags[1,::3] = 1 + flags[2,::4] = 1 + + bin_edges, nbins = self.get_bins(x, bins, ranges) + # numpy binning + binned_signal, binned_signal_sigma, bin_counts_dets = self.numpy_binning(x, signal, weight, bin_edges, nbins, ranges, flags) + + # so3g + binned_signal_so3g, binned_signal_sigma_so3g, bin_counts_so3g = self.make_so3g(nsamps, ndets, nbins, dtype, order) + + so3g.bin_flagged_signal(x, signal, weight, binned_signal_so3g, + binned_signal_sigma_so3g, bin_counts_so3g, + bin_edges, ranges[0], ranges[1], flags) + + tolerance = 1e-10 + np.testing.assert_allclose(binned_signal_so3g, binned_signal, atol=tolerance) + np.testing.assert_allclose(binned_signal_sigma_so3g, binned_signal_sigma, atol=tolerance) + np.testing.assert_allclose(bin_counts_so3g, bin_counts_dets, atol=tolerance) + + if __name__ == "__main__": unittest.main()