diff --git a/src/kbmod/search/kernels.cu b/src/kbmod/search/kernels.cu index 4a798155..50c49775 100644 --- a/src/kbmod/search/kernels.cu +++ b/src/kbmod/search/kernels.cu @@ -23,26 +23,43 @@ namespace search { -extern "C" void device_allocate_psi_phi_array(PsiPhiArray *data) { - if (!data->cpu_array_allocated()) throw std::runtime_error("CPU data is not allocated."); - if (data->gpu_array_allocated()) throw std::runtime_error("GPU data is already allocated."); +extern "C" void device_allocate_psi_phi_arrays(PsiPhiArray *data) { + if (!data->cpu_array_allocated() || !data->cpu_time_array_allocated()) { + throw std::runtime_error("CPU data is not allocated."); + } + if (data->gpu_array_allocated() || data->gpu_time_array_allocated()) { + throw std::runtime_error("GPU data is already allocated."); + } + // Allocate space for the psi/phi data. void *device_array_ptr; checkCudaErrors(cudaMalloc((void **)&device_array_ptr, data->get_total_array_size())); checkCudaErrors(cudaMemcpy(device_array_ptr, data->get_cpu_array_ptr(), data->get_total_array_size(), cudaMemcpyHostToDevice)); data->set_gpu_array_ptr(device_array_ptr); + + // Allocate space for the times data. + float *device_times_ptr; + long unsigned time_bytes = data->get_num_times() * sizeof(float); + checkCudaErrors(cudaMalloc((void **)&device_times_ptr, time_bytes)); + checkCudaErrors( + cudaMemcpy(device_times_ptr, data->get_cpu_time_array_ptr(), time_bytes, cudaMemcpyHostToDevice)); + data->set_gpu_time_array_ptr(device_times_ptr); } -extern "C" void device_free_psi_phi_array(PsiPhiArray *data) { +extern "C" void device_free_psi_phi_arrays(PsiPhiArray *data) { if (data->gpu_array_allocated()) { checkCudaErrors(cudaFree(data->get_gpu_array_ptr())); data->set_gpu_array_ptr(nullptr); } + if (data->gpu_time_array_allocated()) { + checkCudaErrors(cudaFree(data->get_gpu_time_array_ptr())); + data->set_gpu_time_array_ptr(nullptr); + } } __host__ __device__ PsiPhi read_encoded_psi_phi(PsiPhiArrayMeta ¶ms, void *psi_phi_vect, int time, - int row, int col) { + int row, int col) { // Bounds checking. if ((row < 0) || (col < 0) || (row >= params.height) || (col >= params.width)) { return {NO_DATA, NO_DATA}; @@ -124,11 +141,12 @@ extern "C" __device__ __host__ void SigmaGFilteredIndicesCU(float *values, int n /* * Evaluate the likelihood score (as computed with from the psi and phi values) for a single - * given candidate trajectory. Modifies the trajectory in place to update the number of + * given candidate trajectory. Modifies the trajectory in place to update the number of * observations, likelihood, and flux. */ -extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_vect, float *image_times, - SearchParameters params, Trajectory *candidate) { +extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_vect, + float *image_times, SearchParameters params, + Trajectory *candidate) { // Data structures used for filtering. We fill in only what we need. float psi_array[MAX_NUM_IMAGES]; float phi_array[MAX_NUM_IMAGES]; @@ -274,12 +292,10 @@ __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_v } } -extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, float *image_times, SearchParameters params, - int num_trajectories, Trajectory *trj_to_search, int num_results, - Trajectory *best_results) { +extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, SearchParameters params, int num_trajectories, + Trajectory *trj_to_search, int num_results, Trajectory *best_results) { // Allocate Device memory Trajectory *device_tests; - float *device_img_times; Trajectory *device_search_results; // Check the hard coded maximum number of images against the num_images. @@ -288,31 +304,27 @@ extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, float *image_time throw std::runtime_error("Number of images exceeds GPU maximum."); } - // Check that the device psi_phi vector has been allocated. + // Check that the device vectors have already been allocated. if (psi_phi_array.gpu_array_allocated() == false) { throw std::runtime_error("PsiPhi data has not been created."); } + if (psi_phi_array.gpu_time_array_allocated() == false) { + throw std::runtime_error("GPU time data has not been created."); + } // Copy trajectories to search if (params.debug) { printf("Allocating GPU memory for testing grid with %i elements using %lu bytes.\n", num_trajectories, - sizeof(Trajectory) * num_trajectories); + sizeof(Trajectory) * num_trajectories); } checkCudaErrors(cudaMalloc((void **)&device_tests, sizeof(Trajectory) * num_trajectories)); checkCudaErrors(cudaMemcpy(device_tests, trj_to_search, sizeof(Trajectory) * num_trajectories, cudaMemcpyHostToDevice)); - // Copy the time vector. - if (params.debug) { - printf("Allocating GPU memory for time data using %lu bytes.\n", sizeof(float) * num_images); - } - checkCudaErrors(cudaMalloc((void **)&device_img_times, sizeof(float) * num_images)); - checkCudaErrors( - cudaMemcpy(device_img_times, image_times, sizeof(float) * num_images, cudaMemcpyHostToDevice)); - // Allocate space for the results. if (params.debug) { - printf("Allocating GPU memory for %i results using %lu bytes.\n", num_results, sizeof(Trajectory) * num_results); + printf("Allocating GPU memory for %i results using %lu bytes.\n", num_results, + sizeof(Trajectory) * num_results); } checkCudaErrors(cudaMalloc((void **)&device_search_results, sizeof(Trajectory) * num_results)); @@ -326,17 +338,16 @@ extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, float *image_time // Launch Search searchFilterImages<<>>(psi_phi_array.get_meta_data(), psi_phi_array.get_gpu_array_ptr(), - device_img_times, params, num_trajectories, device_tests, - device_search_results); + static_cast(psi_phi_array.get_gpu_time_array_ptr()), + params, num_trajectories, device_tests, device_search_results); cudaDeviceSynchronize(); // Read back results checkCudaErrors(cudaMemcpy(best_results, device_search_results, sizeof(Trajectory) * num_results, cudaMemcpyDeviceToHost)); - // Free the on GPU memory. + // Free the on GPU memory for this specific search. checkCudaErrors(cudaFree(device_search_results)); - checkCudaErrors(cudaFree(device_img_times)); checkCudaErrors(cudaFree(device_tests)); } diff --git a/src/kbmod/search/layered_image.cpp b/src/kbmod/search/layered_image.cpp index 903ec82a..ba7efbb0 100644 --- a/src/kbmod/search/layered_image.cpp +++ b/src/kbmod/search/layered_image.cpp @@ -40,12 +40,12 @@ LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawIm variance = var; } -LayeredImage::LayeredImage(unsigned w, unsigned h, float noise_stdev, float pixel_variance, - double time, const PSF& psf) +LayeredImage::LayeredImage(unsigned w, unsigned h, float noise_stdev, float pixel_variance, double time, + const PSF& psf) : LayeredImage(w, h, noise_stdev, pixel_variance, time, psf, -1) {} -LayeredImage::LayeredImage(unsigned w, unsigned h, float noise_stdev, float pixel_variance, - double time, const PSF& psf, int seed) +LayeredImage::LayeredImage(unsigned w, unsigned h, float noise_stdev, float pixel_variance, double time, + const PSF& psf, int seed) : psf(psf), width(w), height(h) { std::random_device r; std::default_random_engine generator(r()); diff --git a/src/kbmod/search/psi_phi_array.cpp b/src/kbmod/search/psi_phi_array.cpp index 3f265668..e0640e9d 100644 --- a/src/kbmod/search/psi_phi_array.cpp +++ b/src/kbmod/search/psi_phi_array.cpp @@ -6,9 +6,9 @@ namespace search { // Declaration of CUDA functions that will be linked in. #ifdef HAVE_CUDA -extern "C" void device_allocate_psi_phi_array(PsiPhiArray* data); +extern "C" void device_allocate_psi_phi_arrays(PsiPhiArray* data); -extern "C" void device_free_psi_phi_array(PsiPhiArray* data); +extern "C" void device_free_psi_phi_arrays(PsiPhiArray* data); #endif // ------------------------------------------------------- @@ -17,16 +17,7 @@ extern "C" void device_free_psi_phi_array(PsiPhiArray* data); PsiPhiArray::PsiPhiArray() {} -PsiPhiArray::~PsiPhiArray() { - if (cpu_array_ptr != nullptr) { - free(cpu_array_ptr); - } -#ifdef HAVE_CUDA - if (gpu_array_ptr != nullptr) { - device_free_psi_phi_array(this); - } -#endif -} +PsiPhiArray::~PsiPhiArray() { clear(); } void PsiPhiArray::clear() { // Free all used memory on CPU and GPU. @@ -34,10 +25,15 @@ void PsiPhiArray::clear() { free(cpu_array_ptr); cpu_array_ptr = nullptr; } + if (cpu_time_array != nullptr) { + free(cpu_time_array); + cpu_time_array = nullptr; + } #ifdef HAVE_CUDA - if (gpu_array_ptr != nullptr) { - device_free_psi_phi_array(this); + if ((gpu_array_ptr != nullptr) || (gpu_time_array != nullptr)) { + device_free_psi_phi_arrays(this); gpu_array_ptr = nullptr; + gpu_time_array = nullptr; } #endif @@ -140,6 +136,14 @@ PsiPhi PsiPhiArray::read_psi_phi(int time, int row, int col) { return result; } +float PsiPhiArray::read_time(int time_index) { + if (cpu_time_array == nullptr) throw std::runtime_error("Read from unallocated times array."); + if ((time_index < 0) || (time_index >= meta_data.num_times)) { + throw std::runtime_error("Out of bounds read for time step."); + } + return cpu_time_array[time_index]; +} + // ------------------------------------------- // --- Implementation of utility functions --- // ------------------------------------------- @@ -242,7 +246,8 @@ void set_float_cpu_psi_phi_array(PsiPhiArray& data, const std::vector& } void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector& psi_imgs, - const std::vector& phi_imgs, bool debug) { + const std::vector& phi_imgs, const std::vector zeroed_times, + bool debug) { if (result_data.get_cpu_array_ptr() != nullptr) { return; } @@ -251,6 +256,8 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect int num_times = psi_imgs.size(); if (num_times <= 0) throw std::runtime_error("Trying to fill PsiPhi from empty vectors."); if (num_times != phi_imgs.size()) throw std::runtime_error("Size mismatch between psi and phi."); + if (num_times != zeroed_times.size()) + throw std::runtime_error("Size mismatch between psi and zeroed times."); int width = phi_imgs[0].get_width(); int height = phi_imgs[0].get_height(); @@ -287,19 +294,61 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect set_float_cpu_psi_phi_array(result_data, psi_imgs, phi_imgs, debug); } + // Copy the time array. + const long unsigned times_bytes = result_data.get_num_times() * sizeof(float); + if (debug) printf("Allocating %lu bytes on the CPU for times.\n", times_bytes); + + float* times_array = (float*)malloc(times_bytes); + if (times_array == nullptr) throw std::runtime_error("Unable to allocate space for CPU times."); + for (int i = 0; i < result_data.get_num_times(); ++i) { + times_array[i] = zeroed_times[i]; + } + result_data.set_cpu_time_array_ptr(times_array); + #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 times array using %lu bytes.\n", times_bytes); } - device_allocate_psi_phi_array(&result_data); + + device_allocate_psi_phi_arrays(&result_data); if (result_data.get_gpu_array_ptr() == nullptr) { throw std::runtime_error("Unable to allocate GPU PsiPhi array."); } + if (result_data.get_gpu_time_array_ptr() == nullptr) { + throw std::runtime_error("Unable to allocate GPU time array."); + } #endif } +void fill_psi_phi_array_from_image_stack(PsiPhiArray& result_data, ImageStack& stack, int num_bytes, + bool debug) { + // Compute Phi and Psi from convolved images while leaving masked pixels alone + // Reinsert 0s for NO_DATA? + std::vector psi_images; + std::vector phi_images; + const int num_images = stack.img_count(); + if (debug) { + unsigned long num_bytes = 2 * stack.get_height() * stack.get_width() * num_images * sizeof(float); + printf("Building %i temporary %i by %i images (psi and phi), requiring %lu bytes", (num_images * 2), + stack.get_width(), stack.get_height(), num_bytes); + } + + // Build the psi and phi images first. + for (int i = 0; i < num_images; ++i) { + LayeredImage& img = stack.get_single_image(i); + psi_images.push_back(img.generate_psi_image()); + phi_images.push_back(img.generate_phi_image()); + } + + // Convert these into an array form. Needs the full psi and phi computed first so the + // encoding can compute the bounds of each array. + std::vector zeroed_times = stack.build_zeroed_times(); + fill_psi_phi_array(result_data, num_bytes, psi_images, phi_images, zeroed_times, debug); +} + // ------------------------------------------- // --- Python definitions -------------------- // ------------------------------------------- @@ -340,13 +389,20 @@ static void psi_phi_array_binding(py::module& m) { pydocs::DOC_PsiPhiArray_get_cpu_array_allocated) .def_property_readonly("gpu_array_allocated", &ppa::gpu_array_allocated, pydocs::DOC_PsiPhiArray_get_gpu_array_allocated) + .def_property_readonly("cpu_time_array_allocated", &ppa::cpu_time_array_allocated, + pydocs::DOC_PsiPhiArray_get_cpu_time_array_allocated) + .def_property_readonly("gpu_time_array_allocated", &ppa::gpu_time_array_allocated, + pydocs::DOC_PsiPhiArray_get_gpu_time_array_allocated) .def("set_meta_data", &ppa::set_meta_data, pydocs::DOC_PsiPhiArray_set_meta_data) .def("clear", &ppa::clear, pydocs::DOC_PsiPhiArray_clear) - .def("read_psi_phi", &ppa::read_psi_phi, pydocs::DOC_PsiPhiArray_read_psi_phi); + .def("read_psi_phi", &ppa::read_psi_phi, pydocs::DOC_PsiPhiArray_read_psi_phi) + .def("read_time", &ppa::read_time, pydocs::DOC_PsiPhiArray_read_time); m.def("compute_scale_params_from_image_vect", &search::compute_scale_params_from_image_vect); m.def("decode_uint_scalar", &search::decode_uint_scalar); m.def("encode_uint_scalar", &search::encode_uint_scalar); m.def("fill_psi_phi_array", &search::fill_psi_phi_array, pydocs::DOC_PsiPhiArray_fill_psi_phi_array); + m.def("fill_psi_phi_array_from_image_stack", &search::fill_psi_phi_array_from_image_stack, + pydocs::DOC_PsiPhiArray_fill_psi_phi_array_from_image_stack); } #endif diff --git a/src/kbmod/search/psi_phi_array_ds.h b/src/kbmod/search/psi_phi_array_ds.h index 719ec868..477b1704 100644 --- a/src/kbmod/search/psi_phi_array_ds.h +++ b/src/kbmod/search/psi_phi_array_ds.h @@ -1,16 +1,18 @@ /* * psi_phi_array_ds.h * - * The data structure for the interleaved psi/phi array. The the data + * The data structure for the raw data needed for the search algorith, + * including the psi/phi values and the zeroed times. The the data * structure and core functions are included in the header (and separated out * from the rest of the utility functions) to allow the CUDA files to import * only what they need. * * The data structure allocates memory on both the CPU and GPU for the - * interleaved psi/phi array and maintains ownership of the pointers - * until clear() is called or the PsiPhiArray's destructor is called. This allows - * the object to be passed repeatedly to the on-device search without reallocating - * and copying the memory on the GPU. + * arrays and maintains ownership of the pointers until clear() is called + * the object's destructor is called. This allows the object to be passed + * repeatedly to the on-device search without reallocating and copying the + * memory on the GPU. All arrays are stored as pointers (instead of vectors) + * for compatibility with CUDA. * * Created on: Dec 5, 2023 */ @@ -95,9 +97,12 @@ class PsiPhiArray { inline bool cpu_array_allocated() { return cpu_array_ptr != nullptr; } inline bool gpu_array_allocated() { return gpu_array_ptr != nullptr; } + inline bool cpu_time_array_allocated() { return cpu_time_array != nullptr; } + inline bool gpu_time_array_allocated() { return gpu_time_array != nullptr; } - // Primary getter function for interaction (read the data). - PsiPhi read_psi_phi(int time, int row, int col); + // Primary getter functions for interaction (read the data). + PsiPhi read_psi_phi(int time_index, int row, int col); + float read_time(int time_index); // Setters for the utility functions to allocate the data. void set_meta_data(int new_num_bytes, int new_num_times, int new_height, int new_width); @@ -110,12 +115,19 @@ class PsiPhiArray { inline void set_cpu_array_ptr(void* new_ptr) { cpu_array_ptr = new_ptr; } inline void set_gpu_array_ptr(void* new_ptr) { gpu_array_ptr = new_ptr; } + inline float* get_cpu_time_array_ptr() { return cpu_time_array; } + inline float* get_gpu_time_array_ptr() { return gpu_time_array; } + inline void set_cpu_time_array_ptr(float* new_ptr) { cpu_time_array = new_ptr; } + inline void set_gpu_time_array_ptr(float* new_ptr) { gpu_time_array = new_ptr; } + private: PsiPhiArrayMeta meta_data; - // Pointers the array (CPU space). + // Pointers to the arrays void* cpu_array_ptr = nullptr; void* gpu_array_ptr = nullptr; + float* cpu_time_array = nullptr; + float* gpu_time_array = nullptr; }; } /* namespace search */ diff --git a/src/kbmod/search/psi_phi_array_utils.h b/src/kbmod/search/psi_phi_array_utils.h index c363f65d..3573df26 100644 --- a/src/kbmod/search/psi_phi_array_utils.h +++ b/src/kbmod/search/psi_phi_array_utils.h @@ -17,6 +17,8 @@ #include #include "common.h" +#include "image_stack.h" +#include "layered_image.h" #include "psi_phi_array_ds.h" #include "raw_image.h" @@ -26,7 +28,11 @@ namespace search { std::array compute_scale_params_from_image_vect(const std::vector& imgs, int num_bytes); void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector& psi_imgs, - const std::vector& phi_imgs, bool debug = false); + const std::vector& phi_imgs, const std::vector zeroed_times, + bool debug = false); + +void fill_psi_phi_array_from_image_stack(PsiPhiArray& result_data, ImageStack& stack, int num_bytes, + bool debug = false); } /* namespace search */ diff --git a/src/kbmod/search/pydocs/psi_phi_array_docs.h b/src/kbmod/search/pydocs/psi_phi_array_docs.h index ffb7b009..05abe305 100644 --- a/src/kbmod/search/pydocs/psi_phi_array_docs.h +++ b/src/kbmod/search/pydocs/psi_phi_array_docs.h @@ -76,11 +76,19 @@ static const auto DOC_PsiPhiArray_get_phi_scale = R"doc( )doc"; static const auto DOC_PsiPhiArray_get_cpu_array_allocated = R"doc( - A Boolean indicating whether the cpu array exists. + A Boolean indicating whether the cpu data (psi/phi) array exists. )doc"; static const auto DOC_PsiPhiArray_get_gpu_array_allocated = R"doc( - A Boolean indicating whether the gpu array exists. + A Boolean indicating whether the gpu data (psi/phi) array exists. + )doc"; + +static const auto DOC_PsiPhiArray_get_cpu_time_array_allocated = R"doc( + A Boolean indicating whether the cpu time array exists. + )doc"; + +static const auto DOC_PsiPhiArray_get_gpu_time_array_allocated = R"doc( + A Boolean indicating whether the gpu time array exists. )doc"; static const auto DOC_PsiPhiArray_clear = R"doc( @@ -105,6 +113,20 @@ static const auto DOC_PsiPhiArray_read_psi_phi = R"doc( The pixel values. )doc"; +static const auto DOC_PsiPhiArray_read_time = R"doc( + Read a zeroed time value from the CPU array. + + Parameters + ---------- + time : `int` + The timestep to read. + + Returns + ------- + `float` + The time. + )doc"; + static const auto DOC_PsiPhiArray_set_meta_data = R"doc( Set the meta data for the array. Automatically called by fill_psi_phi_array(). @@ -134,6 +156,21 @@ static const auto DOC_PsiPhiArray_fill_psi_phi_array = R"doc( A list of psi images. phi_imgs : `list` A list of phi images. + zeroed_times : `list` + A list of floating point times starting at zero. + )doc"; + +static const auto DOC_PsiPhiArray_fill_psi_phi_array_from_image_stack = R"doc( + Fill the PsiPhiArray an ImageStack. + + Parameters + ---------- + result_data : `PsiPhiArray` + The location to store the data. + num_bytes : `int` + The type of encoding to use (1, 2, or 4). + stack : `ImageStack` + The stack of LayeredImages from which to build the psi and phi images. )doc"; } // namespace pydocs diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index eb782930..cf06968a 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -2,9 +2,8 @@ namespace search { #ifdef HAVE_CUDA -extern "C" void deviceSearchFilter(PsiPhiArray& psi_phi_data, float* image_times, SearchParameters params, - int num_trajectories, Trajectory* trj_to_search, int num_results, - Trajectory* best_results); +extern "C" void deviceSearchFilter(PsiPhiArray& psi_phi_array, SearchParameters params, int num_trajectories, + Trajectory* trj_to_search, int num_results, Trajectory* best_results); #endif StackSearch::StackSearch(ImageStack& imstack) : stack(imstack) { @@ -77,7 +76,8 @@ void StackSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ DebugTimer psi_phi_timer = DebugTimer("Creating psi/phi buffers", debug_info); prepare_psi_phi(); PsiPhiArray psi_phi_data; - fill_psi_phi_array(psi_phi_data, params.encode_num_bytes, psi_images, phi_images, debug_info); + fill_psi_phi_array(psi_phi_data, params.encode_num_bytes, psi_images, phi_images, image_times, + debug_info); psi_phi_timer.stop(); // Allocate a vector for the results. @@ -98,8 +98,8 @@ void StackSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ // Do the actual search on the GPU. DebugTimer search_timer = DebugTimer("Running search", debug_info); #ifdef HAVE_CUDA - deviceSearchFilter(psi_phi_data, image_times.data(), params, search_list.size(), search_list.data(), - max_results, results.data()); + deviceSearchFilter(psi_phi_data, params, search_list.size(), search_list.data(), max_results, + results.data()); #else throw std::runtime_error("Non-GPU search is not implemented."); #endif diff --git a/tests/test_psi_phi_array.py b/tests/test_psi_phi_array.py index 8260fffa..c261e4ed 100644 --- a/tests/test_psi_phi_array.py +++ b/tests/test_psi_phi_array.py @@ -3,7 +3,11 @@ import numpy as np from kbmod.search import ( + HAS_GPU, KB_NO_DATA, + PSF, + ImageStack, + LayeredImage, PsiPhi, PsiPhiArray, RawImage, @@ -11,6 +15,7 @@ decode_uint_scalar, encode_uint_scalar, fill_psi_phi_array, + fill_psi_phi_array_from_image_stack, ) @@ -31,6 +36,8 @@ def setUp(self): self.phi_1 = RawImage(np.full((self.height, self.width), 0.1, dtype=np.single), obs_time=1.0) self.phi_2 = RawImage(np.full((self.height, self.width), 0.2, dtype=np.single), obs_time=2.0) + self.zeroed_times = [0.0, 1.0] + def test_set_meta_data(self): arr = PsiPhiArray() self.assertEqual(arr.num_times, 0) @@ -124,7 +131,9 @@ def test_compute_scale_params_from_image_vect(self): def test_fill_psi_phi_array(self): for num_bytes in [2, 4]: arr = PsiPhiArray() - fill_psi_phi_array(arr, num_bytes, [self.psi_1, self.psi_2], [self.phi_1, self.phi_2], False) + fill_psi_phi_array( + arr, num_bytes, [self.psi_1, self.psi_2], [self.phi_1, self.phi_2], self.zeroed_times, False + ) # Check the meta data. self.assertEqual(arr.num_times, self.num_times) @@ -139,9 +148,16 @@ def test_fill_psi_phi_array(self): self.assertEqual(arr.block_size, num_bytes) self.assertEqual(arr.total_array_size, arr.num_entries * arr.block_size) - # Check that we can correctly read the values from the CPU. + # Check that we allocate the arrays self.assertTrue(arr.cpu_array_allocated) + self.assertTrue(arr.cpu_time_array_allocated) + if HAS_GPU: + self.assertTrue(arr.gpu_array_allocated) + self.assertTrue(arr.gpu_time_array_allocated) + + # Check that we can correctly read the values from the CPU. for time in range(self.num_times): + self.assertAlmostEqual(arr.read_time(time), self.zeroed_times[time]) offset = time * self.width * self.height for row in range(self.height): for col in range(self.width): @@ -152,6 +168,53 @@ def test_fill_psi_phi_array(self): # Check that the arrays are set to NULL after we clear it (memory should be freed too). arr.clear() self.assertFalse(arr.cpu_array_allocated) + self.assertFalse(arr.cpu_time_array_allocated) + if HAS_GPU: + self.assertFalse(arr.gpu_array_allocated) + self.assertFalse(arr.gpu_time_array_allocated) + + def test_fill_psi_phi_array_from_image_stack(self): + # Build a fake image stack. + num_times = 5 + width = 21 + height = 15 + images = [None] * num_times + p = PSF(1.0) + for i in range(num_times): + images[i] = LayeredImage( + width, + height, + 2.0, # noise_level + 4.0, # variance + 2.0 * i + 1.0, # time + p, + ) + im_stack = ImageStack(images) + + # Create the PsiPhiArray from the ImageStack. + arr = PsiPhiArray() + fill_psi_phi_array_from_image_stack(arr, im_stack, 4, False) + + # Check the meta data. + self.assertEqual(arr.num_times, num_times) + self.assertEqual(arr.num_bytes, 4) + self.assertEqual(arr.width, width) + self.assertEqual(arr.height, height) + self.assertEqual(arr.pixels_per_image, width * height) + self.assertEqual(arr.num_entries, 2 * arr.pixels_per_image * num_times) + self.assertEqual(arr.block_size, 4) + self.assertEqual(arr.total_array_size, arr.num_entries * arr.block_size) + + # Check that we allocated the arrays. + self.assertTrue(arr.cpu_array_allocated) + self.assertTrue(arr.cpu_time_array_allocated) + if HAS_GPU: + self.assertTrue(arr.gpu_array_allocated) + self.assertTrue(arr.gpu_time_array_allocated) + + # Since we filled the images with random data, we only test the times. + for time in range(num_times): + self.assertAlmostEqual(arr.read_time(time), 2.0 * time) if __name__ == "__main__":