diff --git a/README.md b/README.md index 31983e377..4b77a270f 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,7 @@ import kbmod.search as kb import numpy as np # Create a point spread function -psf = kb.psf(1.5) +psf = kb.PSF(1.5) # Create fake data with ten 512x512 pixel images. from kbmod.fake_data_creator import * @@ -120,10 +120,10 @@ for im in imgs: ) # Create a new image stack with the inserted object. -stack = kb.image_stack(imgs) +stack = kb.ImageStack(imgs) # Recover the object by searching a set of trajectories. -search = kb.stack_search(stack) +search = kb.StackSearch(stack) search.search( 5, # Number of search angles to try (-0.1, -0.05, 0.0, 0.05, 0.1) 5, # Number of search velocities to try (0, 1, 2, 3, 4) diff --git a/docs/source/project_details/release_notes.rst b/docs/source/project_details/release_notes.rst index 433e7b49d..6d2543d3f 100644 --- a/docs/source/project_details/release_notes.rst +++ b/docs/source/project_details/release_notes.rst @@ -54,8 +54,8 @@ Version 1.0.0 represents a full refactoring and redesign of KBMOD, including maj **Efficiency Improvements** * Reduce internal copies (`115 `_, `119 `_) -* Speed up growMask function (`129 `_) -* Move growMask to GPU (`153 `_) +* Speed up grow_mask function (`129 `_) +* Move grow_mask to GPU (`153 `_) * Improve handling of single/multi-threading in post processing filtering (`155 `_) * Skip masking functions if there are no masking keys (`164 `_) * Move coadded stamp generation to GPU (`179 `_, `189 `_) @@ -68,11 +68,11 @@ Version 1.0.0 represents a full refactoring and redesign of KBMOD, including maj * Fix unwanted append when saving layered images (`136 `_) * Fix median computation with masked images (`137 `_) * Drop masked pixels before conducting sigmaG filtering on GPU (`142 `_, `180 `_, `181 `_) -* Reset global mask when `createGlobalMask` is created ((`164 `_) +* Reset global mask when `create_global_mask` is created ((`164 `_) * Correctly apply `mask_threshold` (`164 `_) * Account for masked pixels in `createAveTemplate` and `simpleDifference` (`164 `_) * Fix bugs in on-GPU stamp generation (`182 `_) -* Add bounds checking in `getResults` (`192 `_) +* Add bounds checking in `get_results` (`192 `_) * Check for divide by zero in clustering function (`215 `_) * Update pybind version (`217 `_) * Ignore invalidly named files during file loading (`233 `_) diff --git a/docs/source/user_manual/index.rst b/docs/source/user_manual/index.rst index afb0a7a48..ee95be9a6 100644 --- a/docs/source/user_manual/index.rst +++ b/docs/source/user_manual/index.rst @@ -77,16 +77,16 @@ KBMOD uses an hierarchy of three nested data structures to store the image data ImageStack __________ -The :py:class:`~kbmod.search.image_stack` holds all of the image data for every time step. The main information stored is the images array, which holds one :py:class:`~kbmod.search.layered_image` structure for each time step. The ImageStack also stores information that applies for all images such as a ``globalMask`` and an ``avgTemplate``. +The :py:class:`~kbmod.search.ImageStack` holds all of the image data for every time step. The main information stored is the images array, which holds one :py:class:`~kbmod.search.LayeredImage` structure for each time step. The ImageStack also stores information that applies for all images such as a ``globalMask`` and an ``avgTemplate``. LayeredImage ____________ -Each layered image holds the data from a single exposure, which consists of multiple layers including: a science image (the flux values), the variance image (representing the noise at each pixel), and MaskImage (representing per-pixel errors). In addition the :py:class:`~kbmod.search.layered_image` tracks per-exposure information such as the PSF for the image and the time at which the image was taken. +Each layered image holds the data from a single exposure, which consists of multiple layers including: a science image (the flux values), the variance image (representing the noise at each pixel), and MaskImage (representing per-pixel errors). In addition the :py:class:`~kbmod.search.LayeredImage` tracks per-exposure information such as the PSF for the image and the time at which the image was taken. RawImages _________ -A :py:class:`~kbmod.search.raw_image` is the lowest level of data storage and effectively consists of a two-dimensional array of floating point values. These values can take on a variety of meanings depending on the use, including flux values, variance values, mask indicators, psi values, and phi values. +A :py:class:`~kbmod.search.RawImage` is the lowest level of data storage and effectively consists of a two-dimensional array of floating point values. These values can take on a variety of meanings depending on the use, including flux values, variance values, mask indicators, psi values, and phi values. diff --git a/notebooks/Kbmod_Reference.ipynb b/notebooks/Kbmod_Reference.ipynb index 3edf933e7..fecab3de6 100644 --- a/notebooks/Kbmod_Reference.ipynb +++ b/notebooks/Kbmod_Reference.ipynb @@ -53,20 +53,20 @@ "\n", "### [psf](#psf) \n", "2D Point Spread Function Array \n", - "### [raw_image](#raw)\n", + "### [RawImage](#raw)\n", "2D Image array \n", "\n", - "### [layered_image](#layered) \n", - "A Complete image represented as 3 raw_image layers (science, mask, variance) \n", + "### [LayeredImage](#layered) \n", + "A Complete image represented as 3 RawImage layers (science, mask, variance) \n", "\n", - "### [image_stack](#stack) \n", - "Stack of layered_images, intended to be the same frame captured at different times\n", + "### [ImageStack](#stack) \n", + "Stack of LayeredImages, intended to be the same frame captured at different times\n", "\n", - "### [stack_search](#search) \n", - "Searches an image_stack for a moving psf\n", + "### [StackSearch](#search) \n", + "Searches an ImageStack for a moving psf\n", "\n", "### [trajectory](#traj)\n", - "Stores an object's position and motion through an image_stack\n" + "Stores an object's position and motion through an ImageStack\n" ] }, { @@ -86,7 +86,7 @@ "metadata": {}, "outputs": [], "source": [ - "p = kb.psf(1.0)" + "p = kb.PSF(1.0)" ] }, { @@ -119,7 +119,7 @@ "outputs": [], "source": [ "arr = np.linspace(0.0, 1.0, 9).reshape(3, 3)\n", - "p2 = kb.psf(arr) # initialized from array\n", + "p2 = kb.PSF(arr) # initialized from array\n", "arr = np.square(arr)\n", "p2.set_array(arr) # set from array\n", "np.array(p2)" @@ -151,13 +151,13 @@ "metadata": {}, "source": [ "\n", - "# layered_image\n", - "Stores the science, mask, and variance image for a single image. The \"layered\" means it contains all of them together. The layered_image also stores auxiliary data, including the time of the image and the image’s PSF.\n", + "# LayeredImage\n", + "Stores the science, mask, and variance image for a single image. The \"layered\" means it contains all of them together. The LayeredImage also stores auxiliary data, including the time of the image and the image’s PSF.\n", "\n", - "A layered_image can be initialized 2 ways: \n", + "A LayeredImage can be initialized 2 ways: \n", "\n", "### A. Load a file for kbmod reference:\n", - "The layered_image is loaded given the path and filename to the FITS file as well as the PSF for the image. " + "The LayeredImage is loaded given the path and filename to the FITS file as well as the PSF for the image. " ] }, { @@ -166,7 +166,7 @@ "metadata": {}, "outputs": [], "source": [ - "im = kb.layered_image(im_path + \"000000.fits\", p)\n", + "im = kb.LayeredImage(im_path + \"000000.fits\", p)\n", "print(f\"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_obstime()}\")" ] }, @@ -185,7 +185,7 @@ "metadata": {}, "outputs": [], "source": [ - "im = kb.layered_image(\"image2\", 100, 100, 5.0, 25.0, 0.0, p)\n", + "im = kb.LayeredImage(\"image2\", 100, 100, 5.0, 25.0, 0.0, p)\n", "# name, width, height, background_noise_sigma, variance, capture_time, PSF" ] }, @@ -193,7 +193,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "You can access a variety of information from the layered_image object" + "You can access a variety of information from the LayeredImage object" ] }, { @@ -229,7 +229,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "A layered_image can have its layers set from any numpy array" + "A LayeredImage can have its layers set from any numpy array" ] }, { @@ -238,7 +238,7 @@ "metadata": {}, "outputs": [], "source": [ - "raw = kb.raw_image(np.ones_like(pixels))" + "raw = kb.RawImage(np.ones_like(pixels))" ] }, { @@ -281,7 +281,7 @@ "source": [ "### Inserting artificial objects\n", "\n", - "Artificial objects can easily be added into a layered_image. The layered_image generates a point observation at the given pixel and applies the image's PSF." + "Artificial objects can easily be added into a LayeredImage. The LayeredImage generates a point observation at the given pixel and applies the image's PSF." ] }, { @@ -338,8 +338,8 @@ "metadata": {}, "source": [ "\n", - "# image_stack\n", - "A collection of layered_images (usually at different times). Used to apply operations to a group of images. " + "# ImageStack\n", + "A collection of LayeredImages (usually at different times). Used to apply operations to a group of images. " ] }, { @@ -350,8 +350,8 @@ "source": [ "# Create a stack with 10 50x50 images with random noise and times ranging from 0 to 1\n", "count = 10\n", - "imlist = [kb.layered_image(\"img\" + str(n), 100, 100, 10.0, 5.0, n / count, p) for n in range(count)]\n", - "stack = kb.image_stack(imlist)" + "imlist = [kb.LayeredImage(\"img\" + str(n), 100, 100, 10.0, 5.0, n / count, p) for n in range(count)]\n", + "stack = kb.ImageStack(imlist)" ] }, { @@ -396,7 +396,7 @@ " all_psfs = [p for _ in range(len(files))]\n", "\n", " # Load the images.\n", - " stack = kb.image_stack(files, all_psfs)\n", + " stack = kb.ImageStack(files, all_psfs)\n", "else:\n", " print(\"Cannot find data directory. Using fake images.\")" ] @@ -442,7 +442,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Most features of the layered_image can be used on the whole stack" + "Most features of the LayeredImage can be used on the whole stack" ] }, { @@ -463,7 +463,7 @@ "print(f\"Height = {stack.get_height()}\")\n", "print(f\"Pixels Per Image = {stack.get_npixels()}\")\n", "\n", - "# Retrieve a list of layered_images back from the stack.\n", + "# Retrieve a list of LayeredImages back from the stack.\n", "stack.get_images()\n", "\n", "# Get the list of image time stamps.\n", @@ -483,17 +483,17 @@ "metadata": {}, "outputs": [], "source": [ - "# Get the individual layered_images.\n", + "# Get the individual LayeredImages.\n", "im_list = stack.get_images()\n", "\n", - "# Create a new list of layered_images with the added object.\n", + "# Create a new list of LayeredImages with the added object.\n", "new_im_list = []\n", "for im, time in zip(im_list, stack.get_times()):\n", " im.add_object(20.0 + (time * 8.0), 35.0 + (time * 0.0), 25000.0)\n", " new_im_list.append(im)\n", "\n", - "# Save these images in a new image_stack.\n", - "stack = kb.image_stack(new_im_list)" + "# Save these images in a new ImageStack.\n", + "stack = kb.ImageStack(new_im_list)" ] }, { @@ -501,7 +501,7 @@ "metadata": {}, "source": [ "\n", - "# stack_search\n", + "# StackSearch\n", "\n", "We can create a search object that will compute auxiliary data for the images and run the search algorithms." ] @@ -512,7 +512,7 @@ "metadata": {}, "outputs": [], "source": [ - "search = kb.stack_search(stack)" + "search = kb.StackSearch(stack)" ] }, { @@ -628,8 +628,8 @@ "print(f\"Likelihood = {best.lh}\")\n", "print(f\"x = {best.x}\")\n", "print(f\"y = {best.y}\")\n", - "print(f\"x_v = {best.x_v}\")\n", - "print(f\"y_v = {best.y_v}\")" + "print(f\"x_v = {best.vx}\")\n", + "print(f\"y_v = {best.vy}\")" ] }, { diff --git a/notebooks/create_fake_data.ipynb b/notebooks/create_fake_data.ipynb index 824ee1f01..6973c51e2 100644 --- a/notebooks/create_fake_data.ipynb +++ b/notebooks/create_fake_data.ipynb @@ -67,7 +67,7 @@ ], "source": [ "trj = ds.insert_random_object(500)\n", - "print(f\"x={trj.x}, y={trj.y}, xv={trj.x_v}, yv={trj.y_v}\")" + "print(f\"x={trj.x}, y={trj.y}, xv={trj.vx}, yv={trj.vy}\")" ] }, { @@ -115,8 +115,8 @@ "for i in range(ds.stack.img_count()):\n", " ti = ds.stack.get_single_image(i).get_obstime()\n", " dt = ti - t0\n", - " px = int(trj.x + dt * trj.x_v + 0.5)\n", - " py = int(trj.y + dt * trj.y_v + 0.5)\n", + " px = int(trj.x + dt * trj.vx + 0.5)\n", + " py = int(trj.y + dt * trj.vy + 0.5)\n", "\n", " print(f\"{i}: t={ti:.3f} at ({px}, {py})\")" ] diff --git a/notebooks/kbmod_visualize.ipynb b/notebooks/kbmod_visualize.ipynb index 7c6881191..f04a56be8 100644 --- a/notebooks/kbmod_visualize.ipynb +++ b/notebooks/kbmod_visualize.ipynb @@ -43,7 +43,7 @@ "# Loading data for visualization demo\n", "\n", "### A. Load a file for visualization:\n", - "The layered_image is loaded given the path and filename to the FITS file as well as the PSF for the image. We use a default psf." + "The LayeredImage is loaded given the path and filename to the FITS file as well as the PSF for the image. We use a default psf." ] }, { @@ -60,8 +60,8 @@ } ], "source": [ - "p = kb.psf(1.0)\n", - "im = kb.layered_image(im_path + \"000000.fits\", p)\n", + "p = kb.PSF(1.0)\n", + "im = kb.LayeredImage(im_path + \"000000.fits\", p)\n", "print(f\"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_obstime()}\")" ] }, @@ -109,7 +109,7 @@ "metadata": {}, "source": [ "### B. Load a stack of images\n", - "A load collection of layered_images at different times. " + "A load collection of LayeredImages at different times. " ] }, { @@ -134,7 +134,7 @@ "all_psfs = [p for _ in range(len(files))]\n", "\n", "# Load the images.\n", - "stack = kb.image_stack(files, all_psfs)\n", + "stack = kb.ImageStack(files, all_psfs)\n", "\n", "num_images = stack.img_count()\n", "print(f\"Loaded {num_images} images.\")" @@ -193,17 +193,17 @@ "outputs": [], "source": [ "# Create the trajectory with a given parameters and then the trajectory result.\n", - "trj = kb.trajectory()\n", + "trj = kb.Trajectory()\n", "trj.x = 11\n", "trj.y = 27\n", - "trj.x_v = 16.0\n", - "trj.y_v = 3.3\n", + "trj.vx = 16.0\n", + "trj.vy = 3.3\n", "\n", "# Create the search stack.\n", - "search = kb.stack_search(stack)\n", + "search = kb.StackSearch(stack)\n", "\n", "# Create the stamps around this trajectory.\n", - "stamps = search.science_viz_stamps(trj, 20)" + "stamps = search.get_stamps(trj, 20)" ] }, { @@ -278,13 +278,13 @@ "source": [ "fig, axs = plt.subplots(1, 3)\n", "\n", - "axs[0].imshow(search.summed_sci_stamp(trj, 10, []), cmap=\"gray\")\n", + "axs[0].imshow(search.get_summed_stamp(trj, 10, []), cmap=\"gray\")\n", "axs[0].set_title(\"Summed\")\n", "\n", - "axs[1].imshow(search.mean_sci_stamp(trj, 10, []), cmap=\"gray\")\n", + "axs[1].imshow(search.get_mean_stamp(trj, 10, []), cmap=\"gray\")\n", "axs[1].set_title(\"Mean\")\n", "\n", - "axs[2].imshow(search.median_sci_stamp(trj, 10, []), cmap=\"gray\")\n", + "axs[2].imshow(search.get_median_stamp(trj, 10, []), cmap=\"gray\")\n", "axs[2].set_title(\"Median\")" ] }, diff --git a/src/kbmod/analysis_utils.py b/src/kbmod/analysis_utils.py index 86c1cc345..153a24ab6 100644 --- a/src/kbmod/analysis_utils.py +++ b/src/kbmod/analysis_utils.py @@ -52,7 +52,7 @@ def load_images( Returns ------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack of images loaded. img_info : `ImageInfo` The information for the images loaded. @@ -124,12 +124,12 @@ def load_images( # Check if the image has a specific PSF. psf = default_psf if header_info.visit_id in image_psf_dict: - psf = kb.psf(image_psf_dict[header_info.visit_id]) + psf = kb.PSF(image_psf_dict[header_info.visit_id]) # Load the image file and set its time. if verbose: print(f"Loading file: {full_file_path}") - img = kb.layered_image(full_file_path, psf) + img = kb.LayeredImage(full_file_path, psf) img.set_obstime(time_stamp) # Save the file, time, and image information. @@ -138,7 +138,7 @@ def load_images( images.append(img) print(f"Loaded {len(images)} images") - stack = kb.image_stack(images) + stack = kb.ImageStack(images) # Create a list of visit times and visit times shifted to 0.0. img_info.set_times_mjd(np.array(visit_times)) @@ -224,8 +224,8 @@ def load_and_filter_results( break if trj.lh < max_lh: row = ResultRow(trj, len(self._mjds)) - psi_curve = np.array(search.psi_curves(trj)) - phi_curve = np.array(search.phi_curves(trj)) + psi_curve = np.array(search.get_psi_curves(trj)) + phi_curve = np.array(search.get_phi_curves(trj)) row.set_psi_phi(psi_curve, phi_curve) result_batch.append_result(row) total_count += 1 @@ -251,14 +251,14 @@ def get_all_stamps(self, result_list, search, stamp_radius): ---------- result_list : `ResultList` The values from trajectories. The stamps are inserted into this data structure. - search : `kbmod.stack_search` + search : `kbmod.StackSearch` The search object stamp_radius : int The radius of the stamps to create. """ stamp_edge = stamp_radius * 2 + 1 for row in result_list.results: - stamps = search.science_viz_stamps(row.trajectory, stamp_radius) + stamps = search.get_stamps(row.trajectory, stamp_radius) row.all_stamps = np.array([np.array(stamp).reshape(stamp_edge, stamp_edge) for stamp in stamps]) def apply_clipped_sigmaG(self, result_list): @@ -402,7 +402,7 @@ def apply_stamp_filter( result_list : `ResultList` The values from trajectories. This data gets modified directly by the filtering. - search : `kbmod.stack_search` + search : `kbmod.StackSearch` The search object. center_thresh : float The fraction of the total flux that must be contained in a single @@ -422,17 +422,17 @@ def apply_stamp_filter( The radius of the stamp. """ # Set the stamp creation and filtering parameters. - params = kb.stamp_parameters() + params = kb.StampParameters() params.radius = stamp_radius params.do_filtering = True params.center_thresh = center_thresh params.peak_offset_x = peak_offset[0] params.peak_offset_y = peak_offset[1] - params.m20 = mom_lims[0] - params.m02 = mom_lims[1] - params.m11 = mom_lims[2] - params.m10 = mom_lims[3] - params.m01 = mom_lims[4] + params.m20_limit = mom_lims[0] + params.m02_limit = mom_lims[1] + params.m11_limit = mom_lims[2] + params.m10_limit = mom_lims[3] + params.m01_limit = mom_lims[4] if stamp_type == "cpp_median" or stamp_type == "median": params.stamp_type = kb.StampType.STAMP_MEDIAN @@ -472,7 +472,7 @@ 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.coadded_stamps( + stamps_slice = search.get_coadded_stamps( trj_slice, bool_slice, params, kb.HAS_GPU and len(trj_slice) > 100 ) for ind, stamp in enumerate(stamps_slice): diff --git a/src/kbmod/fake_data_creator.py b/src/kbmod/fake_data_creator.py index 3d50c7902..9f2cb16ea 100644 --- a/src/kbmod/fake_data_creator.py +++ b/src/kbmod/fake_data_creator.py @@ -31,7 +31,7 @@ def add_fake_object(img, x, y, flux, psf=None): psf : PointSpreadFunc The PSF for the image. """ - if type(img) is layered_image: + if type(img) is LayeredImage: sci = img.get_science() else: sci = img @@ -99,20 +99,20 @@ def __init__(self, width, height, num_times, noise_level=2.0, psf_val=0.5, obs_p day_num += 1 # Make the image stack. - self.stack = self.make_fake_image_stack() + self.stack = self.make_fake_ImageStack() - def make_fake_image_stack(self): + def make_fake_ImageStack(self): """Make a stack of fake layered images. Returns ------- - stack : image_stack + stack : ImageStack """ - p = psf(self.psf_val) + p = PSF(self.psf_val) image_list = [] for i in range(self.num_times): - img = layered_image( + img = LayeredImage( ("%06i" % i), self.width, self.height, @@ -124,7 +124,7 @@ def make_fake_image_stack(self): ) image_list.append(img) - stack = image_stack(image_list) + stack = ImageStack(image_list) return stack def insert_object(self, trj): @@ -139,8 +139,8 @@ def insert_object(self, trj): for i in range(self.num_times): dt = self.times[i] - t0 - px = trj.x + dt * trj.x_v + 0.5 - py = trj.y + dt * trj.y_v + 0.5 + px = trj.x + dt * trj.vx + 0.5 + py = trj.y + dt * trj.vy + 0.5 # Get the image for the timestep, add the object, and # re-set the image. This last step needs to be done @@ -168,13 +168,13 @@ def insert_random_object(self, flux): dt = self.times[-1] - self.times[0] # Create the random trajectory. - t = trajectory() + t = Trajectory() t.x = int(random.random() * self.width) xe = int(random.random() * self.width) - t.x_v = (xe - t.x) / dt + t.vx = (xe - t.x) / dt t.y = int(random.random() * self.height) ye = int(random.random() * self.height) - t.y_v = (ye - t.y) / dt + t.vy = (ye - t.y) / dt t.flux = flux # Insert the object. diff --git a/src/kbmod/file_utils.py b/src/kbmod/file_utils.py index b638944a4..5e0df4fb3 100644 --- a/src/kbmod/file_utils.py +++ b/src/kbmod/file_utils.py @@ -220,11 +220,11 @@ def trajectory_from_np_object(result): trj : trajectory The corresponding trajectory object. """ - trj = kb.trajectory() + trj = kb.Trajectory() trj.x = int(result["x"]) trj.y = int(result["y"]) - trj.x_v = float(result["vx"]) - trj.y_v = float(result["vy"]) + trj.vx = float(result["vx"]) + trj.vy = float(result["vy"]) trj.flux = float(result["flux"]) trj.lh = float(result["lh"]) trj.obs_count = int(result["num_obs"]) diff --git a/src/kbmod/filters/clustering_filters.py b/src/kbmod/filters/clustering_filters.py index e9ba2d7f3..435596126 100644 --- a/src/kbmod/filters/clustering_filters.py +++ b/src/kbmod/filters/clustering_filters.py @@ -72,8 +72,8 @@ def keep_indices(self, result_list: ResultList): # Create arrays of each the trajectories information. x_arr = np.array([row.trajectory.x for row in result_list.results]) y_arr = np.array([row.trajectory.y for row in result_list.results]) - vx_arr = np.array([row.trajectory.x_v for row in result_list.results]) - vy_arr = np.array([row.trajectory.y_v for row in result_list.results]) + vx_arr = np.array([row.trajectory.vx for row in result_list.results]) + vy_arr = np.array([row.trajectory.vy for row in result_list.results]) vel_arr = np.sqrt(np.square(vx_arr) + np.square(vy_arr)) ang_arr = np.arctan2(vy_arr, vx_arr) diff --git a/src/kbmod/filters/stamp_filters.py b/src/kbmod/filters/stamp_filters.py index a1c4b4d59..6d803b99b 100644 --- a/src/kbmod/filters/stamp_filters.py +++ b/src/kbmod/filters/stamp_filters.py @@ -7,7 +7,7 @@ import abc from kbmod.result_list import ResultRow -from kbmod.search import KB_NO_DATA, raw_image +from kbmod.search import KB_NO_DATA, RawImage class BaseStampFilter(abc.ABC): @@ -104,7 +104,7 @@ def keep_row(self, row: ResultRow): # Find the peak in the image. stamp = row.stamp.reshape([self.width, self.width]) - peak_pos = raw_image(stamp).find_peak(True) + peak_pos = RawImage(stamp).find_peak(True) return ( abs(peak_pos.x - self.stamp_radius) < self.x_thresh and abs(peak_pos.y - self.stamp_radius) < self.y_thresh @@ -180,7 +180,7 @@ def keep_row(self, row: ResultRow): # Find the peack in the image. stamp = row.stamp.reshape([self.width, self.width]) - moments = raw_image(stamp).find_central_moments() + moments = RawImage(stamp).find_central_moments() return ( (abs(moments.m01) < self.m01_thresh) and (abs(moments.m10) < self.m10_thresh) diff --git a/src/kbmod/image_info.py b/src/kbmod/image_info.py index 22f78f775..926acb031 100644 --- a/src/kbmod/image_info.py +++ b/src/kbmod/image_info.py @@ -9,7 +9,7 @@ from astropy.wcs import WCS from kbmod.file_utils import FileUtils -from kbmod.search import pixel_pos, layered_image +from kbmod.search import PixelPos, LayeredImage # ImageInfo is a helper class that wraps basic data extracted from a @@ -34,7 +34,7 @@ def populate_from_fits_file(self, filename, load_image=False, p=None): The path and name of the FITS file. load_image : bool Load the image data into a LayeredImage object. - p : `psf` + p : `PSF` The PSF for this layered image. Optional when load `load_image` is False. Otherwise this is required. """ @@ -46,7 +46,7 @@ def populate_from_fits_file(self, filename, load_image=False, p=None): if load_image: if p is None: raise ValueError("Loading image without a PSF.") - self.image = layered_image(filename, p) + self.image = LayeredImage(filename, p) self.filename = filename with fits.open(filename) as hdu_list: @@ -90,12 +90,12 @@ def populate_from_fits_file(self, filename, load_image=False, p=None): # Compute the center of the image in sky coordinates. self.center = self.wcs.pixel_to_world(self.width / 2, self.height / 2) - def set_layered_image(self, image): + def set_LayeredImage(self, image): """Manually set the layered image. Parameters ---------- - image : `layered_image` + image : `LayeredImage` The layered image to use. """ self.image = image @@ -173,10 +173,10 @@ def skycoords_to_pixels(self, pos): Returns ------- - result : `pixel_pos` - A `pixel_pos` object with the (x, y) pixel location. + result : `PixelPos` + A `PixelPos` object with the (x, y) pixel location. """ - result = pixel_pos() + result = PixelPos() result.x, result.y = self.wcs.world_to_pixel(pos) return result @@ -185,8 +185,8 @@ def pixels_to_skycoords(self, pos): Parameters ---------- - pos : `pixel_pos` - A `pixel_pos` object containing the x and y + pos : `PixelPos` + A `PixelPos` object containing the x and y coordinates on the pixel. Returns @@ -337,7 +337,7 @@ def pixels_to_skycoords(self, pos): Parameters ---------- - pos : a list of `pixel_pos` objects + pos : a list of `PixelPos` objects The positions in pixel coordinates. Returns @@ -370,7 +370,7 @@ def trajectory_to_skycoords(self, trj): results = [] for i in range(self.num_images): dt = self.stats[i].get_epoch().mjd - t0 - pos_x = trj.x + dt * trj.x_v - pos_y = trj.y + dt * trj.y_v + pos_x = trj.x + dt * trj.vx + pos_y = trj.y + dt * trj.vy results.append(self.stats[i].wcs.pixel_to_world(pos_x, pos_y)) return results diff --git a/src/kbmod/masking.py b/src/kbmod/masking.py index 89ad4b211..372b8ba0b 100644 --- a/src/kbmod/masking.py +++ b/src/kbmod/masking.py @@ -17,14 +17,14 @@ def apply_mask_operations(stack, mask_list): Parameters ---------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack before the masks have been applied. mask_list : `list` A list of mask_list objects. Returns ------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The same stack object to allow chaining. """ for mask in mask_list: @@ -44,12 +44,12 @@ def apply_mask(self, stack): Parameters ---------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack before the masks have been applied. Returns ------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The same stack object to allow chaining. """ pass @@ -77,12 +77,12 @@ def apply_mask(self, stack): Parameters ---------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack before the masks have been applied. Returns ------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The same stack object to allow chaining. """ if self.flags != 0: @@ -150,12 +150,12 @@ def apply_mask(self, stack): Parameters ---------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack before the masks have been applied. Returns ------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The same stack object to allow chaining. """ if self.global_flags != 0: @@ -181,12 +181,12 @@ def apply_mask(self, stack): Parameters ---------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack before the masks have been applied. Returns ------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The same stack object to allow chaining. """ stack.apply_mask_threshold(self.mask_threshold) @@ -214,12 +214,12 @@ def apply_mask(self, stack): Parameters ---------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack before the masks have been applied. Returns ------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The same stack object to allow chaining. """ stack.grow_mask(self.num_pixels) diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py index 9eb7572fd..da5824137 100644 --- a/src/kbmod/run_search.py +++ b/src/kbmod/run_search.py @@ -62,7 +62,7 @@ def do_masking(self, stack): Parameters ---------- - stack : `kbmod.image_stack` + stack : `kbmod.ImageStack` The stack before the masks have been applied. """ mask_steps = [] @@ -131,13 +131,13 @@ def do_gpu_search(self, search, img_info, suggested_angle, post_process): if self.config["x_pixel_bounds"] and len(self.config["x_pixel_bounds"]) == 2: search.set_start_bounds_x(self.config["x_pixel_bounds"][0], self.config["x_pixel_bounds"][1]) elif self.config["x_pixel_buffer"] and self.config["x_pixel_buffer"] > 0: - width = search.get_image_stack().get_width() + width = search.get_ImageStack().get_width() search.set_start_bounds_x(-self.config["x_pixel_buffer"], width + self.config["x_pixel_buffer"]) if self.config["y_pixel_bounds"] and len(self.config["y_pixel_bounds"]) == 2: search.set_start_bounds_y(self.config["y_pixel_bounds"][0], self.config["y_pixel_bounds"][1]) elif self.config["y_pixel_buffer"] and self.config["y_pixel_buffer"] > 0: - height = search.get_image_stack().get_height() + height = search.get_ImageStack().get_height() search.set_start_bounds_y(-self.config["y_pixel_buffer"], height + self.config["y_pixel_buffer"]) # If we are using barycentric corrections, compute the parameters and @@ -240,7 +240,7 @@ def run_search(self): kb_interface = Interface() # Load the PSF. - default_psf = kb.psf(self.config["psf_val"]) + default_psf = kb.PSF(self.config["psf_val"]) # Load images to search stack, img_info = kb_interface.load_images( @@ -264,7 +264,7 @@ def run_search(self): stack = self.do_masking(stack) # Perform the actual search. - search = kb.stack_search(stack) + search = kb.StackSearch(stack) search, search_params = self.do_gpu_search(search, img_info, suggested_angle, kb_post_process) # Load the KBMOD results into Python and apply a filter based on @@ -328,8 +328,8 @@ def _count_known_matches(self, result_list, search): ---------- result_list : ``kbmod.ResultList`` The result objects found by the search. - search : ``kbmod.search.stack_search`` - A stack_search object containing information about the search. + search : ``kbmod.search.StackSearch`` + A StackSearch object containing information about the search. """ # Get the image metadata im_filepath = self.config["im_filepath"] @@ -341,10 +341,10 @@ def _count_known_matches(self, result_list, search): ps_list = [] for row in result_list.results: - pix_pos_objs = search.get_mult_traj_pos(row.trajectory) - pixel_positions = list(map(lambda p: [p.x, p.y], pix_pos_objs)) + pix_pos_objs = search.get_trajectory_positions(row.trajectory) + PixelPositions = list(map(lambda p: [p.x, p.y], pix_pos_objs)) ps = koffi.PotentialSource() - ps.build_from_images_and_xy_positions(pixel_positions, metadata) + ps.build_from_images_and_xy_positions(PixelPositions, metadata) ps_list.append(ps) print("-----------------") diff --git a/src/kbmod/search/Filtering.h b/src/kbmod/search/Filtering.h deleted file mode 100644 index 17efb4fea..000000000 --- a/src/kbmod/search/Filtering.h +++ /dev/null @@ -1,24 +0,0 @@ -/* - * Filtering.h - * - * Created on: Sept 2, 2022 - * - * Helper functions for filtering results. - */ - -#ifndef FILTERING_H_ -#define FILTERING_H_ - -#include - -namespace search { - -/* Return the list of indices from the values array such that those elements - pass the sigmaG filtering defined by percentiles [sGL0, sGL1] with coefficient - sigmag_coeff and a multiplicative factor of width. */ -std::vector sigmaGFilteredIndices(const std::vector& values, float sgl0, float sgl1, - float sigma_g_coeff, float width); - -} /* namespace search */ - -#endif /* FILTERING_H_ */ diff --git a/src/kbmod/search/ImageStack.cpp b/src/kbmod/search/ImageStack.cpp deleted file mode 100644 index aaef58201..000000000 --- a/src/kbmod/search/ImageStack.cpp +++ /dev/null @@ -1,137 +0,0 @@ -/* - * ImageStack.cpp - * - * Created on: Jun 22, 2017 - * Author: kbmod-usr - */ - -#include "ImageStack.h" - -namespace search { - -ImageStack::ImageStack(const std::vector& filenames, const std::vector& psfs) { - verbose = true; - resetImages(); - loadImages(filenames, psfs); - extractImageTimes(); - setTimeOrigin(); - global_mask = RawImage(getWidth(), getHeight()); - global_mask.setAllPix(0.0); -} - -ImageStack::ImageStack(const std::vector& imgs) { - verbose = true; - images = imgs; - extractImageTimes(); - setTimeOrigin(); - global_mask = RawImage(getWidth(), getHeight()); - global_mask.setAllPix(0.0); -} - -void ImageStack::loadImages(const std::vector& filenames, - const std::vector& psfs) { - const int num_files = filenames.size(); - if (num_files == 0) { - std::cout << "No files provided" - << "\n"; - } - - if (psfs.size() != num_files) throw std::runtime_error("Mismatched PSF array in ImageStack creation."); - - // Load images from file - for (int i = 0; i < num_files; ++i) { - images.push_back(LayeredImage(filenames[i], psfs[i])); - if (verbose) std::cout << "." << std::flush; - } - if (verbose) std::cout << "\n"; -} - -void ImageStack::extractImageTimes() { - // Load image times - image_times = std::vector(); - for (auto& i : images) { - image_times.push_back(float(i.getObstime())); - } -} - -void ImageStack::setTimeOrigin() { - // Set beginning time to 0.0 - double initial_time = image_times[0]; - for (auto& t : image_times) t = t - initial_time; -} - -LayeredImage& ImageStack::getSingleImage(int index) { - if (index < 0 || index > images.size()) throw std::out_of_range("ImageStack index out of bounds."); - return images[index]; -} - -void ImageStack::setSingleImage(int index, LayeredImage& img) { - if (index < 0 || index > images.size()) throw std::out_of_range("ImageStack index out of bounds."); - images[index] = img; -} - -void ImageStack::setTimes(const std::vector& times) { - if (times.size() != imgCount()) - throw std::runtime_error( - "List of times provided" - " does not match the number of images!"); - image_times = times; - setTimeOrigin(); -} - -void ImageStack::resetImages() { images = std::vector(); } - -void ImageStack::convolvePSF() { - for (auto& i : images) i.convolvePSF(); -} - -void ImageStack::saveGlobalMask(const std::string& path) { global_mask.saveToFile(path); } - -void ImageStack::saveImages(const std::string& path) { - for (auto& i : images) i.saveLayers(path); -} - -const RawImage& ImageStack::getGlobalMask() const { return global_mask; } - -void ImageStack::applyMaskFlags(int flags, const std::vector& exceptions) { - for (auto& i : images) { - i.applyMaskFlags(flags, exceptions); - } -} - -void ImageStack::applyGlobalMask(int flags, int threshold) { - createGlobalMask(flags, threshold); - for (auto& i : images) { - i.applyGlobalMask(global_mask); - } -} - -void ImageStack::applyMaskThreshold(float thresh) { - for (auto& i : images) i.applyMaskThreshold(thresh); -} - -void ImageStack::growMask(int steps) { - for (auto& i : images) i.growMask(steps); -} - -void ImageStack::createGlobalMask(int flags, int threshold) { - int npixels = getNPixels(); - - // For each pixel count the number of images where it is masked. - std::vector counts(npixels, 0); - for (unsigned int img = 0; img < images.size(); ++img) { - float* imgMask = images[img].getMDataRef(); - // Count the number of times a pixel has any of the flags - for (unsigned int pixel = 0; pixel < npixels; ++pixel) { - if ((flags & static_cast(imgMask[pixel])) != 0) counts[pixel]++; - } - } - - // Set all pixels below threshold to 0 and all above to 1 - float* global_m = global_mask.getDataRef(); - for (unsigned int p = 0; p < npixels; ++p) { - global_m[p] = counts[p] < threshold ? 0.0 : 1.0; - } -} - -} /* namespace search */ diff --git a/src/kbmod/search/ImageStack.h b/src/kbmod/search/ImageStack.h deleted file mode 100644 index c9e013c9e..000000000 --- a/src/kbmod/search/ImageStack.h +++ /dev/null @@ -1,70 +0,0 @@ -/* - * ImageStack.h - * - * Created on: Jun 22, 2017 - * Author: kbmod-usr - * ImageStack stores a series of LayeredImages from different times. - */ - -#ifndef IMAGESTACK_H_ -#define IMAGESTACK_H_ - -#include -#include -#include -#include -#include -#include -#include "LayeredImage.h" - -namespace search { - -class ImageStack { -public: - ImageStack(const std::vector& filenames, const std::vector& psfs); - ImageStack(const std::vector& imgs); - - // Simple getters. - unsigned imgCount() const { return images.size(); } - unsigned getWidth() const { return images.size() > 0 ? images[0].getWidth() : 0; } - unsigned getHeight() const { return images.size() > 0 ? images[0].getHeight() : 0; } - unsigned getNPixels() const { return images.size() > 0 ? images[0].getNPixels() : 0; } - std::vector& getImages() { return images; } - const std::vector& getTimes() const { return image_times; } - float* getTimesDataRef() { return image_times.data(); } - LayeredImage& getSingleImage(int index); - - // Simple setters. - void setTimes(const std::vector& times); - void resetImages(); - void setSingleImage(int index, LayeredImage& img); - - // Apply makes to all the images. - void applyGlobalMask(int flags, int threshold); - void applyMaskFlags(int flags, const std::vector& exceptions); - void applyMaskThreshold(float thresh); - void growMask(int steps); - const RawImage& getGlobalMask() const; - - void convolvePSF(); - - // Save data to files. - void saveGlobalMask(const std::string& path); - void saveImages(const std::string& path); - - virtual ~ImageStack(){}; - -private: - void loadImages(const std::vector& filenames, const std::vector& psfs); - void extractImageTimes(); - void setTimeOrigin(); - void createGlobalMask(int flags, int threshold); - std::vector images; - RawImage global_mask; - std::vector image_times; - bool verbose; -}; - -} /* namespace search */ - -#endif /* IMAGESTACK_H_ */ diff --git a/src/kbmod/search/KBMOSearch.cpp b/src/kbmod/search/KBMOSearch.cpp deleted file mode 100644 index ae84246c7..000000000 --- a/src/kbmod/search/KBMOSearch.cpp +++ /dev/null @@ -1,614 +0,0 @@ -/* - * KBMOSearch.cpp - * - * Created on: Jun 28, 2017 - * Author: kbmod-usr - */ - -#include "KBMOSearch.h" - -namespace search { - -#ifdef HAVE_CUDA -extern "C" void deviceSearchFilter(int num_images, int width, int height, float* psi_vect, float* phi_vect, - 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 >& use_index_vect, float* results); -#endif - -KBMOSearch::KBMOSearch(ImageStack& imstack) : stack(imstack) { - max_result_count = 100000; - debug_info = false; - psi_phi_generated = false; - - // Default the thresholds. - params.min_observations = 0; - params.min_lh = 0.0; - - // Default filtering arguments. - params.do_sigmag_filter = false; - params.sgl_L = 0.25; - params.sgl_H = 0.75; - params.sigmag_coeff = -1.0; - - // Default the encoding parameters. - params.psi_num_bytes = -1; - params.phi_num_bytes = -1; - - // Default pixel starting bounds. - params.x_start_min = 0; - params.x_start_max = stack.getWidth(); - params.y_start_min = 0; - params.y_start_max = stack.getHeight(); - - // Set default values for the barycentric correction. - bary_corrs = std::vector(stack.imgCount()); - params.use_corr = false; - use_corr = false; - - params.debug = false; -} - -void KBMOSearch::setDebug(bool d) { - debug_info = d; - params.debug = d; -} - -void KBMOSearch::enableCorr(std::vector bary_corr_coeff) { - use_corr = true; - params.use_corr = true; - for (int i = 0; i < stack.imgCount(); i++) { - int j = i * 6; - bary_corrs[i].dx = bary_corr_coeff[j]; - bary_corrs[i].dxdx = bary_corr_coeff[j + 1]; - bary_corrs[i].dxdy = bary_corr_coeff[j + 2]; - bary_corrs[i].dy = bary_corr_coeff[j + 3]; - bary_corrs[i].dydx = bary_corr_coeff[j + 4]; - bary_corrs[i].dydy = bary_corr_coeff[j + 5]; - } -} - -void KBMOSearch::enableGPUSigmaGFilter(std::vector percentiles, float sigmag_coeff, - float min_lh) { - params.do_sigmag_filter = true; - params.sgl_L = percentiles[0]; - params.sgl_H = percentiles[1]; - params.sigmag_coeff = sigmag_coeff; - params.min_lh = min_lh; -} - -void KBMOSearch::enableGPUEncoding(int psi_num_bytes, int phi_num_bytes) { - // Make sure the encoding is one of the supported options. - // Otherwise use default float (aka no encoding). - if (psi_num_bytes == 1 || psi_num_bytes == 2) { - params.psi_num_bytes = psi_num_bytes; - } else { - params.psi_num_bytes = -1; - } - if (phi_num_bytes == 1 || phi_num_bytes == 2) { - params.phi_num_bytes = phi_num_bytes; - } else { - params.phi_num_bytes = -1; - } -} - -void KBMOSearch::setStartBoundsX(int x_min, int x_max) { - params.x_start_min = x_min; - params.x_start_max = x_max; -} - -void KBMOSearch::setStartBoundsY(int y_min, int y_max) { - params.y_start_min = y_min; - params.y_start_max = y_max; -} - -void KBMOSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ang, float min_vel, - float max_vel, int min_observations) { - preparePsiPhi(); - createSearchList(ang_steps, vel_steps, min_ang, max_ang, min_vel, max_vel); - - startTimer("Creating psi/phi buffers"); - std::vector psi_vect; - std::vector phi_vect; - fillPsiAndphi_vects(psi_images, phi_images, &psi_vect, &phi_vect); - endTimer(); - - // Create a data stucture for the per-image data. - PerImageData img_data; - img_data.num_images = stack.imgCount(); - img_data.image_times = stack.getTimesDataRef(); - if (params.use_corr) img_data.bary_corrs = &bary_corrs[0]; - - // Compute the encoding parameters for psi and phi if needed. - // Vectors need to be created outside the if so they stay in scope. - std::vector psi_scale_vect; - std::vector phi_scale_vect; - if (params.psi_num_bytes > 0) { - psi_scale_vect = computeImageScaling(psi_images, params.psi_num_bytes); - img_data.psi_params = psi_scale_vect.data(); - } - if (params.phi_num_bytes > 0) { - phi_scale_vect = computeImageScaling(phi_images, params.phi_num_bytes); - img_data.phi_params = phi_scale_vect.data(); - } - - // Allocate a vector for the results. - int num_search_pixels = - ((params.x_start_max - params.x_start_min) * (params.y_start_max - params.y_start_min)); - int max_results = num_search_pixels * RESULTS_PER_PIXEL; - if (debug_info) { - std::cout << "Searching X=[" << params.x_start_min << ", " << params.x_start_max << "]" - << " Y=[" << params.y_start_min << ", " << params.y_start_max << "]\n"; - std::cout << "Allocating space for " << max_results << " results.\n"; - } - results = std::vector(max_results); - if (debug_info) std::cout << search_list.size() << " trajectories... \n" << std::flush; - - // Set the minimum number of observations. - params.min_observations = min_observations; - - // Do the actual search on the GPU. - startTimer("Searching"); -#ifdef HAVE_CUDA - deviceSearchFilter(stack.imgCount(), stack.getWidth(), stack.getHeight(), psi_vect.data(), phi_vect.data(), - img_data, params, search_list.size(), search_list.data(), max_results, results.data()); -#else - throw std::runtime_error("Non-GPU search is not implemented."); -#endif - endTimer(); - - startTimer("Sorting results"); - sortResults(); - endTimer(); -} - -void KBMOSearch::savePsiPhi(const std::string& path) { - preparePsiPhi(); - saveImages(path); -} - -void KBMOSearch::preparePsiPhi() { - if (!psi_phi_generated) { - psi_images.clear(); - phi_images.clear(); - - // Compute Phi and Psi from convolved images - // while leaving masked pixels alone - // Reinsert 0s for NO_DATA? - const int num_images = stack.imgCount(); - for (int i = 0; i < num_images; ++i) { - LayeredImage& img = stack.getSingleImage(i); - psi_images.push_back(img.generatePsiImage()); - phi_images.push_back(img.generatePhiImage()); - } - - psi_phi_generated = true; - } -} - -std::vector KBMOSearch::computeImageScaling(const std::vector& vect, - int encoding_bytes) const { - std::vector result; - - const int num_images = vect.size(); - for (int i = 0; i < num_images; ++i) { - scaleParameters params; - params.scale = 1.0; - - std::array bnds = vect[i].computeBounds(); - params.min_val = bnds[0]; - params.max_val = bnds[1]; - - // Increase width to avoid divide by zero. - float width = (params.max_val - params.min_val); - if (width < 1e-6) width = 1e-6; - - // Set the scale if we are encoding the values. - if (encoding_bytes == 1 || encoding_bytes == 2) { - long int num_values = (1 << (8 * encoding_bytes)) - 1; - params.scale = width / (double)num_values; - } - - result.push_back(params); - } - - return result; -} - -void KBMOSearch::saveImages(const std::string& path) { - for (int i = 0; i < stack.imgCount(); ++i) { - std::string number = std::to_string(i); - // Add leading zeros - number = std::string(4 - number.length(), '0') + number; - psi_images[i].saveToFile(path + "/psi/PSI" + number + ".fits"); - phi_images[i].saveToFile(path + "/phi/PHI" + number + ".fits"); - } -} - -void KBMOSearch::createSearchList(int angle_steps, int velocity_steps, float min_ang, float max_ang, - float min_vel, float max_vel) { - std::vector angles(angle_steps); - float ang_stepsize = (max_ang - min_ang) / float(angle_steps); - for (int i = 0; i < angle_steps; ++i) { - angles[i] = min_ang + float(i) * ang_stepsize; - } - - std::vector velocities(velocity_steps); - float vel_stepsize = (max_vel - min_vel) / float(velocity_steps); - for (int i = 0; i < velocity_steps; ++i) { - velocities[i] = min_vel + float(i) * vel_stepsize; - } - - int trajCount = angle_steps * velocity_steps; - search_list = std::vector(trajCount); - for (int a = 0; a < angle_steps; ++a) { - for (int v = 0; v < velocity_steps; ++v) { - search_list[a * velocity_steps + v].x_vel = cos(angles[a]) * velocities[v]; - search_list[a * velocity_steps + v].y_vel = sin(angles[a]) * velocities[v]; - } - } -} - -void KBMOSearch::fillPsiAndphi_vects(const std::vector& psi_imgs, - const std::vector& phi_imgs, std::vector* psi_vect, - std::vector* phi_vect) { - assert(psi_vect != NULL); - assert(phi_vect != NULL); - - int num_images = psi_imgs.size(); - assert(num_images > 0); - assert(phi_imgs.size() == num_images); - - int num_pixels = psi_imgs[0].getNPixels(); - for (int i = 0; i < num_images; ++i) { - assert(psi_imgs[i].getNPixels() == num_pixels); - assert(phi_imgs[i].getNPixels() == num_pixels); - } - - psi_vect->clear(); - psi_vect->reserve(num_images * num_pixels); - phi_vect->clear(); - phi_vect->reserve(num_images * num_pixels); - - for (int i = 0; i < num_images; ++i) { - const std::vector& psi_ref = psi_imgs[i].getPixels(); - const std::vector& phi_ref = phi_imgs[i].getPixels(); - for (unsigned p = 0; p < num_pixels; ++p) { - psi_vect->push_back(psi_ref[p]); - phi_vect->push_back(phi_ref[p]); - } - } -} - -std::vector KBMOSearch::scienceStamps(const trajectory& trj, int radius, bool interpolate, - bool keep_no_data, const std::vector& use_index) { - if (use_index.size() > 0 && use_index.size() != stack.imgCount()) { - throw std::runtime_error("Wrong size use_index passed into scienceStamps()"); - } - bool use_all_stamps = use_index.size() == 0; - - std::vector stamps; - int num_times = stack.imgCount(); - for (int i = 0; i < num_times; ++i) { - if (use_all_stamps || use_index[i]) { - PixelPos pos = getTrajPos(trj, i); - RawImage& img = stack.getSingleImage(i).getScience(); - stamps.push_back(img.createStamp(pos.x, pos.y, radius, interpolate, keep_no_data)); - } - } - return stamps; -} - -// For stamps used for visualization we interpolate the pixel values, replace -// NO_DATA tages with zeros, and return all the stamps (regardless of whether -// individual timesteps have been filtered). -std::vector KBMOSearch::scienceStampsForViz(const trajectory& t, int radius) { - std::vector empty_vect; - return scienceStamps(t, radius, true /*=interpolate*/, false /*=keep_no_data*/, empty_vect); -} - -// For creating coadded stamps, we do not interpolate the pixel values and keep -// NO_DATA tagged (so we can filter it out of mean/median). -RawImage KBMOSearch::medianScienceStamp(const trajectory& trj, int radius, - const std::vector& use_index) { - return createMedianImage( - scienceStamps(trj, radius, false /*=interpolate*/, true /*=keep_no_data*/, use_index)); -} - -// For creating coadded stamps, we do not interpolate the pixel values and keep -// NO_DATA tagged (so we can filter it out of mean/median). -RawImage KBMOSearch::meanScienceStamp(const trajectory& trj, int radius, const std::vector& use_index) { - return createMeanImage( - scienceStamps(trj, radius, false /*=interpolate*/, true /*=keep_no_data*/, use_index)); -} - -// For creating summed stamps, we do not interpolate the pixel values and replace NO_DATA -// with zero (which is the same as filtering it out for the sum). -RawImage KBMOSearch::summedScienceStamp(const trajectory& trj, int radius, - const std::vector& use_index) { - return createSummedImage( - scienceStamps(trj, radius, false /*=interpolate*/, false /*=keep_no_data*/, use_index)); -} - -bool KBMOSearch::filterStamp(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& pixels = img.getPixels(); - - // Filter on the peak's position. - PixelPos pos = img.findPeak(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& pixels = img.getPixels(); - 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.findCentralMoments(); - 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 KBMOSearch::coaddedScienceStamps(std::vector& t_array, - std::vector >& use_index_vect, - const StampParameters& params, bool use_gpu) { - if (use_gpu) { -#ifdef HAVE_CUDA - return coaddedScienceStampsGPU(t_array, use_index_vect, params); -#else - print("WARNING: GPU is not enabled. Performing co-adds on the CPU."); -#endif - } - return coaddedScienceStampsCPU(t_array, use_index_vect, params); -} - -std::vector KBMOSearch::coaddedScienceStampsCPU(std::vector& t_array, - std::vector >& use_index_vect, - const StampParameters& params) { - const int num_trajectories = t_array.size(); - std::vector results(num_trajectories); - std::vector empty_pixels(1, NO_DATA); - - for (int i = 0; i < num_trajectories; ++i) { - std::vector stamps = - scienceStamps(t_array[i], params.radius, false, true, use_index_vect[i]); - - RawImage coadd(1, 1); - switch (params.stamp_type) { - case STAMP_MEDIAN: - coadd = createMedianImage(stamps); - break; - case STAMP_MEAN: - coadd = createMeanImage(stamps); - break; - case STAMP_SUM: - coadd = createSummedImage(stamps); - break; - default: - throw std::runtime_error("Invalid stamp coadd type."); - } - - // Do the filtering if needed. - if (params.do_filtering && filterStamp(coadd, params)) { - results[i] = RawImage(1, 1, empty_pixels); - } else { - results[i] = coadd; - } - } - - return results; -} - -std::vector KBMOSearch::coaddedScienceStampsGPU(std::vector& t_array, - std::vector >& 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.imgCount(); - const int width = stack.getWidth(); - const int height = stack.getHeight(); - - // Create a data stucture for the per-image data. - PerImageData img_data; - img_data.num_images = num_images; - img_data.image_times = stack.getTimesDataRef(); - - // 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 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 results(num_trajectories); - std::vector current_pixels(stamp_ppi, 0.0); - std::vector 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 && filterStamp(current_image, params)) { - results[t] = RawImage(1, 1, empty_pixels); - } else { - results[t] = RawImage(stamp_width, stamp_width, current_pixels); - } - } - return results; -} - -std::vector KBMOSearch::createStamps(trajectory t, int radius, const std::vector& imgs, - bool interpolate) { - if (radius < 0) throw std::runtime_error("stamp radius must be at least 0"); - std::vector stamps; - for (int i = 0; i < imgs.size(); ++i) { - PixelPos pos = getTrajPos(t, i); - stamps.push_back(imgs[i]->createStamp(pos.x, pos.y, radius, interpolate, false)); - } - return stamps; -} - -PixelPos KBMOSearch::getTrajPos(const trajectory& t, int i) const { - float time = stack.getTimes()[i]; - if (use_corr) { - return {t.x + time * t.x_vel + bary_corrs[i].dx + t.x * bary_corrs[i].dxdx + t.y * bary_corrs[i].dxdy, - t.y + time * t.y_vel + bary_corrs[i].dy + t.x * bary_corrs[i].dydx + - t.y * bary_corrs[i].dydy}; - } else { - return {t.x + time * t.x_vel, t.y + time * t.y_vel}; - } -} - -std::vector KBMOSearch::getMultTrajPos(trajectory& t) const { - std::vector results; - int num_times = stack.imgCount(); - for (int i = 0; i < num_times; ++i) { - PixelPos pos = getTrajPos(t, i); - results.push_back(pos); - } - return results; -} - -std::vector KBMOSearch::createCurves(trajectory t, const std::vector& imgs) { - /*Create a lightcurve from an image along a trajectory - * - * INPUT- - * trajectory t - The trajectory along which to compute the lightcurve - * std::vector imgs - The image from which to compute the - * trajectory. Most likely a psiImage or a phiImage. - * Output- - * std::vector lightcurve - The computed trajectory - */ - - int img_size = imgs.size(); - std::vector lightcurve; - lightcurve.reserve(img_size); - const std::vector& times = stack.getTimes(); - for (int i = 0; i < img_size; ++i) { - /* Do not use getPixelInterp(), because results from createCurves must - * be able to recover the same likelihoods as the ones reported by the - * gpu search.*/ - float pix_val; - if (use_corr) { - PixelPos pos = getTrajPos(t, i); - pix_val = imgs[i].getPixel(int(pos.x + 0.5), int(pos.y + 0.5)); - } - /* Does not use getTrajPos to be backwards compatible with Hits_Rerun */ - else { - pix_val = imgs[i].getPixel(t.x + int(times[i] * t.x_vel + 0.5), - t.y + int(times[i] * t.y_vel + 0.5)); - } - if (pix_val == NO_DATA) pix_val = 0.0; - lightcurve.push_back(pix_val); - } - return lightcurve; -} - -std::vector KBMOSearch::psiCurves(trajectory& t) { - /*Generate a psi lightcurve for further analysis - * INPUT- - * trajectory& t - The trajectory along which to find the lightcurve - * OUTPUT- - * std::vector - A vector of the lightcurve values - */ - preparePsiPhi(); - return createCurves(t, psi_images); -} - -std::vector KBMOSearch::phiCurves(trajectory& t) { - /*Generate a phi lightcurve for further analysis - * INPUT- - * trajectory& t - The trajectory along which to find the lightcurve - * OUTPUT- - * std::vector - A vector of the lightcurve values - */ - preparePsiPhi(); - return createCurves(t, phi_images); -} - -std::vector& KBMOSearch::getPsiImages() { return psi_images; } - -std::vector& KBMOSearch::getPhiImages() { return phi_images; } - -void KBMOSearch::sortResults() { - __gnu_parallel::sort(results.begin(), results.end(), - [](trajectory a, trajectory b) { return b.lh < a.lh; }); -} - -void KBMOSearch::filterResults(int min_observations) { - results.erase(std::remove_if(results.begin(), results.end(), - std::bind([](trajectory t, int cutoff) { return t.obs_count < cutoff; }, - std::placeholders::_1, min_observations)), - results.end()); -} - -void KBMOSearch::filterResultsLH(float min_lh) { - results.erase(std::remove_if(results.begin(), results.end(), - std::bind([](trajectory t, float cutoff) { return t.lh < cutoff; }, - std::placeholders::_1, min_lh)), - results.end()); -} - -std::vector KBMOSearch::getResults(int start, int count) { - if (start + count >= results.size()) { - count = results.size() - start; - } - if (start < 0) throw std::runtime_error("start must be 0 or greater"); - return std::vector(results.begin() + start, results.begin() + start + count); -} - -// This function is used only for testing by injecting known result trajectories. -void KBMOSearch::setResults(const std::vector& new_results) { results = new_results; } - -void KBMOSearch::startTimer(const std::string& message) { - if (debug_info) { - std::cout << message << "... " << std::flush; - t_start = std::chrono::system_clock::now(); - } -} - -void KBMOSearch::endTimer() { - if (debug_info) { - t_end = std::chrono::system_clock::now(); - t_delta = t_end - t_start; - std::cout << " Took " << t_delta.count() << " seconds.\n" << std::flush; - } -} - -} /* namespace search */ diff --git a/src/kbmod/search/KBMOSearch.h b/src/kbmod/search/KBMOSearch.h deleted file mode 100644 index d42fa83be..000000000 --- a/src/kbmod/search/KBMOSearch.h +++ /dev/null @@ -1,157 +0,0 @@ -/* - * KBMOSearch.h - * - * Created on: Jun 28, 2017 - * Author: kbmod-usr - * - * The KBMOSearch class holds all of the information and functions - * to perform the core stacked search. - */ - -#ifndef KBMODSEARCH_H_ -#define KBMODSEARCH_H_ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "common.h" -#include "ImageStack.h" -#include "PointSpreadFunc.h" - -namespace search { - -class KBMOSearch { -public: - KBMOSearch(ImageStack& imstack); - - int numImages() const { return stack.imgCount(); } - const ImageStack& getImageStack() const { return stack; } - - void setDebug(bool d); - - // The primary search functions. - void enableGPUSigmaGFilter(std::vector percentiles, float sigmag_coeff, float min_lh); - void enableCorr(std::vector bary_corr_coeff); - void enableGPUEncoding(int psi_num_bytes, int phi_num_bytes); - - void setStartBoundsX(int x_min, int x_max); - void setStartBoundsY(int y_min, int y_max); - - void search(int a_steps, int v_steps, float min_angle, float max_angle, float min_velocity, - float max_velocity, int min_observations); - - // Gets the vector of result trajectories. - std::vector getResults(int start, int end); - - // Get the predicted (pixel) positions for a given trajectory. - PixelPos getTrajPos(const trajectory& t, int i) const; - std::vector getMultTrajPos(trajectory& t) const; - - // Filters the results based on various parameters. - void filterResults(int min_observations); - void filterResultsLH(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 indicating whether to use - // each time step. An empty (size=0) vector will use all time steps. - std::vector scienceStamps(const trajectory& trj, int radius, bool interpolate, - bool keep_no_data, const std::vector& use_index); - std::vector scienceStampsForViz(const trajectory& t, int radius); - RawImage medianScienceStamp(const trajectory& trj, int radius, const std::vector& use_index); - RawImage meanScienceStamp(const trajectory& trj, int radius, const std::vector& use_index); - RawImage summedScienceStamp(const trajectory& trj, int radius, const std::vector& 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 coaddedScienceStamps(std::vector& t_array, - std::vector >& use_index_vect, - const StampParameters& params, bool use_cpu); - - // Function to do the actual stamp filtering. - bool filterStamp(const RawImage& img, const StampParameters& params); - - // Getters for the Psi and Phi data. - std::vector& getPsiImages(); - std::vector& getPhiImages(); - std::vector psiCurves(trajectory& t); - std::vector phiCurves(trajectory& t); - - // Save internal data products to a file. - void savePsiPhi(const std::string& path); - - // Helper functions for computing Psi and Phi. - void preparePsiPhi(); - - // Helper functions for testing. - void setResults(const std::vector& new_results); - - virtual ~KBMOSearch(){}; - -protected: - void saveImages(const std::string& path); - void sortResults(); - std::vector createCurves(trajectory t, const std::vector& imgs); - - // Fill an interleaved vector for the GPU functions. - void fillPsiAndphi_vects(const std::vector& psi_imgs, const std::vector& phi_imgs, - std::vector* psi_vect, std::vector* phi_vect); - - // Set the parameter min/max/scale from the psi/phi/other images. - std::vector computeImageScaling(const std::vector& vect, - int encoding_bytes) const; - - // Functions to create and access stamps around proposed trajectories or - // regions. Used to visualize the results. - // This function replaces NO_DATA with a value of 0.0. - std::vector createStamps(trajectory t, int radius, const std::vector& imgs, - bool interpolate); - - // Creates list of trajectories to search. - void createSearchList(int angle_steps, int velocity_steps, float min_ang, float max_ang, - float min_vel, float max_vel); - - std::vector coaddedScienceStampsGPU(std::vector& t_array, - std::vector >& use_index_vect, - const StampParameters& params); - - std::vector coaddedScienceStampsCPU(std::vector& t_array, - std::vector >& use_index_vect, - const StampParameters& params); - // Helper functions for timing operations of the search. - void startTimer(const std::string& message); - void endTimer(); - - unsigned max_result_count; - bool psi_phi_generated; - bool debug_info; - ImageStack stack; - std::vector search_list; - std::vector psi_images; - std::vector phi_images; - std::vector results; - - // Variables for the timer. - std::chrono::time_point t_start, t_end; - std::chrono::duration t_delta; - - // Parameters for the GPU search. - SearchParameters params; - - // Parameters to do barycentric corrections. - bool use_corr; - std::vector bary_corrs; -}; - -} /* namespace search */ - -#endif /* KBMODSEARCH_H_ */ diff --git a/src/kbmod/search/LayeredImage.cpp b/src/kbmod/search/LayeredImage.cpp deleted file mode 100644 index d1cfba0ce..000000000 --- a/src/kbmod/search/LayeredImage.cpp +++ /dev/null @@ -1,233 +0,0 @@ -/* - * LayeredImage.cpp - * - * Created on: Jul 11, 2017 - * Author: kbmod-usr - */ - -#include "LayeredImage.h" - -namespace search { - -LayeredImage::LayeredImage(std::string path, const PointSpreadFunc& psf) : psf(psf) { - int f_begin = path.find_last_of("/"); - int f_end = path.find_last_of(".fits") - 4; - filename = path.substr(f_begin, f_end - f_begin); - - science = RawImage(); - science.loadFromFile(path, 1); - width = science.getWidth(); - height = science.getHeight(); - - mask = RawImage(); - mask.loadFromFile(path, 2); - - variance = RawImage(); - variance.loadFromFile(path, 3); - - if (width != variance.getWidth() or height != variance.getHeight()) - throw std::runtime_error("Science and Variance layers are not the same size."); - if (width != mask.getWidth() or height != mask.getHeight()) - throw std::runtime_error("Science and Mask layers are not the same size."); -} - -LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, - const PointSpreadFunc& psf) - : psf(psf) { - // Get the dimensions of the science layer and check for consistency with - // the other two layers. - width = sci.getWidth(); - height = sci.getHeight(); - if (width != var.getWidth() or height != var.getHeight()) - throw std::runtime_error("Science and Variance layers are not the same size."); - if (width != msk.getWidth() or height != msk.getHeight()) - throw std::runtime_error("Science and Mask layers are not the same size."); - - // Copy the image layers. - science = sci; - mask = msk; - variance = var; -} - -LayeredImage::LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, - const PointSpreadFunc& psf) - : LayeredImage(name, w, h, noise_stdev, pixel_variance, time, psf, -1) {} - -LayeredImage::LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, - const PointSpreadFunc& psf, int seed) - : psf(psf) { - filename = name; - width = w; - height = h; - - std::vector raw_sci(width * height); - std::random_device r; - std::default_random_engine generator(r()); - if (seed >= 0) { - generator.seed(seed); - } - std::normal_distribution distrib(0.0, noise_stdev); - for (float& p : raw_sci) p = distrib(generator); - - science = RawImage(w, h, raw_sci); - science.setObstime(time); - - mask = RawImage(w, h, std::vector(w * h, 0.0)); - variance = RawImage(w, h, std::vector(w * h, pixel_variance)); -} - -void LayeredImage::setPSF(const PointSpreadFunc& new_psf) { - psf = new_psf; -} - -void LayeredImage::growMask(int steps) { - science.growMask(steps); - variance.growMask(steps); -} - -void LayeredImage::convolveGivenPSF(const PointSpreadFunc& given_psf) { - science.convolve(given_psf); - - // Square the PSF use that on the variance image. - PointSpreadFunc psfsq = PointSpreadFunc(given_psf); // Copy - psfsq.squarePSF(); - variance.convolve(psfsq); -} - -void LayeredImage::convolvePSF() { - convolveGivenPSF(psf); -} - -void LayeredImage::applyMaskFlags(int flags, const std::vector& exceptions) { - science.applyMask(flags, exceptions, mask); - variance.applyMask(flags, exceptions, mask); -} - -/* Mask all pixels that are not 0 in global mask */ -void LayeredImage::applyGlobalMask(const RawImage& global_mask) { - science.applyMask(0xFFFFFF, {}, global_mask); - variance.applyMask(0xFFFFFF, {}, global_mask); -} - -void LayeredImage::applyMaskThreshold(float thresh) { - const int num_pixels = getNPixels(); - float* sci_pixels = science.getDataRef(); - float* var_pix = variance.getDataRef(); - for (int i = 0; i < num_pixels; ++i) { - if (sci_pixels[i] > thresh) { - sci_pixels[i] = NO_DATA; - var_pix[i] = NO_DATA; - } - } -} - -void LayeredImage::subtractTemplate(const RawImage& sub_template) { - assert(getHeight() == sub_template.getHeight() && getWidth() == sub_template.getWidth()); - const int num_pixels = getNPixels(); - - float* sci_pixels = science.getDataRef(); - const std::vector& tem_pixels = sub_template.getPixels(); - for (unsigned i = 0; i < num_pixels; ++i) { - if ((sci_pixels[i] != NO_DATA) && (tem_pixels[i] != NO_DATA)) { - sci_pixels[i] -= tem_pixels[i]; - } - } -} - -void LayeredImage::saveLayers(const std::string& path) { - fitsfile* fptr; - int status = 0; - long naxes[2] = {0, 0}; - double obstime = science.getObstime(); - - fits_create_file(&fptr, (path + filename + ".fits").c_str(), &status); - - // If we are unable to create the file, check if it already exists - // and, if so, delete it and retry the create. - if (status == 105) { - status = 0; - fits_open_file(&fptr, (path + filename + ".fits").c_str(), READWRITE, &status); - if (status == 0) { - fits_delete_file(fptr, &status); - fits_create_file(&fptr, (path + filename + ".fits").c_str(), &status); - } - } - - fits_create_img(fptr, SHORT_IMG, 0, naxes, &status); - fits_update_key(fptr, TDOUBLE, "MJD", &obstime, "[d] Generated Image time", &status); - fits_close_file(fptr, &status); - fits_report_error(stderr, status); - - science.appendLayerToFile(path + filename + ".fits"); - mask.appendLayerToFile(path + filename + ".fits"); - variance.appendLayerToFile(path + filename + ".fits"); -} - -void LayeredImage::setScience(RawImage& im) { - checkDims(im); - science = im; -} - -void LayeredImage::setMask(RawImage& im) { - checkDims(im); - mask = im; -} - -void LayeredImage::setVariance(RawImage& im) { - checkDims(im); - variance = im; -} - -void LayeredImage::checkDims(RawImage& im) { - if (im.getWidth() != getWidth()) throw std::runtime_error("Image width does not match"); - if (im.getHeight() != getHeight()) throw std::runtime_error("Image height does not match"); -} - -RawImage LayeredImage::generatePsiImage() { - RawImage result(width, height); - float* result_arr = result.getDataRef(); - float* sci_array = getSDataRef(); - float* var_array = getVDataRef(); - - // Set each of the result pixels. - const int num_pixels = getNPixels(); - for (int p = 0; p < num_pixels; ++p) { - float var_pix = var_array[p]; - if (var_pix != NO_DATA) { - result_arr[p] = sci_array[p] / var_pix; - } else { - result_arr[p] = NO_DATA; - } - } - - // Convolve with the PSF. - result.convolve(psf); - - return result; -} - -RawImage LayeredImage::generatePhiImage() { - RawImage result(width, height); - float* result_arr = result.getDataRef(); - float* var_array = getVDataRef(); - - // Set each of the result pixels. - const int num_pixels = getNPixels(); - for (int p = 0; p < num_pixels; ++p) { - float var_pix = var_array[p]; - if (var_pix != NO_DATA) { - result_arr[p] = 1.0 / var_pix; - } else { - result_arr[p] = NO_DATA; - } - } - - // Convolve with the PSF squared. - PointSpreadFunc psfsq = PointSpreadFunc(psf); // Copy - psfsq.squarePSF(); - result.convolve(psfsq); - - return result; -} - -} /* namespace search */ diff --git a/src/kbmod/search/LayeredImage.h b/src/kbmod/search/LayeredImage.h deleted file mode 100644 index 4a4cb2781..000000000 --- a/src/kbmod/search/LayeredImage.h +++ /dev/null @@ -1,100 +0,0 @@ -/* - * LayeredImage.h - * - * Created on: Jul 11, 2017 - * Author: kbmod-usr - * - * LayeredImage stores an image from a single time with different layers of - * data, such as science pixels, variance pixels, and mask pixels. - */ - -#ifndef LAYEREDIMAGE_H_ -#define LAYEREDIMAGE_H_ - -#include -#include -#include -#include -#include -#include -#include -#include "RawImage.h" -#include "common.h" - -namespace search { - -class LayeredImage { -public: - explicit LayeredImage(std::string path, const PointSpreadFunc& psf); - explicit LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, - const PointSpreadFunc& psf); - explicit LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, - const PointSpreadFunc& psf); - explicit LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, - const PointSpreadFunc& psf, int seed); - - // Set an image specific point spread function. - void setPSF(const PointSpreadFunc& psf); - const PointSpreadFunc& getPSF() const { return psf; } - - // Basic getter functions for image data. - std::string getName() const { return filename; } - unsigned getWidth() const { return width; } - unsigned getHeight() const { return height; } - unsigned getNPixels() const { return width * height; } - double getObstime() const { return science.getObstime(); } - void setObstime(double obstime) { science.setObstime(obstime); } - - // Getter functions for the data in the individual layers. - RawImage& getScience() { return science; } - RawImage& getMask() { return mask; } - RawImage& getVariance() { return variance; } - - // Get pointers to the raw pixel arrays. - float* getSDataRef() { return science.getDataRef(); } - float* getVDataRef() { return variance.getDataRef(); } - float* getMDataRef() { return mask.getDataRef(); } - - // Applies the mask functions to each of the science and variance layers. - void applyMaskFlags(int flag, const std::vector& exceptions); - void applyGlobalMask(const RawImage& global_mask); - void applyMaskThreshold(float thresh); - void growMask(int steps); - - // Subtracts a template image from the science layer. - void subtractTemplate(const RawImage& sub_template); - - // Saves the data in each later to a file. - void saveLayers(const std::string& path); - - // Setter functions for the individual layers. - void setScience(RawImage& im); - void setMask(RawImage& im); - void setVariance(RawImage& im); - - // Convolve with a given PSF or the default one. - void convolvePSF(); - void convolveGivenPSF(const PointSpreadFunc& psf); - - virtual ~LayeredImage(){}; - - // Generate psi and phi images from the science and variance layers. - RawImage generatePsiImage(); - RawImage generatePhiImage(); - -private: - void checkDims(RawImage& im); - - std::string filename; - unsigned width; - unsigned height; - - PointSpreadFunc psf; - RawImage science; - RawImage mask; - RawImage variance; -}; - -} /* namespace search */ - -#endif /* LAYEREDIMAGE_H_ */ diff --git a/src/kbmod/search/PointSpreadFunc.cpp b/src/kbmod/search/PointSpreadFunc.cpp deleted file mode 100644 index dd786aa94..000000000 --- a/src/kbmod/search/PointSpreadFunc.cpp +++ /dev/null @@ -1,150 +0,0 @@ -/* - * GeneratorPSF.cpp - * - * Created on: Nov 12, 2016 - * Author: peter - */ - -#include "PointSpreadFunc.h" - -namespace search { - -PointSpreadFunc::PointSpreadFunc() : kernel(1, 1.0) { - dim = 1; - radius = 0; - width = 1e-12; - sum = 1.0; -} - -PointSpreadFunc::PointSpreadFunc(float stdev) { - width = stdev; - float simple_gauss[MAX_KERNEL_RADIUS]; - double psf_coverage = 0.0; - double norm_factor = stdev * sqrt(2.0); - int i = 0; - - // Create 1D gaussian array - while (psf_coverage < 0.98 && i < MAX_KERNEL_RADIUS) { - float current_bin = - 0.5 * (std::erf((float(i) + 0.5) / norm_factor) - std::erf((float(i) - 0.5) / norm_factor)); - simple_gauss[i] = current_bin; - if (i == 0) { - psf_coverage += current_bin; - } else { - psf_coverage += 2.0 * current_bin; - } - i++; - } - - radius = i - 1; // This value is good for - dim = 2 * radius + 1; - - // Create 2D gaussain by multiplying with itself - kernel = std::vector(); - for (int ii = 0; ii < dim; ++ii) { - for (int jj = 0; jj < dim; ++jj) { - float current = simple_gauss[abs(radius - ii)] * simple_gauss[abs(radius - jj)]; - kernel.push_back(current); - } - } - calcSum(); -} - -// Copy constructor. -PointSpreadFunc::PointSpreadFunc(const PointSpreadFunc& other) { - kernel = other.kernel; - dim = other.dim; - radius = other.radius; - width = other.width; - sum = other.sum; -} - -// Copy assignment. -PointSpreadFunc& PointSpreadFunc::operator=(const PointSpreadFunc& other) { - kernel = other.kernel; - dim = other.dim; - radius = other.radius; - width = other.width; - sum = other.sum; - return *this; -} - -// Move constructor. -PointSpreadFunc::PointSpreadFunc(PointSpreadFunc&& other) - : kernel(std::move(other.kernel)), - dim(other.dim), - radius(other.radius), - width(other.width), - sum(other.sum) {} - -// Move assignment. -PointSpreadFunc& PointSpreadFunc::operator=(PointSpreadFunc&& other) { - if (this != &other) { - kernel = std::move(other.kernel); - dim = other.dim; - radius = other.radius; - width = other.width; - sum = other.sum; - } - return *this; -} - -#ifdef Py_PYTHON_H -PointSpreadFunc::PointSpreadFunc(pybind11::array_t arr) { setArray(arr); } - -void PointSpreadFunc::setArray(pybind11::array_t arr) { - pybind11::buffer_info info = arr.request(); - - if (info.ndim != 2) - throw std::runtime_error( - "Array must have 2 dimensions. (It " - " must also be a square with odd dimensions)"); - - if (info.shape[0] != info.shape[1]) - throw std::runtime_error( - "Array must be square (x-dimension == y-dimension)." - "It also must have odd dimensions."); - float* pix = static_cast(info.ptr); - dim = info.shape[0]; - if (dim % 2 == 0) - throw std::runtime_error( - "Array dimension must be odd. The " - "middle of an even numbered array is ambiguous."); - radius = dim / 2; // Rounds down - sum = 0.0; - kernel = std::vector(pix, pix + dim * dim); - calcSum(); - width = 0.0; -} -#endif - -void PointSpreadFunc::calcSum() { - sum = 0.0; - for (auto& i : kernel) sum += i; -} - -void PointSpreadFunc::squarePSF() { - for (float& i : kernel) { - i = i * i; - } - calcSum(); -} - -std::string PointSpreadFunc::printPSF() { - std::stringstream ss; - ss.setf(std::ios::fixed, std::ios::floatfield); - ss.precision(3); - for (int row = 0; row < dim; ++row) { - ss << "| "; - for (int col = 0; col < dim; ++col) { - ss << kernel[row * dim + col] << " | "; - } - ss << "\n "; - for (int space = 0; space < dim * 8 - 1; ++space) ss << "-"; - ss << "\n"; - } - ss << 100.0 * sum << "% of PSF contained within kernel\n"; - return ss.str(); -} - -} /* namespace search */ diff --git a/src/kbmod/search/PointSpreadFunc.h b/src/kbmod/search/PointSpreadFunc.h deleted file mode 100644 index d270ae6ca..000000000 --- a/src/kbmod/search/PointSpreadFunc.h +++ /dev/null @@ -1,68 +0,0 @@ -/* - * PointSpreadFunc.h - * - * Created on: Nov 12, 2016 - * Author: peter - * - * A class for working with point spread functions. - */ - -#ifndef POINTSPREADFUNC_H_ -#define POINTSPREADFUNC_H_ - -#include -#include -#include -#include -#include -#ifdef Py_PYTHON_H -#include -#include -#include -#endif -#include "common.h" - -namespace search { - -class PointSpreadFunc { -public: - PointSpreadFunc(); // Create a no-op PSF. - PointSpreadFunc(float stdev); - PointSpreadFunc(const PointSpreadFunc& other); // Copy constructor - PointSpreadFunc(PointSpreadFunc&& other); // Move constructor -#ifdef Py_PYTHON_H - PointSpreadFunc(pybind11::array_t arr); - void setArray(pybind11::array_t arr); -#endif - virtual ~PointSpreadFunc(){}; - - // Assignment functions. - PointSpreadFunc& operator=(const PointSpreadFunc& other); // Copy assignment - PointSpreadFunc& operator=(PointSpreadFunc&& other); // Move assignment - - // Getter functions (inlined) - float getStdev() const { return width; } - float getSum() const { return sum; } - float getValue(int x, int y) const { return kernel[y * dim + x]; } - int getDim() const { return dim; } - int getRadius() const { return radius; } - int getSize() const { return kernel.size(); } - const std::vector& getKernel() const { return kernel; }; - float* kernelData() { return kernel.data(); } - - // Computation functions. - void calcSum(); - void squarePSF(); - std::string printPSF(); - -private: - std::vector kernel; - float width; - float sum; - int dim; - int radius; -}; - -} /* namespace search */ - -#endif /* POINTSPREADFUNC_H_ */ diff --git a/src/kbmod/search/RawImage.cpp b/src/kbmod/search/RawImage.cpp deleted file mode 100644 index 637d69de8..000000000 --- a/src/kbmod/search/RawImage.cpp +++ /dev/null @@ -1,633 +0,0 @@ -/* - * RawImage.cpp - * - * Created on: Jun 22, 2017 - * Author: kbmod-usr - */ - -#include "RawImage.h" - -namespace search { - -#ifdef HAVE_CUDA -// Performs convolution between an image represented as an array of floats -// and a PSF on a GPU device. -extern "C" void deviceConvolve(float* source_img, float* result_img, int width, int height, float* psf_kernel, - int psf_size, int psf_dim, int psf_radius, float psf_sum); -#endif - -RawImage::RawImage() : width(0), height(0), obstime(-1.0) { pixels = std::vector(); } - -// Copy constructor -RawImage::RawImage(const RawImage& old) { - width = old.getWidth(); - height = old.getHeight(); - pixels = old.getPixels(); - obstime = old.getObstime(); -} - -// Copy assignment -RawImage& RawImage::operator=(const RawImage& source) { - width = source.width; - height = source.height; - pixels = source.pixels; - obstime = source.obstime; - return *this; -} - -// Move constructor -RawImage::RawImage(RawImage&& source) - : width(source.width), - height(source.height), - obstime(source.obstime), - pixels(std::move(source.pixels)) {} - -// Move assignment -RawImage& RawImage::operator=(RawImage&& source) { - if (this != &source) { - width = source.width; - height = source.height; - pixels = std::move(source.pixels); - obstime = source.obstime; - } - return *this; -} - -RawImage::RawImage(unsigned w, unsigned h) : height(h), width(w), obstime(-1.0), pixels(w * h) {} - -RawImage::RawImage(unsigned w, unsigned h, const std::vector& pix) - : width(w), height(h), obstime(-1.0), pixels(pix) { - assert(w * h == pix.size()); -} - -#ifdef Py_PYTHON_H -RawImage::RawImage(pybind11::array_t arr) { - obstime = -1.0; - setArray(arr); -} - -void RawImage::setArray(pybind11::array_t& arr) { - pybind11::buffer_info info = arr.request(); - - if (info.ndim != 2) throw std::runtime_error("Array must have 2 dimensions."); - - width = info.shape[1]; - height = info.shape[0]; - float* pix = static_cast(info.ptr); - - pixels = std::vector(pix, pix + getNPixels()); -} -#endif - -bool RawImage::approxEqual(const RawImage& img_b, float atol) const { - if ((width != img_b.width) || (height != img_b.height)) return false; - - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - float p1 = getPixel(x, y); - float p2 = img_b.getPixel(x, y); - - // NO_DATA values must match exactly. - if ((p1 == NO_DATA) && (p2 != NO_DATA)) return false; - if ((p1 != NO_DATA) && (p2 == NO_DATA)) return false; - - // Other values match within an absolute tolerance. - if (fabs(p1 - p2) > atol) return false; - } - } - return true; -} - -// Load the image data from a specific layer of a FITS file. -void RawImage::loadFromFile(const std::string& file_path, int layer_num) { - // Open the file's header and read in the obstime and the dimensions. - fitsfile* fptr; - int status = 0; - int mjdStatus = 0; - int file_not_found; - int nullval = 0; - int anynull = 0; - - // Open the correct layer to extract the RawImage. - std::string layerPath = file_path + "[" + std::to_string(layer_num) + "]"; - if (fits_open_file(&fptr, layerPath.c_str(), READONLY, &status)) { - fits_report_error(stderr, status); - throw std::runtime_error("Could not open FITS file to read RawImage"); - } - - // Read image dimensions. - long dimensions[2]; - if (fits_read_keys_lng(fptr, "NAXIS", 1, 2, dimensions, &file_not_found, &status)) - fits_report_error(stderr, status); - width = dimensions[0]; - height = dimensions[1]; - - // Read in the image. - pixels = std::vector(width * height); - if (fits_read_img(fptr, TFLOAT, 1, getNPixels(), &nullval, pixels.data(), &anynull, &status)) - fits_report_error(stderr, status); - - // Read image observation time, ignore error if does not exist - obstime = -1.0; - if (fits_read_key(fptr, TDOUBLE, "MJD", &obstime, NULL, &mjdStatus)) obstime = -1.0; - if (fits_close_file(fptr, &status)) fits_report_error(stderr, status); - - // If we are reading from a sublayer and did not find a time, try the overall header. - if (obstime < 0.0) { - if (fits_open_file(&fptr, file_path.c_str(), READONLY, &status)) - throw std::runtime_error("Could not open FITS file to read RawImage"); - fits_read_key(fptr, TDOUBLE, "MJD", &obstime, NULL, &mjdStatus); - if (fits_close_file(fptr, &status)) fits_report_error(stderr, status); - } -} - -void RawImage::saveToFile(const std::string& filename) { - fitsfile* fptr; - int status = 0; - long naxes[2] = {0, 0}; - - fits_create_file(&fptr, filename.c_str(), &status); - - // If we are unable to create the file, check if it already exists - // and, if so, delete it and retry the create. - if (status == 105) { - status = 0; - fits_open_file(&fptr, filename.c_str(), READWRITE, &status); - if (status == 0) { - fits_delete_file(fptr, &status); - fits_create_file(&fptr, filename.c_str(), &status); - } - } - - // Create the primary array image (32-bit float pixels) - long dimensions[2]; - dimensions[0] = width; - dimensions[1] = height; - fits_create_img(fptr, FLOAT_IMG, 2 /*naxis*/, dimensions, &status); - fits_report_error(stderr, status); - - /* Write the array of floats to the image */ - fits_write_img(fptr, TFLOAT, 1, getNPixels(), pixels.data(), &status); - fits_report_error(stderr, status); - - // Add the basic header data. - fits_update_key(fptr, TDOUBLE, "MJD", &obstime, "[d] Generated Image time", &status); - fits_report_error(stderr, status); - - fits_close_file(fptr, &status); - fits_report_error(stderr, status); -} - -void RawImage::appendLayerToFile(const std::string& filename) { - int status = 0; - fitsfile* f; - - // Check that we can open the file. - if (fits_open_file(&f, filename.c_str(), READWRITE, &status)) { - fits_report_error(stderr, status); - throw std::runtime_error("Unable to open FITS file for appending."); - } - - // This appends a layer (extension) if the file exists) - /* Create the primary array image (32-bit float pixels) */ - long dimensions[2]; - dimensions[0] = width; - dimensions[1] = height; - fits_create_img(f, FLOAT_IMG, 2 /*naxis*/, dimensions, &status); - fits_report_error(stderr, status); - - /* Write the array of floats to the image */ - fits_write_img(f, TFLOAT, 1, getNPixels(), pixels.data(), &status); - fits_report_error(stderr, status); - - // Save the image time in the header. - fits_update_key(f, TDOUBLE, "MJD", &obstime, "[d] Generated Image time", &status); - fits_report_error(stderr, status); - - fits_close_file(f, &status); - fits_report_error(stderr, status); -} - -RawImage RawImage::createStamp(float x, float y, int radius, bool interpolate, bool keep_no_data) const { - if (radius < 0) throw std::runtime_error("stamp radius must be at least 0"); - - int dim = radius * 2 + 1; - RawImage stamp(dim, dim); - for (int yoff = 0; yoff < dim; ++yoff) { - for (int xoff = 0; xoff < dim; ++xoff) { - float pix_val; - if (interpolate) - pix_val = getPixelInterp(x + static_cast(xoff - radius), - y + static_cast(yoff - radius)); - else - pix_val = getPixel(static_cast(x) + xoff - radius, static_cast(y) + yoff - radius); - if ((pix_val == NO_DATA) && !keep_no_data) pix_val = 0.0; - stamp.setPixel(xoff, yoff, pix_val); - } - } - - stamp.setObstime(obstime); - return stamp; -} - -void RawImage::convolve_cpu(const PointSpreadFunc& psf) { - std::vector result(width * height, 0.0); - const int psf_rad = psf.getRadius(); - const float psf_total = psf.getSum(); - - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - // Pixels with NO_DATA remain NO_DATA. - if (pixels[y * width + x] == NO_DATA) { - result[y * width + x] = NO_DATA; - continue; - } - - float sum = 0.0; - float psf_portion = 0.0; - for (int j = -psf_rad; j <= psf_rad; j++) { - for (int i = -psf_rad; i <= psf_rad; i++) { - if ((x + i >= 0) && (x + i < width) && (y + j >= 0) && (y + j < height)) { - float current_pixel = pixels[(y + j) * width + (x + i)]; - if (current_pixel != NO_DATA) { - float current_psf = psf.getValue(i + psf_rad, j + psf_rad); - psf_portion += current_psf; - sum += current_pixel * current_psf; - } - } - } - } - result[y * width + x] = (sum * psf_total) / psf_portion; - } - } - - // Copy the data into the pixels vector. - const int npixels = getNPixels(); - for (int i = 0; i < npixels; ++i) { - pixels[i] = result[i]; - } -} - -void RawImage::convolve(PointSpreadFunc psf) { -#ifdef HAVE_CUDA - deviceConvolve(pixels.data(), pixels.data(), getWidth(), getHeight(), psf.kernelData(), psf.getSize(), - psf.getDim(), psf.getRadius(), psf.getSum()); -#else - convolve_cpu(psf); -#endif -} - -void RawImage::applyMask(int flags, const std::vector& exceptions, const RawImage& mask) { - const std::vector& mask_pix = mask.getPixels(); - const int num_pixels = getNPixels(); - assert(num_pixels == mask.getNPixels()); - for (unsigned int p = 0; p < num_pixels; ++p) { - int pix_flags = static_cast(mask_pix[p]); - bool is_exception = false; - for (auto& e : exceptions) is_exception = is_exception || e == pix_flags; - if (!is_exception && ((flags & pix_flags) != 0)) pixels[p] = NO_DATA; - } -} - -/* This implementation of growMask is optimized for steps > 1 - (which is how the code is generally used. If you are only - growing the mask by 1, the extra copy will be a little slower. -*/ -void RawImage::growMask(int steps) { - const int num_pixels = width * height; - - // Set up the initial masked vector that stores the number of steps - // each pixel is from a masked pixel in the original image. - std::vector masked(num_pixels, -1); - for (int i = 0; i < num_pixels; ++i) { - if (pixels[i] == NO_DATA) masked[i] = 0; - } - - // Grow out the mask one for each step. - for (int itr = 1; itr <= steps; ++itr) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - int center = width * y + x; - if (masked[center] == -1) { - // Mask pixels that are adjacent to a pixel masked during - // the last iteration only. - if ((x + 1 < width && masked[center + 1] == itr - 1) || - (x - 1 >= 0 && masked[center - 1] == itr - 1) || - (y + 1 < height && masked[center + width] == itr - 1) || - (y - 1 >= 0 && masked[center - width] == itr - 1)) { - masked[center] = itr; - } - } - } - } - } - - // Mask the pixels in the image. - for (std::size_t i = 0; i < num_pixels; ++i) { - if (masked[i] > -1) { - pixels[i] = NO_DATA; - } - } -} - -std::vector RawImage::bilinearInterp(float x, float y) const { - // Linear interpolation - // Find the 4 pixels (aPix, bPix, cPix, dPix) - // that the corners (a, b, c, d) of the - // new pixel land in, and blend into those - - // Returns a vector with 4 pixel locations - // and their interpolation value - - // Top right - float ax = x + 0.5; - float ay = y + 0.5; - float a_px = floor(ax); - float a_py = floor(ay); - float a_amt = (ax - a_px) * (ay - a_py); - - // Bottom right - float bx = x + 0.5; - float by = y - 0.5; - float b_px = floor(bx); - float b_py = floor(by); - float b_amt = (bx - b_px) * (b_py + 1.0 - by); - - // Bottom left - float cx = x - 0.5; - float cy = y - 0.5; - float c_px = floor(cx); - float c_py = floor(cy); - float c_amt = (c_px + 1.0 - cx) * (c_py + 1.0 - cy); - - // Top left - float dx = x - 0.5; - float dy = y + 0.5; - float d_px = floor(dx); - float d_py = floor(dy); - float d_amt = (d_px + 1.0 - dx) * (dy - d_py); - - // make sure the right amount has been distributed - float diff = std::abs(a_amt + b_amt + c_amt + d_amt - 1.0); - if (diff > 0.01) std::cout << "warning: bilinearInterpSum == " << diff << "\n"; - return {a_px, a_py, a_amt, b_px, b_py, b_amt, c_px, c_py, c_amt, d_px, d_py, d_amt}; -} - -void RawImage::addPixelInterp(float x, float y, float value) { - // Interpolation values - std::vector iv = bilinearInterp(x, y); - - addToPixel(iv[0], iv[1], value * iv[2]); - addToPixel(iv[3], iv[4], value * iv[5]); - addToPixel(iv[6], iv[7], value * iv[8]); - addToPixel(iv[9], iv[10], value * iv[11]); -} - -void RawImage::addToPixel(float fx, float fy, float value) { - assert(fx - floor(fx) == 0.0 && fy - floor(fy) == 0.0); - int x = static_cast(fx); - int y = static_cast(fy); - if (x >= 0 && x < width && y >= 0 && y < height) pixels[y * width + x] += value; -} - -float RawImage::getPixelInterp(float x, float y) const { - if ((x < 0.0 || y < 0.0) || (x > static_cast(width) || y > static_cast(height))) - return NO_DATA; - std::vector iv = bilinearInterp(x, y); - float a = getPixel(iv[0], iv[1]); - float b = getPixel(iv[3], iv[4]); - float c = getPixel(iv[6], iv[7]); - float d = getPixel(iv[9], iv[10]); - float interpSum = 0.0; - float total = 0.0; - if (a != NO_DATA) { - interpSum += iv[2]; - total += a * iv[2]; - } - if (b != NO_DATA) { - interpSum += iv[5]; - total += b * iv[5]; - } - if (c != NO_DATA) { - interpSum += iv[8]; - total += c * iv[8]; - } - if (d != NO_DATA) { - interpSum += iv[11]; - total += d * iv[11]; - } - if (interpSum == 0.0) { - return NO_DATA; - } else { - return total / interpSum; - } -} - -void RawImage::setAllPix(float value) { - for (auto& p : pixels) p = value; -} - -std::array RawImage::computeBounds() const { - const int num_pixels = getNPixels(); - float min_val = FLT_MAX; - float max_val = -FLT_MAX; - for (unsigned p = 0; p < num_pixels; ++p) { - if (pixels[p] != NO_DATA) { - min_val = std::min(min_val, pixels[p]); - max_val = std::max(max_val, pixels[p]); - } - } - - // Assert that we have seen at least some valid data. - assert(max_val != -FLT_MAX); - assert(min_val != FLT_MAX); - - // Set and return the result array. - std::array res; - res[0] = min_val; - res[1] = max_val; - return res; -} - -// The maximum value of the image and return the coordinates. -PixelPos RawImage::findPeak(bool furthest_from_center) const { - int c_x = width / 2; - int c_y = height / 2; - - // Initialize the variables for tracking the peak's location. - PixelPos result = {0, 0}; - float max_val = pixels[0]; - float dist2 = c_x * c_x + c_y * c_y; - - // Search each pixel for the peak. - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - float pix_val = pixels[y * width + x]; - if (pix_val > max_val) { - max_val = pix_val; - result.x = x; - result.y = y; - dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y); - } else if (pix_val == max_val) { - int new_dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y); - if ((furthest_from_center && (new_dist2 > dist2)) || - (!furthest_from_center && (new_dist2 < dist2))) { - max_val = pix_val; - result.x = x; - result.y = y; - dist2 = new_dist2; - } - } - } - } - - return result; -} - -// Find the basic image moments in order to test if stamps have a gaussian shape. -// It computes the moments on the "normalized" image where the minimum -// value has been shifted to zero and the sum of all elements is 1.0. -// Elements with NO_DATA are treated as zero. -ImageMoments RawImage::findCentralMoments() const { - const int num_pixels = width * height; - const int c_x = width / 2; - const int c_y = height / 2; - - // Set all the moments to zero initially. - ImageMoments res = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0}; - - // Find the min (non-NO_DATA) value to subtract off. - float min_val = FLT_MAX; - for (int p = 0; p < num_pixels; ++p) { - min_val = ((pixels[p] != NO_DATA) && (pixels[p] < min_val)) ? pixels[p] : min_val; - } - - // Find the sum of the zero-shifted (non-NO_DATA) pixels. - double sum = 0.0; - for (int p = 0; p < num_pixels; ++p) { - sum += (pixels[p] != NO_DATA) ? (pixels[p] - min_val) : 0.0; - } - if (sum == 0.0) return res; - - // Compute the rest of the moments. - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - int ind = y * width + x; - float pix_val = (pixels[ind] != NO_DATA) ? (pixels[ind] - min_val) / sum : 0.0; - - res.m00 += pix_val; - res.m10 += (x - c_x) * pix_val; - res.m20 += (x - c_x) * (x - c_x) * pix_val; - res.m01 += (y - c_y) * pix_val; - res.m02 += (y - c_y) * (y - c_y) * pix_val; - res.m11 += (x - c_x) * (y - c_y) * pix_val; - } - } - - return res; -} - -RawImage createMedianImage(const std::vector& images) { - int num_images = images.size(); - assert(num_images > 0); - - int width = images[0].getWidth(); - int height = images[0].getHeight(); - for (auto& img : images) assert(img.getWidth() == width and img.getHeight() == height); - - RawImage result = RawImage(width, height); - std::vector pix_array(num_images); - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - int num_unmasked = 0; - for (int i = 0; i < num_images; ++i) { - // Only used the unmasked pixels. - float pix_val = images[i].getPixel(x, y); - if ((pix_val != NO_DATA) && (!std::isnan(pix_val))) { - pix_array[num_unmasked] = pix_val; - num_unmasked += 1; - } - } - - if (num_unmasked > 0) { - std::sort(pix_array.begin(), pix_array.begin() + num_unmasked); - - // If we have an even number of elements, take the mean of the two - // middle ones. - int median_ind = num_unmasked / 2; - if (num_unmasked % 2 == 0) { - float ave_middle = (pix_array[median_ind] + pix_array[median_ind - 1]) / 2.0; - result.setPixel(x, y, ave_middle); - } else { - result.setPixel(x, y, pix_array[median_ind]); - } - } else { - // We use a 0.0 value if there is no data to allow for visualization - // and value based filtering. - result.setPixel(x, y, 0.0); - } - } - } - - return result; -} - -RawImage createSummedImage(const std::vector& images) { - int num_images = images.size(); - assert(num_images > 0); - - int width = images[0].getWidth(); - int height = images[0].getHeight(); - for (auto& img : images) assert(img.getWidth() == width and img.getHeight() == height); - - RawImage result = RawImage(width, height); - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - float sum = 0.0; - for (int i = 0; i < num_images; ++i) { - float pix_val = images[i].getPixel(x, y); - if ((pix_val == NO_DATA) || (std::isnan(pix_val))) pix_val = 0.0; - sum += pix_val; - } - result.setPixel(x, y, sum); - } - } - - return result; -} - -RawImage createMeanImage(const std::vector& images) { - int num_images = images.size(); - assert(num_images > 0); - - int width = images[0].getWidth(); - int height = images[0].getHeight(); - for (auto& img : images) assert(img.getWidth() == width and img.getHeight() == height); - - RawImage result = RawImage(width, height); - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - float sum = 0.0; - float count = 0.0; - for (int i = 0; i < num_images; ++i) { - float pix_val = images[i].getPixel(x, y); - if ((pix_val != NO_DATA) && (!std::isnan(pix_val))) { - count += 1.0; - sum += pix_val; - } - } - - if (count > 0.0) { - result.setPixel(x, y, sum / count); - } else { - // We use a 0.0 value if there is no data to allow for visualization - // and value based filtering. - result.setPixel(x, y, 0.0); - } - } - } - - return result; -} - -} /* namespace search */ diff --git a/src/kbmod/search/bindings.cpp b/src/kbmod/search/bindings.cpp index ebf6c49b1..314ff8871 100644 --- a/src/kbmod/search/bindings.cpp +++ b/src/kbmod/search/bindings.cpp @@ -2,288 +2,43 @@ #include #include -#include "PointSpreadFunc.cpp" -#include "RawImage.cpp" -#include "LayeredImage.cpp" -#include "ImageStack.cpp" -#include "KBMOSearch.cpp" -#include "Filtering.cpp" +#include "psf.cpp" +#include "raw_image.cpp" +#include "layered_image.cpp" +#include "image_stack.cpp" +#include "stack_search.cpp" +#include "filtering.cpp" +#include "common.h" -namespace py = pybind11; -using pf = search::PointSpreadFunc; -using ri = search::RawImage; -using li = search::LayeredImage; -using is = search::ImageStack; -using ks = search::KBMOSearch; -using tj = search::trajectory; using bc = search::BaryCorrection; using pp = search::PixelPos; - using std::to_string; -PYBIND11_MODULE(search, m) { - m.attr("KB_NO_DATA") = pybind11::float_(search::NO_DATA); - m.attr("HAS_GPU") = pybind11::bool_(search::HAVE_GPU); - py::enum_(m, "StampType") - .value("STAMP_SUM", search::StampType::STAMP_SUM) - .value("STAMP_MEAN", search::StampType::STAMP_MEAN) - .value("STAMP_MEDIAN", search::StampType::STAMP_MEDIAN) - .export_values(); - py::class_(m, "psf", py::buffer_protocol(), R"pbdoc( - Point Spread Function. - - Parameters - ---------- - arg : `float`, `numpy.array` or `psf` - Given value represents one of: - - * standard deviation of a Gaussian PSF (`float`) - * kernel representing the PSF (array-like) - * another `psf` object. - - - Notes - ----- - When instantiated with another `psf` object, returns its copy. - When instantiated with an array-like object, that array must be - a square matrix and have an odd number of dimensions - )pbdoc") - .def_buffer([](pf &m) -> py::buffer_info { - return py::buffer_info(m.kernelData(), sizeof(float), py::format_descriptor::format(), - 2, {m.getDim(), m.getDim()}, - {sizeof(float) * m.getDim(), sizeof(float)}); - }) - .def(py::init<>()) - .def(py::init()) - .def(py::init>()) - .def(py::init()) - .def("set_array", &pf::setArray, R"pbdoc( - Sets the PSF kernel. - - Parameters - ---------- - arr : `numpy.array` - Numpy array representing the PSF. - )pbdoc") - .def("get_stdev", &pf::getStdev, "Returns the PSF's standard deviation.") - .def("get_sum", &pf::getSum, "Returns the sum of PSFs kernel elements.") - .def("get_dim", &pf::getDim, "Returns the PSF kernel dimensions.") - .def("get_radius", &pf::getRadius, "Returns the radius of the PSF") - .def("get_size", &pf::getSize, "Returns the number of elements in the PSFs kernel.") - .def("get_kernel", &pf::getKernel, "Returns the PSF kernel.") - .def("get_value", &pf::getValue, "Returns the PSF kernel value at a specific point.") - .def("square_psf", &pf::squarePSF, - "Squares, raises to the power of two, the elements of the PSF kernel.") - .def("print_psf", &pf::printPSF, "Pretty-prints the PSF."); - - py::class_(m, "raw_image", py::buffer_protocol()) - .def_buffer([](ri &m) -> py::buffer_info { - return py::buffer_info(m.getDataRef(), sizeof(float), py::format_descriptor::format(), - 2, {m.getHeight(), m.getWidth()}, - {sizeof(float) * m.getWidth(), sizeof(float)}); - }) - .def(py::init()) - .def(py::init()) - .def(py::init>()) - .def("get_height", &ri::getHeight, "Returns the image's height in pixels.") - .def("get_width", &ri::getWidth, "Returns the image's width in pixels.") - .def("get_npixels", &ri::getNPixels, "Returns the image's total number of pixels.") - .def("get_all_pixels", &ri::getPixels, "Returns a list of the images pixels.") - .def("set_array", &ri::setArray, "Sets all image pixels given an array of values.") - .def("get_obstime", &ri::getObstime, "Get the observation time of the image.") - .def("set_obstime", &ri::setObstime, "Set the observation time of the image.") - .def("approx_equal", &ri::approxEqual, "Checks if two images are approximately equal.") - .def("compute_bounds", &ri::computeBounds, "Returns min and max pixel values.") - .def("find_peak", &ri::findPeak, "Returns the pixel coordinates of the maximum value.") - .def("find_central_moments", &ri::findCentralMoments, "Returns the central moments of the image.") - .def("create_stamp", &ri::createStamp) - .def("set_pixel", &ri::setPixel, "Set the value of a given pixel.") - .def("add_pixel", &ri::addToPixel, "Add to the value of a given pixel.") - .def("add_pixel_interp", &ri::addPixelInterp, "Add to the interpolated value of a given pixel.") - .def("apply_mask", &ri::applyMask) - .def("grow_mask", &ri::growMask) - .def("pixel_has_data", &ri::pixelHasData, - "Returns a Boolean indicating whether the pixel has data.") - .def("set_all", &ri::setAllPix, "Set all pixel values given an array.") - .def("get_pixel", &ri::getPixel, "Returns the value of a pixel.") - .def("get_pixel_interp", &ri::getPixelInterp, "Get the interoplated value of a pixel.") - .def("convolve", &ri::convolve, "Convolve the image with a PSF.") - .def("convolve_cpu", &ri::convolve_cpu, "Convolve the image with a PSF.") - .def("load_fits", &ri::loadFromFile, "Load the image data from a FITS file.") - .def("save_fits", &ri::saveToFile, "Save the image to a FITS file.") - .def("append_fits_layer", &ri::appendLayerToFile, "Append the image as a layer in a FITS file."); - m.def("create_median_image", &search::createMedianImage); - m.def("create_summed_image", &search::createSummedImage); - m.def("create_mean_image", &search::createMeanImage); - py::class_
  • (m, "layered_image") - .def(py::init()) - .def(py::init(), R"pbdoc( - Creates a layered_image out of individual `raw_image` layers. - - Parameters - ---------- - sci : `raw_image` - The `raw_image` for the science layer. - var : `raw_image` - The `raw_image` for the cariance layer. - msk : `raw_image` - The `raw_image` for the mask layer. - p : `psf` - The PSF for the image. - - Raises - ------ - Raises an exception if the layers are not the same size. - )pbdoc") - .def(py::init()) - .def(py::init()) - .def("set_psf", &li::setPSF, "Sets the PSF object.") - .def("get_psf", &li::getPSF, "Returns the PSF object.") - .def("apply_mask_flags", &li::applyMaskFlags) - .def("apply_mask_threshold", &li::applyMaskThreshold) - .def("sub_template", &li::subtractTemplate) - .def("save_layers", &li::saveLayers) - .def("get_science", &li::getScience, "Returns the science layer raw_image.") - .def("get_mask", &li::getMask, "Returns the mask layer raw_image.") - .def("get_variance", &li::getVariance, "Returns the variance layer raw_image.") - .def("set_science", &li::setScience) - .def("set_mask", &li::setMask) - .def("set_variance", &li::setVariance) - .def("convolve_psf", &li::convolvePSF, "Convolve each layer with the object's PSF.") - .def("convolve_given_psf", &li::convolveGivenPSF, "Convolve each layer with a given PSF.") - .def("grow_mask", &li::growMask) - .def("get_name", &li::getName, "Returns the name of the layered image.") - .def("get_width", &li::getWidth, "Returns the image's width in pixels.") - .def("get_height", &li::getHeight, "Returns the image's height in pixels.") - .def("get_npixels", &li::getNPixels, "Returns the image's total number of pixels.") - .def("get_obstime", &li::getObstime, "Get the image's observation time.") - .def("set_obstime", &li::setObstime, "Set the image's observation time.") - .def("generate_psi_image", &li::generatePsiImage) - .def("generate_phi_image", &li::generatePhiImage); - py::class_(m, "image_stack") - .def(py::init, std::vector>()) - .def(py::init>()) - .def("get_images", &is::getImages) - .def("get_single_image", &is::getSingleImage) - .def("set_single_image", &is::setSingleImage) - .def("get_times", &is::getTimes) - .def("set_times", &is::setTimes) - .def("img_count", &is::imgCount) - .def("apply_mask_flags", &is::applyMaskFlags) - .def("apply_mask_threshold", &is::applyMaskThreshold) - .def("apply_global_mask", &is::applyGlobalMask) - .def("grow_mask", &is::growMask) - .def("save_global_mask", &is::saveGlobalMask) - .def("save_images", &is::saveImages) - .def("get_global_mask", &is::getGlobalMask) - .def("convolve_psf", &is::convolvePSF) - .def("get_width", &is::getWidth) - .def("get_height", &is::getHeight) - .def("get_npixels", &is::getNPixels); - py::class_(m, "stack_search") - .def(py::init()) - .def("save_psi_phi", &ks::savePsiPhi) - .def("search", &ks::search) - .def("enable_gpu_sigmag_filter", &ks::enableGPUSigmaGFilter) - .def("enable_gpu_encoding", &ks::enableGPUEncoding) - .def("enable_corr", &ks::enableCorr) - .def("set_start_bounds_x", &ks::setStartBoundsX) - .def("set_start_bounds_y", &ks::setStartBoundsY) - .def("set_debug", &ks::setDebug) - .def("filter_min_obs", &ks::filterResults) - .def("get_num_images", &ks::numImages) - .def("get_image_stack", &ks::getImageStack) - // Science Stamp Functions - .def("science_viz_stamps", &ks::scienceStampsForViz) - .def("median_sci_stamp", &ks::medianScienceStamp) - .def("mean_sci_stamp", &ks::meanScienceStamp) - .def("summed_sci_stamp", &ks::summedScienceStamp) - .def("coadded_stamps", - (std::vector(ks::*)(std::vector &, std::vector> &, - const search::StampParameters &, bool)) & - ks::coaddedScienceStamps) - // For testing - .def("filter_stamp", &ks::filterStamp) - .def("get_traj_pos", &ks::getTrajPos) - .def("get_mult_traj_pos", &ks::getMultTrajPos) - .def("psi_curves", (std::vector(ks::*)(tj &)) & ks::psiCurves) - .def("phi_curves", (std::vector(ks::*)(tj &)) & ks::phiCurves) - .def("prepare_psi_phi", &ks::preparePsiPhi) - .def("get_psi_images", &ks::getPsiImages) - .def("get_phi_images", &ks::getPhiImages) - .def("get_results", &ks::getResults) - .def("set_results", &ks::setResults); - py::class_(m, "trajectory", R"pbdoc( - A trajectory structure holding basic information about potential results. - )pbdoc") - .def(py::init<>()) - .def_readwrite("x_v", &tj::x_vel) - .def_readwrite("y_v", &tj::y_vel) - .def_readwrite("lh", &tj::lh) - .def_readwrite("flux", &tj::flux) - .def_readwrite("x", &tj::x) - .def_readwrite("y", &tj::y) - .def_readwrite("obs_count", &tj::obs_count) - .def("__repr__", - [](const tj &t) { - return "lh: " + to_string(t.lh) + " flux: " + to_string(t.flux) + - " x: " + to_string(t.x) + " y: " + to_string(t.y) + - " x_v: " + to_string(t.x_vel) + " y_v: " + to_string(t.y_vel) + - " obs_count: " + to_string(t.obs_count); - }) - .def(py::pickle( - [](const tj &p) { // __getstate__ - return py::make_tuple(p.x_vel, p.y_vel, p.lh, p.flux, p.x, p.y, p.obs_count); - }, - [](py::tuple t) { // __setstate__ - if (t.size() != 7) throw std::runtime_error("Invalid state!"); - tj trj = {t[0].cast(), t[1].cast(), t[2].cast(), - t[3].cast(), t[4].cast(), t[5].cast(), - t[6].cast()}; - return trj; - })); - py::class_(m, "pixel_pos") - .def(py::init<>()) - .def_readwrite("x", &pp::x) - .def_readwrite("y", &pp::y) - .def("__repr__", [](const pp &p) { return "x: " + to_string(p.x) + " y: " + to_string(p.y); }); - py::class_(m, "image_moments") - .def(py::init<>()) - .def_readwrite("m00", &search::ImageMoments::m00) - .def_readwrite("m01", &search::ImageMoments::m01) - .def_readwrite("m10", &search::ImageMoments::m10) - .def_readwrite("m11", &search::ImageMoments::m11) - .def_readwrite("m02", &search::ImageMoments::m02) - .def_readwrite("m20", &search::ImageMoments::m20); - py::class_(m, "stamp_parameters") - .def(py::init<>()) - .def_readwrite("radius", &search::StampParameters::radius) - .def_readwrite("stamp_type", &search::StampParameters::stamp_type) - .def_readwrite("do_filtering", &search::StampParameters::do_filtering) - .def_readwrite("center_thresh", &search::StampParameters::center_thresh) - .def_readwrite("peak_offset_x", &search::StampParameters::peak_offset_x) - .def_readwrite("peak_offset_y", &search::StampParameters::peak_offset_y) - .def_readwrite("m01", &search::StampParameters::m01_limit) - .def_readwrite("m10", &search::StampParameters::m10_limit) - .def_readwrite("m11", &search::StampParameters::m11_limit) - .def_readwrite("m02", &search::StampParameters::m02_limit) - .def_readwrite("m20", &search::StampParameters::m20_limit); - py::class_(m, "BaryCorrection") - .def(py::init<>()) - .def_readwrite("dx", &bc::dx) - .def_readwrite("dxdx", &bc::dxdx) - .def_readwrite("dxdy", &bc::dxdy) - .def_readwrite("dy", &bc::dy) - .def_readwrite("dydx", &bc::dydx) - .def_readwrite("dydy", &bc::dydy) - .def("__repr__", [](const bc &b) { - return "dx = " + to_string(b.dx) + " + " + to_string(b.dxdx) + " x + " + to_string(b.dxdy) + - " y; " + " dy = " + to_string(b.dy) + " + " + to_string(b.dydx) + " x + " + - to_string(b.dydy) + " y"; - }); - // Functions from Filtering.cpp - m.def("sigmag_filtered_indices", &search::sigmaGFilteredIndices); - m.def("calculate_likelihood_psi_phi", &search::calculateLikelihoodFromPsiPhi); +PYBIND11_MODULE(search, m) { + m.attr("KB_NO_DATA") = pybind11::float_(search::NO_DATA); + m.attr("HAS_GPU") = pybind11::bool_(search::HAVE_GPU); + py::enum_(m, "StampType") + .value("STAMP_SUM", search::StampType::STAMP_SUM) + .value("STAMP_MEAN", search::StampType::STAMP_MEAN) + .value("STAMP_MEDIAN", search::StampType::STAMP_MEDIAN) + .export_values(); + search::psf_bindings(m); + search::raw_image_bindings(m); + search::layered_image_bindings(m); + search::image_stack_bindings(m); + search::stack_search_bindings(m); + search::trajectory_bindings(m); + search::pixel_pos_bindings(m); + search::image_moments_bindings(m); + search::stamp_parameters_bindings(m); + search::bary_correction_bindings(m); + // Functions from raw_image.cpp + m.def("create_median_image", &search::create_median_image); + m.def("create_summed_image", &search::create_summed_image); + m.def("create_mean_image", &search::create_mean_image); + // Functions from filtering.cpp + m.def("sigmag_filtered_indices", &search::sigmaGFilteredIndices); + m.def("calculate_likelihood_psi_phi", &search::calculateLikelihoodFromPsiPhi); } diff --git a/src/kbmod/search/common.h b/src/kbmod/search/common.h index 529cba5a8..c2fa124a8 100644 --- a/src/kbmod/search/common.h +++ b/src/kbmod/search/common.h @@ -1,40 +1,37 @@ -/* - * common.h - * - * Created on: Jun 20, 2017 - * Author: kbmod-usr - */ - #ifndef COMMON_H_ #define COMMON_H_ -namespace search { +#include +#include "pydocs/common_docs.h" + + +namespace search { #ifdef HAVE_CUDA -constexpr bool HAVE_GPU = true; + constexpr bool HAVE_GPU = true; #else -constexpr bool HAVE_GPU = false; + constexpr bool HAVE_GPU = false; #endif -constexpr unsigned int MAX_KERNEL_RADIUS = 15; -constexpr unsigned short MAX_STAMP_EDGE = 64; -constexpr unsigned short CONV_THREAD_DIM = 32; -constexpr unsigned short THREAD_DIM_X = 128; -constexpr unsigned short THREAD_DIM_Y = 2; -constexpr unsigned short RESULTS_PER_PIXEL = 8; -constexpr float NO_DATA = -9999.0; - -enum StampType { STAMP_SUM = 0, STAMP_MEAN, STAMP_MEDIAN }; - -/* - * Data structure to represent an objects trajectory - * through a stack of images - */ -struct trajectory { + constexpr unsigned int MAX_KERNEL_RADIUS = 15; + constexpr unsigned short MAX_STAMP_EDGE = 64; + constexpr unsigned short CONV_THREAD_DIM = 32; + constexpr unsigned short THREAD_DIM_X = 128; + constexpr unsigned short THREAD_DIM_Y = 2; + constexpr unsigned short RESULTS_PER_PIXEL = 8; + constexpr float NO_DATA = -9999.0; + + enum StampType { STAMP_SUM = 0, STAMP_MEAN, STAMP_MEDIAN }; + + /* + * Data structure to represent an objects trajectory + * through a stack of images + */ + struct Trajectory { // Trajectory velocities - float x_vel; - float y_vel; - // Likelyhood + float vx; + float vy; + // Likelihood float lh; // Est. Flux float flux; @@ -43,20 +40,51 @@ struct trajectory { short y; // Number of images summed short obs_count; -}; -// The position (in pixels) of a trajectory. -struct PixelPos { + // I can't believe string::format is not a thing until C++ 20 + const std::string to_string() const { + return "lh: " + std::to_string(lh) + + " flux: " + std::to_string(flux) + + " x: " + std::to_string(x) + + " y: " + std::to_string(y) + + " vx: " + std::to_string(vx) + + " vy: " + std::to_string(vy) + + " obs_count: " + std::to_string(obs_count); + } + + // returns a yaml-compliant string + const std::string to_yaml() const { + return "{lh: " + std::to_string(lh) + + ", flux: " + std::to_string(flux) + + ", x: " + std::to_string(x) + + ", y: " + std::to_string(y) + + ", vx: " + std::to_string(vx) + + ", vy: " + std::to_string(vy) + + ", obs_count: " + std::to_string(obs_count) + +"}"; + } + }; + + // The position (in pixels) of a trajectory. + struct PixelPos { float x; float y; -}; - -/* - * Linear approximation to the barycentric correction needed to transform a - * pixel in the first image to a pixel in a consequent image. One struct needed - * per image. Correction calculated in higher level code. - */ -struct BaryCorrection { + + const std::string to_string() const { + return "x: " + std::to_string(x) + " y: " + std::to_string(y); + } + + const std::string to_yaml() const { + return "{x: " + std::to_string(x) + " y: " + std::to_string(y) + "}"; + } + }; + + /* + * Linear approximation to the barycentric correction needed to transform a + * pixel in the first image to a pixel in a consequent image. One struct needed + * per image. Correction calculated in higher level code. + */ + struct BaryCorrection { // linear coefficients of linear fit of pixel dependence float dx; float dxdx; @@ -64,11 +92,30 @@ struct BaryCorrection { float dy; float dydx; float dydy; -}; -/* The parameters to use for the on device search. */ + const std::string to_string() const { + return "dx: " + std::to_string(dx) + + " dxdx: " + std::to_string(dxdx) + + " dxdy: " + std::to_string(dxdy) + + " dy: " + std::to_string(dy) + + " dydx: " + std::to_string(dydx) + + " dydy: " + std::to_string(dydy); + } + + const std::string to_yaml() const { + return "{dx: " + std::to_string(dx) + + " dxdx: " + std::to_string(dxdx) + + " dxdy: " + std::to_string(dxdy) + + " dy: " + std::to_string(dy) + + " dydx: " + std::to_string(dydx) + + " dydy: " + std::to_string(dydy) + + "}"; + } + }; + + /* The parameters to use for the on device search. */ -struct SearchParameters { + struct SearchParameters { // Basic filtering paramets. int min_observations; float min_lh; @@ -95,16 +142,16 @@ struct SearchParameters { // Provide debugging output. bool debug; -}; + }; -struct scaleParameters { + struct scaleParameters { float min_val; float max_val; float scale; -}; + }; -// Search data on a per-image basis. -struct PerImageData { + // Search data on a per-image basis. + struct PerImageData { int num_images = 0; float* image_times = nullptr; @@ -112,9 +159,9 @@ struct PerImageData { scaleParameters* psi_params = nullptr; scaleParameters* phi_params = nullptr; -}; + }; -struct StampParameters { + struct StampParameters { int radius = 10; StampType stamp_type = STAMP_SUM; bool do_filtering = false; @@ -130,17 +177,106 @@ struct StampParameters { float m11_limit; float m02_limit; float m20_limit; -}; + }; -// Basic image moments use for analysis. -struct ImageMoments { + // Basic image moments use for analysis. + struct ImageMoments { float m00; float m01; float m10; float m11; float m02; float m20; -}; + }; + +#ifdef Py_PYTHON_H + namespace py = pybind11; + + static void trajectory_bindings(py::module &m) { + using tj = Trajectory; + + py::class_(m, "Trajectory", pydocs::DOC_Trajectory) + .def(py::init<>()) + .def_readwrite("vx", &tj::vx) + .def_readwrite("vy", &tj::vy) + .def_readwrite("lh", &tj::lh) + .def_readwrite("flux", &tj::flux) + .def_readwrite("x", &tj::x) + .def_readwrite("y", &tj::y) + .def_readwrite("obs_count", &tj::obs_count) + .def("__repr__", [](const tj &t) { return "Trajectory(" + t.to_string() + ")"; }) + .def("__str__", &tj::to_string) + .def(py::pickle( + [] (const tj &p) { // __getstate__ + return py::make_tuple(p.vx, p.vy, p.lh, p.flux, p.x, p.y, p.obs_count); + }, + [] (py::tuple t) { // __setstate__ + if (t.size() != 7) + throw std::runtime_error("Invalid state!"); + tj trj = { + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), t[5].cast(), + t[6].cast() + }; + return trj; + }) + ); + } + + static void pixel_pos_bindings(py::module &m) { + py::class_(m, "PixelPos", pydocs::DOC_PixelPos) + .def(py::init<>()) + .def_readwrite("x", &PixelPos::x) + .def_readwrite("y", &PixelPos::y) + .def("__repr__", [] (const PixelPos &p) { + return "PixelPos(" + p.to_string() + ")"; + }) + .def("__str__", &PixelPos::to_string); + } + + static void image_moments_bindings(py::module &m) { + py::class_(m, "ImageMoments", pydocs::DOC_ImageMoments) + .def(py::init<>()) + .def_readwrite("m00", &ImageMoments::m00) + .def_readwrite("m01", &ImageMoments::m01) + .def_readwrite("m10", &ImageMoments::m10) + .def_readwrite("m11", &ImageMoments::m11) + .def_readwrite("m02", &ImageMoments::m02) + .def_readwrite("m20", &ImageMoments::m20); + } + + static void stamp_parameters_bindings(py::module &m) { + py::class_(m, "StampParameters", pydocs::DOC_StampParameters) + .def(py::init<>()) + .def_readwrite("radius", &StampParameters::radius) + .def_readwrite("stamp_type", &StampParameters::stamp_type) + .def_readwrite("do_filtering", &StampParameters::do_filtering) + .def_readwrite("center_thresh", &StampParameters::center_thresh) + .def_readwrite("peak_offset_x", &StampParameters::peak_offset_x) + .def_readwrite("peak_offset_y", &StampParameters::peak_offset_y) + .def_readwrite("m01_limit", &StampParameters::m01_limit) + .def_readwrite("m10_limit", &StampParameters::m10_limit) + .def_readwrite("m11_limit", &StampParameters::m11_limit) + .def_readwrite("m02_limit", &StampParameters::m02_limit) + .def_readwrite("m20_limit", &StampParameters::m20_limit); + } + + static void bary_correction_bindings(py::module &m) { + py::class_(m, "BaryCorrection", pydocs::DOC_BaryCorrection) + .def(py::init<>()) + .def_readwrite("dx", &BaryCorrection::dx) + .def_readwrite("dxdx", &BaryCorrection::dxdx) + .def_readwrite("dxdy", &BaryCorrection::dxdy) + .def_readwrite("dy", &BaryCorrection::dy) + .def_readwrite("dydx", &BaryCorrection::dydx) + .def_readwrite("dydy", &BaryCorrection::dydy) + .def("__repr__", [](const BaryCorrection &b) { + return "BaryCorrection(" + b.to_string() + ")"; + }) + .def("__str__", &BaryCorrection::to_string); + } + +#endif /* Py_PYTHON_H */ } /* namespace search */ diff --git a/src/kbmod/search/Filtering.cpp b/src/kbmod/search/filtering.cpp similarity index 50% rename from src/kbmod/search/Filtering.cpp rename to src/kbmod/search/filtering.cpp index 0cc893d52..0bc8f4f3c 100644 --- a/src/kbmod/search/Filtering.cpp +++ b/src/kbmod/search/filtering.cpp @@ -1,27 +1,19 @@ -/* - * Filtering.cpp - * - * Created on: Sept 2, 2022 - * - * Helper functions for filtering results. - */ - -#include "Filtering.h" +#include "filtering.h" #include -namespace search { +namespace search { #ifdef HAVE_CUDA -/* The filter_kenerls.cu functions. */ -extern "C" void SigmaGFilteredIndicesCU(float *values, int num_values, float sgl0, float sgl1, float sg_coeff, - float width, int *idx_array, int *min_keep_idx, int *max_keep_idx); + /* The filter_kenerls.cu functions. */ + extern "C" void SigmaGFilteredIndicesCU(float *values, int num_values, float sgl0, float sgl1, float sg_coeff, + float width, int *idx_array, int *min_keep_idx, int *max_keep_idx); #endif -/* Return the list of indices from the values array such that those elements - pass the sigmaG filtering defined by percentiles [sgl0, sgl1] with coefficient - sigma_g_coeff and a multiplicative factor of width. */ -std::vector sigmaGFilteredIndices(const std::vector &values, float sgl0, float sgl1, - float sigma_g_coeff, float width) { + /* Return the list of indices from the values array such that those elements + pass the sigmaG filtering defined by percentiles [sgl0, sgl1] with coefficient + sigma_g_coeff and a multiplicative factor of width. */ + std::vector sigmaGFilteredIndices(const std::vector &values, float sgl0, float sgl1, + float sigma_g_coeff, float width) { // Bounds check the percentile values. assert(sgl0 > 0.0); assert(sgl1 < 1.0); @@ -31,7 +23,7 @@ std::vector sigmaGFilteredIndices(const std::vector &values, float s float values_arr[num_values]; int idx_array[num_values]; for (int i = 0; i < num_values; ++i) { - values_arr[i] = values[i]; + values_arr[i] = values[i]; } int min_keep_idx = 0; @@ -47,28 +39,28 @@ std::vector sigmaGFilteredIndices(const std::vector &values, float s // Copy the result into a vector and return it. std::vector result; for (int i = min_keep_idx; i <= max_keep_idx; ++i) { - result.push_back(idx_array[i]); + result.push_back(idx_array[i]); } return result; -} + } -/* Given a set of psi and phi values, - return a likelihood value */ -double calculateLikelihoodFromPsiPhi(std::vector psi_values, std::vector phi_values) { + /* Given a set of psi and phi values, + return a likelihood value */ + double calculateLikelihoodFromPsiPhi(std::vector psi_values, std::vector phi_values) { assert(psi_values.size() == phi_values.size()); double psi_sum = 0.0; double phi_sum = 0.0; for (int i = 0; i < psi_values.size(); i++) { - psi_sum += psi_values[i]; - phi_sum += phi_values[i]; + psi_sum += psi_values[i]; + phi_sum += phi_values[i]; } if (psi_sum == 0.0 || phi_sum <= 0.0) { - return 0.0; + return 0.0; } return psi_sum / sqrt(phi_sum); -} + } } /* namespace search */ diff --git a/src/kbmod/search/filtering.h b/src/kbmod/search/filtering.h new file mode 100644 index 000000000..3f4243742 --- /dev/null +++ b/src/kbmod/search/filtering.h @@ -0,0 +1,16 @@ +#ifndef FILTERING_H_ +#define FILTERING_H_ + +#include + + +namespace search { + /* Return the list of indices from the values array such that those elements + pass the sigmaG filtering defined by percentiles [sGL0, sGL1] with coefficient + sigmag_coeff and a multiplicative factor of width. */ + std::vector sigmaGFilteredIndices(const std::vector& values, float sgl0, float sgl1, + float sigma_g_coeff, float width); + +} /* namespace search */ + +#endif /* FILTERING_H_ */ diff --git a/src/kbmod/search/image_kernels.cu b/src/kbmod/search/image_kernels.cu index c3ec078cf..bd368ec15 100644 --- a/src/kbmod/search/image_kernels.cu +++ b/src/kbmod/search/image_kernels.cu @@ -19,7 +19,7 @@ namespace search { /* * Device kernel that convolves the provided image with the psf */ -__global__ void convolvePSF(int width, int height, float *source_img, float *result_img, float *psf, +__global__ void convolve_psf(int width, int height, float *source_img, float *result_img, float *psf, int psf_radius, int psf_dim, float psf_sum) { // Find bounds of convolution area const int x = blockIdx.x * CONV_THREAD_DIM + threadIdx.x; @@ -73,7 +73,7 @@ extern "C" void deviceConvolve(float *source_img, float *result_img, int width, checkCudaErrors( cudaMemcpy(devicesource_img, source_img, sizeof(float) * n_pixels, cudaMemcpyHostToDevice)); - convolvePSF<<>>(width, height, devicesource_img, deviceresult_img, device_kernel, + convolve_psf<<>>(width, height, devicesource_img, deviceresult_img, device_kernel, psf_radiusius, psf_dim, psf_sum); checkCudaErrors( diff --git a/src/kbmod/search/image_stack.cpp b/src/kbmod/search/image_stack.cpp new file mode 100644 index 000000000..6f20219b7 --- /dev/null +++ b/src/kbmod/search/image_stack.cpp @@ -0,0 +1,161 @@ +#include "image_stack.h" + + +namespace py = pybind11; + + +namespace search { + ImageStack::ImageStack(const std::vector& filenames, const std::vector& psfs) { + verbose = true; + reset_images(); + load_images(filenames, psfs); + extract_image_times(); + set_time_origin(); + global_mask = RawImage(get_width(), get_height()); + global_mask.set_all_pix(0.0); + } + + ImageStack::ImageStack(const std::vector& imgs) { + verbose = true; + images = imgs; + extract_image_times(); + set_time_origin(); + global_mask = RawImage(get_width(), get_height()); + global_mask.set_all_pix(0.0); + } + + void ImageStack::load_images(const std::vector& filenames, + const std::vector& psfs) { + const int num_files = filenames.size(); + if (num_files == 0) { + std::cout << "No files provided" + << "\n"; + } + + if (psfs.size() != num_files) throw std::runtime_error("Mismatched PSF array in ImageStack creation."); + + // Load images from file + for (int i = 0; i < num_files; ++i) { + images.push_back(LayeredImage(filenames[i], psfs[i])); + if (verbose) std::cout << "." << std::flush; + } + if (verbose) std::cout << "\n"; + } + + void ImageStack::extract_image_times() { + // Load image times + image_times = std::vector(); + for (auto& i : images) { + image_times.push_back(float(i.get_obstime())); + } + } + + void ImageStack::set_time_origin() { + // Set beginning time to 0.0 + double initial_time = image_times[0]; + for (auto& t : image_times) t = t - initial_time; + } + + LayeredImage& ImageStack::get_single_image(int index) { + if (index < 0 || index > images.size()) throw std::out_of_range("ImageStack index out of bounds."); + return images[index]; + } + + void ImageStack::set_single_image(int index, LayeredImage& img) { + if (index < 0 || index > images.size()) throw std::out_of_range("ImageStack index out of bounds."); + images[index] = img; + } + + void ImageStack::set_times(const std::vector& times) { + if (times.size() != img_count()) + throw std::runtime_error("List of times provided does not match the number of images!"); + image_times = times; + set_time_origin(); + } + + void ImageStack::reset_images() { images = std::vector(); } + + void ImageStack::convolve_psf() { + for (auto& i : images) i.convolve_psf(); + } + + void ImageStack::save_global_mask(const std::string& path) { global_mask.save_to_file(path); } + + void ImageStack::save_images(const std::string& path) { + for (auto& i : images) i.save_layers(path); + } + + const RawImage& ImageStack::get_global_mask() const { return global_mask; } + + void ImageStack::apply_mask_flags(int flags, const std::vector& exceptions) { + for (auto& i : images) { + i.apply_mask_flags(flags, exceptions); + } + } + + void ImageStack::apply_global_mask(int flags, int threshold) { + create_global_mask(flags, threshold); + for (auto& i : images) { + i.apply_global_mask(global_mask); + } + } + + void ImageStack::apply_mask_threshold(float thresh) { + for (auto& i : images) i.apply_mask_threshold(thresh); + } + + void ImageStack::grow_mask(int steps) { + for (auto& i : images) i.grow_mask(steps); + } + + void ImageStack::create_global_mask(int flags, int threshold) { + int npixels = get_npixels(); + + // For each pixel count the number of images where it is masked. + std::vector counts(npixels, 0); + for (unsigned int img = 0; img < images.size(); ++img) { + float* imgMask = images[img].get_mask().data(); + // Count the number of times a pixel has any of the flags + for (unsigned int pixel = 0; pixel < npixels; ++pixel) { + if ((flags & static_cast(imgMask[pixel])) != 0) counts[pixel]++; + } + } + + // Set all pixels below threshold to 0 and all above to 1 + float* global_m = global_mask.data(); + for (unsigned int p = 0; p < npixels; ++p) { + global_m[p] = counts[p] < threshold ? 0.0 : 1.0; + } + } + +#ifdef Py_PYTHON_H + static void image_stack_bindings(py::module &m) { + using is = search::ImageStack; + using li = search::LayeredImage; + using pf = search::PSF; + + py::class_(m, "ImageStack", pydocs::DOC_ImageStack) + .def(py::init, std::vector>()) + .def(py::init>()) + .def("get_images", &is::get_images, pydocs::DOC_ImageStack_get_images) + .def("get_single_image", &is::get_single_image, pydocs::DOC_ImageStack_get_single_image) + .def("set_single_image", &is::set_single_image, pydocs::DOC_ImageStack_set_single_image) + .def("get_times", &is::get_times, pydocs::DOC_ImageStack_get_times) + .def("set_times", &is::set_times, pydocs::DOC_ImageStack_set_times ) + .def("img_count", &is::img_count, pydocs::DOC_ImageStack_img_count) + .def("apply_mask_flags", &is::apply_mask_flags, pydocs::DOC_ImageStack_apply_mask_flags) + .def("apply_mask_threshold", &is::apply_mask_threshold, pydocs::DOC_ImageStack_apply_mask_threshold) + .def("apply_global_mask", &is::apply_global_mask, pydocs::DOC_ImageStack_apply_global_mask) + .def("grow_mask", &is::grow_mask, pydocs::DOC_ImageStack_grow_mask) + .def("save_global_mask", &is::save_global_mask, pydocs::DOC_ImageStack_save_global_mask) + .def("save_images", &is::save_images, pydocs::DOC_ImageStack_save_images) + .def("get_global_mask", &is::get_global_mask, pydocs::DOC_ImageStack_get_global_mask) + .def("convolve_psf", &is::convolve_psf, pydocs::DOC_ImageStack_convolve_psf) + .def("get_width", &is::get_width, pydocs::DOC_ImageStack_get_width) + .def("get_height", &is::get_height, pydocs::DOC_ImageStack_get_height) + .def("get_npixels", &is::get_npixels, pydocs::DOC_ImageStack_get_npixels); + } + +#endif /* Py_PYTHON_H */ + +} /* namespace search */ diff --git a/src/kbmod/search/image_stack.h b/src/kbmod/search/image_stack.h new file mode 100644 index 000000000..12d56affe --- /dev/null +++ b/src/kbmod/search/image_stack.h @@ -0,0 +1,63 @@ +#ifndef IMAGESTACK_H_ +#define IMAGESTACK_H_ + +#include +#include +#include +#include +#include +#include +#include "layered_image.h" +#include "pydocs/image_stack_docs.h" + + +namespace search { + class ImageStack { + public: + ImageStack(const std::vector& filenames, const std::vector& psfs); + ImageStack(const std::vector& imgs); + + // Simple getters. + unsigned img_count() const { return images.size(); } + unsigned get_width() const { return images.size() > 0 ? images[0].get_width() : 0; } + unsigned get_height() const { return images.size() > 0 ? images[0].get_height() : 0; } + unsigned get_npixels() const { return images.size() > 0 ? images[0].get_npixels() : 0; } + std::vector& get_images() { return images; } + const std::vector& get_times() const { return image_times; } + float* get_timesDataRef() { return image_times.data(); } + LayeredImage& get_single_image(int index); + + // Simple setters. + void set_times(const std::vector& times); + void reset_images(); + void set_single_image(int index, LayeredImage& img); + + // Apply makes to all the images. + void apply_global_mask(int flags, int threshold); + void apply_mask_flags(int flags, const std::vector& exceptions); + void apply_mask_threshold(float thresh); + void grow_mask(int steps); + const RawImage& get_global_mask() const; + + void convolve_psf(); + + // Save data to files. + void save_global_mask(const std::string& path); + void save_images(const std::string& path); + + virtual ~ImageStack(){}; + + private: + void load_images(const std::vector& filenames, const std::vector& psfs); + void extract_image_times(); + void set_time_origin(); + void create_global_mask(int flags, int threshold); + std::vector images; + RawImage global_mask; + std::vector image_times; + bool verbose; + }; + +} /* namespace search */ + +#endif /* IMAGESTACK_H_ */ diff --git a/src/kbmod/search/kernels.cu b/src/kbmod/search/kernels.cu index bfd5efa97..2b2ee9c13 100644 --- a/src/kbmod/search/kernels.cu +++ b/src/kbmod/search/kernels.cu @@ -18,7 +18,7 @@ #include #include -#include "ImageStack.h" +#include "image_stack.h" namespace search { @@ -88,7 +88,7 @@ __device__ float ReadEncodedPixel(void *image_vect, int index, int n_bytes, cons */ __global__ void searchFilterImages(int num_images, int width, int height, void *psi_vect, void *phi_vect, PerImageData image_data, SearchParameters params, int num_trajectories, - trajectory *trajectories, trajectory *results) { + Trajectory *trajectories, Trajectory *results) { // Get the x and y coordinates within the search space. const int x_i = blockIdx.x * THREAD_DIM_X + threadIdx.x; const int y_i = blockIdx.y * THREAD_DIM_Y + threadIdx.y; @@ -114,7 +114,7 @@ __global__ void searchFilterImages(int num_images, int width, int height, void * // Create an initial set of best results with likelihood -1.0. // We also set (x, y) because they are used in the later python // functions. - trajectory best[RESULTS_PER_PIXEL]; + Trajectory best[RESULTS_PER_PIXEL]; for (int r = 0; r < RESULTS_PER_PIXEL; ++r) { best[r].x = x; best[r].y = y; @@ -124,11 +124,11 @@ __global__ void searchFilterImages(int num_images, int width, int height, void * // For each trajectory we'd like to search for (int t = 0; t < num_trajectories; ++t) { // Create a trajectory for this search. - trajectory curr_trj; + Trajectory curr_trj; curr_trj.x = x; curr_trj.y = y; - curr_trj.x_vel = trajectories[t].x_vel; - curr_trj.y_vel = trajectories[t].y_vel; + curr_trj.vx = trajectories[t].vx; + curr_trj.vy = trajectories[t].vy; curr_trj.obs_count = 0; float psi_sum = 0.0; @@ -147,15 +147,15 @@ __global__ void searchFilterImages(int num_images, int width, int height, void * for (int i = 0; i < num_images; ++i) { // Predict the trajectory's position. float curr_time = image_data.image_times[i]; - int current_x = x + int(curr_trj.x_vel * curr_time + 0.5); - int current_y = y + int(curr_trj.y_vel * curr_time + 0.5); + int current_x = x + int(curr_trj.vx * curr_time + 0.5); + int current_y = y + int(curr_trj.vy * curr_time + 0.5); // If using barycentric correction, apply it. // Must be before out of bounds check if (params.use_corr && (image_data.bary_corrs != nullptr)) { BaryCorrection bc = image_data.bary_corrs[i]; - current_x = int(x + curr_trj.x_vel * curr_time + bc.dx + x * bc.dxdx + y * bc.dxdy + 0.5); - current_y = int(y + curr_trj.y_vel * curr_time + bc.dy + x * bc.dydx + y * bc.dydy + 0.5); + current_x = int(x + curr_trj.vx * curr_time + bc.dx + x * bc.dxdx + y * bc.dxdy + 0.5); + current_y = int(y + curr_trj.vy * curr_time + bc.dy + x * bc.dydx + y * bc.dydy + 0.5); } // Test if trajectory goes out of the image, in which case we do not @@ -221,7 +221,7 @@ __global__ void searchFilterImages(int num_images, int width, int height, void * // Insert the new trajectory into the sorted list of results. // Only sort the values with valid likelihoods. - trajectory temp; + Trajectory temp; for (int r = 0; r < RESULTS_PER_PIXEL; ++r) { if (curr_trj.lh > best[r].lh && curr_trj.lh > -1.0) { temp = best[r]; @@ -293,13 +293,13 @@ void *encodeImageFloat(float *image_vect, unsigned int vectLength, bool debug) { extern "C" void deviceSearchFilter(int num_images, int width, int height, float *psi_vect, float *phi_vect, PerImageData img_data, SearchParameters params, int num_trajectories, - trajectory *trj_to_search, int num_results, trajectory *best_results) { + Trajectory *trj_to_search, int num_results, Trajectory *best_results) { // Allocate Device memory - trajectory *device_tests; + Trajectory *device_tests; float *device_img_times; void *device_psi; void *device_phi; - trajectory *device_search_results; + Trajectory *device_search_results; BaryCorrection *device_bary_corrs = nullptr; scaleParameters *device_psi_params = nullptr; scaleParameters *device_phi_params = nullptr; @@ -310,9 +310,9 @@ extern "C" void deviceSearchFilter(int num_images, int width, int height, float } if (params.debug) { - printf("Allocating %lu bytes for testing grid.\n", sizeof(trajectory) * num_trajectories); + printf("Allocating %lu bytes for testing grid.\n", sizeof(Trajectory) * num_trajectories); } - checkCudaErrors(cudaMalloc((void **)&device_tests, sizeof(trajectory) * num_trajectories)); + checkCudaErrors(cudaMalloc((void **)&device_tests, sizeof(Trajectory) * num_trajectories)); if (params.debug) { printf("Allocating %lu bytes for time data.\n", sizeof(float) * num_images); @@ -320,12 +320,12 @@ extern "C" void deviceSearchFilter(int num_images, int width, int height, float checkCudaErrors(cudaMalloc((void **)&device_img_times, sizeof(float) * num_images)); if (params.debug) { - printf("Allocating %lu bytes for testing grid.\n", sizeof(trajectory) * num_trajectories); + printf("Allocating %lu bytes for testing grid.\n", sizeof(Trajectory) * num_trajectories); } - checkCudaErrors(cudaMalloc((void **)&device_search_results, sizeof(trajectory) * num_results)); + checkCudaErrors(cudaMalloc((void **)&device_search_results, sizeof(Trajectory) * num_results)); // Copy trajectories to search - checkCudaErrors(cudaMemcpy(device_tests, trj_to_search, sizeof(trajectory) * num_trajectories, + checkCudaErrors(cudaMemcpy(device_tests, trj_to_search, sizeof(Trajectory) * num_trajectories, cudaMemcpyHostToDevice)); // Copy image times @@ -397,7 +397,7 @@ extern "C" void deviceSearchFilter(int num_images, int width, int height, float device_search_results); // Read back results - checkCudaErrors(cudaMemcpy(best_results, device_search_results, sizeof(trajectory) * num_results, + checkCudaErrors(cudaMemcpy(best_results, device_search_results, sizeof(Trajectory) * num_results, cudaMemcpyDeviceToHost)); // Free the on GPU memory. @@ -412,12 +412,12 @@ extern "C" void deviceSearchFilter(int num_images, int width, int height, float } __global__ void deviceGetCoaddStamp(int num_images, int width, int height, float *image_vect, - PerImageData image_data, int num_trajectories, trajectory *trajectories, + PerImageData image_data, int num_trajectories, Trajectory *trajectories, StampParameters params, int *use_index_vect, float *results) { // Get the trajectory that we are going to be using. const int trj_index = blockIdx.x * blockDim.x + threadIdx.x; if (trj_index < 0 || trj_index >= num_trajectories) return; - trajectory trj = trajectories[trj_index]; + Trajectory trj = trajectories[trj_index]; // Get the pixel coordinates within the stamp to use. const int stamp_width = 2 * params.radius + 1; @@ -446,12 +446,12 @@ __global__ void deviceGetCoaddStamp(int num_images, int width, int height, float // Predict the trajectory's position including the barycentric correction if needed. float curr_time = image_data.image_times[t]; - int current_x = int(trj.x + trj.x_vel * curr_time); - int current_y = int(trj.y + trj.y_vel * curr_time); + int current_x = int(trj.x + trj.vx * curr_time); + int current_y = int(trj.y + trj.vy * curr_time); if (image_data.bary_corrs != nullptr) { BaryCorrection bc = image_data.bary_corrs[t]; - current_x = int(trj.x + trj.x_vel * curr_time + bc.dx + trj.x * bc.dxdx + trj.y * bc.dxdy); - current_y = int(trj.y + trj.y_vel * curr_time + bc.dy + trj.x * bc.dydx + trj.y * bc.dydy); + current_x = int(trj.x + trj.vx * curr_time + bc.dx + trj.x * bc.dxdx + trj.y * bc.dxdy); + current_y = int(trj.y + trj.vy * curr_time + bc.dy + trj.x * bc.dydx + trj.y * bc.dydy); } // Get the stamp and add it to the list of values. @@ -513,10 +513,10 @@ __global__ void deviceGetCoaddStamp(int num_images, int width, int height, float } void deviceGetCoadds(ImageStack &stack, PerImageData image_data, int num_trajectories, - trajectory *trajectories, StampParameters params, + Trajectory *trajectories, StampParameters params, std::vector> &use_index_vect, float *results) { // Allocate Device memory - trajectory *device_trjs; + Trajectory *device_trjs; int *device_use_index = nullptr; float *device_times; float *device_img; @@ -524,17 +524,17 @@ void deviceGetCoadds(ImageStack &stack, PerImageData image_data, int num_traject BaryCorrection *device_bary_corrs = nullptr; // Compute the dimensions for the data. - const unsigned int num_images = stack.imgCount(); - const unsigned int width = stack.getWidth(); - const unsigned int height = stack.getHeight(); + const unsigned int num_images = stack.img_count(); + const unsigned int width = stack.get_width(); + const unsigned int height = stack.get_height(); const unsigned int num_image_pixels = num_images * width * height; const unsigned int stamp_width = 2 * params.radius + 1; const unsigned int stamp_ppi = (2 * params.radius + 1) * (2 * params.radius + 1); const unsigned int num_stamp_pixels = num_trajectories * stamp_ppi; // Allocate and copy the trajectories. - checkCudaErrors(cudaMalloc((void **)&device_trjs, sizeof(trajectory) * num_trajectories)); - checkCudaErrors(cudaMemcpy(device_trjs, trajectories, sizeof(trajectory) * num_trajectories, + checkCudaErrors(cudaMalloc((void **)&device_trjs, sizeof(Trajectory) * num_trajectories)); + checkCudaErrors(cudaMemcpy(device_trjs, trajectories, sizeof(Trajectory) * num_trajectories, cudaMemcpyHostToDevice)); // Check if we need to create a vector of per-trajectory, per-image use. @@ -565,7 +565,7 @@ void deviceGetCoadds(ImageStack &stack, PerImageData image_data, int num_traject checkCudaErrors(cudaMalloc((void **)&device_img, sizeof(float) * num_image_pixels)); float *next_ptr = device_img; for (unsigned t = 0; t < num_images; ++t) { - const std::vector &data_ref = stack.getSingleImage(t).getScience().getPixels(); + const std::vector &data_ref = stack.get_single_image(t).get_science().get_pixels(); assert(data_ref.size() == width * height); checkCudaErrors(cudaMemcpy(next_ptr, data_ref.data(), sizeof(float) * width * height, diff --git a/src/kbmod/search/layered_image.cpp b/src/kbmod/search/layered_image.cpp new file mode 100644 index 000000000..f741957e8 --- /dev/null +++ b/src/kbmod/search/layered_image.cpp @@ -0,0 +1,265 @@ +#include "layered_image.h" + + +namespace search { + LayeredImage::LayeredImage(std::string path, const PSF& psf) : psf(psf) { + int f_begin = path.find_last_of("/"); + int f_end = path.find_last_of(".fits") - 4; + filename = path.substr(f_begin, f_end - f_begin); + + science = RawImage(); + science.load_from_file(path, 1); + width = science.get_width(); + height = science.get_height(); + + mask = RawImage(); + mask.load_from_file(path, 2); + + variance = RawImage(); + variance.load_from_file(path, 3); + + if (width != variance.get_width() or height != variance.get_height()) + throw std::runtime_error("Science and Variance layers are not the same size."); + if (width != mask.get_width() or height != mask.get_height()) + throw std::runtime_error("Science and Mask layers are not the same size."); + } + + LayeredImage::LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, + const PSF& psf) + : psf(psf) { + // Get the dimensions of the science layer and check for consistency with + // the other two layers. + width = sci.get_width(); + height = sci.get_height(); + if (width != var.get_width() or height != var.get_height()) + throw std::runtime_error("Science and Variance layers are not the same size."); + if (width != msk.get_width() or height != msk.get_height()) + throw std::runtime_error("Science and Mask layers are not the same size."); + + // Copy the image layers. + science = sci; + mask = msk; + variance = var; + } + + LayeredImage::LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, + const PSF& psf) + : LayeredImage(name, w, h, noise_stdev, pixel_variance, time, psf, -1) {} + + LayeredImage::LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, + const PSF& psf, int seed) + : psf(psf) { + filename = name; + width = w; + height = h; + + std::vector raw_sci(width * height); + std::random_device r; + std::default_random_engine generator(r()); + if (seed >= 0) { + generator.seed(seed); + } + std::normal_distribution distrib(0.0, noise_stdev); + for (float& p : raw_sci) p = distrib(generator); + + science = RawImage(w, h, raw_sci); + science.set_obstime(time); + + mask = RawImage(w, h, std::vector(w * h, 0.0)); + variance = RawImage(w, h, std::vector(w * h, pixel_variance)); + } + + void LayeredImage::set_psf(const PSF& new_psf) { + psf = new_psf; + } + + void LayeredImage::grow_mask(int steps) { + science.grow_mask(steps); + variance.grow_mask(steps); + } + + void LayeredImage::convolve_given_psf(const PSF& given_psf) { + science.convolve(given_psf); + + // Square the PSF use that on the variance image. + PSF psfsq = PSF(given_psf); // Copy + psfsq.square_psf(); + variance.convolve(psfsq); + } + + void LayeredImage::convolve_psf() { + convolve_given_psf(psf); + } + + void LayeredImage::apply_mask_flags(int flags, const std::vector& exceptions) { + science.apply_mask(flags, exceptions, mask); + variance.apply_mask(flags, exceptions, mask); + } + + /* Mask all pixels that are not 0 in global mask */ + void LayeredImage::apply_global_mask(const RawImage& global_mask) { + science.apply_mask(0xFFFFFF, {}, global_mask); + variance.apply_mask(0xFFFFFF, {}, global_mask); + } + + void LayeredImage::apply_mask_threshold(float thresh) { + const int num_pixels = get_npixels(); + float* sci_pixels = science.data(); + float* var_pix = variance.data(); + for (int i = 0; i < num_pixels; ++i) { + if (sci_pixels[i] > thresh) { + sci_pixels[i] = NO_DATA; + var_pix[i] = NO_DATA; + } + } + } + + void LayeredImage::subtract_template(const RawImage& sub_template) { + assert(get_height() == sub_template.get_height() && get_width() == sub_template.get_width()); + const int num_pixels = get_npixels(); + + float* sci_pixels = science.data(); + const std::vector& tem_pixels = sub_template.get_pixels(); + for (unsigned i = 0; i < num_pixels; ++i) { + if ((sci_pixels[i] != NO_DATA) && (tem_pixels[i] != NO_DATA)) { + sci_pixels[i] -= tem_pixels[i]; + } + } + } + + void LayeredImage::save_layers(const std::string& path) { + fitsfile* fptr; + int status = 0; + long naxes[2] = {0, 0}; + double obstime = science.get_obstime(); + + fits_create_file(&fptr, (path + filename + ".fits").c_str(), &status); + + // If we are unable to create the file, check if it already exists + // and, if so, delete it and retry the create. + if (status == 105) { + status = 0; + fits_open_file(&fptr, (path + filename + ".fits").c_str(), READWRITE, &status); + if (status == 0) { + fits_delete_file(fptr, &status); + fits_create_file(&fptr, (path + filename + ".fits").c_str(), &status); + } + } + + fits_create_img(fptr, SHORT_IMG, 0, naxes, &status); + fits_update_key(fptr, TDOUBLE, "MJD", &obstime, "[d] Generated Image time", &status); + fits_close_file(fptr, &status); + fits_report_error(stderr, status); + + science.append_layer_to_file(path + filename + ".fits"); + mask.append_layer_to_file(path + filename + ".fits"); + variance.append_layer_to_file(path + filename + ".fits"); + } + + void LayeredImage::set_science(RawImage& im) { + check_dims(im); + science = im; + } + + void LayeredImage::set_mask(RawImage& im) { + check_dims(im); + mask = im; + } + + void LayeredImage::set_variance(RawImage& im) { + check_dims(im); + variance = im; + } + + void LayeredImage::check_dims(RawImage& im) { + if (im.get_width() != get_width()) throw std::runtime_error("Image width does not match"); + if (im.get_height() != get_height()) throw std::runtime_error("Image height does not match"); + } + + RawImage LayeredImage::generate_psi_image() { + RawImage result(width, height); + float* result_arr = result.data(); + float* sci_array = science.data(); + float* var_array = variance.data(); + + // Set each of the result pixels. + const int num_pixels = get_npixels(); + for (int p = 0; p < num_pixels; ++p) { + float var_pix = var_array[p]; + if (var_pix != NO_DATA) { + result_arr[p] = sci_array[p] / var_pix; + } else { + result_arr[p] = NO_DATA; + } + } + + // Convolve with the PSF. + result.convolve(psf); + + return result; + } + + RawImage LayeredImage::generate_phi_image() { + RawImage result(width, height); + float* result_arr = result.data(); + float* var_array = variance.data(); + + // Set each of the result pixels. + const int num_pixels = get_npixels(); + for (int p = 0; p < num_pixels; ++p) { + float var_pix = var_array[p]; + if (var_pix != NO_DATA) { + result_arr[p] = 1.0 / var_pix; + } else { + result_arr[p] = NO_DATA; + } + } + + // Convolve with the PSF squared. + PSF psfsq = PSF(psf); // Copy + psfsq.square_psf(); + result.convolve(psfsq); + + return result; + } + +#ifdef Py_PYTHON_H + static void layered_image_bindings(py::module &m) { + using li = search::LayeredImage; + using ri = search::RawImage; + using pf = search::PSF; + + py::class_
  • (m, "LayeredImage", pydocs::DOC_LayeredImage) + .def(py::init()) + .def(py::init()) + .def(py::init()) + .def(py::init()) + .def("set_psf", &li::set_psf, pydocs::DOC_LayeredImage_set_psf) + .def("get_psf", &li::get_psf, pydocs::DOC_LayeredImage_get_psf) + .def("apply_mask_flags", &li::apply_mask_flags, pydocs::DOC_LayeredImage_apply_mask_flags) + .def("apply_mask_threshold", &li::apply_mask_threshold, pydocs::DOC_LayeredImage_apply_mask_threshold) + .def("sub_template", &li::subtract_template, pydocs::DOC_LayeredImage_sub_template) + .def("save_layers", &li::save_layers, pydocs::DOC_LayeredImage_save_layers) + .def("get_science", &li::get_science, pydocs::DOC_LayeredImage_get_science) + .def("get_mask", &li::get_mask, pydocs::DOC_LayeredImage_get_mask) + .def("get_variance", &li::get_variance, pydocs::DOC_LayeredImage_get_variance) + .def("set_science", &li::set_science, pydocs::DOC_LayeredImage_set_science) + .def("set_mask", &li::set_mask, pydocs::DOC_LayeredImage_set_mask) + .def("set_variance", &li::set_variance, pydocs::DOC_LayeredImage_set_variance) + .def("convolve_psf", &li::convolve_psf, pydocs::DOC_LayeredImage_convolve_psf) + .def("convolve_given_psf", &li::convolve_given_psf, pydocs::DOC_LayeredImage_convolve_given_psf) + .def("grow_mask", &li::grow_mask, pydocs::DOC_LayeredImage_grow_mask) + .def("get_name", &li::get_name, pydocs::DOC_LayeredImage_get_name) + .def("get_width", &li::get_width, pydocs::DOC_LayeredImage_get_width) + .def("get_height", &li::get_height, pydocs::DOC_LayeredImage_get_height) + .def("get_npixels", &li::get_npixels, pydocs::DOC_LayeredImage_get_npixels) + .def("get_obstime", &li::get_obstime, pydocs::DOC_LayeredImage_get_obstime) + .def("set_obstime", &li::set_obstime, pydocs::DOC_LayeredImage_set_obstime) + .def("generate_psi_image", &li::generate_psi_image, pydocs::DOC_LayeredImage_generate_psi_image) + .def("generate_phi_image", &li::generate_phi_image, pydocs::DOC_LayeredImage_generate_phi_image); + } + +#endif /* Py_PYTHON_H */ + +} /* namespace search */ + diff --git a/src/kbmod/search/layered_image.h b/src/kbmod/search/layered_image.h new file mode 100644 index 000000000..b4b3abb41 --- /dev/null +++ b/src/kbmod/search/layered_image.h @@ -0,0 +1,86 @@ +#ifndef LAYEREDIMAGE_H_ +#define LAYEREDIMAGE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include "raw_image.h" +#include "common.h" +#include "pydocs/layered_image_docs.h" + + +namespace search { + class LayeredImage { + public: + explicit LayeredImage(std::string path, const PSF& psf); + explicit LayeredImage(const RawImage& sci, const RawImage& var, const RawImage& msk, + const PSF& psf); + explicit LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, + const PSF& psf); + explicit LayeredImage(std::string name, int w, int h, float noise_stdev, float pixel_variance, double time, + const PSF& psf, int seed); + + // Set an image specific point spread function. + void set_psf(const PSF& psf); + const PSF& get_psf() const { return psf; } + + // Basic getter functions for image data. + std::string get_name() const { return filename; } + unsigned get_width() const { return width; } + unsigned get_height() const { return height; } + unsigned get_npixels() const { return width * height; } + double get_obstime() const { return science.get_obstime(); } + void set_obstime(double obstime) { science.set_obstime(obstime); } + + // Getter functions for the data in the individual layers. + RawImage& get_science() { return science; } + RawImage& get_mask() { return mask; } + RawImage& get_variance() { return variance; } + + // Applies the mask functions to each of the science and variance layers. + void apply_mask_flags(int flag, const std::vector& exceptions); + void apply_global_mask(const RawImage& global_mask); + void apply_mask_threshold(float thresh); + void grow_mask(int steps); + + // Subtracts a template image from the science layer. + void subtract_template(const RawImage& sub_template); + + // Saves the data in each later to a file. + void save_layers(const std::string& path); + + // Setter functions for the individual layers. + void set_science(RawImage& im); + void set_mask(RawImage& im); + void set_variance(RawImage& im); + + // Convolve with a given PSF or the default one. + void convolve_psf(); + void convolve_given_psf(const PSF& psf); + + virtual ~LayeredImage(){}; + + // Generate psi and phi images from the science and variance layers. + RawImage generate_psi_image(); + RawImage generate_phi_image(); + + private: + void check_dims(RawImage& im); + + std::string filename; + unsigned width; + unsigned height; + + PSF psf; + RawImage science; + RawImage mask; + RawImage variance; + }; + +} /* namespace search */ + +#endif /* LAYEREDIMAGE_H_ */ diff --git a/src/kbmod/search/psf.cpp b/src/kbmod/search/psf.cpp new file mode 100644 index 000000000..c75107e4d --- /dev/null +++ b/src/kbmod/search/psf.cpp @@ -0,0 +1,175 @@ +#include "psf.h" + + +namespace py = pybind11; + + + +namespace search { + PSF::PSF() : kernel(1, 1.0) { + dim = 1; + radius = 0; + width = 1e-12; + sum = 1.0; + } + + PSF::PSF(float stdev) { + width = stdev; + float simple_gauss[MAX_KERNEL_RADIUS]; + double psf_coverage = 0.0; + double norm_factor = stdev * sqrt(2.0); + int i = 0; + + // Create 1D gaussian array + while (psf_coverage < 0.98 && i < MAX_KERNEL_RADIUS) { + float current_bin = + 0.5 * (std::erf((float(i) + 0.5) / norm_factor) - std::erf((float(i) - 0.5) / norm_factor)); + simple_gauss[i] = current_bin; + if (i == 0) { + psf_coverage += current_bin; + } else { + psf_coverage += 2.0 * current_bin; + } + i++; + } + + radius = i - 1; // This value is good for + dim = 2 * radius + 1; + + // Create 2D gaussain by multiplying with itself + kernel = std::vector(); + for (int ii = 0; ii < dim; ++ii) { + for (int jj = 0; jj < dim; ++jj) { + float current = simple_gauss[abs(radius - ii)] * simple_gauss[abs(radius - jj)]; + kernel.push_back(current); + } + } + calc_sum(); + } + + // Copy constructor. + PSF::PSF(const PSF& other) { + kernel = other.kernel; + dim = other.dim; + radius = other.radius; + width = other.width; + sum = other.sum; + } + + // Copy assignment. + PSF& PSF::operator=(const PSF& other) { + kernel = other.kernel; + dim = other.dim; + radius = other.radius; + width = other.width; + sum = other.sum; + return *this; + } + + // Move constructor. + PSF::PSF(PSF&& other) + : kernel(std::move(other.kernel)), + dim(other.dim), + radius(other.radius), + width(other.width), + sum(other.sum) {} + + // Move assignment. + PSF& PSF::operator=(PSF&& other) { + if (this != &other) { + kernel = std::move(other.kernel); + dim = other.dim; + radius = other.radius; + width = other.width; + sum = other.sum; + } + return *this; + } + + void PSF::calc_sum() { + sum = 0.0; + for (auto& i : kernel) sum += i; + } + + void PSF::square_psf() { + for (float& i : kernel) { + i = i * i; + } + calc_sum(); + } + + std::string PSF::print() { + std::stringstream ss; + ss.setf(std::ios::fixed, std::ios::floatfield); + ss.precision(3); + for (int row = 0; row < dim; ++row) { + ss << "| "; + for (int col = 0; col < dim; ++col) { + ss << kernel[row * dim + col] << " | "; + } + ss << "\n "; + for (int space = 0; space < dim * 8 - 1; ++space) ss << "-"; + ss << "\n"; + } + ss << 100.0 * sum << "% of PSF contained within kernel\n"; + return ss.str(); + } + +#ifdef Py_PYTHON_H + PSF::PSF(pybind11::array_t arr) { set_array(arr); } + + void PSF::set_array(pybind11::array_t arr) { + pybind11::buffer_info info = arr.request(); + + if (info.ndim != 2) + throw std::runtime_error( + "Array must have 2 dimensions. (It " + " must also be a square with odd dimensions)"); + + if (info.shape[0] != info.shape[1]) + throw std::runtime_error( + "Array must be square (x-dimension == y-dimension)." + "It also must have odd dimensions."); + float* pix = static_cast(info.ptr); + dim = info.shape[0]; + if (dim % 2 == 0) + throw std::runtime_error( + "Array dimension must be odd. The " + "middle of an even numbered array is ambiguous."); + radius = dim / 2; // Rounds down + sum = 0.0; + kernel = std::vector(pix, pix + dim * dim); + calc_sum(); + width = 0.0; + } + + static void psf_bindings(py::module &m) { + using psf = search::PSF; + + py::class_(m, "PSF", py::buffer_protocol(), pydocs::DOC_PSF) + .def_buffer([](psf &m) -> py::buffer_info { + return py::buffer_info(m.data(), // void *ptr; + sizeof(float), // py::ssize_t itemsize; + py::format_descriptor::format(), // std::string format; + 2, // py::ssize_t ndim; + {m.get_dim(), m.get_dim()}, // std::vector shape; + {sizeof(float) * m.get_dim(), sizeof(float)}); // std::vector strides; + }) + .def(py::init<>()) + .def(py::init()) + .def(py::init>()) + .def(py::init()) + .def("set_array", &psf::set_array, pydocs::DOC_PSF_set_array) + .def("get_std", &psf::get_std, pydocs::DOC_PSF_get_std) + .def("get_sum", &psf::get_sum, pydocs::DOC_PSF_get_sum) + .def("get_dim", &psf::get_dim, pydocs::DOC_PSF_get_dim) + .def("get_radius", &psf::get_radius, pydocs::DOC_PSF_get_radius) + .def("get_size", &psf::get_size, pydocs::DOC_PSF_get_size) + .def("get_kernel", &psf::get_kernel, pydocs::DOC_PSF_get_kernel) + .def("get_value", &psf::get_value, pydocs::DOC_PSF_get_value) + .def("square_psf", &psf::square_psf, pydocs::DOC_PSF_square_psf) + .def("print", &psf::print, pydocs::DOC_PSF_print); + } +#endif + +} /* namespace search */ diff --git a/src/kbmod/search/psf.h b/src/kbmod/search/psf.h new file mode 100644 index 000000000..0a38c873c --- /dev/null +++ b/src/kbmod/search/psf.h @@ -0,0 +1,56 @@ +#ifndef PSF_H_ +#define PSF_H_ + +#include +#include +#include +#include +#include +#include "common.h" +#include "pydocs/psf_docs.h" + + +namespace search { + class PSF { + public: + PSF(); // Create a no-op PSF. + PSF(float stdev); + PSF(const PSF& other); // Copy constructor + PSF(PSF&& other); // Move constructor +#ifdef Py_PYTHON_H + PSF(pybind11::array_t arr); + void set_array(pybind11::array_t arr); +#endif + + virtual ~PSF(){}; + + // Assignment functions. + PSF& operator=(const PSF& other); // Copy assignment + PSF& operator=(PSF&& other); // Move assignment + + // Getter functions (inlined) + float get_std() const { return width; } + float get_sum() const { return sum; } + float get_value(int x, int y) const { return kernel[y * dim + x]; } + int get_dim() const { return dim; } + int get_radius() const { return radius; } + int get_size() const { return kernel.size(); } + const std::vector& get_kernel() const { return kernel; }; + float* data() { return kernel.data(); } + + // Computation functions. + void calc_sum(); + void square_psf(); + std::string print(); + + private: + std::vector kernel; + float width; + float sum; + int dim; + int radius; + }; + +} /* namespace search */ + +#endif /* PSF_H_ */ diff --git a/src/kbmod/search/pydocs/common_docs.h b/src/kbmod/search/pydocs/common_docs.h new file mode 100644 index 000000000..ae76909b7 --- /dev/null +++ b/src/kbmod/search/pydocs/common_docs.h @@ -0,0 +1,44 @@ +#ifndef COMMON_DOCS +#define COMMON_DOCS + +namespace pydocs { + static const auto DOC_Trajectory = R"doc( + A trajectory structure holding basic information about potential results. + + Attributes + ---------- + x : `float` + x coordinate of the origin (what?) + y : `float` + y coordinate of the origin (what?) + vx : `float` + x component of the velocity, as projected on the image + vy : `float` + y component of the velocity, as projected on the image + lh : `float` + Likelihood (accumulated?) + flux : `float` + Flux (accumulated?) + obs_count : `int` + Number of observations trajectory was seen in. + )doc"; + + static const auto DOC_PixelPos = R"doc( + todo + )doc"; + + static const auto DOC_ImageMoments = R"doc( + todo + )doc"; + + static const auto DOC_StampParameters = R"doc( + todo + ")doc"; + + static const auto DOC_BaryCorrection = R"doc( + todo + ")doc"; + +} // namespace pydocs + +#endif /* COMMON_DOCS */ diff --git a/src/kbmod/search/pydocs/image_stack_docs.h b/src/kbmod/search/pydocs/image_stack_docs.h new file mode 100644 index 000000000..ed21bf377 --- /dev/null +++ b/src/kbmod/search/pydocs/image_stack_docs.h @@ -0,0 +1,79 @@ +#ifndef IMAGESTACK_DOCS +#define IMAGESTACK_DOCS + +namespace pydocs { + static const auto DOC_ImageStack = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_get_images = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_get_single_image = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_set_single_image = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_get_times = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_set_times = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_img_count = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_apply_mask_flags = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_apply_mask_threshold = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_apply_global_mask = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_grow_mask = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_save_global_mask = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_save_images = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_get_global_mask = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_convolve_psf = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_get_width = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_get_height = R"doc( + todo + )doc"; + + static const auto DOC_ImageStack_get_npixels = R"doc( + todo + )doc"; + +} /* pydocs */ + +#endif /* IMAGESTACK_DOCS */ diff --git a/src/kbmod/search/pydocs/layered_image_docs.h b/src/kbmod/search/pydocs/layered_image_docs.h new file mode 100644 index 000000000..3d9a55f66 --- /dev/null +++ b/src/kbmod/search/pydocs/layered_image_docs.h @@ -0,0 +1,143 @@ +#ifndef LAYEREDIMAGE_DOCS +#define LAYEREDIMAGE_DOCS + +namespace pydocs { + static const auto DOC_LayeredImage = R"doc( + Creates a layered_image out of individual `RawImage` layers. + + Parameters + ---------- + path : `str`, optional + Path to a FITS file containing ``science``, ``mask`` and ``variance`` + extensions. + sci : `RawImage`, optional + The `RawImage` for the science layer. + var : `RawImage`, optional + The `RawImage` for the cariance layer. + msk : `RawImage`, optional + The `RawImage` for the mask layer. + name : `str`, optional + File/layered image name. + width : `int`, optional + Width of the images (in pixels). + height : `int`, optional + Height of the images (in pixels). + std: `float`, optional + Standard deviation of the image. + var: `float`, optional + Variance of the pixels, assumed uniform. + time : `float`, optional + Observation time. + seed : `int`, optional + Pseudorandom number generator. + + psf : `PSF` + The PSF for the image. + + Raises + ------ + RuntimeError: + If the science, variance and mask are not the same size. + + Notes + ----- + Class can be instantiated from a file, or from 3 `RawImage` (science, mask, + variance) objects, or by providing the name, dimensions, standard deviation, + variance, observation time to which, additionally, a random seed generator can + be provided. PSF is always required. + )doc"; + + static const auto DOC_LayeredImage_set_psf = R"doc( + Sets the PSF object. + )doc"; + + static const auto DOC_LayeredImage_get_psf = R"doc( + Returns the PSF object. + )doc"; + + static const auto DOC_LayeredImage_apply_mask_flags = R"doc( + No idea + )doc"; + + static const auto DOC_LayeredImage_apply_mask_threshold = R"doc( + No idea + )doc"; + + static const auto DOC_LayeredImage_sub_template = R"doc( + Subtract given image template + )doc"; + + static const auto DOC_LayeredImage_save_layers = R"doc( + Save image? + )doc"; + + static const auto DOC_LayeredImage_get_science = R"doc( + Returns the science layer raw_image. + )doc"; + + static const auto DOC_LayeredImage_get_mask = R"doc( + Returns the mask layer raw_image. + )doc"; + + static const auto DOC_LayeredImage_get_variance = R"doc( + Returns the variance layer raw_image. + )doc"; + + static const auto DOC_LayeredImage_set_science = R"doc( + Returns the science layer raw_image. + )doc"; + + static const auto DOC_LayeredImage_set_mask = R"doc( + Returns the mask layer raw_image. + )doc"; + + static const auto DOC_LayeredImage_set_variance = R"doc( + Returns the science layer raw_image. + )doc"; + + static const auto DOC_LayeredImage_convolve_psf = R"doc( + todo + )doc"; + + static const auto DOC_LayeredImage_convolve_given_psf = R"doc( + todo + )doc"; + + static const auto DOC_LayeredImage_grow_mask = R"doc( + todo + )doc"; + + static const auto DOC_LayeredImage_get_name = R"doc( + Returns the name of the layered image. + )doc"; + + static const auto DOC_LayeredImage_get_width = R"doc( + Returns the image's width in pixels. + )doc"; + + static const auto DOC_LayeredImage_get_height = R"doc( + Returns the image's height in pixels. + )doc"; + + static const auto DOC_LayeredImage_get_npixels = R"doc( + Returns the image's total number of pixels. + )doc"; + + static const auto DOC_LayeredImage_get_obstime = R"doc( + Get the image's observation time. + )doc"; + + static const auto DOC_LayeredImage_set_obstime = R"doc( + Set the image's observation time. + )doc"; + + static const auto DOC_LayeredImage_generate_psi_image = R"doc( + todo + )doc"; + + static const auto DOC_LayeredImage_generate_phi_image = R"doc( + todo + )doc"; +} /* pydocs */ + +#endif /* LAYEREDIMAGE_DOCS */ diff --git a/src/kbmod/search/pydocs/psf_docs.h b/src/kbmod/search/pydocs/psf_docs.h new file mode 100644 index 000000000..a3d679abc --- /dev/null +++ b/src/kbmod/search/pydocs/psf_docs.h @@ -0,0 +1,76 @@ +#ifndef PSF_DOCS +#define PSF_DOCS + +namespace pydocs { + static const auto DOC_PSF = R"doc( + Point Spread Function. + + Parameters + ---------- + stdev : `float`, optional + Standard deviation of the Gaussian PSF. + psf : `PSF`, optional + Another PSF object. + arr : `numpy.array`, optional + A realization of the PSF. + + Notes + ----- + When instantiated with another `psf` object, returns its copy. + When instantiated with an array-like object, that array must be + a square matrix and have an odd number of dimensions. Only one + of the arguments is required. + )doc"; + + static const auto DOC_PSF_set_array = R"doc( + Set the kernel values of a realized PSF. + + Parameters + ---------- + arr : `numpy.array` + A realization of the PSF. + + Notes + ----- + Given realization of a PSF has to be an odd-dimensional square + matrix. + )doc"; + + static const auto DOC_PSF_get_std = R"doc( + "Returns the PSF's standard deviation." + )doc"; + + static const auto DOC_PSF_get_sum = R"doc( + "Returns the sum of PSFs kernel elements. + ")doc"; + + static const auto DOC_PSF_get_dim = R"doc( + "Returns the PSF kernel dimensions. + ")doc"; + + static const auto DOC_PSF_get_radius = R"doc( + "Returns the radius of the PSF + ")doc"; + + static const auto DOC_PSF_get_size = R"doc( + "Returns the number of elements in the PSFs kernel. + ")doc"; + + static const auto DOC_PSF_get_kernel = R"doc( + "Returns the PSF kernel. + ")doc"; + + static const auto DOC_PSF_get_value = R"doc( + "Returns the PSF kernel value at a specific point. + ")doc"; + + static const auto DOC_PSF_square_psf = R"doc( + "Squares, raises to the power of two, the elements of the PSF kernel. + ")doc"; + + static const auto DOC_PSF_print = R"doc( + "Pretty-prints the PSF. + ")doc"; +} // namespace pydocs + +#endif /* PSF_DOCS */ diff --git a/src/kbmod/search/pydocs/raw_image_docs.h b/src/kbmod/search/pydocs/raw_image_docs.h new file mode 100644 index 000000000..dead10120 --- /dev/null +++ b/src/kbmod/search/pydocs/raw_image_docs.h @@ -0,0 +1,116 @@ +#ifndef RAWIMAGEDOCS +#define RAWIMAGEDOCS + +namespace pydocs{ + static const auto DOC_RawImage = R"doc( + Raw Image, a row-ordered 2D array unraveled into a vector. + )doc"; + + static const auto DOC_RawImage_get_height = R"doc( + Returns the image's height in pixels. + )doc"; + + static const auto DOC_RawImage_get_width = R"doc( + Returns the image's width in pixels. + )doc"; + + static const auto DOC_RawImage_get_npixels = R"doc( + Returns the image's total number of pixels. + )doc"; + + static const auto DOC_RawImage_get_all_pixels = R"doc( + Returns a list of the images pixels. + )doc"; + + static const auto DOC_RawImage_set_array = R"doc( + Sets all image pixels given an array of values. + )doc"; + + static const auto DOC_RawImage_get_obstime = R"doc( + Get the observation time of the image. + )doc"; + + static const auto DOC_RawImage_set_obstime = R"doc( + Set the observation time of the image. + )doc"; + + static const auto DOC_RawImage_approx_equal = R"doc( + Checks if two images are approximately equal. + )doc"; + + static const auto DOC_RawImage_compute_bounds = R"doc( + Returns min and max pixel values. + )doc"; + + static const auto DOC_RawImage_find_peak = R"doc( + Returns the pixel coordinates of the maximum value. + )doc"; + + static const auto DOC_RawImage_find_central_moments = R"doc( + Returns the central moments of the image. + )doc"; + + static const auto DOC_RawImage_create_stamp = R"doc( + Create stamp. + )doc"; + + static const auto DOC_RawImage_set_pixel = R"doc( + Set the value of a given pixel. + )doc"; + + static const auto DOC_RawImage_add_pixel = R"doc( + Add to the raw value of a given pixel. + )doc"; + + static const auto DOC_RawImage_add_pixel_interp = R"doc( + Add to the value calculated by bilinear interpolation + of the neighborhood of the given pixel position. + )doc"; + + static const auto DOC_RawImage_apply_mask = R"doc( + applies mask + )doc"; + + static const auto DOC_RawImage_grow_mask = R"doc( + grows mask + )doc"; + + static const auto DOC_RawImage_pixel_has_data = R"doc( + Returns a Boolean indicating whether the pixel has data. + )doc"; + + static const auto DOC_RawImage_set_all = R"doc( + Set all pixel values given an array. + )doc"; + + static const auto DOC_RawImage_get_pixel = R"doc( + Returns the value of a pixel. + )doc"; + + static const auto DOC_RawImage_get_pixel_interp = R"doc( + Get the interoplated value of a pixel. + )doc"; + + static const auto DOC_RawImage_convolve = R"doc( + Convolve the image with a PSF. + )doc"; + + static const auto DOC_RawImage_convolve_cpu = R"doc( + Convolve the image with a PSF. + )doc"; + + static const auto DOC_RawImage_load_fits = R"doc( + Load the image data from a FITS file. + )doc"; + + static const auto DOC_RawImage_save_fits = R"doc( + Save the image to a FITS file. + )doc"; + + static const auto DOC_RawImage_append_fits_layer = R"doc( + Append the image as a layer in a FITS file. + )doc"; + +} /* namespace pydocs */ + +#endif /* RAWIMAGE_DOCS */ diff --git a/src/kbmod/search/pydocs/stack_search_docs.h b/src/kbmod/search/pydocs/stack_search_docs.h new file mode 100644 index 000000000..5fe0eb71f --- /dev/null +++ b/src/kbmod/search/pydocs/stack_search_docs.h @@ -0,0 +1,114 @@ +#ifndef STACKSEARCH_DOCS +#define STACHSEARCH_DOCS + +namespace pydocs { + static const auto DOC_StackSearch = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_save_psi_phi = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_search = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_enable_gpu_sigmag_filter = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_enable_gpu_encoding = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_enable_corr = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_set_start_bounds_x = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_set_start_bounds_y = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_set_debug = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_filter_min_obs = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_num_images = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_imagestack = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_stamps = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_median_stamp = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_mean_stamp = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_summed_stamp = R"doc( + 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"; + + static const auto DOC_StackSearch_get_trajectory_positions = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_psi_curves = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_phi_curves= R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_prepare_psi_phi = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_psi_images = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_phi_images = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_get_results = R"doc( + todo + )doc"; + + static const auto DOC_StackSearch_set_results = R"doc( + todo + )doc"; + +} /* pydocs */ +#endif /* STACKSEARCH_DOCS */ diff --git a/src/kbmod/search/raw_image.cpp b/src/kbmod/search/raw_image.cpp new file mode 100644 index 000000000..af4ef2e5b --- /dev/null +++ b/src/kbmod/search/raw_image.cpp @@ -0,0 +1,671 @@ +#include "raw_image.h" + + +namespace py = pybind11; + + +namespace search { +#ifdef HAVE_CUDA + // Performs convolution between an image represented as an array of floats + // and a PSF on a GPU device. + extern "C" void deviceConvolve(float* source_img, float* result_img, int width, int height, float* psf_kernel, + int psf_size, int psf_dim, int psf_radius, float psf_sum); +#endif + + RawImage::RawImage() : width(0), height(0), obstime(-1.0) { pixels = std::vector(); } + + // Copy constructor + RawImage::RawImage(const RawImage& old) { + width = old.get_width(); + height = old.get_height(); + pixels = old.get_pixels(); + obstime = old.get_obstime(); + } + + // Copy assignment + RawImage& RawImage::operator=(const RawImage& source) { + width = source.width; + height = source.height; + pixels = source.pixels; + obstime = source.obstime; + return *this; + } + + // Move constructor + RawImage::RawImage(RawImage&& source) + : width(source.width), + height(source.height), + obstime(source.obstime), + pixels(std::move(source.pixels)) {} + + // Move assignment + RawImage& RawImage::operator=(RawImage&& source) { + if (this != &source) { + width = source.width; + height = source.height; + pixels = std::move(source.pixels); + obstime = source.obstime; + } + return *this; + } + + RawImage::RawImage(unsigned w, unsigned h) : height(h), width(w), obstime(-1.0), pixels(w * h) {} + + RawImage::RawImage(unsigned w, unsigned h, const std::vector& pix) + : width(w), height(h), obstime(-1.0), pixels(pix) { + assert(w * h == pix.size()); + } + + bool RawImage::approx_equal(const RawImage& img_b, float atol) const { + if ((width != img_b.width) || (height != img_b.height)) return false; + + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + float p1 = get_pixel(x, y); + float p2 = img_b.get_pixel(x, y); + + // NO_DATA values must match exactly. + if ((p1 == NO_DATA) && (p2 != NO_DATA)) return false; + if ((p1 != NO_DATA) && (p2 == NO_DATA)) return false; + + // Other values match within an absolute tolerance. + if (fabs(p1 - p2) > atol) return false; + } + } + return true; + } + + // Load the image data from a specific layer of a FITS file. + void RawImage::load_from_file(const std::string& file_path, int layer_num) { + // Open the file's header and read in the obstime and the dimensions. + fitsfile* fptr; + int status = 0; + int mjdStatus = 0; + int file_not_found; + int nullval = 0; + int anynull = 0; + + // Open the correct layer to extract the RawImage. + std::string layerPath = file_path + "[" + std::to_string(layer_num) + "]"; + if (fits_open_file(&fptr, layerPath.c_str(), READONLY, &status)) { + fits_report_error(stderr, status); + throw std::runtime_error("Could not open FITS file to read RawImage"); + } + + // Read image dimensions. + long dimensions[2]; + if (fits_read_keys_lng(fptr, "NAXIS", 1, 2, dimensions, &file_not_found, &status)) + fits_report_error(stderr, status); + width = dimensions[0]; + height = dimensions[1]; + + // Read in the image. + pixels = std::vector(width * height); + if (fits_read_img(fptr, TFLOAT, 1, get_npixels(), &nullval, pixels.data(), &anynull, &status)) + fits_report_error(stderr, status); + + // Read image observation time, ignore error if does not exist + obstime = -1.0; + if (fits_read_key(fptr, TDOUBLE, "MJD", &obstime, NULL, &mjdStatus)) obstime = -1.0; + if (fits_close_file(fptr, &status)) fits_report_error(stderr, status); + + // If we are reading from a sublayer and did not find a time, try the overall header. + if (obstime < 0.0) { + if (fits_open_file(&fptr, file_path.c_str(), READONLY, &status)) + throw std::runtime_error("Could not open FITS file to read RawImage"); + fits_read_key(fptr, TDOUBLE, "MJD", &obstime, NULL, &mjdStatus); + if (fits_close_file(fptr, &status)) fits_report_error(stderr, status); + } + } + + void RawImage::save_to_file(const std::string& filename) { + fitsfile* fptr; + int status = 0; + long naxes[2] = {0, 0}; + + fits_create_file(&fptr, filename.c_str(), &status); + + // If we are unable to create the file, check if it already exists + // and, if so, delete it and retry the create. + if (status == 105) { + status = 0; + fits_open_file(&fptr, filename.c_str(), READWRITE, &status); + if (status == 0) { + fits_delete_file(fptr, &status); + fits_create_file(&fptr, filename.c_str(), &status); + } + } + + // Create the primary array image (32-bit float pixels) + long dimensions[2]; + dimensions[0] = width; + dimensions[1] = height; + fits_create_img(fptr, FLOAT_IMG, 2 /*naxis*/, dimensions, &status); + fits_report_error(stderr, status); + + /* Write the array of floats to the image */ + fits_write_img(fptr, TFLOAT, 1, get_npixels(), pixels.data(), &status); + fits_report_error(stderr, status); + + // Add the basic header data. + fits_update_key(fptr, TDOUBLE, "MJD", &obstime, "[d] Generated Image time", &status); + fits_report_error(stderr, status); + + fits_close_file(fptr, &status); + fits_report_error(stderr, status); + } + + void RawImage::append_layer_to_file(const std::string& filename) { + int status = 0; + fitsfile* f; + + // Check that we can open the file. + if (fits_open_file(&f, filename.c_str(), READWRITE, &status)) { + fits_report_error(stderr, status); + throw std::runtime_error("Unable to open FITS file for appending."); + } + + // This appends a layer (extension) if the file exists) + /* Create the primary array image (32-bit float pixels) */ + long dimensions[2]; + dimensions[0] = width; + dimensions[1] = height; + fits_create_img(f, FLOAT_IMG, 2 /*naxis*/, dimensions, &status); + fits_report_error(stderr, status); + + /* Write the array of floats to the image */ + fits_write_img(f, TFLOAT, 1, get_npixels(), pixels.data(), &status); + fits_report_error(stderr, status); + + // Save the image time in the header. + fits_update_key(f, TDOUBLE, "MJD", &obstime, "[d] Generated Image time", &status); + fits_report_error(stderr, status); + + fits_close_file(f, &status); + fits_report_error(stderr, status); + } + + RawImage RawImage::create_stamp(float x, float y, int radius, bool interpolate, bool keep_no_data) const { + if (radius < 0) throw std::runtime_error("stamp radius must be at least 0"); + + int dim = radius * 2 + 1; + RawImage stamp(dim, dim); + for (int yoff = 0; yoff < dim; ++yoff) { + for (int xoff = 0; xoff < dim; ++xoff) { + float pix_val; + if (interpolate) + pix_val = get_pixel_interp(x + static_cast(xoff - radius), + y + static_cast(yoff - radius)); + else + pix_val = get_pixel(static_cast(x) + xoff - radius, static_cast(y) + yoff - radius); + if ((pix_val == NO_DATA) && !keep_no_data) pix_val = 0.0; + stamp.set_pixel(xoff, yoff, pix_val); + } + } + + stamp.set_obstime(obstime); + return stamp; + } + + void RawImage::convolve_cpu(const PSF& psf) { + std::vector result(width * height, 0.0); + const int psf_rad = psf.get_radius(); + const float psf_total = psf.get_sum(); + + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + // Pixels with NO_DATA remain NO_DATA. + if (pixels[y * width + x] == NO_DATA) { + result[y * width + x] = NO_DATA; + continue; + } + + float sum = 0.0; + float psf_portion = 0.0; + for (int j = -psf_rad; j <= psf_rad; j++) { + for (int i = -psf_rad; i <= psf_rad; i++) { + if ((x + i >= 0) && (x + i < width) && (y + j >= 0) && (y + j < height)) { + float current_pixel = pixels[(y + j) * width + (x + i)]; + if (current_pixel != NO_DATA) { + float current_psf = psf.get_value(i + psf_rad, j + psf_rad); + psf_portion += current_psf; + sum += current_pixel * current_psf; + } + } + } + } + result[y * width + x] = (sum * psf_total) / psf_portion; + } + } + + // Copy the data into the pixels vector. + const int npixels = get_npixels(); + for (int i = 0; i < npixels; ++i) { + pixels[i] = result[i]; + } + } + + void RawImage::convolve(PSF psf) { +#ifdef HAVE_CUDA + deviceConvolve(pixels.data(), pixels.data(), get_width(), get_height(), psf.data(), psf.get_size(), + psf.get_dim(), psf.get_radius(), psf.get_sum()); +#else + convolve_cpu(psf); +#endif + } + + void RawImage::apply_mask(int flags, const std::vector& exceptions, const RawImage& mask) { + const std::vector& mask_pix = mask.get_pixels(); + const int num_pixels = get_npixels(); + assert(num_pixels == mask.get_npixels()); + for (unsigned int p = 0; p < num_pixels; ++p) { + int pix_flags = static_cast(mask_pix[p]); + bool is_exception = false; + for (auto& e : exceptions) is_exception = is_exception || e == pix_flags; + if (!is_exception && ((flags & pix_flags) != 0)) pixels[p] = NO_DATA; + } + } + + /* This implementation of grow_mask is optimized for steps > 1 + (which is how the code is generally used. If you are only + growing the mask by 1, the extra copy will be a little slower. + */ + void RawImage::grow_mask(int steps) { + const int num_pixels = width * height; + + // Set up the initial masked vector that stores the number of steps + // each pixel is from a masked pixel in the original image. + std::vector masked(num_pixels, -1); + for (int i = 0; i < num_pixels; ++i) { + if (pixels[i] == NO_DATA) masked[i] = 0; + } + + // Grow out the mask one for each step. + for (int itr = 1; itr <= steps; ++itr) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + int center = width * y + x; + if (masked[center] == -1) { + // Mask pixels that are adjacent to a pixel masked during + // the last iteration only. + if ((x + 1 < width && masked[center + 1] == itr - 1) || + (x - 1 >= 0 && masked[center - 1] == itr - 1) || + (y + 1 < height && masked[center + width] == itr - 1) || + (y - 1 >= 0 && masked[center - width] == itr - 1)) { + masked[center] = itr; + } + } + } + } + } + + // Mask the pixels in the image. + for (std::size_t i = 0; i < num_pixels; ++i) { + if (masked[i] > -1) { + pixels[i] = NO_DATA; + } + } + } + + std::vector RawImage::bilinear_interp(float x, float y) const { + // Linear interpolation + // Find the 4 pixels (aPix, bPix, cPix, dPix) + // that the corners (a, b, c, d) of the + // new pixel land in, and blend into those + + // Returns a vector with 4 pixel locations + // and their interpolation value + + // Top right + float ax = x + 0.5; + float ay = y + 0.5; + float a_px = floor(ax); + float a_py = floor(ay); + float a_amt = (ax - a_px) * (ay - a_py); + + // Bottom right + float bx = x + 0.5; + float by = y - 0.5; + float b_px = floor(bx); + float b_py = floor(by); + float b_amt = (bx - b_px) * (b_py + 1.0 - by); + + // Bottom left + float cx = x - 0.5; + float cy = y - 0.5; + float c_px = floor(cx); + float c_py = floor(cy); + float c_amt = (c_px + 1.0 - cx) * (c_py + 1.0 - cy); + + // Top left + float dx = x - 0.5; + float dy = y + 0.5; + float d_px = floor(dx); + float d_py = floor(dy); + float d_amt = (d_px + 1.0 - dx) * (dy - d_py); + + // make sure the right amount has been distributed + float diff = std::abs(a_amt + b_amt + c_amt + d_amt - 1.0); + if (diff > 0.01) std::cout << "warning: bilinear_interpSum == " << diff << "\n"; + return {a_px, a_py, a_amt, b_px, b_py, b_amt, c_px, c_py, c_amt, d_px, d_py, d_amt}; + } + + void RawImage::add_pixel_interp(float x, float y, float value) { + // Interpolation values + std::vector iv = bilinear_interp(x, y); + + add_to_pixel(iv[0], iv[1], value * iv[2]); + add_to_pixel(iv[3], iv[4], value * iv[5]); + add_to_pixel(iv[6], iv[7], value * iv[8]); + add_to_pixel(iv[9], iv[10], value * iv[11]); + } + + void RawImage::add_to_pixel(float fx, float fy, float value) { + assert(fx - floor(fx) == 0.0 && fy - floor(fy) == 0.0); + int x = static_cast(fx); + int y = static_cast(fy); + if (x >= 0 && x < width && y >= 0 && y < height) pixels[y * width + x] += value; + } + + float RawImage::get_pixel_interp(float x, float y) const { + if ((x < 0.0 || y < 0.0) || (x > static_cast(width) || y > static_cast(height))) + return NO_DATA; + std::vector iv = bilinear_interp(x, y); + float a = get_pixel(iv[0], iv[1]); + float b = get_pixel(iv[3], iv[4]); + float c = get_pixel(iv[6], iv[7]); + float d = get_pixel(iv[9], iv[10]); + float interpSum = 0.0; + float total = 0.0; + if (a != NO_DATA) { + interpSum += iv[2]; + total += a * iv[2]; + } + if (b != NO_DATA) { + interpSum += iv[5]; + total += b * iv[5]; + } + if (c != NO_DATA) { + interpSum += iv[8]; + total += c * iv[8]; + } + if (d != NO_DATA) { + interpSum += iv[11]; + total += d * iv[11]; + } + if (interpSum == 0.0) { + return NO_DATA; + } else { + return total / interpSum; + } + } + + void RawImage::set_all_pix(float value) { + for (auto& p : pixels) p = value; + } + + std::array RawImage::compute_bounds() const { + const int num_pixels = get_npixels(); + float min_val = FLT_MAX; + float max_val = -FLT_MAX; + for (unsigned p = 0; p < num_pixels; ++p) { + if (pixels[p] != NO_DATA) { + min_val = std::min(min_val, pixels[p]); + max_val = std::max(max_val, pixels[p]); + } + } + + // Assert that we have seen at least some valid data. + assert(max_val != -FLT_MAX); + assert(min_val != FLT_MAX); + + // Set and return the result array. + std::array res; + res[0] = min_val; + res[1] = max_val; + return res; + } + + // The maximum value of the image and return the coordinates. + PixelPos RawImage::find_peak(bool furthest_from_center) const { + int c_x = width / 2; + int c_y = height / 2; + + // Initialize the variables for tracking the peak's location. + PixelPos result = {0, 0}; + float max_val = pixels[0]; + float dist2 = c_x * c_x + c_y * c_y; + + // Search each pixel for the peak. + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + float pix_val = pixels[y * width + x]; + if (pix_val > max_val) { + max_val = pix_val; + result.x = x; + result.y = y; + dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y); + } else if (pix_val == max_val) { + int new_dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y); + if ((furthest_from_center && (new_dist2 > dist2)) || + (!furthest_from_center && (new_dist2 < dist2))) { + max_val = pix_val; + result.x = x; + result.y = y; + dist2 = new_dist2; + } + } + } + } + + return result; + } + + // Find the basic image moments in order to test if stamps have a gaussian shape. + // It computes the moments on the "normalized" image where the minimum + // value has been shifted to zero and the sum of all elements is 1.0. + // Elements with NO_DATA are treated as zero. + ImageMoments RawImage::find_central_moments() const { + const int num_pixels = width * height; + const int c_x = width / 2; + const int c_y = height / 2; + + // Set all the moments to zero initially. + ImageMoments res = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0}; + + // Find the min (non-NO_DATA) value to subtract off. + float min_val = FLT_MAX; + for (int p = 0; p < num_pixels; ++p) { + min_val = ((pixels[p] != NO_DATA) && (pixels[p] < min_val)) ? pixels[p] : min_val; + } + + // Find the sum of the zero-shifted (non-NO_DATA) pixels. + double sum = 0.0; + for (int p = 0; p < num_pixels; ++p) { + sum += (pixels[p] != NO_DATA) ? (pixels[p] - min_val) : 0.0; + } + if (sum == 0.0) return res; + + // Compute the rest of the moments. + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + int ind = y * width + x; + float pix_val = (pixels[ind] != NO_DATA) ? (pixels[ind] - min_val) / sum : 0.0; + + res.m00 += pix_val; + res.m10 += (x - c_x) * pix_val; + res.m20 += (x - c_x) * (x - c_x) * pix_val; + res.m01 += (y - c_y) * pix_val; + res.m02 += (y - c_y) * (y - c_y) * pix_val; + res.m11 += (x - c_x) * (y - c_y) * pix_val; + } + } + + return res; + } + + RawImage create_median_image(const std::vector& images) { + int num_images = images.size(); + assert(num_images > 0); + + int width = images[0].get_width(); + int height = images[0].get_height(); + for (auto& img : images) assert(img.get_width() == width and img.get_height() == height); + + RawImage result = RawImage(width, height); + std::vector pix_array(num_images); + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + int num_unmasked = 0; + for (int i = 0; i < num_images; ++i) { + // Only used the unmasked pixels. + float pix_val = images[i].get_pixel(x, y); + if ((pix_val != NO_DATA) && (!std::isnan(pix_val))) { + pix_array[num_unmasked] = pix_val; + num_unmasked += 1; + } + } + + if (num_unmasked > 0) { + std::sort(pix_array.begin(), pix_array.begin() + num_unmasked); + + // If we have an even number of elements, take the mean of the two + // middle ones. + int median_ind = num_unmasked / 2; + if (num_unmasked % 2 == 0) { + float ave_middle = (pix_array[median_ind] + pix_array[median_ind - 1]) / 2.0; + result.set_pixel(x, y, ave_middle); + } else { + result.set_pixel(x, y, pix_array[median_ind]); + } + } else { + // We use a 0.0 value if there is no data to allow for visualization + // and value based filtering. + result.set_pixel(x, y, 0.0); + } + } + } + + return result; + } + + RawImage create_summed_image(const std::vector& images) { + int num_images = images.size(); + assert(num_images > 0); + + int width = images[0].get_width(); + int height = images[0].get_height(); + for (auto& img : images) assert(img.get_width() == width and img.get_height() == height); + + RawImage result = RawImage(width, height); + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + float sum = 0.0; + for (int i = 0; i < num_images; ++i) { + float pix_val = images[i].get_pixel(x, y); + if ((pix_val == NO_DATA) || (std::isnan(pix_val))) pix_val = 0.0; + sum += pix_val; + } + result.set_pixel(x, y, sum); + } + } + + return result; + } + + RawImage create_mean_image(const std::vector& images) { + int num_images = images.size(); + assert(num_images > 0); + + int width = images[0].get_width(); + int height = images[0].get_height(); + for (auto& img : images) assert(img.get_width() == width and img.get_height() == height); + + RawImage result = RawImage(width, height); + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + float sum = 0.0; + float count = 0.0; + for (int i = 0; i < num_images; ++i) { + float pix_val = images[i].get_pixel(x, y); + if ((pix_val != NO_DATA) && (!std::isnan(pix_val))) { + count += 1.0; + sum += pix_val; + } + } + + if (count > 0.0) { + result.set_pixel(x, y, sum / count); + } else { + // We use a 0.0 value if there is no data to allow for visualization + // and value based filtering. + result.set_pixel(x, y, 0.0); + } + } + } + + return result; + } + + +#ifdef Py_PYTHON_H + RawImage::RawImage(pybind11::array_t arr) { + obstime = -1.0; + set_array(arr); + } + + void RawImage::set_array(pybind11::array_t& arr) { + pybind11::buffer_info info = arr.request(); + + if (info.ndim != 2) throw std::runtime_error("Array must have 2 dimensions."); + + width = info.shape[1]; + height = info.shape[0]; + float* pix = static_cast(info.ptr); + + pixels = std::vector(pix, pix + get_npixels()); + } + + + static void raw_image_bindings(py::module &m) { + using ri = search::RawImage; + + py::class_(m, "RawImage", py::buffer_protocol(), pydocs::DOC_RawImage) + .def_buffer([](ri &m) -> py::buffer_info { + return py::buffer_info(m.data(), sizeof(float), py::format_descriptor::format(), + 2, {m.get_height(), m.get_width()}, + {sizeof(float) * m.get_width(), sizeof(float)}); + }) + .def(py::init()) + .def(py::init()) + .def(py::init>()) + .def("get_height", &ri::get_height, pydocs::DOC_RawImage_get_height ) + .def("get_width", &ri::get_width, pydocs:: DOC_RawImage_get_width) + .def("get_npixels", &ri::get_npixels, pydocs:: DOC_RawImage_get_npixels) + .def("get_all_pixels", &ri::get_pixels, pydocs:: DOC_RawImage_get_all_pixels) + .def("set_array", &ri::set_array, pydocs:: DOC_RawImage_set_array) + .def("get_obstime", &ri::get_obstime, pydocs:: DOC_RawImage_get_obstime) + .def("set_obstime", &ri::set_obstime, pydocs:: DOC_RawImage_set_obstime) + .def("approx_equal", &ri::approx_equal, pydocs:: DOC_RawImage_approx_equal) + .def("compute_bounds", &ri::compute_bounds, pydocs:: DOC_RawImage_compute_bounds) + .def("find_peak", &ri::find_peak, pydocs:: DOC_RawImage_find_peak) + .def("find_central_moments", &ri::find_central_moments, pydocs:: DOC_RawImage_find_central_moments) + .def("create_stamp", &ri::create_stamp, pydocs:: DOC_RawImage_create_stamp) + .def("set_pixel", &ri::set_pixel, pydocs:: DOC_RawImage_set_pixel) + .def("add_pixel", &ri::add_to_pixel, pydocs:: DOC_RawImage_add_pixel) + .def("add_pixel_interp", &ri::add_pixel_interp, pydocs:: DOC_RawImage_add_pixel_interp) + .def("apply_mask", &ri::apply_mask, pydocs:: DOC_RawImage_apply_mask) + .def("grow_mask", &ri::grow_mask, pydocs:: DOC_RawImage_grow_mask) + .def("pixel_has_data", &ri::pixel_has_data, pydocs:: DOC_RawImage_pixel_has_data) + .def("set_all", &ri::set_all_pix, pydocs:: DOC_RawImage_set_all) + .def("get_pixel", &ri::get_pixel, pydocs:: DOC_RawImage_get_pixel) + .def("get_pixel_interp", &ri::get_pixel_interp, pydocs:: DOC_RawImage_get_pixel_interp) + .def("convolve", &ri::convolve, pydocs:: DOC_RawImage_convolve) + .def("convolve_cpu", &ri::convolve_cpu, pydocs:: DOC_RawImage_convolve_cpu) + .def("load_fits", &ri::load_from_file, pydocs:: DOC_RawImage_load_fits) + .def("save_fits", &ri::save_to_file, pydocs:: DOC_RawImage_save_fits) + .def("append_fits_layer", &ri::append_layer_to_file, pydocs:: DOC_RawImage_append_fits_layer); + } +#endif + +} /* namespace search */ diff --git a/src/kbmod/search/RawImage.h b/src/kbmod/search/raw_image.h similarity index 51% rename from src/kbmod/search/RawImage.h rename to src/kbmod/search/raw_image.h index 46d576a2f..cf43d0cf9 100644 --- a/src/kbmod/search/RawImage.h +++ b/src/kbmod/search/raw_image.h @@ -1,12 +1,3 @@ -/* - * RawImage.h - * - * Created on: Jun 22, 2017 - * Author: kbmod-usr - * - * RawImage stores pixel level data for a single image. - */ - #ifndef RAWIMAGE_H_ #define RAWIMAGE_H_ @@ -19,18 +10,14 @@ #include #include #include -#ifdef Py_PYTHON_H -#include -#include -#include -#endif #include "common.h" -#include "PointSpreadFunc.h" +#include "psf.h" +#include "pydocs/raw_image_docs.h" -namespace search { -class RawImage { -public: +namespace search { + class RawImage { + public: RawImage(); RawImage(const RawImage& old); // Copy constructor RawImage(RawImage&& source); // Move constructor @@ -39,102 +26,102 @@ class RawImage { #ifdef Py_PYTHON_H explicit RawImage(pybind11::array_t arr); - void setArray(pybind11::array_t& arr); + void set_array(pybind11::array_t& arr); #endif RawImage& operator=(const RawImage& source); // Copy assignment RawImage& operator=(RawImage&& source); // Move assignment // Basic getter functions for image data. - unsigned getWidth() const { return width; } - unsigned getHeight() const { return height; } - unsigned getNPixels() const { return width * height; } + unsigned get_width() const { return width; } + unsigned get_height() const { return height; } + unsigned get_npixels() const { return width * height; } // Inline pixel functions. - float getPixel(int x, int y) const { - return (x >= 0 && x < width && y >= 0 && y < height) ? pixels[y * width + x] : NO_DATA; + float get_pixel(int x, int y) const { + return (x >= 0 && x < width && y >= 0 && y < height) ? pixels[y * width + x] : NO_DATA; } - bool pixelHasData(int x, int y) const { - return (x >= 0 && x < width && y >= 0 && y < height) ? pixels[y * width + x] != NO_DATA : false; + bool pixel_has_data(int x, int y) const { + return (x >= 0 && x < width && y >= 0 && y < height) ? pixels[y * width + x] != NO_DATA : false; } - void setPixel(int x, int y, float value) { - if (x >= 0 && x < width && y >= 0 && y < height) pixels[y * width + x] = value; + void set_pixel(int x, int y, float value) { + if (x >= 0 && x < width && y >= 0 && y < height) pixels[y * width + x] = value; } - const std::vector& getPixels() const { return pixels; } - float* getDataRef() { return pixels.data(); } // Get pointer to pixels + const std::vector& get_pixels() const { return pixels; } + float* data() { return pixels.data(); } // Get pointer to pixels // Get the interpolated brightness of a real values point // using the four neighboring pixels. - float getPixelInterp(float x, float y) const; + float get_pixel_interp(float x, float y) const; // Check if two raw images are approximately equal. - bool approxEqual(const RawImage& imgB, float atol) const; + bool approx_equal(const RawImage& imgB, float atol) const; // Functions for locally storing the image time. - double getObstime() const { return obstime; } - void setObstime(double new_time) { obstime = new_time; } + double get_obstime() const { return obstime; } + void set_obstime(double new_time) { obstime = new_time; } // Compute the min and max bounds of values in the image. - std::array computeBounds() const; + std::array compute_bounds() const; // Masks out the pixels of the image where: // flags a bit vector of mask flags to apply // (use 0xFFFFFF to apply all flags) // exceptions is a vector of pixel flags to ignore // mask is an image of bit vector mask flags - void applyMask(int flags, const std::vector& exceptions, const RawImage& mask); + void apply_mask(int flags, const std::vector& exceptions, const RawImage& mask); - void setAllPix(float value); - void addToPixel(float fx, float fy, float value); - void addPixelInterp(float x, float y, float value); - std::vector bilinearInterp(float x, float y) const; + void set_all_pix(float value); + void add_to_pixel(float fx, float fy, float value); + void add_pixel_interp(float x, float y, float value); + std::vector bilinear_interp(float x, float y) const; // Grow the area of masked pixels. - void growMask(int steps); + void grow_mask(int steps); // Load the image data from a specific layer of a FITS file. // Overwrites the current image data. - void loadFromFile(const std::string& file_path, int layer_num); + void load_from_file(const std::string& file_path, int layer_num); // Save the RawImage to a file (single layer) or append the layer to an existing file. - void saveToFile(const std::string& filename); - void appendLayerToFile(const std::string& filename); + void save_to_file(const std::string& filename); + void append_layer_to_file(const std::string& filename); // Convolve the image with a point spread function. - void convolve(PointSpreadFunc psf); - void convolve_cpu(const PointSpreadFunc& psf); + void convolve(PSF psf); + void convolve_cpu(const PSF& psf); // Create a "stamp" image of a give radius (width=2*radius+1) // about the given point. // keep_no_data indicates whether to use the NO_DATA flag or replace with 0.0. - RawImage createStamp(float x, float y, int radius, bool interpolate, bool keep_no_data) const; + RawImage create_stamp(float x, float y, int radius, bool interpolate, bool keep_no_data) const; // The maximum value of the image and return the coordinates. The parameter // furthest_from_center indicates whether to break ties using the peak further // or closer to the center of the image. - PixelPos findPeak(bool furthest_from_center) const; + PixelPos find_peak(bool furthest_from_center) const; // Find the basic image moments in order to test if stamps have a gaussian shape. // It computes the moments on the "normalized" image where the minimum // value has been shifted to zero and the sum of all elements is 1.0. // Elements with NO_DATA are treated as zero. - ImageMoments findCentralMoments() const; + ImageMoments find_central_moments() const; virtual ~RawImage(){}; -private: + private: unsigned width; unsigned height; std::vector pixels; double obstime; -}; + }; -// Helper functions for creating composite images. -RawImage createMedianImage(const std::vector& images); -RawImage createSummedImage(const std::vector& images); -RawImage createMeanImage(const std::vector& images); + // Helper functions for creating composite images. + RawImage create_median_image(const std::vector& images); + RawImage create_summed_image(const std::vector& images); + RawImage create_mean_image(const std::vector& images); } /* namespace search */ diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp new file mode 100644 index 000000000..71dcdca21 --- /dev/null +++ b/src/kbmod/search/stack_search.cpp @@ -0,0 +1,653 @@ +#include "stack_search.h" + + +namespace search { +#ifdef HAVE_CUDA + extern "C" void deviceSearchFilter(int num_images, int width, int height, float* psi_vect, float* phi_vect, + 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 >& use_index_vect, float* results); +#endif + + StackSearch::StackSearch(ImageStack& imstack) : stack(imstack) { + max_result_count = 100000; + debug_info = false; + psi_phi_generated = false; + + // Default the thresholds. + params.min_observations = 0; + params.min_lh = 0.0; + + // Default filtering arguments. + params.do_sigmag_filter = false; + params.sgl_L = 0.25; + params.sgl_H = 0.75; + params.sigmag_coeff = -1.0; + + // Default the encoding parameters. + params.psi_num_bytes = -1; + params.phi_num_bytes = -1; + + // Default pixel starting bounds. + params.x_start_min = 0; + params.x_start_max = stack.get_width(); + params.y_start_min = 0; + params.y_start_max = stack.get_height(); + + // Set default values for the barycentric correction. + bary_corrs = std::vector(stack.img_count()); + params.use_corr = false; + use_corr = false; + + params.debug = false; + } + + void StackSearch::set_debug(bool d) { + debug_info = d; + params.debug = d; + } + + void StackSearch::enable_corr(std::vector bary_corr_coeff) { + use_corr = true; + params.use_corr = true; + for (int i = 0; i < stack.img_count(); i++) { + int j = i * 6; + bary_corrs[i].dx = bary_corr_coeff[j]; + bary_corrs[i].dxdx = bary_corr_coeff[j + 1]; + bary_corrs[i].dxdy = bary_corr_coeff[j + 2]; + bary_corrs[i].dy = bary_corr_coeff[j + 3]; + bary_corrs[i].dydx = bary_corr_coeff[j + 4]; + bary_corrs[i].dydy = bary_corr_coeff[j + 5]; + } + } + + void StackSearch::enable_gpu_sigmag_filter(std::vector percentiles, float sigmag_coeff, + float min_lh) { + params.do_sigmag_filter = true; + params.sgl_L = percentiles[0]; + params.sgl_H = percentiles[1]; + params.sigmag_coeff = sigmag_coeff; + params.min_lh = min_lh; + } + + void StackSearch::enable_gpu_encoding(int psi_num_bytes, int phi_num_bytes) { + // Make sure the encoding is one of the supported options. + // Otherwise use default float (aka no encoding). + if (psi_num_bytes == 1 || psi_num_bytes == 2) { + params.psi_num_bytes = psi_num_bytes; + } else { + params.psi_num_bytes = -1; + } + if (phi_num_bytes == 1 || phi_num_bytes == 2) { + params.phi_num_bytes = phi_num_bytes; + } else { + params.phi_num_bytes = -1; + } + } + + void StackSearch::set_start_bounds_x(int x_min, int x_max) { + params.x_start_min = x_min; + params.x_start_max = x_max; + } + + void StackSearch::set_start_bounds_y(int y_min, int y_max) { + params.y_start_min = y_min; + params.y_start_max = y_max; + } + + void StackSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ang, float min_vel, + float mavx, int min_observations) { + prepare_psi_phi(); + create_search_list(ang_steps, vel_steps, min_ang, max_ang, min_vel, mavx); + + start_timer("Creating psi/phi buffers"); + std::vector psi_vect; + std::vector phi_vect; + fill_psi_phi(psi_images, phi_images, &psi_vect, &phi_vect); + end_timer(); + + // Create a data stucture for the per-image data. + PerImageData img_data; + img_data.num_images = stack.img_count(); + img_data.image_times = stack.get_timesDataRef(); + if (params.use_corr) img_data.bary_corrs = &bary_corrs[0]; + + // Compute the encoding parameters for psi and phi if needed. + // Vectors need to be created outside the if so they stay in scope. + std::vector psi_scale_vect; + std::vector phi_scale_vect; + if (params.psi_num_bytes > 0) { + psi_scale_vect = compute_image_scaling(psi_images, params.psi_num_bytes); + img_data.psi_params = psi_scale_vect.data(); + } + if (params.phi_num_bytes > 0) { + phi_scale_vect = compute_image_scaling(phi_images, params.phi_num_bytes); + img_data.phi_params = phi_scale_vect.data(); + } + + // Allocate a vector for the results. + int num_search_pixels = + ((params.x_start_max - params.x_start_min) * (params.y_start_max - params.y_start_min)); + int max_results = num_search_pixels * RESULTS_PER_PIXEL; + if (debug_info) { + std::cout << "Searching X=[" << params.x_start_min << ", " << params.x_start_max << "]" + << " Y=[" << params.y_start_min << ", " << params.y_start_max << "]\n"; + std::cout << "Allocating space for " << max_results << " results.\n"; + } + results = std::vector(max_results); + if (debug_info) std::cout << search_list.size() << " trajectories... \n" << std::flush; + + // Set the minimum number of observations. + params.min_observations = min_observations; + + // Do the actual search on the GPU. + start_timer("Searching"); +#ifdef HAVE_CUDA + deviceSearchFilter(stack.img_count(), stack.get_width(), stack.get_height(), psi_vect.data(), phi_vect.data(), + img_data, params, search_list.size(), search_list.data(), max_results, results.data()); +#else + throw std::runtime_error("Non-GPU search is not implemented."); +#endif + end_timer(); + + start_timer("Sorting results"); + sort_results(); + end_timer(); + } + + void StackSearch::save_psiphi(const std::string& path) { + prepare_psi_phi(); + save_images(path); + } + + void StackSearch::prepare_psi_phi() { + if (!psi_phi_generated) { + psi_images.clear(); + phi_images.clear(); + + // Compute Phi and Psi from convolved images + // while leaving masked pixels alone + // Reinsert 0s for NO_DATA? + const int num_images = stack.img_count(); + 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()); + } + + psi_phi_generated = true; + } + } + + std::vector StackSearch::compute_image_scaling(const std::vector& vect, + int encoding_bytes) const { + std::vector result; + + const int num_images = vect.size(); + for (int i = 0; i < num_images; ++i) { + scaleParameters params; + params.scale = 1.0; + + std::array bnds = vect[i].compute_bounds(); + params.min_val = bnds[0]; + params.max_val = bnds[1]; + + // Increase width to avoid divide by zero. + float width = (params.max_val - params.min_val); + if (width < 1e-6) width = 1e-6; + + // Set the scale if we are encoding the values. + if (encoding_bytes == 1 || encoding_bytes == 2) { + long int num_values = (1 << (8 * encoding_bytes)) - 1; + params.scale = width / (double)num_values; + } + + result.push_back(params); + } + + return result; + } + + void StackSearch::save_images(const std::string& path) { + for (int i = 0; i < stack.img_count(); ++i) { + std::string number = std::to_string(i); + // Add leading zeros + number = std::string(4 - number.length(), '0') + number; + psi_images[i].save_to_file(path + "/psi/PSI" + number + ".fits"); + phi_images[i].save_to_file(path + "/phi/PHI" + number + ".fits"); + } + } + + void StackSearch::create_search_list(int angle_steps, int velocity_steps, float min_ang, float max_ang, + float min_vel, float mavx) { + std::vector angles(angle_steps); + float ang_stepsize = (max_ang - min_ang) / float(angle_steps); + for (int i = 0; i < angle_steps; ++i) { + angles[i] = min_ang + float(i) * ang_stepsize; + } + + std::vector velocities(velocity_steps); + float vel_stepsize = (mavx - min_vel) / float(velocity_steps); + for (int i = 0; i < velocity_steps; ++i) { + velocities[i] = min_vel + float(i) * vel_stepsize; + } + + int trajCount = angle_steps * velocity_steps; + search_list = std::vector(trajCount); + for (int a = 0; a < angle_steps; ++a) { + for (int v = 0; v < velocity_steps; ++v) { + search_list[a * velocity_steps + v].vx = cos(angles[a]) * velocities[v]; + search_list[a * velocity_steps + v].vy = sin(angles[a]) * velocities[v]; + } + } + } + + void StackSearch::fill_psi_phi(const std::vector& psi_imgs, + const std::vector& phi_imgs, std::vector* psi_vect, + std::vector* phi_vect) { + assert(psi_vect != NULL); + assert(phi_vect != NULL); + + int num_images = psi_imgs.size(); + assert(num_images > 0); + assert(phi_imgs.size() == num_images); + + int num_pixels = psi_imgs[0].get_npixels(); + for (int i = 0; i < num_images; ++i) { + assert(psi_imgs[i].get_npixels() == num_pixels); + assert(phi_imgs[i].get_npixels() == num_pixels); + } + + psi_vect->clear(); + psi_vect->reserve(num_images * num_pixels); + phi_vect->clear(); + phi_vect->reserve(num_images * num_pixels); + + for (int i = 0; i < num_images; ++i) { + const std::vector& psi_ref = psi_imgs[i].get_pixels(); + const std::vector& phi_ref = phi_imgs[i].get_pixels(); + for (unsigned p = 0; p < num_pixels; ++p) { + psi_vect->push_back(psi_ref[p]); + phi_vect->push_back(phi_ref[p]); + } + } + } + + std::vector StackSearch::create_stamps(const Trajectory& trj, int radius, bool interpolate, + bool keep_no_data, const std::vector& use_index) { + if (use_index.size() > 0 && use_index.size() != stack.img_count()) { + throw std::runtime_error("Wrong size use_index passed into create_stamps()"); + } + bool use_all_stamps = use_index.size() == 0; + + std::vector stamps; + int num_times = stack.img_count(); + for (int i = 0; i < num_times; ++i) { + if (use_all_stamps || use_index[i]) { + PixelPos pos = get_trajectory_position(trj, i); + RawImage& img = stack.get_single_image(i).get_science(); + stamps.push_back(img.create_stamp(pos.x, pos.y, radius, interpolate, keep_no_data)); + } + } + return stamps; + } + + // For stamps used for visualization we interpolate the pixel values, replace + // NO_DATA tages with zeros, and return all the stamps (regardless of whether + // individual timesteps have been filtered). + std::vector StackSearch::get_stamps(const Trajectory& t, int radius) { + std::vector empty_vect; + return create_stamps(t, radius, true /*=interpolate*/, false /*=keep_no_data*/, empty_vect); + } + + // For creating coadded stamps, we do not interpolate the pixel values and keep + // NO_DATA tagged (so we can filter it out of mean/median). + RawImage StackSearch::get_median_stamp(const Trajectory& trj, int radius, + const std::vector& use_index) { + return create_median_image( + create_stamps(trj, radius, false /*=interpolate*/, true /*=keep_no_data*/, use_index)); + } + + // For creating coadded stamps, we do not interpolate the pixel values and keep + // NO_DATA tagged (so we can filter it out of mean/median). + RawImage StackSearch::get_mean_stamp(const Trajectory& trj, int radius, const std::vector& use_index) { + return create_mean_image( + create_stamps(trj, radius, false /*=interpolate*/, true /*=keep_no_data*/, use_index)); + } + + // For creating summed stamps, we do not interpolate the pixel values and replace NO_DATA + // with zero (which is the same as filtering it out for the sum). + RawImage StackSearch::get_summed_stamp(const Trajectory& trj, int radius, + const std::vector& use_index) { + return create_summed_image( + create_stamps(trj, radius, false /*=interpolate*/, false /*=keep_no_data*/, use_index)); + } + + 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& 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& 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 StackSearch::get_coadded_stamps(std::vector& t_array, + std::vector >& 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 StackSearch::get_coadded_stamps_cpu(std::vector& t_array, + std::vector >& use_index_vect, + const StampParameters& params) { + const int num_trajectories = t_array.size(); + std::vector results(num_trajectories); + std::vector empty_pixels(1, NO_DATA); + + for (int i = 0; i < num_trajectories; ++i) { + std::vector stamps = + create_stamps(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 StackSearch::get_coadded_stamps_gpu(std::vector& t_array, + std::vector >& 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. + PerImageData img_data; + img_data.num_images = num_images; + img_data.image_times = stack.get_timesDataRef(); + + // 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 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 results(num_trajectories); + std::vector current_pixels(stamp_ppi, 0.0); + std::vector 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 StackSearch::create_stamps(Trajectory t, int radius, const std::vector& imgs, + bool interpolate) { + if (radius < 0) throw std::runtime_error("stamp radius must be at least 0"); + std::vector stamps; + for (int i = 0; i < imgs.size(); ++i) { + PixelPos pos = get_trajectory_position(t, i); + stamps.push_back(imgs[i]->create_stamp(pos.x, pos.y, radius, interpolate, false)); + } + return stamps; + } + + PixelPos StackSearch::get_trajectory_position(const Trajectory& t, int i) const { + float time = stack.get_times()[i]; + if (use_corr) { + return {t.x + time * t.vx + bary_corrs[i].dx + t.x * bary_corrs[i].dxdx + t.y * bary_corrs[i].dxdy, + t.y + time * t.vy + bary_corrs[i].dy + t.x * bary_corrs[i].dydx + + t.y * bary_corrs[i].dydy}; + } else { + return {t.x + time * t.vx, t.y + time * t.vy}; + } + } + + std::vector StackSearch::get_trajectory_positions(Trajectory& t) const { + std::vector results; + int num_times = stack.img_count(); + for (int i = 0; i < num_times; ++i) { + PixelPos pos = get_trajectory_position(t, i); + results.push_back(pos); + } + return results; + } + + std::vector StackSearch::create_curves(Trajectory t, const std::vector& imgs) { + /*Create a lightcurve from an image along a trajectory + * + * INPUT- + * Trajectory t - The trajectory along which to compute the lightcurve + * std::vector imgs - The image from which to compute the + * trajectory. Most likely a psiImage or a phiImage. + * Output- + * std::vector lightcurve - The computed trajectory + */ + + int img_size = imgs.size(); + std::vector lightcurve; + lightcurve.reserve(img_size); + const std::vector& times = stack.get_times(); + for (int i = 0; i < img_size; ++i) { + /* Do not use get_pixel_interp(), because results from create_curves must + * be able to recover the same likelihoods as the ones reported by the + * gpu search.*/ + float pix_val; + if (use_corr) { + PixelPos pos = get_trajectory_position(t, i); + pix_val = imgs[i].get_pixel(int(pos.x + 0.5), int(pos.y + 0.5)); + } + /* Does not use get_trajectory_position to be backwards compatible with Hits_Rerun */ + else { + pix_val = imgs[i].get_pixel(t.x + int(times[i] * t.vx + 0.5), + t.y + int(times[i] * t.vy + 0.5)); + } + if (pix_val == NO_DATA) pix_val = 0.0; + lightcurve.push_back(pix_val); + } + return lightcurve; + } + + std::vector StackSearch::get_psi_curves(Trajectory& t) { + /*Generate a psi lightcurve for further analysis + * INPUT- + * Trajectory& t - The trajectory along which to find the lightcurve + * OUTPUT- + * std::vector - A vector of the lightcurve values + */ + prepare_psi_phi(); + return create_curves(t, psi_images); + } + + std::vector StackSearch::get_phi_curves(Trajectory& t) { + /*Generate a phi lightcurve for further analysis + * INPUT- + * Trajectory& t - The trajectory along which to find the lightcurve + * OUTPUT- + * std::vector - A vector of the lightcurve values + */ + prepare_psi_phi(); + return create_curves(t, phi_images); + } + + std::vector& StackSearch::get_psi_images() { return psi_images; } + + std::vector& StackSearch::getPhiImages() { return phi_images; } + + void StackSearch::sort_results() { + __gnu_parallel::sort(results.begin(), results.end(), + [](Trajectory a, Trajectory b) { return b.lh < a.lh; }); + } + + void StackSearch::filter_results(int min_observations) { + results.erase(std::remove_if(results.begin(), results.end(), + std::bind([](Trajectory t, int cutoff) { return t.obs_count < cutoff; }, + std::placeholders::_1, min_observations)), + results.end()); + } + + void StackSearch::filter_results_lh(float min_lh) { + results.erase(std::remove_if(results.begin(), results.end(), + std::bind([](Trajectory t, float cutoff) { return t.lh < cutoff; }, + std::placeholders::_1, min_lh)), + results.end()); + } + + std::vector StackSearch::get_results(int start, int count) { + if (start + count >= results.size()) { + count = results.size() - start; + } + if (start < 0) throw std::runtime_error("start must be 0 or greater"); + return std::vector(results.begin() + start, results.begin() + start + count); + } + + // This function is used only for testing by injecting known result trajectories. + void StackSearch::set_results(const std::vector& new_results) { results = new_results; } + + void StackSearch::start_timer(const std::string& message) { + if (debug_info) { + std::cout << message << "... " << std::flush; + t_start = std::chrono::system_clock::now(); + } + } + + void StackSearch::end_timer() { + if (debug_info) { + t_end = std::chrono::system_clock::now(); + t_delta = t_end - t_start; + std::cout << " Took " << t_delta.count() << " seconds.\n" << std::flush; + } + } + +#ifdef Py_PYTHON_H + static void stack_search_bindings(py::module &m) { + using tj = search::Trajectory; + using pf = search::PSF; + using ri = search::RawImage; + using is = search::ImageStack; + using ks = search::StackSearch; + + py::class_(m, "StackSearch", pydocs::DOC_StackSearch) + .def(py::init()) + .def("save_psi_phi", &ks::save_psiphi, pydocs::DOC_StackSearch_save_psi_phi) + .def("search", &ks::search, pydocs::DOC_StackSearch_search) + .def("enable_gpu_sigmag_filter", &ks::enable_gpu_sigmag_filter, pydocs::DOC_StackSearch_enable_gpu_sigmag_filter) + .def("enable_gpu_encoding", &ks::enable_gpu_encoding, pydocs::DOC_StackSearch_enable_gpu_encoding) + .def("enable_corr", &ks::enable_corr, pydocs::DOC_StackSearch_enable_corr) + .def("set_start_bounds_x", &ks::set_start_bounds_x, pydocs::DOC_StackSearch_set_start_bounds_x) + .def("set_start_bounds_y", &ks::set_start_bounds_y, pydocs::DOC_StackSearch_set_start_bounds_y) + .def("set_debug", &ks::set_debug, pydocs::DOC_StackSearch_set_debug) + .def("filter_min_obs", &ks::filter_results, pydocs::DOC_StackSearch_filter_min_obs) + .def("get_num_images", &ks::num_images, pydocs::DOC_StackSearch_get_num_images) + .def("get_imagestack", &ks::get_imagestack, pydocs::DOC_StackSearch_get_imagestack) + // Science Stamp Functions + .def("get_stamps", &ks::get_stamps, pydocs::DOC_StackSearch_get_stamps) + .def("get_median_stamp", &ks::get_median_stamp, pydocs::DOC_StackSearch_get_median_stamp) + .def("get_mean_stamp", &ks::get_mean_stamp, pydocs::DOC_StackSearch_get_mean_stamp) + .def("get_summed_stamp", &ks::get_summed_stamp, pydocs::DOC_StackSearch_get_summed_stamp) + .def("get_coadded_stamps", //wth is happening here + (std::vector(ks::*)(std::vector &, std::vector> &, + 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) + .def("get_trajectory_position", &ks::get_trajectory_position, pydocs::DOC_StackSearch_get_trajectory_position) + .def("get_trajectory_positions", &ks::get_trajectory_positions, pydocs::DOC_StackSearch_get_trajectory_positions) + .def("get_psi_curves", (std::vector(ks::*)(tj &)) & ks::get_psi_curves, pydocs::DOC_StackSearch_get_psi_curves) + .def("get_phi_curves", (std::vector(ks::*)(tj &)) & ks::get_phi_curves, pydocs::DOC_StackSearch_get_phi_curves) + .def("prepare_psi_phi", &ks::prepare_psi_phi, pydocs::DOC_StackSearch_prepare_psi_phi) + .def("get_psi_images", &ks::get_psi_images, pydocs::DOC_StackSearch_get_psi_images) + .def("get_phi_images", &ks::getPhiImages, pydocs::DOC_StackSearch_get_phi_images) + .def("get_results", &ks::get_results, pydocs::DOC_StackSearch_get_results) + .def("set_results", &ks::set_results, pydocs::DOC_StackSearch_set_results); + } + +#endif /* Py_PYTHON_H */ + +} /* namespace search */ diff --git a/src/kbmod/search/stack_search.h b/src/kbmod/search/stack_search.h new file mode 100644 index 000000000..2e39069a7 --- /dev/null +++ b/src/kbmod/search/stack_search.h @@ -0,0 +1,148 @@ +#ifndef KBMODSEARCH_H_ +#define KBMODSEARCH_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "common.h" +#include "image_stack.h" +#include "psf.h" +#include "pydocs/stack_search_docs.h" + + +namespace search { + class StackSearch { + public: + StackSearch(ImageStack& imstack); + + int num_images() const { return stack.img_count(); } + const ImageStack& get_imagestack() const { return stack; } + + void set_debug(bool d); + + // The primary search functions. + void enable_gpu_sigmag_filter(std::vector percentiles, float sigmag_coeff, float min_lh); + void enable_corr(std::vector bary_corr_coeff); + void enable_gpu_encoding(int psi_num_bytes, int phi_num_bytes); + + void set_start_bounds_x(int x_min, int x_max); + void set_start_bounds_y(int y_min, int y_max); + + void search(int a_steps, int v_steps, float min_angle, float max_angle, float min_velocity, + float max_velocity, int min_observations); + + // Gets the vector of result trajectories. + std::vector get_results(int start, int end); + + // Get the predicted (pixel) positions for a given trajectory. + PixelPos get_trajectory_position(const Trajectory& t, int i) const; + std::vector get_trajectory_positions(Trajectory& t) const; + + // Filters the results based on various parameters. + 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 indicating whether to use + // each time step. An empty (size=0) vector will use all time steps. + std::vector create_stamps(const Trajectory& trj, int radius, bool interpolate, + bool keep_no_data, const std::vector& use_index); + std::vector get_stamps(const Trajectory& t, int radius); + RawImage get_median_stamp(const Trajectory& trj, int radius, const std::vector& use_index); + RawImage get_mean_stamp(const Trajectory& trj, int radius, const std::vector& use_index); + RawImage get_summed_stamp(const Trajectory& trj, int radius, const std::vector& 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 get_coadded_stamps(std::vector& t_array, + std::vector >& 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& get_psi_images(); + std::vector& getPhiImages(); + std::vector get_psi_curves(Trajectory& t); + std::vector get_phi_curves(Trajectory& t); + + // Save internal data products to a file. + void save_psiphi(const std::string& path); + + // Helper functions for computing Psi and Phi. + void prepare_psi_phi(); + + // Helper functions for testing. + void set_results(const std::vector& new_results); + + virtual ~StackSearch(){}; + + protected: + void save_images(const std::string& path); + void sort_results(); + std::vector create_curves(Trajectory t, const std::vector& imgs); + + // Fill an interleaved vector for the GPU functions. + void fill_psi_phi(const std::vector& psi_imgs, const std::vector& phi_imgs, + std::vector* psi_vect, std::vector* phi_vect); + + // Set the parameter min/max/scale from the psi/phi/other images. + std::vector compute_image_scaling(const std::vector& vect, + int encoding_bytes) const; + + // Functions to create and access stamps around proposed trajectories or + // regions. Used to visualize the results. + // This function replaces NO_DATA with a value of 0.0. + std::vector create_stamps(Trajectory t, int radius, const std::vector& imgs, + bool interpolate); + + // Creates list of trajectories to search. + void create_search_list(int angle_steps, int velocity_steps, float min_ang, float max_ang, + float min_vel, float max_vel); + + std::vector get_coadded_stamps_gpu(std::vector& t_array, + std::vector >& use_index_vect, + const StampParameters& params); + + std::vector get_coadded_stamps_cpu(std::vector& t_array, + std::vector >& use_index_vect, + const StampParameters& params); + // Helper functions for timing operations of the search. + void start_timer(const std::string& message); + void end_timer(); + + unsigned max_result_count; + bool psi_phi_generated; + bool debug_info; + ImageStack stack; + std::vector search_list; + std::vector psi_images; + std::vector phi_images; + std::vector results; + + // Variables for the timer. + std::chrono::time_point t_start, t_end; + std::chrono::duration t_delta; + + // Parameters for the GPU search. + SearchParameters params; + + // Parameters to do barycentric corrections. + bool use_corr; + std::vector bary_corrs; + }; + +} /* namespace search */ + +#endif /* KBMODSEARCH_H_ */ diff --git a/tests/benchmark.py b/tests/benchmark.py index 15ab58a8f..05998c8bc 100644 --- a/tests/benchmark.py +++ b/tests/benchmark.py @@ -51,7 +51,7 @@ def test_benchmark(benchmark): run_search = kbmod.run_search.run_search(input_parameters) # Load the PSF. kb_interface = kbmod.analysis_utils.Interface() - default_psf = kbmod.search.psf(run_search.config["psf_val"]) + default_psf = kbmod.search.PSF(run_search.config["psf_val"]) # Load images to search _, img_info = kb_interface.load_images( diff --git a/tests/data/fake_results.txt b/tests/data/fake_results.txt index 72923b067..766b3c2a7 100644 --- a/tests/data/fake_results.txt +++ b/tests/data/fake_results.txt @@ -1,2 +1,2 @@ -lh: 300.0 flux: 750.0 x: 106 y: 44 x_v: 9.52 y_v: -0.5 obs_count: 10 -lh: 250.0 flux: 500.00 x: 55 y: 60 x_v: 10.5 y_v: -1.7 obs_count: 9 \ No newline at end of file +lh: 300.0 flux: 750.0 x: 106 y: 44 vx: 9.52 vy: -0.5 obs_count: 10 +lh: 250.0 flux: 500.00 x: 55 y: 60 vx: 10.5 vy: -1.7 obs_count: 9 diff --git a/tests/regression_test.py b/tests/regression_test.py index 1bc8c6479..fa204b3ed 100644 --- a/tests/regression_test.py +++ b/tests/regression_test.py @@ -25,9 +25,9 @@ def ave_trajectory_distance(trjA, trjB, times=[0.0]): Parameters ---------- trjA : `trajectory` - The first trajectory to evaluate. + The first Trajectory to evaluate. trjB : `trajectory` - The second trajectory to evaluate. + The second Trajectory to evaluate. times : list The list of zero-shifted times at which to evaluate the matches. The average of the distances at these times @@ -40,8 +40,8 @@ def ave_trajectory_distance(trjA, trjB, times=[0.0]): """ total = 0.0 for t in times: - dx = (trjA.x + t * trjA.x_v) - (trjB.x + t * trjB.x_v) - dy = (trjA.y + t * trjA.y_v) - (trjB.y + t * trjB.y_v) + dx = (trjA.x + t * trjA.vx) - (trjB.x + t * trjB.vx) + dy = (trjA.y + t * trjA.vy) - (trjB.y + t * trjB.vy) total += math.sqrt(dx * dx + dy * dy) ave_dist = total / len(times) @@ -50,7 +50,7 @@ def ave_trajectory_distance(trjA, trjB, times=[0.0]): def find_unique_overlap(traj_query, traj_base, threshold, times=[0.0]): """Finds the set of trajectories in traj_query that are 'close' to - trajectories in traj_base such that each trajectory in traj_base + trajectories in traj_base such that each Trajectory in traj_base is used at most once. Used to evaluate the performance of algorithms. @@ -72,7 +72,7 @@ def find_unique_overlap(traj_query, traj_base, threshold, times=[0.0]): ------- results : list The list of trajectories that appear in both traj1 and traj2 - where each trajectory in each set is only used once. + where each Trajectory in each set is only used once. """ num_times = len(times) size_base = len(traj_base) @@ -100,7 +100,7 @@ def find_unique_overlap(traj_query, traj_base, threshold, times=[0.0]): def find_set_difference(traj_query, traj_base, threshold, times=[0.0]): """Finds the set of trajectories in traj_query that are NOT 'close' to - any trajectories in traj_base such that each trajectory in traj_base + any trajectories in traj_base such that each Trajectory in traj_base is used at most once. Used to evaluate the performance of algorithms. @@ -122,7 +122,7 @@ def find_set_difference(traj_query, traj_base, threshold, times=[0.0]): ------- results : list A list of trajectories that appear in traj_query but not - in traj_base where each trajectory in each set is only + in traj_base where each Trajectory in each set is only used once. """ num_times = len(times) @@ -151,7 +151,7 @@ def find_set_difference(traj_query, traj_base, threshold, times=[0.0]): def make_trajectory(x, y, vx, vy, flux): - """Create a fake trajectory given the parameters. + """Create a fake Trajectory given the parameters. Arguments: x : int @@ -166,18 +166,18 @@ def make_trajectory(x, y, vx, vy, flux): The flux of the object. Returns: - A trajectory object. + A Trajectory object. """ - t = trajectory() + t = Trajectory() t.x = x t.y = y - t.x_v = vx - t.y_v = vy + t.vx = vx + t.vy = vy t.flux = flux return t -def make_fake_image_stack(times, trjs, psf_vals): +def make_fake_ImageStack(times, trjs, psf_vals): """ Make a stack of fake layered images. @@ -190,7 +190,7 @@ def make_fake_image_stack(times, trjs, psf_vals): A list of PSF variances. Returns: - A image_stack + A ImageStack """ imCount = len(times) t0 = times[0] @@ -201,7 +201,7 @@ def make_fake_image_stack(times, trjs, psf_vals): imlist = [] for i in range(imCount): - p = psf(psf_vals[i]) + p = PSF(psf_vals[i]) time = times[i] - t0 # For each odd case, don't save the time. These will be provided by the time file. @@ -209,15 +209,15 @@ def make_fake_image_stack(times, trjs, psf_vals): if i % 2 == 1: saved_time = 0.0 - img = layered_image(("%06i" % i), dim_x, dim_y, noise_level, variance, saved_time, p, i) + img = LayeredImage(("%06i" % i), dim_x, dim_y, noise_level, variance, saved_time, p, i) for trj in trjs: - px = trj.x + time * trj.x_v + 0.5 - py = trj.y + time * trj.y_v + 0.5 + px = trj.x + time * trj.vx + 0.5 + py = trj.y + time * trj.vy + 0.5 add_fake_object(img, px, py, trj.flux, p) imlist.append(img) - stack = image_stack(imlist) + stack = ImageStack(imlist) return stack @@ -468,7 +468,7 @@ def perform_search(im_filepath, time_file, psf_file, res_filepath, results_suffi times.append(67130.2 + i) psf_vals.append(default_psf + 0.01) - stack = make_fake_image_stack(times, trjs, psf_vals) + stack = make_fake_ImageStack(times, trjs, psf_vals) save_fake_data(dir_name, stack, times, psf_vals, default_psf) # Do the search. diff --git a/tests/test_analysis_utils.py b/tests/test_analysis_utils.py index 4f4bfade4..930670035 100644 --- a/tests/test_analysis_utils.py +++ b/tests/test_analysis_utils.py @@ -8,11 +8,11 @@ class test_analysis_utils(unittest.TestCase): def _make_trajectory(self, x0, y0, xv, yv, lh): - t = trajectory() + t = Trajectory() t.x = x0 t.y = y0 - t.x_v = xv - t.y_v = yv + t.vx = xv + t.vy = yv t.lh = lh return t @@ -97,7 +97,7 @@ def setUp(self): self.dim_y = 20 self.noise_level = 1.0 self.variance = self.noise_level**2 - self.p = psf(0.5) + self.p = PSF(0.5) # create image set with single moving object self.imlist = [] @@ -105,20 +105,20 @@ def setUp(self): for i in range(self.img_count): time = i / self.img_count self.time_list.append(time) - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) self.imlist.append(im) - self.stack = image_stack(self.imlist) + self.stack = ImageStack(self.imlist) # Set up old_results object for analysis_utils.PostProcess self.num_curves = 4 curve_num_times = 20 # First 3 passing indices - psi_curves = [ + get_psi_curves = [ np.array([1.0 + (x / 100) for x in range(curve_num_times)]) for _ in range(self.num_curves - 1) ] - phi_curves = [ + get_phi_curves = [ np.array([1.0 + (y / 100) for y in range(curve_num_times)]) for _ in range(self.num_curves - 1) ] # Failing index @@ -127,8 +127,8 @@ def setUp(self): failing_psi = [0.0 + (z / 100) for z in range(curve_num_times)] failing_psi[14] = -100.0 failing_psi[2] = 100.0 - psi_curves.append(np.array(failing_psi)) - phi_curves.append(np.array([1.0 for _ in range(curve_num_times)])) + get_psi_curves.append(np.array(failing_psi)) + get_phi_curves.append(np.array([1.0 for _ in range(curve_num_times)])) self.good_indices = [z for z in range(curve_num_times)] self.good_indices.remove(14) @@ -137,8 +137,8 @@ def setUp(self): self.curve_time_list = [i for i in range(curve_num_times)] self.curve_result_set = ResultList(self.curve_time_list) for i in range(self.num_curves): - row = ResultRow(trajectory(), curve_num_times) - row.set_psi_phi(psi_curves[i], phi_curves[i]) + row = ResultRow(Trajectory(), curve_num_times) + row.set_psi_phi(get_psi_curves[i], get_phi_curves[i]) self.curve_result_set.append_result(row) def test_apply_clipped_sigmaG_single_thread(self): @@ -170,26 +170,27 @@ def test_apply_clipped_sigmaG_multi_thread(self): self.assertEqual(self.curve_result_set.results[2].valid_indices, all_indices) self.assertEqual(self.curve_result_set.results[3].valid_indices, self.good_indices) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_apply_stamp_filter(self): # object properties self.object_flux = 250.0 self.start_x = 4 self.start_y = 3 - self.x_vel = 2.0 - self.y_vel = 1.0 + self.vxel = 2.0 + self.vyel = 1.0 for i in range(self.img_count): time = i / self.img_count add_fake_object( self.imlist[i], - self.start_x + time * self.x_vel + 0.5, - self.start_y + time * self.y_vel + 0.5, + self.start_x + time * self.vxel + 0.5, + self.start_y + time * self.vyel + 0.5, self.object_flux, self.p, ) - stack = image_stack(self.imlist) - search = stack_search(stack) + stack = ImageStack(self.imlist) + search = StackSearch(stack) search.search( self.angle_steps, self.velocity_steps, @@ -226,49 +227,49 @@ def test_apply_stamp_filter_2(self): self.object_flux = 250.0 self.start_x = 4 self.start_y = 3 - self.x_vel = 2.0 - self.y_vel = 1.0 + self.vxel = 2.0 + self.vyel = 1.0 for i in range(self.img_count): time = i / self.img_count add_fake_object( self.imlist[i], - self.start_x + time * self.x_vel, - self.start_y + time * self.y_vel, + self.start_x + time * self.vxel, + self.start_y + time * self.vyel, self.object_flux, self.p, ) - stack = image_stack(self.imlist) - search = stack_search(stack) + stack = ImageStack(self.imlist) + search = StackSearch(stack) - # Create a first trajectory that matches perfectly. - trj = trajectory() + # Create a first Trajectory that matches perfectly. + trj = Trajectory() trj.x = self.start_x trj.y = self.start_y - trj.x_v = self.x_vel - trj.y_v = self.y_vel + trj.vx = self.vxel + trj.vy = self.vyel - # Create a second trajectory that isn't any good. - trj2 = trajectory() + # Create a second Trajectory that isn't any good. + trj2 = Trajectory() trj2.x = 1 trj2.y = 1 - trj2.x_v = 0 - trj2.y_v = 0 + trj2.vx = 0 + trj2.vy = 0 - # Create a third trajectory that is close to good, but offset. - trj3 = trajectory() + # Create a third Trajectory that is close to good, but offset. + trj3 = Trajectory() trj3.x = trj.x + 2 trj3.y = trj.y + 2 - trj3.x_v = trj.x_v - trj3.y_v = trj.y_v + trj3.vx = trj.vx + trj3.vy = trj.vy - # Create a fourth trajectory that is just close enough - trj4 = trajectory() + # Create a fourth Trajectory that is just close enough + trj4 = Trajectory() trj4.x = trj.x + 1 trj4.y = trj.y + 1 - trj4.x_v = trj.x_v - trj4.y_v = trj.y_v + trj4.vx = trj.vx + trj4.vy = trj.vy # Create the ResultList. keep = ResultList(self.time_list) @@ -347,7 +348,7 @@ def test_load_and_filter_results_lh(self): imlist = [] for i in range(self.img_count): t = self.time_list[i] - im = layered_image(str(i), 100, 100, self.noise_level, self.variance, t, self.p, i) + im = LayeredImage(str(i), 100, 100, self.noise_level, self.variance, t, self.p, i) # Add the objects. for j, trj in enumerate(trjs): @@ -357,7 +358,7 @@ def test_load_and_filter_results_lh(self): imlist.append(im) # Create the stack search and insert the fake results. - search = stack_search(image_stack(imlist)) + search = StackSearch(ImageStack(imlist)) search.set_results(trjs) # Do the filtering. @@ -381,7 +382,7 @@ def test_file_load_basic(self): None, None, [0, 157130.2], - psf(1.0), + PSF(1.0), verbose=False, ) self.assertEqual(stack.img_count(), 4) @@ -393,7 +394,7 @@ def test_file_load_basic(self): self.assertEqual(img.get_width(), 64) self.assertEqual(img.get_height(), 64) self.assertAlmostEqual(img.get_obstime(), true_times[i], delta=0.005) - self.assertAlmostEqual(1.0, img.get_psf().get_stdev()) + self.assertAlmostEqual(1.0, img.get_psf().get_std()) # Check that visit IDs and times were extracted for each file in img_info. true_visit_ids = ["000000", "000001", "000002", "000003"] @@ -405,7 +406,7 @@ def test_file_load_basic(self): self.assertAlmostEqual(time_obj.mjd, true_times[i], delta=0.005) def test_file_load_extra(self): - p = psf(1.0) + p = PSF(1.0) loader = Interface() stack, img_info = loader.load_images( @@ -426,7 +427,7 @@ def test_file_load_extra(self): self.assertEqual(img.get_width(), 64) self.assertEqual(img.get_height(), 64) self.assertAlmostEqual(img.get_obstime(), true_times[i], delta=0.005) - self.assertAlmostEqual(psfs_std[i], img.get_psf().get_stdev()) + self.assertAlmostEqual(psfs_std[i], img.get_psf().get_std()) # Check that visit IDs and times were extracted for each file in img_info. true_visit_ids = ["000000", "000001", "000002", "000003"] diff --git a/tests/test_bilinear_interp.py b/tests/test_bilinear_interp.py index c4dd8dae0..07520b02f 100644 --- a/tests/test_bilinear_interp.py +++ b/tests/test_bilinear_interp.py @@ -9,10 +9,10 @@ class test_bilinear_interp(unittest.TestCase): def setUp(self): self.im_count = 5 - p = kb.psf(0.05) + p = kb.PSF(0.05) self.images = [] for c in range(self.im_count): - im = kb.layered_image(str(c), 10, 10, 0.0, 1.0, c, p) + im = kb.LayeredImage(str(c), 10, 10, 0.0, 1.0, c, p) add_fake_object(im, 2 + c * 0.5 + 0.5, 2 + c * 0.5 + 0.5, 1, p) self.images.append(im) @@ -35,7 +35,7 @@ def test_pixels(self): def test_pixel_interp(self): pixels = numpy.array([[0.0, 1.2, 0.0], [1.0, 2.0, 1.0]]) - im = kb.raw_image(pixels) + im = kb.RawImage(pixels) self.assertEqual(im.get_width(), 3) self.assertEqual(im.get_height(), 2) self.assertEqual(im.get_npixels(), 6) diff --git a/tests/test_calc_barycentric_corr.py b/tests/test_calc_barycentric_corr.py index d84a5dd98..eb4eada79 100644 --- a/tests/test_calc_barycentric_corr.py +++ b/tests/test_calc_barycentric_corr.py @@ -116,13 +116,13 @@ def _synthetic_wcs(self, params): baryCoeff = run_search._calc_barycentric_corr(img_info, params.dist) return baryCoeff - def test_image_stack(self): + def test_ImageStack(self): # Test the calc_barycentric function of run_search run_search = kbmod.run_search.run_search(self.input_parameters) self.assertIsNotNone(run_search) # Load the PSF. kb_interface = kbmod.analysis_utils.Interface() - default_psf = kbmod.search.psf(run_search.config["psf_val"]) + default_psf = kbmod.search.PSF(run_search.config["psf_val"]) # Load images to search stack, img_info = kb_interface.load_images( @@ -158,7 +158,7 @@ def test_single_image(self): self.assertIsNotNone(run_search) # Load the PSF. kb_interface = kbmod.analysis_utils.Interface() - default_psf = kbmod.search.psf(run_search.config["psf_val"]) + default_psf = kbmod.search.PSF(run_search.config["psf_val"]) full_file_path = f"{self.input_parameters['im_filepath']}/000000.fits" header_info = kbmod.analysis_utils.ImageInfo() header_info.populate_from_fits_file(full_file_path) diff --git a/tests/test_clustering_filters.py b/tests/test_clustering_filters.py index e2d95c623..b2361a15e 100644 --- a/tests/test_clustering_filters.py +++ b/tests/test_clustering_filters.py @@ -12,7 +12,7 @@ def _make_data(self, objs): Parameters ---------- obj : list of lists - A list where each element specifies a trajectory + A list where each element specifies a Trajectory as [x, y, xv, yv]. Returns @@ -21,11 +21,11 @@ def _make_data(self, objs): """ rs = ResultList(self.times, track_filtered=True) for x in objs: - t = trajectory() + t = Trajectory() t.x = x[0] t.y = x[1] - t.x_v = x[2] - t.y_v = x[3] + t.vx = x[2] + t.vy = x[3] t.lh = 100.0 row = ResultRow(t, self.num_times) diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py index 80ae83580..713733674 100644 --- a/tests/test_end_to_end.py +++ b/tests/test_end_to_end.py @@ -45,18 +45,21 @@ def setUp(self): "average_angle": 0.0, } + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_demo_defaults(self): rs = run_search(self.input_parameters) keep = rs.run_search() self.assertGreaterEqual(keep.num_results(), 1) self.assertEqual(keep.results[0].stamp.size, 441) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_demo_config_file(self): rs = run_search({"im_filepath": "../data/demo"}, config_file="../data/demo_config.yml") keep = rs.run_search() self.assertGreaterEqual(keep.num_results(), 1) self.assertEqual(keep.results[0].stamp.size, 441) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_demo_stamp_size(self): self.input_parameters["stamp_radius"] = 15 self.input_parameters["mom_lims"] = [80.0, 80.0, 50.0, 20.0, 20.0] diff --git a/tests/test_fake_data_creator.py b/tests/test_fake_data_creator.py index 319a4af0b..ee21d71b4 100644 --- a/tests/test_fake_data_creator.py +++ b/tests/test_fake_data_creator.py @@ -34,8 +34,8 @@ def test_insert_object(self): t0 = ds.stack.get_single_image(0).get_obstime() for i in range(ds.stack.img_count()): dt = ds.stack.get_single_image(i).get_obstime() - t0 - px = int(trj.x + dt * trj.x_v + 0.5) - py = int(trj.y + dt * trj.y_v + 0.5) + px = int(trj.x + dt * trj.vx + 0.5) + py = int(trj.y + dt * trj.vy + 0.5) # Check the trajectory stays in the image. self.assertGreaterEqual(px, 0) diff --git a/tests/test_file_utils.py b/tests/test_file_utils.py index ce9ed4e49..200a19bd7 100644 --- a/tests/test_file_utils.py +++ b/tests/test_file_utils.py @@ -117,30 +117,30 @@ def test_load_results_trajectories(self): trj_results = FileUtils.load_results_file_as_trajectories("./data/fake_results.txt") self.assertEqual(len(trj_results), 2) - self.assertTrue(isinstance(trj_results[0], trajectory)) + self.assertTrue(isinstance(trj_results[0], Trajectory)) self.assertEqual(trj_results[0].x, 106) self.assertEqual(trj_results[0].y, 44) - self.assertAlmostEqual(trj_results[0].x_v, 9.52, delta=1e-6) - self.assertAlmostEqual(trj_results[0].y_v, -0.5, delta=1e-6) + self.assertAlmostEqual(trj_results[0].vx, 9.52, delta=1e-6) + self.assertAlmostEqual(trj_results[0].vy, -0.5, delta=1e-6) self.assertAlmostEqual(trj_results[0].lh, 300.0, delta=1e-6) self.assertAlmostEqual(trj_results[0].flux, 750.0, delta=1e-6) self.assertEqual(trj_results[0].obs_count, 10) - self.assertTrue(isinstance(trj_results[1], trajectory)) + self.assertTrue(isinstance(trj_results[1], Trajectory)) self.assertEqual(trj_results[1].x, 55) self.assertEqual(trj_results[1].y, 60) - self.assertAlmostEqual(trj_results[1].x_v, 10.5, delta=1e-6) - self.assertAlmostEqual(trj_results[1].y_v, -1.7, delta=1e-6) + self.assertAlmostEqual(trj_results[1].vx, 10.5, delta=1e-6) + self.assertAlmostEqual(trj_results[1].vy, -1.7, delta=1e-6) self.assertAlmostEqual(trj_results[1].lh, 250.0, delta=1e-6) self.assertAlmostEqual(trj_results[1].flux, 500.0, delta=1e-6) self.assertEqual(trj_results[1].obs_count, 9) def test_save_and_load_single_result(self): - trj = trajectory() + trj = Trajectory() trj.x = 1 trj.y = 2 - trj.x_v = 3.0 - trj.y_v = 4.0 + trj.vx = 3.0 + trj.vy = 4.0 with tempfile.TemporaryDirectory() as dir_name: filename = f"{dir_name}/results_tmp.txt" @@ -150,8 +150,8 @@ def test_save_and_load_single_result(self): self.assertEqual(len(loaded_trjs), 1) self.assertEqual(loaded_trjs[0].x, trj.x) self.assertEqual(loaded_trjs[0].y, trj.y) - self.assertEqual(loaded_trjs[0].x_v, trj.x_v) - self.assertEqual(loaded_trjs[0].y_v, trj.y_v) + self.assertEqual(loaded_trjs[0].vx, trj.vx) + self.assertEqual(loaded_trjs[0].vy, trj.vy) def test_load_mpc(self): coords, obs_times = FileUtils.mpc_reader("./data/mpcs.txt") diff --git a/tests/test_filtering.py b/tests/test_filtering.py index 8eca1342c..3c87f99f3 100644 --- a/tests/test_filtering.py +++ b/tests/test_filtering.py @@ -8,12 +8,14 @@ class test_kernels_wrappers(unittest.TestCase): + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_sigmag_filtered_indices_same(self): # With everything the same, nothing should be filtered. values = [1.0 for _ in range(20)] inds = sigmag_filtered_indices(values, 0.25, 0.75, 0.7413, 2.0) self.assertEqual(len(inds), 20) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_sigmag_filtered_indices_no_outliers(self): # Try with a median of 1.0 and a percentile range of 3.0 (2.0 - -1.0). # It should filter any values outside [-3.45, 5.45] @@ -21,6 +23,7 @@ def test_sigmag_filtered_indices_no_outliers(self): inds = sigmag_filtered_indices(values, 0.25, 0.75, 0.7413, 2.0) self.assertEqual(len(inds), len(values)) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_sigmag_filtered_indices_one_outlier(self): # Try with a median of 1.0 and a percentile range of 3.0 (2.0 - -1.0). # It should filter any values outside [-3.45, 5.45] @@ -37,6 +40,7 @@ def test_sigmag_filtered_indices_one_outlier(self): inds = sigmag_filtered_indices(values, 0.25, 0.75, 0.7413, 3.0) self.assertEqual(len(inds), len(values)) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_sigmag_filtered_indices_other_bounds(self): # Do the filtering of test_sigmag_filtered_indices_one_outlier # with wider bounds [-1.8944, 3.8944]. @@ -58,6 +62,7 @@ def test_sigmag_filtered_indices_other_bounds(self): for i in range(1, 9): self.assertTrue(i in inds) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_sigmag_filtered_indices_two_outliers(self): # Try with a median of 0.0 and a percentile range of 1.1 (1.0 - -0.1). # It should filter any values outside [-1.631, 1.631]. @@ -72,6 +77,7 @@ def test_sigmag_filtered_indices_two_outliers(self): inds = sigmag_filtered_indices(values, 0.25, 0.75, 0.7413, 20.0) self.assertEqual(len(inds), len(values) - 1) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_sigmag_filtered_indices_three_outliers(self): # Try with a median of 5.0 and a percentile range of 4.0 (7.0-3.0). # It should filter any values outside [-0.93, 10.93]. diff --git a/tests/test_image_info.py b/tests/test_image_info.py index 56da37ebd..7cf11652b 100644 --- a/tests/test_image_info.py +++ b/tests/test_image_info.py @@ -81,12 +81,12 @@ def test_set_obs_position(self): def test_load_image(self): img_info = ImageInfo() - img_info.populate_from_fits_file("./data/fake_images/000000.fits", load_image=True, p=psf(1.0)) + img_info.populate_from_fits_file("./data/fake_images/000000.fits", load_image=True, p=PSF(1.0)) self.assertIsNotNone(img_info.image) self.assertEqual(img_info.image.get_width(), 64) self.assertEqual(img_info.image.get_height(), 64) - def test_load_image_no_psf(self): + def test_load_image_no_PSF(self): img_info = ImageInfo() with self.assertRaises(ValueError): img_info.populate_from_fits_file("./data/fake_images/000000.fits", load_image=True) @@ -137,7 +137,7 @@ def test_load_files(self): # The (0, 0) pixel should be the same as the RA, DEC # provided in the fits header. - pos00 = pixel_pos() + pos00 = PixelPos() pos00.x = 0.0 pos00.y = 0.0 sky_pos00 = img_info.stats[0].pixels_to_skycoords(pos00) @@ -145,7 +145,7 @@ def test_load_files(self): self.assertAlmostEqual(sky_pos00.dec.degree, -10.788) # The (10, 20) pixel should be moved by 0.001 per pixel. - pos2 = pixel_pos() + pos2 = PixelPos() pos2.x = 10.0 pos2.y = 20.0 sky_pos2 = img_info.stats[0].pixels_to_skycoords(pos2) @@ -153,21 +153,21 @@ def test_load_files(self): self.assertAlmostEqual(sky_pos2.dec.degree, -10.768) # Test that we can map the sky positions back to the coordinates. - pixel_pos00 = img_info.stats[0].skycoords_to_pixels(sky_pos00) - self.assertAlmostEqual(pixel_pos00.x, 0.0) - self.assertAlmostEqual(pixel_pos00.y, 0.0) + PixelPos00 = img_info.stats[0].skycoords_to_pixels(sky_pos00) + self.assertAlmostEqual(PixelPos00.x, 0.0) + self.assertAlmostEqual(PixelPos00.y, 0.0) - pixel_pos2 = img_info.stats[0].skycoords_to_pixels(sky_pos2) - self.assertAlmostEqual(pixel_pos2.x, 10.0) - self.assertAlmostEqual(pixel_pos2.y, 20.0) + PixelPos2 = img_info.stats[0].skycoords_to_pixels(sky_pos2) + self.assertAlmostEqual(PixelPos2.x, 10.0) + self.assertAlmostEqual(PixelPos2.y, 20.0) - # A trajectory of x=0, y=0, x_v=5.0, y_v=10.0 should produce + # A Trajectory of x=0, y=0, x_v=5.0, y_v=10.0 should produce # the same results as above. - trj = trajectory() + trj = Trajectory() trj.x = 0 trj.y = 0 - trj.x_v = 5.0 - trj.y_v = 10.0 + trj.vx = 5.0 + trj.vy = 10.0 sky_pos_mult = img_info.trajectory_to_skycoords(trj) self.assertAlmostEqual(sky_pos_mult[0].ra.degree, 201.614) self.assertAlmostEqual(sky_pos_mult[0].dec.degree, -10.788) diff --git a/tests/test_image_stack.py b/tests/test_image_stack.py index 2183c8dc9..66b0b4f26 100644 --- a/tests/test_image_stack.py +++ b/tests/test_image_stack.py @@ -5,15 +5,15 @@ from kbmod.search import * -class test_image_stack(unittest.TestCase): +class test_ImageStack(unittest.TestCase): def setUp(self): # Create multiple fake layered images to use. self.num_images = 5 self.images = [None] * self.num_images self.p = [None] * self.num_images for i in range(self.num_images): - self.p[i] = psf(5.0 / float(2 * i + 1)) - self.images[i] = layered_image( + self.p[i] = PSF(5.0 / float(2 * i + 1)) + self.images[i] = LayeredImage( ("layered_test_%i" % i), 80, # dim_x = 80 pixels, 60, # dim_y = 60 pixels, @@ -28,7 +28,7 @@ def setUp(self): mask.set_pixel(10, 10 + i, 1) self.images[i].set_mask(mask) - self.im_stack = image_stack(self.images) + self.im_stack = ImageStack(self.images) def test_create(self): self.assertEqual(self.num_images, self.im_stack.img_count()) diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py index 334233e4f..8dbcd702c 100644 --- a/tests/test_layered_image.py +++ b/tests/test_layered_image.py @@ -7,12 +7,12 @@ from kbmod.search import * -class test_layered_image(unittest.TestCase): +class test_LayeredImage(unittest.TestCase): def setUp(self): - self.p = psf(1.0) + self.p = PSF(1.0) # Create a fake layered image to use. - self.image = layered_image( + self.image = LayeredImage( "layered_test", 80, # dim_x = 80 pixels, 60, # dim_y = 60 pixels, @@ -30,7 +30,7 @@ def test_create(self): self.assertEqual(self.image.get_obstime(), 10.0) self.assertEqual(self.image.get_name(), "layered_test") - # Create a fake layered_image. + # Create a fake LayeredImage. science = self.image.get_science() variance = self.image.get_variance() mask = self.image.get_mask() @@ -45,19 +45,19 @@ def test_create(self): self.assertLessEqual(science.get_pixel(x, y), 100.0) def test_create_from_layers(self): - sci = raw_image(30, 40) + sci = RawImage(30, 40) for y in range(40): for x in range(30): sci.set_pixel(x, y, x + 30.0 * y) - var = raw_image(30, 40) + var = RawImage(30, 40) var.set_all(1.0) - mask = raw_image(30, 40) + mask = RawImage(30, 40) mask.set_all(0.0) # Create the layered image. - img2 = layered_image(sci, var, mask, psf(2.0)) + img2 = LayeredImage(sci, var, mask, PSF(2.0)) self.assertEqual(img2.get_width(), 30) self.assertEqual(img2.get_height(), 40) self.assertEqual(img2.get_npixels(), 30.0 * 40.0) @@ -79,10 +79,10 @@ def test_convolve_psf(self): msk0 = self.image.get_mask() # Create a copy of the image. - img_b = layered_image(sci0, var0, msk0, self.p) + img_b = LayeredImage(sci0, var0, msk0, self.p) # A no-op PSF does not change the image. - img_b.convolve_given_psf(psf()) + img_b.convolve_given_psf(PSF()) sci1 = img_b.get_science() var1 = img_b.get_variance() for y in range(img_b.get_height()): @@ -99,7 +99,7 @@ def test_convolve_psf(self): self.assertNotAlmostEqual(sci0.get_pixel(x, y), sci1.get_pixel(x, y)) self.assertNotAlmostEqual(var0.get_pixel(x, y), var1.get_pixel(x, y)) - def test_overwrite_psf(self): + def test_overwrite_PSF(self): p1 = self.image.get_psf() self.assertEqual(p1.get_size(), 25) self.assertEqual(p1.get_dim(), 5) @@ -111,7 +111,7 @@ def test_overwrite_psf(self): science_pixel_psf1 = self.image.get_science().get_pixel(50, 50) # Change the PSF to a no-op. - self.image.set_psf(psf()) + self.image.set_psf(PSF()) # Check that we retrieve the correct PSF. p2 = self.image.get_psf() @@ -239,8 +239,8 @@ def test_grow_mask_mult(self): self.assertEqual(science.pixel_has_data(x, y), dx + dy > 3) def test_psi_and_phi_image(self): - p = psf(0.00000001) # A point function. - img = layered_image("small_test", 6, 5, 2.0, 4.0, 10.0, p) + p = PSF(0.00000001) # A point function. + img = LayeredImage("small_test", 6, 5, 2.0, 4.0, 10.0, p) # Create fake science and variance images. sci = img.get_science() @@ -281,7 +281,7 @@ def test_subtract_template(self): old_science.set_pixel(10, 21, KB_NO_DATA) self.image.set_science(old_science) - template = raw_image(self.image.get_width(), self.image.get_height()) + template = RawImage(self.image.get_width(), self.image.get_height()) template.set_all(0.0) for h in range(old_science.get_height()): template.set_pixel(10, h, 0.01 * h) @@ -301,7 +301,7 @@ def test_read_write_files(self): with tempfile.TemporaryDirectory() as dir_name: file_name = "tmp_layered_test_data" full_path = "%s/%s.fits" % (dir_name, file_name) - im1 = layered_image( + im1 = LayeredImage( file_name, 15, # dim_x = 15 pixels, 20, # dim_y = 20 pixels, @@ -322,7 +322,7 @@ def test_read_write_files(self): im1.save_layers(dir_name + "/") # Reload the test data and check that it matches. - im2 = layered_image(full_path, self.p) + im2 = LayeredImage(full_path, self.p) self.assertEqual(im1.get_height(), im2.get_height()) self.assertEqual(im1.get_width(), im2.get_width()) self.assertEqual(im1.get_npixels(), im2.get_npixels()) @@ -348,7 +348,7 @@ def test_overwrite_files(self): full_path = "%s/%s.fits" % (dir_name, file_name) # Save the test image. - img1 = layered_image(file_name, 15, 20, 2.0, 4.0, 10.0, self.p) + img1 = LayeredImage(file_name, 15, 20, 2.0, 4.0, 10.0, self.p) img1.save_layers(dir_name + "/") with fits.open(full_path) as hdulist: self.assertEqual(len(hdulist), 4) @@ -357,7 +357,7 @@ def test_overwrite_files(self): # Save a new test image over the first and check # that it replaces it. - img2 = layered_image(file_name, 25, 40, 2.0, 4.0, 10.0, self.p) + img2 = LayeredImage(file_name, 25, 40, 2.0, 4.0, 10.0, self.p) img2.save_layers(dir_name + "/") with fits.open(full_path) as hdulist2: self.assertEqual(len(hdulist2), 4) diff --git a/tests/test_masking.py b/tests/test_masking.py index 978a803db..b7ebc248f 100644 --- a/tests/test_masking.py +++ b/tests/test_masking.py @@ -41,17 +41,17 @@ def setUp(self): self.dim_y = 20 self.noise_level = 0.1 self.variance = self.noise_level**2 - self.p = psf(1.0) + self.p = PSF(1.0) self.imlist = [] self.time_list = [] for i in range(self.img_count): time = i / self.img_count self.time_list.append(time) - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) self.imlist.append(im) - self.stack = image_stack(self.imlist) + self.stack = ImageStack(self.imlist) def test_threshold_masker(self): # Set one science pixel per image above the threshold @@ -177,17 +177,17 @@ def setUp(self): self.dim_y = 50 self.noise_level = 0.1 self.variance = self.noise_level**2 - self.p = psf(1.0) + self.p = PSF(1.0) self.imlist = [] self.time_list = [] for i in range(self.img_count): time = i / self.img_count self.time_list.append(time) - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) self.imlist.append(im) - self.stack = image_stack(self.imlist) + self.stack = ImageStack(self.imlist) def test_apply_masks(self): overrides = { diff --git a/tests/test_psf.py b/tests/test_psf.py index 4d088c90d..f8cecb58e 100644 --- a/tests/test_psf.py +++ b/tests/test_psf.py @@ -1,16 +1,16 @@ import unittest -from kbmod.search import psf +from kbmod.search import PSF -class test_psf(unittest.TestCase): +class test_PSF(unittest.TestCase): def setUp(self): self.psf_count = 12 sigma_list = range(self.psf_count) - self.psf_list = [psf(x / 5 + 0.2) for x in sigma_list] + self.psf_list = [PSF(x / 5 + 0.2) for x in sigma_list] def test_make_noop(self): - psf0 = psf() + psf0 = PSF() self.assertEqual(psf0.get_size(), 1) self.assertEqual(psf0.get_dim(), 1) self.assertEqual(psf0.get_radius(), 0) @@ -20,13 +20,13 @@ def test_make_noop(self): self.assertEqual(kernel0[0], 1.0) def test_make_and_copy(self): - psf1 = psf(1.0) + psf1 = PSF(1.0) self.assertEqual(psf1.get_size(), 25) self.assertEqual(psf1.get_dim(), 5) self.assertEqual(psf1.get_radius(), 2) # Make a copy. - psf2 = psf(psf1) + psf2 = PSF(psf1) self.assertEqual(psf2.get_size(), 25) self.assertEqual(psf2.get_dim(), 5) self.assertEqual(psf2.get_radius(), 2) @@ -38,7 +38,7 @@ def test_make_and_copy(self): # Test the creation of a delta function (no-op) PSF. def test_no_op(self): - psf1 = psf(0.000001) + psf1 = PSF(0.000001) self.assertEqual(psf1.get_size(), 1) self.assertEqual(psf1.get_dim(), 1) self.assertEqual(psf1.get_radius(), 0) @@ -55,7 +55,7 @@ def test_square(self): x_sum = p.get_sum() # Make a copy and confirm that the copy preserves the sum. - x = psf(p) + x = PSF(p) self.assertEqual(x.get_sum(), x_sum) # Since each pixel value is squared the sum of the diff --git a/tests/test_raw_image.py b/tests/test_raw_image.py index 3b2ce688b..afdb95842 100644 --- a/tests/test_raw_image.py +++ b/tests/test_raw_image.py @@ -6,11 +6,11 @@ from kbmod.search import * -class test_raw_image(unittest.TestCase): +class test_RawImage(unittest.TestCase): def setUp(self): self.width = 10 self.height = 12 - self.img = raw_image(self.width, self.height) + self.img = RawImage(self.width, self.height) for x in range(self.width): for y in range(self.height): self.img.set_pixel(x, y, float(x + y * self.width)) @@ -39,7 +39,7 @@ def test_create(self): def test_approx_equal(self): # Make approximate copy. - img2 = raw_image(self.width, self.height) + img2 = RawImage(self.width, self.height) for x in range(self.width): for y in range(self.height): img2.set_pixel(x, y, self.img.get_pixel(x, y) + 0.0001) @@ -63,7 +63,7 @@ def test_approx_equal(self): def test_copy(self): # Copy the image. - img2 = raw_image(self.img) + img2 = RawImage(self.img) self.assertEqual(img2.get_width(), self.width) self.assertEqual(img2.get_height(), self.height) self.assertEqual(img2.get_npixels(), self.width * self.height) @@ -131,7 +131,7 @@ def test_find_peak_duplicate(self): self.assertEqual(int(peak.y), 1) def test_find_central_moments(self): - img = raw_image(5, 5) + img = RawImage(5, 5) # Try something mostly symmetric and centered. img.set_all(0.1) @@ -178,9 +178,9 @@ def test_find_central_moments(self): def test_convolve_psf_identity_cpu(self): psf_data = [[0.0 for _ in range(3)] for _ in range(3)] psf_data[1][1] = 1.0 - p = psf(np.array(psf_data)) + p = PSF(np.array(psf_data)) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve_cpu(p) # Check that the image is unchanged. @@ -190,23 +190,23 @@ def test_convolve_psf_identity_cpu(self): def test_convolve_psf_identity_gpu(self): psf_data = [[0.0 for _ in range(3)] for _ in range(3)] psf_data[1][1] = 1.0 - p = psf(np.array(psf_data)) + p = PSF(np.array(psf_data)) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve(p) # Check that the image is unchanged. self.assertTrue(self.img.approx_equal(img2, 0.0001)) def test_convolve_psf_mask_cpu(self): - p = psf(1.0) + p = PSF(1.0) # Mask out three pixels. self.img.set_pixel(5, 6, KB_NO_DATA) self.img.set_pixel(0, 3, KB_NO_DATA) self.img.set_pixel(5, 7, KB_NO_DATA) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve_cpu(p) # Check that the same pixels are masked. @@ -219,14 +219,14 @@ def test_convolve_psf_mask_cpu(self): @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_convolve_psf_mask_gpu(self): - p = psf(1.0) + p = PSF(1.0) # Mask out three pixels. self.img.set_pixel(5, 6, KB_NO_DATA) self.img.set_pixel(0, 3, KB_NO_DATA) self.img.set_pixel(5, 7, KB_NO_DATA) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve(p) # Check that the same pixels are masked. @@ -246,10 +246,10 @@ def test_convolve_psf_average_cpu(self): for x in range(1, 4): for y in range(1, 4): psf_data[x][y] = 0.1111111 - p = psf(np.array(psf_data)) + p = PSF(np.array(psf_data)) self.assertAlmostEqual(p.get_sum(), 1.0, delta=0.00001) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve_cpu(p) for x in range(self.width): @@ -286,10 +286,10 @@ def test_convolve_psf_average_gpu(self): for x in range(1, 4): for y in range(1, 4): psf_data[x][y] = 0.1111111 - p = psf(np.array(psf_data)) + p = PSF(np.array(psf_data)) self.assertAlmostEqual(p.get_sum(), 1.0, delta=0.00001) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve(p) for x in range(self.width): @@ -319,9 +319,9 @@ def test_convolve_psf_average_gpu(self): def test_convolve_psf_orientation_cpu(self): # Set up a non-symmetric psf where orientation matters. psf_data = [[0.0, 0.0, 0.0], [0.0, 0.5, 0.4], [0.0, 0.1, 0.0]] - p = psf(np.array(psf_data)) + p = PSF(np.array(psf_data)) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve_cpu(p) for x in range(self.width): @@ -343,9 +343,9 @@ def test_convolve_psf_orientation_cpu(self): def test_convolve_psf_orientation_gpu(self): # Set up a non-symmetric psf where orientation matters. psf_data = [[0.0, 0.0, 0.0], [0.0, 0.5, 0.4], [0.0, 0.1, 0.0]] - p = psf(np.array(psf_data)) + p = PSF(np.array(psf_data)) - img2 = raw_image(self.img) + img2 = RawImage(self.img) img2.convolve(p) for x in range(self.width): @@ -405,13 +405,13 @@ def test_make_stamp(self): def test_read_write_file(self): with tempfile.TemporaryDirectory() as dir_name: - file_name = "tmp_raw_image" + file_name = "tmp_RawImage" full_path = "%s/%s.fits" % (dir_name, file_name) self.img.save_fits(full_path) # Reload the file. - img2 = raw_image(0, 0) + img2 = RawImage(0, 0) img2.load_fits(full_path, 0) self.assertEqual(img2.get_width(), self.width) self.assertEqual(img2.get_height(), self.height) @@ -421,7 +421,7 @@ def test_read_write_file(self): def test_stack_file(self): with tempfile.TemporaryDirectory() as dir_name: - file_name = "tmp_raw_image" + file_name = "tmp_RawImage" full_path = "%s/%s.fits" % (dir_name, file_name) # Save the image and create a file. @@ -433,7 +433,7 @@ def test_stack_file(self): self.img.append_fits_layer(full_path) # Check that we get 5 layers with the correct times. - img2 = raw_image(0, 0) + img2 = RawImage(0, 0) for i in range(5): img2.load_fits(full_path, i) @@ -444,9 +444,9 @@ def test_stack_file(self): self.assertTrue(self.img.approx_equal(img2, 1e-5)) def test_create_median_image(self): - img1 = raw_image(np.array([[0.0, -1.0], [2.0, 1.0], [0.7, 3.1]])) - img2 = raw_image(np.array([[1.0, 0.0], [1.0, 3.5], [4.0, 3.0]])) - img3 = raw_image(np.array([[-1.0, -2.0], [3.0, 5.0], [4.1, 3.3]])) + img1 = RawImage(np.array([[0.0, -1.0], [2.0, 1.0], [0.7, 3.1]])) + img2 = RawImage(np.array([[1.0, 0.0], [1.0, 3.5], [4.0, 3.0]])) + img3 = RawImage(np.array([[-1.0, -2.0], [3.0, 5.0], [4.1, 3.3]])) vect = [img1, img2, img3] median_image = create_median_image(vect) @@ -460,8 +460,8 @@ def test_create_median_image(self): self.assertAlmostEqual(median_image.get_pixel(1, 2), 3.1, delta=1e-6) # Apply masks to images 1 and 3. - img1.apply_mask(1, [], raw_image(np.array([[0, 1], [0, 1], [0, 1]]))) - img3.apply_mask(1, [], raw_image(np.array([[0, 0], [1, 1], [1, 0]]))) + img1.apply_mask(1, [], RawImage(np.array([[0, 1], [0, 1], [0, 1]]))) + img3.apply_mask(1, [], RawImage(np.array([[0, 0], [1, 1], [1, 0]]))) median_image2 = create_median_image([img1, img2, img3]) self.assertEqual(median_image2.get_width(), 2) @@ -474,21 +474,21 @@ def test_create_median_image(self): self.assertAlmostEqual(median_image2.get_pixel(1, 2), 3.15, delta=1e-6) def test_create_median_image_more(self): - img1 = raw_image(np.array([[1.0, -1.0], [-1.0, 1.0], [1.0, 0.1]])) - img2 = raw_image(np.array([[2.0, 0.0], [0.0, 2.0], [2.0, 0.0]])) - img3 = raw_image(np.array([[3.0, -2.0], [-2.0, 5.0], [4.0, 0.3]])) - img4 = raw_image(np.array([[4.0, 3.0], [3.0, 6.0], [5.0, 0.1]])) - img5 = raw_image(np.array([[5.0, -3.0], [-3.0, 7.0], [7.0, 0.0]])) - img6 = raw_image(np.array([[6.0, 2.0], [2.0, 4.0], [6.0, 0.1]])) - img7 = raw_image(np.array([[7.0, 3.0], [3.0, 3.0], [3.0, 0.0]])) - - img1.apply_mask(1, [], raw_image(np.array([[0, 0], [1, 1], [0, 0]]))) - img2.apply_mask(1, [], raw_image(np.array([[0, 0], [1, 1], [1, 0]]))) - img3.apply_mask(1, [], raw_image(np.array([[0, 0], [0, 1], [0, 0]]))) - img4.apply_mask(1, [], raw_image(np.array([[0, 0], [0, 1], [0, 0]]))) - img5.apply_mask(1, [], raw_image(np.array([[0, 1], [0, 1], [0, 0]]))) - img6.apply_mask(1, [], raw_image(np.array([[0, 1], [1, 1], [0, 0]]))) - img7.apply_mask(1, [], raw_image(np.array([[0, 0], [1, 1], [0, 0]]))) + img1 = RawImage(np.array([[1.0, -1.0], [-1.0, 1.0], [1.0, 0.1]])) + img2 = RawImage(np.array([[2.0, 0.0], [0.0, 2.0], [2.0, 0.0]])) + img3 = RawImage(np.array([[3.0, -2.0], [-2.0, 5.0], [4.0, 0.3]])) + img4 = RawImage(np.array([[4.0, 3.0], [3.0, 6.0], [5.0, 0.1]])) + img5 = RawImage(np.array([[5.0, -3.0], [-3.0, 7.0], [7.0, 0.0]])) + img6 = RawImage(np.array([[6.0, 2.0], [2.0, 4.0], [6.0, 0.1]])) + img7 = RawImage(np.array([[7.0, 3.0], [3.0, 3.0], [3.0, 0.0]])) + + img1.apply_mask(1, [], RawImage(np.array([[0, 0], [1, 1], [0, 0]]))) + img2.apply_mask(1, [], RawImage(np.array([[0, 0], [1, 1], [1, 0]]))) + img3.apply_mask(1, [], RawImage(np.array([[0, 0], [0, 1], [0, 0]]))) + img4.apply_mask(1, [], RawImage(np.array([[0, 0], [0, 1], [0, 0]]))) + img5.apply_mask(1, [], RawImage(np.array([[0, 1], [0, 1], [0, 0]]))) + img6.apply_mask(1, [], RawImage(np.array([[0, 1], [1, 1], [0, 0]]))) + img7.apply_mask(1, [], RawImage(np.array([[0, 0], [1, 1], [0, 0]]))) vect = [img1, img2, img3, img4, img5, img6, img7] median_image = create_median_image(vect) @@ -503,9 +503,9 @@ def test_create_median_image_more(self): self.assertAlmostEqual(median_image.get_pixel(1, 2), 0.1, delta=1e-6) def test_create_summed_image(self): - img1 = raw_image(np.array([[0.0, -1.0], [2.0, 1.0], [0.7, 3.1]])) - img2 = raw_image(np.array([[1.0, 0.0], [1.0, 3.5], [4.0, 3.0]])) - img3 = raw_image(np.array([[-1.0, -2.0], [3.0, 5.0], [4.1, 3.3]])) + img1 = RawImage(np.array([[0.0, -1.0], [2.0, 1.0], [0.7, 3.1]])) + img2 = RawImage(np.array([[1.0, 0.0], [1.0, 3.5], [4.0, 3.0]])) + img3 = RawImage(np.array([[-1.0, -2.0], [3.0, 5.0], [4.1, 3.3]])) vect = [img1, img2, img3] summed_image = create_summed_image(vect) @@ -519,8 +519,8 @@ def test_create_summed_image(self): self.assertAlmostEqual(summed_image.get_pixel(1, 2), 9.4, delta=1e-6) # Apply masks to images 1 and 3. - img1.apply_mask(1, [], raw_image(np.array([[0, 1], [0, 1], [0, 1]]))) - img3.apply_mask(1, [], raw_image(np.array([[0, 0], [1, 1], [1, 0]]))) + img1.apply_mask(1, [], RawImage(np.array([[0, 1], [0, 1], [0, 1]]))) + img3.apply_mask(1, [], RawImage(np.array([[0, 0], [1, 1], [1, 0]]))) summed_image2 = create_summed_image([img1, img2, img3]) self.assertEqual(summed_image2.get_width(), 2) @@ -533,9 +533,9 @@ def test_create_summed_image(self): self.assertAlmostEqual(summed_image2.get_pixel(1, 2), 6.3, delta=1e-6) def test_create_mean_image(self): - img1 = raw_image(np.array([[0.0, -1.0], [2.0, 1.0], [0.7, 3.1]])) - img2 = raw_image(np.array([[1.0, 0.0], [1.0, 3.5], [4.0, 3.0]])) - img3 = raw_image(np.array([[-1.0, -2.0], [3.0, 5.0], [4.1, 3.3]])) + img1 = RawImage(np.array([[0.0, -1.0], [2.0, 1.0], [0.7, 3.1]])) + img2 = RawImage(np.array([[1.0, 0.0], [1.0, 3.5], [4.0, 3.0]])) + img3 = RawImage(np.array([[-1.0, -2.0], [3.0, 5.0], [4.1, 3.3]])) mean_image = create_mean_image([img1, img2, img3]) self.assertEqual(mean_image.get_width(), 2) @@ -548,9 +548,9 @@ def test_create_mean_image(self): self.assertAlmostEqual(mean_image.get_pixel(1, 2), 9.4 / 3.0, delta=1e-6) # Apply masks to images 1, 2, and 3. - img1.apply_mask(1, [], raw_image(np.array([[0, 1], [0, 1], [0, 1]]))) - img2.apply_mask(1, [], raw_image(np.array([[0, 0], [0, 0], [0, 1]]))) - img3.apply_mask(1, [], raw_image(np.array([[0, 0], [1, 1], [1, 1]]))) + img1.apply_mask(1, [], RawImage(np.array([[0, 1], [0, 1], [0, 1]]))) + img2.apply_mask(1, [], RawImage(np.array([[0, 0], [0, 0], [0, 1]]))) + img3.apply_mask(1, [], RawImage(np.array([[0, 0], [1, 1], [1, 1]]))) mean_image2 = create_mean_image([img1, img2, img3]) self.assertEqual(mean_image2.get_width(), 2) diff --git a/tests/test_readme_example.py b/tests/test_readme_example.py index dc428545e..a3338cd76 100644 --- a/tests/test_readme_example.py +++ b/tests/test_readme_example.py @@ -7,9 +7,10 @@ class test_readme_example(unittest.TestCase): + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_make_and_copy(self): # Create a point spread function - psf = kb.psf(1.5) + psf = kb.PSF(1.5) # Create fake data with ten 512x512 pixel images. ds = FakeDataSet(512, 512, 10) @@ -35,10 +36,10 @@ def test_make_and_copy(self): ) # Create a new image stack with the inserted object. - stack = kb.image_stack(imgs) + stack = kb.ImageStack(imgs) # Recover the object by searching a set of trajectories. - search = kb.stack_search(stack) + search = kb.StackSearch(stack) search.search( 5, # Number of search angles to try (-0.1, -0.05, 0.0, 0.05, 0.1) 5, # Number of search velocities to try (0, 1, 2, 3, 4) diff --git a/tests/test_result_list.py b/tests/test_result_list.py index 595143d7f..6fe952dd1 100644 --- a/tests/test_result_list.py +++ b/tests/test_result_list.py @@ -11,7 +11,7 @@ class test_result_data_row(unittest.TestCase): def setUp(self): - self.trj = trajectory() + self.trj = Trajectory() self.trj.obs_count = 4 self.times = [1.0, 2.0, 3.0, 4.0] @@ -37,7 +37,7 @@ def test_filter(self): self.assertEqual(self.rdr.valid_indices, [0, 2, 3]) self.assertEqual(self.rdr.valid_times(self.times), [1.0, 3.0, 4.0]) - # The values within the trajectory object *should* change. + # The values within the Trajectory object *should* change. self.assertEqual(self.rdr.trajectory.obs_count, 3) self.assertAlmostEqual(self.rdr.trajectory.flux, 1.1666667, delta=1e-5) self.assertAlmostEqual(self.rdr.trajectory.lh, 2.020725, delta=1e-5) @@ -69,7 +69,7 @@ def test_append_single(self): self.assertEqual(rs.num_results(), 0) for i in range(5): - t = trajectory() + t = Trajectory() t.lh = float(i) rs.append_result(ResultRow(t, self.num_times)) self.assertEqual(rs.num_results(), 5) @@ -86,13 +86,13 @@ def test_extend(self): # Fill the first ResultList with 5 rows. for i in range(5): - t = trajectory() + t = Trajectory() t.lh = float(i) rs1.append_result(ResultRow(t, self.num_times)) # Fill a second Result set with 5 different rows. for i in range(5): - t = trajectory() + t = Trajectory() t.lh = float(i) + 5.0 rs2.append_result(ResultRow(t, self.num_times)) @@ -111,7 +111,7 @@ def test_extend(self): def test_clear(self): rs = ResultList(self.times) for i in range(3): - t = trajectory() + t = Trajectory() rs.append_result(ResultRow(t, self.num_times)) self.assertEqual(rs.num_results(), 3) @@ -121,7 +121,7 @@ def test_clear(self): def test_filter(self): rs = ResultList(self.times) for i in range(10): - t = trajectory() + t = Trajectory() t.lh = float(i) rs.append_result(ResultRow(t, self.num_times)) self.assertEqual(rs.num_results(), 10) @@ -141,7 +141,7 @@ def test_filter(self): def test_filter_dups(self): rs = ResultList(self.times, track_filtered=False) for i in range(10): - t = trajectory() + t = Trajectory() t.lh = float(i) rs.append_result(ResultRow(t, self.num_times)) self.assertEqual(rs.num_results(), 10) @@ -162,7 +162,7 @@ def test_filter_dups(self): def test_filter_track(self): rs = ResultList(self.times, track_filtered=True) for i in range(10): - t = trajectory() + t = Trajectory() t.x = i rs.append_result(ResultRow(t, self.num_times)) self.assertEqual(rs.num_results(), 10) @@ -201,7 +201,7 @@ def test_save_results(self): # Fill the ResultList with 3 fake rows. rs = ResultList(times) for i in range(3): - row = ResultRow(trajectory(), 3) + row = ResultRow(Trajectory(), 3) row.set_psi_phi([0.1, 0.2, 0.3], [1.0, 1.0, 0.5]) row.filter_indices([t for t in range(i + 1)]) rs.append_result(row) @@ -270,7 +270,7 @@ def test_save_and_load_results(self): # Fill the ResultList with 3 fake rows. rs = ResultList(times) for i in range(3): - row = ResultRow(trajectory(), num_times) + row = ResultRow(Trajectory(), num_times) row.set_psi_phi([0.1, 0.6, 0.2, float(i)], [2.0, 0.5, float(i), 1.0]) row.filter_indices([t for t in range(num_times - i)]) row.stamp = np.array([[float(i), float(i) / 3.0], [1.0, 0.5]]) @@ -313,7 +313,7 @@ def test_save_and_load_results_filtered(self): # Fill the ResultList with 5 fake rows. rs = ResultList(times, track_filtered=True) for i in range(5): - trj = trajectory() + trj = Trajectory() trj.x = 10 * i row = ResultRow(trj, num_times) row.set_psi_phi([0.1, 0.6, 0.2, float(i)], [2.0, 0.5, float(i), 1.0]) diff --git a/tests/test_search.py b/tests/test_search.py index 411d86a65..e06a76d8b 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -19,21 +19,21 @@ def setUp(self): self.dim_y = 60 self.noise_level = 4.0 self.variance = self.noise_level**2 - self.p = psf(1.0) + self.p = PSF(1.0) # object properties self.object_flux = 250.0 self.start_x = 17 self.start_y = 12 - self.x_vel = 21.0 - self.y_vel = 16.0 + self.vxel = 21.0 + self.vyel = 16.0 - # create a trajectory for the object - self.trj = trajectory() + # create a Trajectory for the object + self.trj = Trajectory() self.trj.x = self.start_x self.trj.y = self.start_y - self.trj.x_v = self.x_vel - self.trj.y_v = self.y_vel + self.trj.vx = self.vxel + self.trj.vy = self.vyel # Add convenience array of all true bools for the stamp tests. self.all_valid = [True] * self.imCount @@ -54,13 +54,13 @@ def setUp(self): self.imlist = [] for i in range(self.imCount): time = i / self.imCount - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) add_fake_object( im, - self.start_x + time * self.x_vel + 0.5, - self.start_y + time * self.y_vel + 0.5, + self.start_x + time * self.vxel + 0.5, + self.start_y + time * self.vyel + 0.5, self.object_flux, self.p, ) @@ -73,41 +73,42 @@ def setUp(self): im.apply_mask_flags(1, []) self.imlist.append(im) - self.stack = image_stack(self.imlist) - self.search = stack_search(self.stack) + self.stack = ImageStack(self.imlist) + self.search = StackSearch(self.stack) # Set the filtering parameters. - self.params = stamp_parameters() + self.params = StampParameters() self.params.radius = 5 self.params.do_filtering = True self.params.stamp_type = StampType.STAMP_MEAN self.params.center_thresh = 0.03 self.params.peak_offset_x = 1.5 self.params.peak_offset_y = 1.5 - self.params.m01 = 0.6 - self.params.m10 = 0.6 - self.params.m11 = 2.0 - self.params.m02 = 35.5 - self.params.m20 = 35.5 + self.params.m01_limit = 0.6 + self.params.m10_limit = 0.6 + self.params.m11_limit = 2.0 + self.params.m02_limit = 35.5 + self.params.m20_limit = 35.5 def test_psiphi(self): - p = psf(0.00001) + p = PSF(0.00001) # Image1 has a single object. - image1 = layered_image("test1", 5, 10, 2.0, 4.0, 1.0, p) + image1 = LayeredImage("test1", 5, 10, 2.0, 4.0, 1.0, p) add_fake_object(image1, 3.5, 2.5, 400.0, p) # Image2 has a single object and a masked pixel. - image2 = layered_image("test2", 5, 10, 2.0, 4.0, 2.0, p) + image2 = LayeredImage("test2", 5, 10, 2.0, 4.0, 2.0, p) add_fake_object(image2, 2.5, 4.5, 400.0, p) + mask = image2.get_mask() mask.set_pixel(4, 9, 1) image2.set_mask(mask) image2.apply_mask_flags(1, []) # Create a stack from the two objects. - stack = image_stack([image1, image2]) - search = stack_search(stack) + stack = ImageStack([image1, image2]) + search = StackSearch(stack) # Generate psi and phi. search.prepare_psi_phi() @@ -138,6 +139,7 @@ def test_psiphi(self): ) self.assertAlmostEqual(phi[1].get_pixel(x, y), 1.0 / var.get_pixel(x, y), delta=1e-6) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results(self): self.search.search( self.angle_steps, @@ -153,10 +155,11 @@ def test_results(self): best = results[0] self.assertAlmostEqual(best.x, self.start_x, delta=self.pixel_error) self.assertAlmostEqual(best.y, self.start_y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / self.x_vel, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / self.y_vel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / self.vxel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / self.vyel, 1, delta=self.velocity_error) self.assertAlmostEqual(best.flux / self.object_flux, 1, delta=self.flux_error) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results_extended_bounds(self): self.search.set_start_bounds_x(-10, self.dim_x + 10) self.search.set_start_bounds_y(-10, self.dim_y + 10) @@ -175,10 +178,11 @@ def test_results_extended_bounds(self): best = results[0] self.assertAlmostEqual(best.x, self.start_x, delta=self.pixel_error) self.assertAlmostEqual(best.y, self.start_y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / self.x_vel, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / self.y_vel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / self.vxel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / self.vyel, 1, delta=self.velocity_error) self.assertAlmostEqual(best.flux / self.object_flux, 1, delta=self.flux_error) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results_reduced_bounds(self): self.search.set_start_bounds_x(5, self.dim_x - 5) self.search.set_start_bounds_y(5, self.dim_y - 5) @@ -197,34 +201,35 @@ def test_results_reduced_bounds(self): best = results[0] self.assertAlmostEqual(best.x, self.start_x, delta=self.pixel_error) self.assertAlmostEqual(best.y, self.start_y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / self.x_vel, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / self.y_vel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / self.vxel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / self.vyel, 1, delta=self.velocity_error) self.assertAlmostEqual(best.flux / self.object_flux, 1, delta=self.flux_error) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results_off_chip(self): - trj = trajectory() + trj = Trajectory() trj.x = -3 trj.y = 12 - trj.x_v = 25.0 - trj.y_v = 10.0 + trj.vx = 25.0 + trj.vy = 10.0 # Create images with this fake object. imlist = [] for i in range(self.imCount): time = i / self.imCount - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) add_fake_object( im, - trj.x + time * trj.x_v + 0.5, - trj.y + time * trj.y_v + 0.5, + trj.x + time * trj.vx + 0.5, + trj.y + time * trj.vy + 0.5, self.object_flux, self.p, ) imlist.append(im) - stack = image_stack(imlist) - search = stack_search(stack) + stack = ImageStack(imlist) + search = StackSearch(stack) # Do the extended search. search.set_start_bounds_x(-10, self.dim_x + 10) @@ -244,11 +249,11 @@ def test_results_off_chip(self): best = results[0] self.assertAlmostEqual(best.x, trj.x, delta=self.pixel_error) self.assertAlmostEqual(best.y, trj.y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / trj.x_v, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / trj.y_v, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / trj.vx, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / trj.vy, 1, delta=self.velocity_error) def test_sci_viz_stamps(self): - sci_stamps = self.search.science_viz_stamps(self.trj, 2) + sci_stamps = self.search.get_stamps(self.trj, 2) self.assertEqual(len(sci_stamps), self.imCount) times = self.stack.get_times() @@ -258,8 +263,8 @@ def test_sci_viz_stamps(self): # Compute the interpolated pixel value at the projected location. t = times[i] - x = float(self.trj.x) + self.trj.x_v * t - y = float(self.trj.y) + self.trj.y_v * t + x = float(self.trj.x) + self.trj.vx * t + y = float(self.trj.y) + self.trj.vy * t pixVal = self.imlist[i].get_science().get_pixel_interp(x, y) if pixVal == KB_NO_DATA: pivVal = 0.0 @@ -269,8 +274,8 @@ def test_sci_viz_stamps(self): self.assertAlmostEqual(sci_stamps[i].get_pixel(2, 2), pixVal, delta=0.001) def test_stacked_sci(self): - # Compute the stacked science from a single trajectory. - sci = self.search.summed_sci_stamp(self.trj, 2, []) + # Compute the stacked science from a single Trajectory. + sci = self.search.get_summed_stamp(self.trj, 2, []) self.assertEqual(sci.get_width(), 5) self.assertEqual(sci.get_height(), 5) @@ -279,8 +284,8 @@ def test_stacked_sci(self): sum_middle = 0.0 for i in range(self.imCount): t = times[i] - x = int(self.trj.x + self.trj.x_v * t) - y = int(self.trj.y + self.trj.y_v * t) + x = int(self.trj.x + self.trj.vx * t) + y = int(self.trj.y + self.trj.vy * t) pixVal = self.imlist[i].get_science().get_pixel(x, y) if pixVal == KB_NO_DATA: pivVal = 0.0 @@ -297,11 +302,11 @@ def test_median_stamps_trj(self): goodIdx[1][5] = 0 goodIdx[1][9] = 0 - medianStamps0 = self.search.median_sci_stamp(self.trj, 2, goodIdx[0]) + medianStamps0 = self.search.get_median_stamp(self.trj, 2, goodIdx[0]) self.assertEqual(medianStamps0.get_width(), 5) self.assertEqual(medianStamps0.get_height(), 5) - medianStamps1 = self.search.median_sci_stamp(self.trj, 2, goodIdx[1]) + medianStamps1 = self.search.get_median_stamp(self.trj, 2, goodIdx[1]) self.assertEqual(medianStamps1.get_width(), 5) self.assertEqual(medianStamps1.get_height(), 5) @@ -311,8 +316,8 @@ def test_median_stamps_trj(self): pix_values1 = [] for i in range(self.imCount): t = times[i] - x = int(self.trj.x + self.trj.x_v * t) - y = int(self.trj.y + self.trj.y_v * t) + x = int(self.trj.x + self.trj.vx * t) + y = int(self.trj.y + self.trj.vy * t) pixVal = self.imlist[i].get_science().get_pixel(x, y) if pixVal != KB_NO_DATA and goodIdx[0][i] == 1: pix_values0.append(pixVal) @@ -326,15 +331,15 @@ def test_median_stamps_trj(self): self.assertAlmostEqual(np.median(pix_values1), medianStamps1.get_pixel(2, 2), delta=1e-5) def test_median_stamps_no_data(self): - # Create a trajectory that goes through the masked pixels. - trj = trajectory() + # Create a Trajectory that goes through the masked pixels. + trj = Trajectory() trj.x = self.masked_x trj.y = self.masked_y - trj.x_v = 0 - trj.y_v = 0 + trj.vx = 0 + trj.vy = 0 - # Compute the stacked science from a single trajectory. - medianStamp = self.search.median_sci_stamp(trj, 2, self.all_valid) + # Compute the stacked science from a single Trajectory. + medianStamp = self.search.get_median_stamp(trj, 2, self.all_valid) self.assertEqual(medianStamp.get_width(), 5) self.assertEqual(medianStamp.get_height(), 5) @@ -356,11 +361,11 @@ def test_mean_stamps_trj(self): goodIdx[1][5] = 0 goodIdx[1][9] = 0 - meanStamp0 = self.search.mean_sci_stamp(self.trj, 2, goodIdx[0]) + meanStamp0 = self.search.get_mean_stamp(self.trj, 2, goodIdx[0]) self.assertEqual(meanStamp0.get_width(), 5) self.assertEqual(meanStamp0.get_height(), 5) - meanStamp1 = self.search.mean_sci_stamp(self.trj, 2, goodIdx[1]) + meanStamp1 = self.search.get_mean_stamp(self.trj, 2, goodIdx[1]) self.assertEqual(meanStamp1.get_width(), 5) self.assertEqual(meanStamp1.get_height(), 5) @@ -372,8 +377,8 @@ def test_mean_stamps_trj(self): pix_count1 = 0.0 for i in range(self.imCount): t = times[i] - x = int(self.trj.x + self.trj.x_v * t) - y = int(self.trj.y + self.trj.y_v * t) + x = int(self.trj.x + self.trj.vx * t) + y = int(self.trj.y + self.trj.vy * t) pixVal = self.imlist[i].get_science().get_pixel(x, y) if pixVal != KB_NO_DATA and goodIdx[0][i] == 1: pix_sum0 += pixVal @@ -389,15 +394,15 @@ def test_mean_stamps_trj(self): self.assertAlmostEqual(pix_sum1 / pix_count1, meanStamp1.get_pixel(2, 2), delta=1e-5) def test_mean_stamps_no_data(self): - # Create a trajectory that goes through the masked pixels. - trj = trajectory() + # Create a Trajectory that goes through the masked pixels. + trj = Trajectory() trj.x = self.masked_x trj.y = self.masked_y - trj.x_v = 0 - trj.y_v = 0 + trj.vx = 0 + trj.vy = 0 - # Compute the stacked science from a single trajectory. - meanStamp = self.search.mean_sci_stamp(trj, 2, self.all_valid) + # Compute the stacked science from a single Trajectory. + meanStamp = self.search.get_mean_stamp(trj, 2, self.all_valid) self.assertEqual(meanStamp.get_width(), 5) self.assertEqual(meanStamp.get_height(), 5) @@ -418,7 +423,7 @@ def test_filter_stamp(self): stamp_width = 2 * self.params.radius + 1 # Test a stamp with nothing in it. - stamp = raw_image(stamp_width, stamp_width) + stamp = RawImage(stamp_width, stamp_width) stamp.set_all(1.0) self.assertTrue(self.search.filter_stamp(stamp, self.params)) @@ -466,7 +471,7 @@ def test_coadd_cpu_simple(self): imlist = [] for i in range(3): time = i - im = layered_image(str(i), 3, 3, 0.1, 0.01, i, self.p, i) + im = LayeredImage(str(i), 3, 3, 0.1, 0.01, i, self.p, i) # Overwrite the middle row to be i + 1. sci = im.get_science() @@ -485,39 +490,39 @@ def test_coadd_cpu_simple(self): im.apply_mask_flags(1, []) imlist.append(im) - stack = image_stack(imlist) - search = stack_search(stack) + stack = ImageStack(imlist) + search = StackSearch(stack) all_valid = [True, True, True] # convenience array - # One trajectory right in the image's middle. - trj = trajectory() + # One Trajectory right in the image's middle. + trj = Trajectory() trj.x = 1 trj.y = 1 - trj.x_v = 0 - trj.y_v = 0 + trj.vx = 0 + trj.vy = 0 # Basic Stamp parameters. - params = stamp_parameters() + params = StampParameters() params.radius = 1 params.do_filtering = False # Test summed. params.stamp_type = StampType.STAMP_SUM - stamps = search.coadded_stamps([trj], [all_valid], params, False) + stamps = search.get_coadded_stamps([trj], [all_valid], params, False) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 5.0) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 6.0) # Test mean. params.stamp_type = StampType.STAMP_MEAN - stamps = search.coadded_stamps([trj], [all_valid], params, False) + stamps = search.get_coadded_stamps([trj], [all_valid], params, False) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0) # Test median. params.stamp_type = StampType.STAMP_MEDIAN - stamps = search.coadded_stamps([trj], [all_valid], params, False) + stamps = search.get_coadded_stamps([trj], [all_valid], params, False) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0) @@ -528,7 +533,7 @@ def test_coadd_gpu_simple(self): imlist = [] for i in range(3): time = i - im = layered_image(str(i), 3, 3, 0.1, 0.01, i, self.p, i) + im = LayeredImage(str(i), 3, 3, 0.1, 0.01, i, self.p, i) # Overwrite the middle row to be i + 1. sci = im.get_science() @@ -547,61 +552,61 @@ def test_coadd_gpu_simple(self): im.apply_mask_flags(1, []) imlist.append(im) - stack = image_stack(imlist) - search = stack_search(stack) + stack = ImageStack(imlist) + search = StackSearch(stack) all_valid = [True, True, True] # convenience array - # One trajectory right in the image's middle. - trj = trajectory() + # One Trajectory right in the image's middle. + trj = Trajectory() trj.x = 1 trj.y = 1 - trj.x_v = 0 - trj.y_v = 0 + trj.vx = 0 + trj.vy = 0 # Basic Stamp parameters. - params = stamp_parameters() + params = StampParameters() params.radius = 1 params.do_filtering = False # Test summed. params.stamp_type = StampType.STAMP_SUM - stamps = search.coadded_stamps([trj], [all_valid], params, True) + stamps = search.get_coadded_stamps([trj], [all_valid], params, True) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 5.0) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 6.0) # Test mean. params.stamp_type = StampType.STAMP_MEAN - stamps = search.coadded_stamps([trj], [all_valid], params, True) + stamps = search.get_coadded_stamps([trj], [all_valid], params, True) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0) # Test median. params.stamp_type = StampType.STAMP_MEDIAN - stamps = search.coadded_stamps([trj], [all_valid], params, True) + stamps = search.get_coadded_stamps([trj], [all_valid], params, True) self.assertAlmostEqual(stamps[0].get_pixel(0, 1), 3.0) self.assertAlmostEqual(stamps[0].get_pixel(1, 1), 2.5) self.assertAlmostEqual(stamps[0].get_pixel(2, 1), 2.0) def test_coadd_cpu(self): - params = stamp_parameters() + params = StampParameters() params.radius = 3 params.do_filtering = False - # Compute the stacked science (summed and mean) from a single trajectory. + # Compute the stacked science (summed and mean) from a single Trajectory. params.stamp_type = StampType.STAMP_SUM - summedStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, False) + summedStamps = self.search.get_coadded_stamps([self.trj], [self.all_valid], params, False) self.assertEqual(summedStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(summedStamps[0].get_height(), 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEAN - meanStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, False) + meanStamps = self.search.get_coadded_stamps([self.trj], [self.all_valid], params, False) self.assertEqual(meanStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(meanStamps[0].get_height(), 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEDIAN - medianStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, False) + medianStamps = self.search.get_coadded_stamps([self.trj], [self.all_valid], params, False) self.assertEqual(medianStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(medianStamps[0].get_height(), 2 * params.radius + 1) @@ -617,8 +622,8 @@ def test_coadd_cpu(self): pix_vals = [] for i in range(self.imCount): t = times[i] - x = int(self.trj.x + self.trj.x_v * t) + x_offset - y = int(self.trj.y + self.trj.y_v * t) + y_offset + x = int(self.trj.x + self.trj.vx * t) + x_offset + y = int(self.trj.y + self.trj.vy * t) + y_offset pixVal = self.imlist[i].get_science().get_pixel(x, y) if pixVal != KB_NO_DATA: pix_sum += pixVal @@ -636,23 +641,23 @@ def test_coadd_cpu(self): @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_gpu(self): - params = stamp_parameters() + params = StampParameters() params.radius = 3 params.do_filtering = False - # Compute the stacked science (summed and mean) from a single trajectory. + # Compute the stacked science (summed and mean) from a single Trajectory. params.stamp_type = StampType.STAMP_SUM - summedStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, True) + summedStamps = self.search.get_coadded_stamps([self.trj], [self.all_valid], params, True) self.assertEqual(summedStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(summedStamps[0].get_height(), 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEAN - meanStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, True) + meanStamps = self.search.get_coadded_stamps([self.trj], [self.all_valid], params, True) self.assertEqual(meanStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(meanStamps[0].get_height(), 2 * params.radius + 1) params.stamp_type = StampType.STAMP_MEDIAN - medianStamps = self.search.coadded_stamps([self.trj], [self.all_valid], params, True) + medianStamps = self.search.get_coadded_stamps([self.trj], [self.all_valid], params, True) self.assertEqual(medianStamps[0].get_width(), 2 * params.radius + 1) self.assertEqual(medianStamps[0].get_height(), 2 * params.radius + 1) @@ -668,8 +673,8 @@ def test_coadd_gpu(self): pix_vals = [] for i in range(self.imCount): t = times[i] - x = int(self.trj.x + self.trj.x_v * t) + x_offset - y = int(self.trj.y + self.trj.y_v * t) + y_offset + x = int(self.trj.x + self.trj.vx * t) + x_offset + y = int(self.trj.y + self.trj.vy * t) + y_offset pixVal = self.imlist[i].get_science().get_pixel(x, y) if pixVal != KB_NO_DATA: pix_sum += pixVal @@ -686,7 +691,7 @@ def test_coadd_gpu(self): ) def test_coadd_cpu_use_inds(self): - params = stamp_parameters() + params = StampParameters() params.radius = 1 params.do_filtering = False params.stamp_type = StampType.STAMP_MEAN @@ -699,8 +704,8 @@ def test_coadd_cpu_use_inds(self): inds[1][7] = False inds[1][11] = False - # Compute the stacked science (summed and mean) from a single trajectory. - meanStamps = self.search.coadded_stamps([self.trj, self.trj], inds, params, False) + # Compute the stacked science (summed and mean) from a single Trajectory. + meanStamps = self.search.get_coadded_stamps([self.trj, self.trj], inds, params, False) # Compute the true summed and mean pixels for all of the pixels in the stamp. times = self.stack.get_times() @@ -715,8 +720,8 @@ def test_coadd_cpu_use_inds(self): count_1 = 0.0 for i in range(self.imCount): t = times[i] - x = int(self.trj.x + self.trj.x_v * t) + x_offset - y = int(self.trj.y + self.trj.y_v * t) + y_offset + x = int(self.trj.x + self.trj.vx * t) + x_offset + y = int(self.trj.y + self.trj.vy * t) + y_offset pixVal = self.imlist[i].get_science().get_pixel(x, y) if pixVal != KB_NO_DATA and inds[0][i] > 0: @@ -735,7 +740,7 @@ def test_coadd_cpu_use_inds(self): @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_gpu_use_inds(self): - params = stamp_parameters() + params = StampParameters() params.radius = 1 params.do_filtering = False params.stamp_type = StampType.STAMP_MEAN @@ -748,8 +753,8 @@ def test_coadd_gpu_use_inds(self): inds[1][7] = False inds[1][11] = False - # Compute the stacked science (summed and mean) from a single trajectory. - meanStamps = self.search.coadded_stamps([self.trj, self.trj], inds, params, True) + # Compute the stacked science (summed and mean) from a single Trajectory. + meanStamps = self.search.get_coadded_stamps([self.trj, self.trj], inds, params, True) # Compute the true summed and mean pixels for all of the pixels in the stamp. times = self.stack.get_times() @@ -764,8 +769,8 @@ def test_coadd_gpu_use_inds(self): count_1 = 0.0 for i in range(self.imCount): t = times[i] - x = int(self.trj.x + self.trj.x_v * t) + x_offset - y = int(self.trj.y + self.trj.y_v * t) + y_offset + x = int(self.trj.x + self.trj.vx * t) + x_offset + y = int(self.trj.y + self.trj.vy * t) + y_offset pixVal = self.imlist[i].get_science().get_pixel(x, y) if pixVal != KB_NO_DATA and inds[0][i] > 0: @@ -783,30 +788,30 @@ def test_coadd_gpu_use_inds(self): self.assertAlmostEqual(sum_1 / count_1, meanStamps[1].get_pixel(stamp_x, stamp_y), delta=1e-3) def test_coadd_filter_cpu(self): - # Create a second trajectory that isn't any good. - trj2 = trajectory() + # Create a second Trajectory that isn't any good. + trj2 = Trajectory() trj2.x = 1 trj2.y = 1 - trj2.x_v = 0 - trj2.y_v = 0 + trj2.vx = 0 + trj2.vy = 0 - # Create a third trajectory that is close to good, but offset. - trj3 = trajectory() + # Create a third Trajectory that is close to good, but offset. + trj3 = Trajectory() trj3.x = self.trj.x + 2 trj3.y = self.trj.y + 2 - trj3.x_v = self.trj.x_v - trj3.y_v = self.trj.y_v + trj3.vx = self.trj.vx + trj3.vy = self.trj.vy - # Create a fourth trajectory that is close enough - trj4 = trajectory() + # Create a fourth Trajectory that is close enough + trj4 = Trajectory() trj4.x = self.trj.x + 1 trj4.y = self.trj.y + 1 - trj4.x_v = self.trj.x_v - trj4.y_v = self.trj.y_v + trj4.vx = self.trj.vx + trj4.vy = self.trj.vy - # Compute the stacked science from a single trajectory. + # Compute the stacked science from a single Trajectory. all_valid_vect = [(self.all_valid) for i in range(4)] - meanStamps = self.search.coadded_stamps( + meanStamps = self.search.get_coadded_stamps( [self.trj, trj2, trj3, trj4], all_valid_vect, self.params, False ) @@ -824,30 +829,30 @@ def test_coadd_filter_cpu(self): @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_filter_gpu(self): - # Create a second trajectory that isn't any good. - trj2 = trajectory() + # Create a second Trajectory that isn't any good. + trj2 = Trajectory() trj2.x = 1 trj2.y = 1 - trj2.x_v = 0 - trj2.y_v = 0 + trj2.vx = 0 + trj2.vy = 0 - # Create a third trajectory that is close to good, but offset. - trj3 = trajectory() + # Create a third Trajectory that is close to good, but offset. + trj3 = Trajectory() trj3.x = self.trj.x + 2 trj3.y = self.trj.y + 2 - trj3.x_v = self.trj.x_v - trj3.y_v = self.trj.y_v + trj3.vx = self.trj.vx + trj3.vy = self.trj.vy - # Create a fourth trajectory that is close enough - trj4 = trajectory() + # Create a fourth Trajectory that is close enough + trj4 = Trajectory() trj4.x = self.trj.x + 1 trj4.y = self.trj.y + 1 - trj4.x_v = self.trj.x_v - trj4.y_v = self.trj.y_v + trj4.vx = self.trj.vx + trj4.vy = self.trj.vy - # Compute the stacked science from a single trajectory. + # Compute the stacked science from a single Trajectory. all_valid_vect = [(self.all_valid) for i in range(4)] - meanStamps = self.search.coadded_stamps( + meanStamps = self.search.get_coadded_stamps( [self.trj, trj2, trj3, trj4], all_valid_vect, self.params, True ) diff --git a/tests/test_search_encode.py b/tests/test_search_encode.py index 4e1f0941c..10ebe54e5 100644 --- a/tests/test_search_encode.py +++ b/tests/test_search_encode.py @@ -19,21 +19,21 @@ def setUp(self): self.dim_y = 110 self.noise_level = 4.0 self.variance = self.noise_level**2 - self.p = psf(1.0) + self.p = PSF(1.0) # object properties self.object_flux = 250.0 self.start_x = 33 self.start_y = 5 - self.x_vel = 12.0 - self.y_vel = 19.0 + self.vxel = 12.0 + self.vyel = 19.0 - # create a trajectory for the object - self.trj = trajectory() + # create a Trajectory for the object + self.trj = Trajectory() self.trj.x = self.start_x self.trj.y = self.start_y - self.trj.x_v = self.x_vel - self.trj.y_v = self.y_vel + self.trj.vx = self.vxel + self.trj.vy = self.vyel # search parameters self.angle_steps = 150 @@ -52,21 +52,22 @@ def setUp(self): self.imlist = [] for i in range(self.imCount): time = i / self.imCount - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) add_fake_object( im, - self.start_x + time * self.x_vel + 0.5, - self.start_y + time * self.y_vel + 0.5, + self.start_x + time * self.vxel + 0.5, + self.start_y + time * self.vyel + 0.5, self.object_flux, self.p, ) self.imlist.append(im) - self.stack = image_stack(self.imlist) + self.stack = ImageStack(self.imlist) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_two_bytes(self): - search = stack_search(self.stack) + search = StackSearch(self.stack) search.enable_gpu_encoding(2, 2) search.search( self.angle_steps, @@ -82,12 +83,13 @@ def test_two_bytes(self): best = results[0] self.assertAlmostEqual(best.x, self.start_x, delta=self.pixel_error) self.assertAlmostEqual(best.y, self.start_y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / self.x_vel, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / self.y_vel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / self.vxel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / self.vyel, 1, delta=self.velocity_error) self.assertAlmostEqual(best.flux / self.object_flux, 1, delta=self.flux_error) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_one_byte(self): - search = stack_search(self.stack) + search = StackSearch(self.stack) search.enable_gpu_encoding(1, 1) search.search( self.angle_steps, @@ -103,12 +105,13 @@ def test_one_byte(self): best = results[0] self.assertAlmostEqual(best.x, self.start_x, delta=self.pixel_error) self.assertAlmostEqual(best.y, self.start_y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / self.x_vel, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / self.y_vel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / self.vxel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / self.vyel, 1, delta=self.velocity_error) self.assertAlmostEqual(best.flux / self.object_flux, 1, delta=self.flux_error) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_different_encodings(self): - search = stack_search(self.stack) + search = StackSearch(self.stack) # Encode phi to 2 bytes, but leave psi as a 4 byte float. search.enable_gpu_encoding(-1, 2) @@ -126,8 +129,8 @@ def test_different_encodings(self): best = results[0] self.assertAlmostEqual(best.x, self.start_x, delta=self.pixel_error) self.assertAlmostEqual(best.y, self.start_y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / self.x_vel, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / self.y_vel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / self.vxel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / self.vyel, 1, delta=self.velocity_error) self.assertAlmostEqual(best.flux / self.object_flux, 1, delta=self.flux_error) diff --git a/tests/test_search_filter.py b/tests/test_search_filter.py index 1c8eb9aca..62e4da1dd 100644 --- a/tests/test_search_filter.py +++ b/tests/test_search_filter.py @@ -19,21 +19,21 @@ def setUp(self): self.dim_y = 60 self.noise_level = 4.0 self.variance = self.noise_level**2 - self.p = psf(1.0) + self.p = PSF(1.0) # object properties self.object_flux = 250.0 self.start_x = 17 self.start_y = 12 - self.x_vel = 21.0 - self.y_vel = 16.0 + self.vxel = 21.0 + self.vyel = 16.0 - # create a trajectory for the object - self.trj = trajectory() + # create a Trajectory for the object + self.trj = Trajectory() self.trj.x = self.start_x self.trj.y = self.start_y - self.trj.x_v = self.x_vel - self.trj.y_v = self.y_vel + self.trj.vx = self.vxel + self.trj.vy = self.vyel # search parameters self.angle_steps = 150 @@ -52,20 +52,20 @@ def setUp(self): self.imlist = [] for i in range(self.imCount): time = i / self.imCount - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) add_fake_object( im, - self.start_x + time * self.x_vel + 0.5, - self.start_y + time * self.y_vel + 0.5, + self.start_x + time * self.vxel + 0.5, + self.start_y + time * self.vyel + 0.5, self.object_flux, self.p, ) self.imlist.append(im) - self.stack = image_stack(self.imlist) + self.stack = ImageStack(self.imlist) - self.search = stack_search(self.stack) + self.search = StackSearch(self.stack) self.search.enable_gpu_sigmag_filter(self.sigmaG_lims, self.sigmaG_coeff, self.lh_level) self.search.search( self.angle_steps, @@ -77,13 +77,14 @@ def setUp(self): int(self.imCount / 2), ) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results(self): results = self.search.get_results(0, 10) best = results[0] self.assertAlmostEqual(best.x, self.start_x, delta=self.pixel_error) self.assertAlmostEqual(best.y, self.start_y, delta=self.pixel_error) - self.assertAlmostEqual(best.x_v / self.x_vel, 1, delta=self.velocity_error) - self.assertAlmostEqual(best.y_v / self.y_vel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vx / self.vxel, 1, delta=self.velocity_error) + self.assertAlmostEqual(best.vy / self.vyel, 1, delta=self.velocity_error) self.assertAlmostEqual(best.flux / self.object_flux, 1, delta=self.flux_error) diff --git a/tests/test_stamp_filters.py b/tests/test_stamp_filters.py index 8f9a68096..ee2b2f2bc 100644 --- a/tests/test_stamp_filters.py +++ b/tests/test_stamp_filters.py @@ -11,7 +11,7 @@ def setUp(self): self.num_times = len(self.times) def _create_row(self, stamp): - row = ResultRow(trajectory(), self.num_times) + row = ResultRow(Trajectory(), self.num_times) row.stamp = np.array(stamp.get_all_pixels()) return row @@ -20,16 +20,16 @@ def test_peak_filtering_name(self): def test_skip_invalid_stamp(self): # No stamp - row = ResultRow(trajectory(), self.num_times) + row = ResultRow(Trajectory(), self.num_times) self.assertFalse(StampPeakFilter(5, 100, 100).keep_row(row)) # Wrong sized stamp - stamp = raw_image(5, 5) + stamp = RawImage(5, 5) row = self._create_row(stamp) self.assertFalse(StampPeakFilter(5, 100, 100).keep_row(row)) def test_peak_filtering(self): - stamp = raw_image(11, 11) + stamp = RawImage(11, 11) stamp.set_all(1.0) stamp.set_pixel(3, 4, 10.0) row = self._create_row(stamp) @@ -41,7 +41,7 @@ def test_peak_filtering(self): self.assertFalse(StampPeakFilter(5, 2, 2).keep_row(row)) def test_peak_filtering_furthest(self): - stamp = raw_image(9, 9) + stamp = RawImage(9, 9) stamp.set_all(1.0) stamp.set_pixel(3, 4, 10.0) stamp.set_pixel(8, 1, 10.0) # Use furthest from center. @@ -61,7 +61,7 @@ def test_central_moments_filtering_name(self): ) def test_central_moments_filtering(self): - stamp = raw_image(11, 11) + stamp = RawImage(11, 11) stamp.set_all(0.0) row = self._create_row(stamp) @@ -97,7 +97,7 @@ def test_center_filtering_name(self): ) def test_center_filtering(self): - stamp = raw_image(7, 7) + stamp = RawImage(7, 7) stamp.set_all(0.0) row = self._create_row(stamp) @@ -105,13 +105,13 @@ def test_center_filtering(self): self.assertFalse(StampCenterFilter(3, True, 0.01).keep_row(row)) # No local maxima (should fail). - stamp = raw_image(7, 7) + stamp = RawImage(7, 7) stamp.set_all(0.01) row = self._create_row(stamp) self.assertFalse(StampCenterFilter(3, True, 0.01).keep_row(row)) # Single strong central peak - stamp = raw_image(7, 7) + stamp = RawImage(7, 7) stamp.set_all(0.05) stamp.set_pixel(3, 3, 10.0) row = self._create_row(stamp) @@ -119,7 +119,7 @@ def test_center_filtering(self): self.assertFalse(StampCenterFilter(3, True, 1.0).keep_row(row)) # Less than 50% in the center pixel. - stamp = raw_image(7, 7) + stamp = RawImage(7, 7) stamp.set_all(0.05) stamp.set_pixel(3, 3, 10.0) stamp.set_pixel(3, 4, 9.0) @@ -128,7 +128,7 @@ def test_center_filtering(self): self.assertTrue(StampCenterFilter(3, True, 0.4).keep_row(row)) # Center is not the maximum value. - stamp = raw_image(7, 7) + stamp = RawImage(7, 7) stamp.set_all(0.05) stamp.set_pixel(1, 2, 10.0) stamp.set_pixel(3, 3, 9.0) diff --git a/tests/test_stamp_parity.py b/tests/test_stamp_parity.py index a3d245880..e1286ec28 100644 --- a/tests/test_stamp_parity.py +++ b/tests/test_stamp_parity.py @@ -25,21 +25,21 @@ def setUp(self): self.dim_y = 60 self.noise_level = 4.0 self.variance = self.noise_level**2 - self.p = psf(1.0) + self.p = PSF(1.0) # object properties self.object_flux = 250.0 self.start_x = 17 self.start_y = 12 - self.x_vel = 21.0 - self.y_vel = 16.0 + self.vxel = 21.0 + self.vyel = 16.0 - # create a trajectory for the object - self.trj = trajectory() + # create a Trajectory for the object + self.trj = Trajectory() self.trj.x = self.start_x self.trj.y = self.start_y - self.trj.x_v = self.x_vel - self.trj.y_v = self.y_vel + self.trj.vx = self.vxel + self.trj.vy = self.vyel # search parameters self.angle_steps = 150 @@ -57,13 +57,13 @@ def setUp(self): self.imlist = [] for i in range(self.imCount): time = i / self.imCount - im = layered_image( + im = LayeredImage( str(i), self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i ) add_fake_object( im, - self.start_x + time * self.x_vel + 0.5, - self.start_y + time * self.y_vel + 0.5, + self.start_x + time * self.vxel + 0.5, + self.start_y + time * self.vyel + 0.5, self.object_flux, self.p, ) @@ -76,14 +76,14 @@ def setUp(self): im.apply_mask_flags(1, []) self.imlist.append(im) - self.stack = image_stack(self.imlist) - self.search = stack_search(self.stack) + self.stack = ImageStack(self.imlist) + self.search = StackSearch(self.stack) @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_gpu_parity(self): radius = 2 width = 2 * radius + 1 - params = stamp_parameters() + params = StampParameters() params.radius = radius params.do_filtering = False @@ -97,11 +97,11 @@ def test_coadd_gpu_parity(self): # Check the summed stamps. Note summed stamp does not use goodIdx. params.stamp_type = StampType.STAMP_SUM stamps_old = [ - self.search.summed_sci_stamp(self.trj, radius, all_valid), - self.search.summed_sci_stamp(self.trj, radius, all_valid), + self.search.get_summed_stamp(self.trj, radius, all_valid), + self.search.get_summed_stamp(self.trj, radius, all_valid), ] - stamps_gpu = self.search.coadded_stamps(results, [all_valid, all_valid], params, True) - stamps_cpu = self.search.coadded_stamps(results, [all_valid, all_valid], params, False) + stamps_gpu = self.search.get_coadded_stamps(results, [all_valid, all_valid], params, True) + stamps_cpu = self.search.get_coadded_stamps(results, [all_valid, all_valid], params, False) for r in range(2): self.assertTrue(stamps_old[r].approx_equal(stamps_gpu[r], 1e-5)) self.assertTrue(stamps_old[r].approx_equal(stamps_cpu[r], 1e-5)) @@ -109,11 +109,11 @@ def test_coadd_gpu_parity(self): # Check the mean stamps. params.stamp_type = StampType.STAMP_MEAN stamps_old = [ - self.search.mean_sci_stamp(self.trj, radius, goodIdx[0]), - self.search.mean_sci_stamp(self.trj, radius, goodIdx[1]), + self.search.get_mean_stamp(self.trj, radius, goodIdx[0]), + self.search.get_mean_stamp(self.trj, radius, goodIdx[1]), ] - stamps_gpu = self.search.coadded_stamps(results, goodIdx, params, True) - stamps_cpu = self.search.coadded_stamps(results, goodIdx, params, False) + stamps_gpu = self.search.get_coadded_stamps(results, goodIdx, params, True) + stamps_cpu = self.search.get_coadded_stamps(results, goodIdx, params, False) for r in range(2): self.assertTrue(stamps_old[r].approx_equal(stamps_gpu[r], 1e-5)) self.assertTrue(stamps_old[r].approx_equal(stamps_cpu[r], 1e-5)) @@ -121,11 +121,11 @@ def test_coadd_gpu_parity(self): # Check the median stamps. params.stamp_type = StampType.STAMP_MEDIAN stamps_old = [ - self.search.median_sci_stamp(self.trj, radius, goodIdx[0]), - self.search.median_sci_stamp(self.trj, radius, goodIdx[1]), + self.search.get_median_stamp(self.trj, radius, goodIdx[0]), + self.search.get_median_stamp(self.trj, radius, goodIdx[1]), ] - stamps_gpu = self.search.coadded_stamps(results, goodIdx, params, True) - stamps_cpu = self.search.coadded_stamps(results, goodIdx, params, False) + stamps_gpu = self.search.get_coadded_stamps(results, goodIdx, params, True) + stamps_cpu = self.search.get_coadded_stamps(results, goodIdx, params, False) for r in range(2): self.assertTrue(stamps_old[r].approx_equal(stamps_gpu[r], 1e-5)) self.assertTrue(stamps_old[r].approx_equal(stamps_cpu[r], 1e-5)) diff --git a/tests/test_stats_filters.py b/tests/test_stats_filters.py index 0f1b7665b..83348a1ec 100644 --- a/tests/test_stats_filters.py +++ b/tests/test_stats_filters.py @@ -12,7 +12,7 @@ def setUp(self): self.rs = ResultList(self.times, track_filtered=True) for i in range(10): - t = trajectory() + t = Trajectory() row = ResultRow(t, self.num_times) row.filter_indices([k for k in range(i)]) row.final_likelihood = float(i) @@ -67,7 +67,7 @@ def test_filter_likelihood_mp(self): # Create a lot more results. rs = ResultList(self.times, track_filtered=True) for i in range(1000): - row = ResultRow(trajectory(), self.num_times) + row = ResultRow(Trajectory(), self.num_times) row.final_likelihood = 0.01 * float(i) rs.append_result(row) self.assertEqual(rs.num_results(), 1000)