diff --git a/docs/source/user_manual/search_params.rst b/docs/source/user_manual/search_params.rst index fd640aeb3..1077bd3a1 100644 --- a/docs/source/user_manual/search_params.rst +++ b/docs/source/user_manual/search_params.rst @@ -8,14 +8,6 @@ This document serves to provide a quick overview of the existing parameters and +------------------------+-----------------------------+----------------------------------------+ | **Parameter** | **Default Value** | **Interpretation** | +------------------------+-----------------------------+----------------------------------------+ -| ``ang_arr`` | [np.pi/15, np.pi/15, 128] | Minimum, maximum and number of angles | -| | | to search through (in radians) | -+------------------------+-----------------------------+----------------------------------------+ -| ``average_angle`` | None | Overrides the ecliptic angle | -| | | calculation and instead centers the | -| | | average search around average_angle | -| | | (in radians). | -+------------------------+-----------------------------+----------------------------------------+ | ``center_thresh`` | 0.00 | The minimum fraction of total flux | | | | within a stamp that must be contained | | | | in the central pixel | @@ -149,17 +141,12 @@ This document serves to provide a quick overview of the existing parameters and | | | if: | | | | * ``sum`` - (default) Per pixel sum | | | | * ``median`` - A per pixel median | -| | | * ``mean`` - A per pixel mean |\ +| | | * ``mean`` - A per pixel mean | +------------------------+-----------------------------+----------------------------------------+ | ``track_filtered`` | False | A Boolean indicating whether to track | | | | the filtered trajectories. Warning | | | | can use a lot of memory. | +------------------------+-----------------------------+----------------------------------------+ -| ``v_arr`` | [92.0, 526.0, 256] | Minimum, maximum and number of | -| | | velocities to search through. The | -| | | minimum and maximum velocities are | -| | | specified in pixels per day. | -+------------------------+-----------------------------+----------------------------------------+ | ``x_pixel_bounds`` | None | A length two list giving the starting | | | | and ending x pixels to use for the | | | | search. `None` uses the image bounds. | diff --git a/docs/source/user_manual/search_space.rst b/docs/source/user_manual/search_space.rst index 939f1a1cd..218aee193 100644 --- a/docs/source/user_manual/search_space.rst +++ b/docs/source/user_manual/search_space.rst @@ -33,7 +33,8 @@ Choosing Velocities Perhaps the most complex aspect of the KBMOD algorithm is how it defines the grid of search velocities. KBMOD allows you to define custom search strategies to best match the data. These include: * ``SingleVelocitySearch`` - A single predefined x and y velocity * ``VelocityGridSearch`` - An evenly spaced grid of x and y velocities -* ``KBMODV1SearchConfig`` - An evenly spaced grid of velocity magnitudes and angles (this was the only option in v1.1 and before) +* ``EclipticCenteredSearch`` - An evenly spaced grid of velocity magnitudes and angles (using a current parameterization) centered on a given or computed ecliptic angle. +* ``KBMODV1SearchConfig`` - An evenly spaced grid of velocity magnitudes and angles (using the legacy parameterization). * ``RandomVelocitySearch`` - Randomly sampled x and y velocities Additional search strategies can be defined by overriding the ``TrajectoryGenerator`` class in trajectory_generator.py. @@ -77,6 +78,52 @@ The ``VelocityGridSearch`` strategy searches a uniform grid of x and y velocitie | ``max_vy`` | The maximum velocity in the y-dimension (pixels per day). | +------------------------+-----------------------------------------------------------+ +EclipticCenteredSearch +---------------------- + +The grid is defined by two sets of parameters: a sampling of absolute velocities in pixels per day and a sampling of the velocities' angles in degrees or radians. Each sampling consists of values defining the range and number of sampling steps. + +Given the linear sampling for both velocities and angles, the full set of candidate trajectories is computed as:: + + + for (int a = 0; a < angleSteps; ++a) { + for (int v = 0; v < velocitySteps; ++v) { + searchList[a * velocitySteps + v].xVel = cos(sampled_angles[a]) * sampled_velocities[v]; + searchList[a * velocitySteps + v].yVel = sin(sampled_angles[a]) * sampled_velocities[v]; + } + } + +where ``sampled_angles`` contains the list of angles to test and ``sampled_velocities`` contains the list of velocities. + +The list of velocities is created from the given bounds list ``velocities=[min_vel, max_vel, vel_steps]``. The range is inclusive of both bounds. + +Each angle in the list is computed as an **offset** from the ecliptic angle. KBMOD uses the following ordering for extracting the ecliptic. +1. If ``given_ecliptic`` is provided (is not ``None``) in the generator’s configuration that value is used directly. +2. If the first image has a WCS, the ecliptic is estimated from that WCS. +3. A default ecliptic of 0.0 is used. +The angles used are defined from the list ``angles=[min_offset, max_offset, angle_steps]`` and will span ``[ecliptic + min_offset, ecliptic + max_offset]`` inclusive of both bounds. Angles can be specified in degrees or radians (as noted by the ``angle_units`` parameter) but must be consistent among all angles. + + ++------------------------+------------------------------------------------------+ +| **Parameter** | **Interpretation** | ++------------------------+------------------------------------------------------+ +| ``angles`` | A length 3 list with the minimum angle offset, | +| | the maximum offset, and the number of angles to | +| | to search through (angles specified in units given | +| | by ``angle_units``). | ++------------------------+------------------------------------------------------+ +| ``angle_units`` | The units to use for angles, such as "rad" or "deg". | ++------------------------+------------------------------------------------------+ +| ``given_ecliptic`` | The given value of the ecliptic angle (specified in | +| | units given by ``angle_units``). | ++------------------------+------------------------------------------------------+ +| ``velocities`` | A length 3 list with the minimum velocity (in | +| | pixels per day), the maximum velocity (in pixels | +| | per day), and number of velocities to test. | ++------------------------+------------------------------------------------------+ +| ``velocity_units`` | The units to use for velocities (e.g. "pix / d") | ++------------------------+------------------------------------------------------+ + KBMODV1SearchConfig ------------------- diff --git a/notebooks/KBMOD_Demo.ipynb b/notebooks/KBMOD_Demo.ipynb index c21cdea39..6797f9917 100644 --- a/notebooks/KBMOD_Demo.ipynb +++ b/notebooks/KBMOD_Demo.ipynb @@ -122,9 +122,15 @@ "ang_arr = [ang_below, ang_above, ang_steps]\n", "\n", "input_parameters = {\n", - " # Grid search parameters\n", - " \"v_arr\": v_arr,\n", - " \"ang_arr\": ang_arr,\n", + " # Use search parameters (including a force ecliptic angle of 0.0)\n", + " # to match what we know is in the demo data.\n", + " \"generator_config\": {\n", + " \"name\": \"EclipticCenteredSearch\",\n", + " \"angles\": [-0.5, 0.5, 11],\n", + " \"velocities\": [0.0, 20.0, 21],\n", + " \"angle_units\": \"radian\",\n", + " \"force_ecliptic\": 0.0,\n", + " },\n", " # Output parameters\n", " \"res_filepath\": res_filepath,\n", " \"output_suffix\": results_suffix,\n", @@ -142,9 +148,6 @@ " # Some basic stamp filtering limits.\n", " \"mom_lims\": [37.5, 37.5, 1.5, 1.0, 1.0],\n", " \"peak_offset\": [3.0, 3.0],\n", - " # Override the ecliptic angle for the demo data since we\n", - " # know the true angle in pixel space.\n", - " \"average_angle\": 0.0,\n", "}\n", "config = SearchConfiguration.from_dict(input_parameters)" ] diff --git a/notebooks/create_fake_data.ipynb b/notebooks/create_fake_data.ipynb index 489789c42..3835c31ae 100644 --- a/notebooks/create_fake_data.ipynb +++ b/notebooks/create_fake_data.ipynb @@ -149,9 +149,13 @@ "\n", "settings = {\n", " # Override the search data to match the known object.\n", - " \"average_angle\": 0.0,\n", - " \"v_arr\": [0, 20, 20],\n", - " \"ang_arr\": [0.5, 0.5, 10],\n", + " \"generator_config\": {\n", + " \"name\": \"EclipticCenteredSearch\",\n", + " \"angles\": [-0.5, 0.5, 11],\n", + " \"velocities\": [0.0, 20.0, 21],\n", + " \"angle_units\": \"radian\",\n", + " \"force_ecliptic\": 0.0,\n", + " },\n", " # Loosen the other filtering parameters.\n", " \"clip_negative\": True,\n", " \"sigmaG_lims\": [15, 60],\n", diff --git a/src/kbmod/configuration.py b/src/kbmod/configuration.py index 7ce545589..b678593c5 100644 --- a/src/kbmod/configuration.py +++ b/src/kbmod/configuration.py @@ -19,8 +19,6 @@ def __init__(self): self._required_params = set() self._params = { - "ang_arr": [math.pi / 15, math.pi / 15, 128], - "average_angle": None, "center_thresh": 0.00, "chunk_size": 500000, "clip_negative": False, @@ -33,7 +31,14 @@ def __init__(self): "do_mask": True, "do_stamp_filter": True, "encode_num_bytes": -1, - "generator_config": None, + "generator_config": { + "name": "EclipticCenteredSearch", + "velocity": [92.0, 526.0, 257], + "angles": [-math.pi / 15, math.pi / 15, 129], + "angle_units": "radian", + "velocity_units": "pix / d", + "given_ecliptic": None, + }, "gpu_filter": False, "ind_output_files": True, "im_filepath": None, @@ -52,7 +57,6 @@ def __init__(self): "stamp_radius": 10, "stamp_type": "sum", "track_filtered": False, - "v_arr": [92.0, 526.0, 256], "x_pixel_bounds": None, "x_pixel_buffer": None, "y_pixel_bounds": None, diff --git a/src/kbmod/fake_data/demo_helper.py b/src/kbmod/fake_data/demo_helper.py index 62c7f9fc5..57ff09d8c 100644 --- a/src/kbmod/fake_data/demo_helper.py +++ b/src/kbmod/fake_data/demo_helper.py @@ -38,9 +38,13 @@ def make_demo_data(filename=None): # Create configuraiton settings that match the object inserted. settings = { # Override the search data to match the known object. - "average_angle": 0.0, - "v_arr": [0, 20, 20], - "ang_arr": [0.5, 0.5, 10], + "generator_config": { + "name": "EclipticCenteredSearch", + "velocities": [0, 20.0, 21], + "angles": [-0.5, 0.5, 11], + "angle_units": "radian", + "given_ecliptic": 0.0, + }, # Loosen the other filtering parameters. "clip_negative": True, "sigmaG_lims": [15, 60], diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py index b5bbd6111..b7f2cd65e 100644 --- a/src/kbmod/run_search.py +++ b/src/kbmod/run_search.py @@ -11,8 +11,7 @@ from .filters.stamp_filters import append_all_stamps, append_coadds, get_coadds_and_filter_results from .results import Results -from .trajectory_generator import create_trajectory_generator, KBMODV1SearchConfig -from .wcs_utils import calc_ecliptic_angle +from .trajectory_generator import create_trajectory_generator from .work_unit import WorkUnit @@ -200,7 +199,7 @@ def run_search(self, config, stack, trj_generator=None): The stack before the masks have been applied. Modified in-place. trj_generator : `TrajectoryGenerator`, optional The object to generate the candidate trajectories for each pixel. - If None uses the default KBMODv1 grid search + If None uses the default EclipticCenteredSearch Returns ------- @@ -219,7 +218,7 @@ def run_search(self, config, stack, trj_generator=None): # Perform the actual search. if trj_generator is None: - trj_generator = create_trajectory_generator(config) + trj_generator = create_trajectory_generator(config, work_unit=None) keep = self.do_gpu_search(config, stack, trj_generator) if config["do_stamp_filter"]: @@ -288,14 +287,7 @@ def run_search_from_work_unit(self, work): keep : `Results` The results. """ - # Set the average angle if it is not set. - if work.config["average_angle"] is None: - center_pixel = (work.im_stack.get_width() / 2, work.im_stack.get_height() / 2) - if work.get_wcs(0) is not None: - work.config.set("average_angle", calc_ecliptic_angle(work.get_wcs(0), center_pixel)) - else: - logger.warning("Average angle not set and no WCS provided. Setting average_angle=0.0") - work.config.set("average_angle", 0.0) + trj_generator = create_trajectory_generator(work.config, work_unit=work) # Run the search. return self.run_search(work.config, work.im_stack) diff --git a/src/kbmod/trajectory_generator.py b/src/kbmod/trajectory_generator.py index 604874770..9b0ede75e 100644 --- a/src/kbmod/trajectory_generator.py +++ b/src/kbmod/trajectory_generator.py @@ -1,15 +1,21 @@ import abc +import copy import logging import math import random +import astropy.units as u + from astropy.table import Table from kbmod.configuration import SearchConfiguration from kbmod.search import Trajectory -def create_trajectory_generator(config): +logger = logging.getLogger(__name__) + + +def create_trajectory_generator(config, work_unit=None, **kwargs): """Create a TrajectoryGenerator object given a dictionary of configuration options. The generator class is specified by the 'name' entry, which must exist and match the class name of one @@ -19,6 +25,9 @@ def create_trajectory_generator(config): ---------- config : `dict` or `SearchConfiguration` The dictionary of generator parameters. + work_unit : `WorkUnit`, optional + A WorkUnit to provide additional information about the data that + can be used to derive parameters that depend on the input. Returns ------- @@ -33,16 +42,8 @@ def create_trajectory_generator(config): # Check if we are dealing with a top level configuration. if type(config) is SearchConfiguration: if config["generator_config"] is None: - # We are dealing with a legacy configuration file. - gen = KBMODV1SearchConfig( - v_arr=config["v_arr"], - ang_arr=config["ang_arr"], - average_angle=config["average_angle"], - ) - return gen - else: - # We are dealing with a top level configuration. - config = config["generator_config"] + raise ValueError("Missing generator_config parameter.") + config = config["generator_config"] if "name" not in config: raise KeyError("The trajectory generator configuration must contain a name field.") @@ -50,10 +51,14 @@ def create_trajectory_generator(config): name = config["name"] if name not in TrajectoryGenerator.generators: raise KeyError("Trajectory generator {name} is undefined.") - logger = logging.getLogger(__name__) logger.info(f"Creating trajectory generator of type {name}") - return TrajectoryGenerator.generators[name](**config) + # Add any keyword arguments to the params, overriding current values. + params = copy.deepcopy(config) + params.update(kwargs) + logger.debug(str(params)) + + return TrajectoryGenerator.generators[name](**params, work_unit=work_unit) class TrajectoryGenerator(abc.ABC): @@ -67,7 +72,7 @@ class TrajectoryGenerator(abc.ABC): generators = {} # A mapping of class name to class object for subclasses. - def __init__(self, *args, **kwargs): + def __init__(self, work_unit=None, **kwargs): pass def __init_subclass__(cls, **kwargs): @@ -131,7 +136,7 @@ def to_table(self): class SingleVelocitySearch(TrajectoryGenerator): """Search a single velocity from each pixel.""" - def __init__(self, vx, vy, *args, **kwargs): + def __init__(self, vx, vy, **kwargs): """Create a class SingleVelocitySearch. Parameters @@ -141,7 +146,7 @@ def __init__(self, vx, vy, *args, **kwargs): vy : `float` The velocity in y pixels (pixels per day). """ - super().__init__(*args, **kwargs) + super().__init__(**kwargs) self.vx = vx self.vy = vy @@ -165,7 +170,7 @@ def generate(self, *args, **kwargs): class VelocityGridSearch(TrajectoryGenerator): """Search a grid defined by steps in velocity space.""" - def __init__(self, vx_steps, min_vx, max_vx, vy_steps, min_vy, max_vy, *args, **kwargs): + def __init__(self, vx_steps, min_vx, max_vx, vy_steps, min_vy, max_vy, **kwargs): """Create a class VelocityGridSearch. Parameters @@ -183,7 +188,7 @@ def __init__(self, vx_steps, min_vx, max_vx, vy_steps, min_vy, max_vy, *args, ** max_vy : `float` The maximum velocity in the y-dimension (pixels per day). """ - super().__init__(*args, **kwargs) + super().__init__(**kwargs) if vx_steps < 2 or vy_steps < 2: raise ValueError("VelocityGridSearch requires at least 2 steps in each dimension") @@ -232,7 +237,7 @@ def generate(self, *args, **kwargs): class KBMODV1Search(TrajectoryGenerator): """Search a grid defined by velocities and angles.""" - def __init__(self, vel_steps, min_vel, max_vel, ang_steps, min_ang, max_ang, *args, **kwargs): + def __init__(self, vel_steps, min_vel, max_vel, ang_steps, min_ang, max_ang, **kwargs): """Create a class KBMODV1Search. Parameters @@ -250,7 +255,7 @@ def __init__(self, vel_steps, min_vel, max_vel, ang_steps, min_ang, max_ang, *ar max_ang : `float` The maximum angle (in radians) """ - super().__init__(*args, **kwargs) + super().__init__(**kwargs) if vel_steps < 1 or ang_steps < 1: raise ValueError("KBMODV1Search requires at least 1 step in each dimension") @@ -301,9 +306,9 @@ def generate(self, *args, **kwargs): class KBMODV1SearchConfig(KBMODV1Search): - """Search a grid defined by velocities and angles in the format of the configuration file.""" + """Search a grid defined by velocities and angles in the format of the legacy configuration file.""" - def __init__(self, v_arr, ang_arr, average_angle, *args, **kwargs): + def __init__(self, v_arr, ang_arr, average_angle=None, work_unit=None, **kwargs): """Create a class KBMODV1SearchConfig. Parameters @@ -314,25 +319,157 @@ def __init__(self, v_arr, ang_arr, average_angle, *args, **kwargs): ang_arr : `list` A triplet of the minimum angle offset (in radians), and the maximum angle offset (in radians), and the number of angles to try. - average_angle : `float` + average_angle : `float`, optional The central angle to search around. Should align with the ecliptic in most cases. + work_unit : `WorkUnit`, optional + A WorkUnit to provide additional information about the data that + can be used to derive parameters that depend on the input. """ if len(v_arr) != 3: raise ValueError("KBMODV1SearchConfig requires v_arr to be length 3") if len(ang_arr) != 3: raise ValueError("KBMODV1SearchConfig requires ang_arr to be length 3") if average_angle is None: - raise ValueError("KBMODV1SearchConfig requires a valid average_angle.") + if work_unit is None: + raise ValueError( + "KBMODV1SearchConfig requires a valid average_angle or a WorkUnit with a WCS." + ) + average_angle = work_unit.compute_ecliptic_angle() ang_min = average_angle - ang_arr[0] ang_max = average_angle + ang_arr[1] - super().__init__(v_arr[2], v_arr[0], v_arr[1], ang_arr[2], ang_min, ang_max, *args, **kwargs) + super().__init__(v_arr[2], v_arr[0], v_arr[1], ang_arr[2], ang_min, ang_max, **kwargs) + + +class EclipticCenteredSearch(TrajectoryGenerator): + """Search a grid defined by velocities and angles relative to the ecliptic. + + Attributes + ---------- + ecliptic_angle : `float` + The angle to use for the ecliptic in the image (in the units defined in ``angle_units``). + velocities : `list` + A triplet of the minimum velocity to use (in pixels per day), and the maximum velocity + magnitude (in pixels per day), and the number of velocity steps to try. + angles : `list` + A triplet of the minimum angle offset (in radians), and the maximum angle offset (in radians), + and the number of angles to try. + min_ang : `float` + The minimum search angle in the image (ecliptic_angle + min_ang_offset) in radians. + max_ang : `float` + The maximum search angle in the image (ecliptic_angle + max_ang_offset) in radians. + """ + + def __init__( + self, + velocities=[0.0, 0.0, 0], + angles=[0.0, 0.0, 0], + angle_units="radian", + velocity_units="pix / d", + given_ecliptic=None, + work_unit=None, + **kwargs, + ): + """Create a class EclipticCenteredSearch. + + Parameters + ---------- + velocities : `list` + A triplet of the minimum velocity to use (in pixels per day), and the maximum velocity + magnitude (in pixels per day), and the number of velocity steps to try. + angles : `list` + A triplet of the minimum angle offset (in the units defined in ``angle_units``), the maximum + angle offset (in the units defined in ``angle_units``), and the number of angles to try. + angle_units : `str` + The units for the angle. + Default: 'radian' + velocity_units : `str` + The units for the angle. + Default: 'pix / d' + given_ecliptic : `float`, optional + An override for the ecliptic as given in the config (in the units defined in + ``angle_units``). This angle takes precedence over ``computed_ecliptic``. + work_unit : `WorkUnit`, optional + A WorkUnit to provide additional information about the data that + can be used to derive parameters that depend on the input. + """ + super().__init__(**kwargs) + ang_units = u.Unit(angle_units) + vel_units = u.Unit(velocity_units) + + if given_ecliptic is not None: + self.ecliptic_angle = (given_ecliptic * ang_units).to(u.rad).value + elif work_unit is not None: + # compute_ecliptic_angle() always produces radians. + self.ecliptic_angle = work_unit.compute_ecliptic_angle() + else: + logger.warning("No ecliptic angle provided. Using 0.0.") + self.ecliptic_angle = 0.0 + + if len(angles) != 3: + raise ValueError("Invalid angles parameter. Expected a length 3 list.") + if len(velocities) != 3: + raise ValueError("Invalid velocity parameter. Expected a length 3 list.") + if velocities[2] < 1 or angles[2] < 1: + raise ValueError("EclipticCenteredSearch requires at least 1 step in each dimension") + if velocities[1] < velocities[0]: + raise ValueError(f"Invalid EclipticCenteredSearch bounds: {velocities}") + + self.velocities = [ + (velocities[0] * vel_units).to(u.pixel / u.day).value, + (velocities[1] * vel_units).to(u.pixel / u.day).value, + velocities[2], + ] + self.vel_stepsize = (velocities[1] - velocities[0]) / float(velocities[2] - 1) + + # Compute the angle bounds and step size in radians. + self.angles = [ + (angles[0] * ang_units).to(u.rad).value, + (angles[1] * ang_units).to(u.rad).value, + angles[2], + ] + self.min_ang = self.ecliptic_angle + self.angles[0] + self.max_ang = self.ecliptic_angle + self.angles[1] + self.ang_stepsize = (self.max_ang - self.min_ang) / float(self.angles[2] - 1) + + def __repr__(self): + return ( + "EclipticSearch:" + f" v=[{self.velocities[0]}, {self.velocities[1]}], {self.velocities[2]}" + f" a=[{self.min_ang}, {self.max_ang}], {self.ang_steps}" + ) + + def __str__(self): + return f"""EclipticSearch: + Vel: [{self.velocities[0]}, {self.velocities[1]}] in {self.velocities[2]} steps. + Ang: + Ecliptic = {self.ecliptic_angle} + Offsets = {self.angles[0]} to {self.angles[1]} + [{self.min_ang}, {self.max_ang}] in {self.angles[2]} steps.""" + + def generate(self, *args, **kwargs): + """Produces a single candidate trajectory to test. + + Returns + ------- + candidate : `Trajectory` + A ``Trajectory`` to test at each pixel. + """ + for ang_i in range(self.angles[2]): + for vel_i in range(self.velocities[2]): + curr_ang = self.min_ang + ang_i * self.ang_stepsize + curr_vel = self.velocities[0] + vel_i * self.vel_stepsize + + vx = math.cos(curr_ang) * curr_vel + vy = math.sin(curr_ang) * curr_vel + + yield Trajectory(vx=vx, vy=vy) class RandomVelocitySearch(TrajectoryGenerator): """Search a grid defined by min/max bounds on pixel velocities.""" - def __init__(self, min_vx, max_vx, min_vy, max_vy, max_samples=1_000_000, *args, **kwargs): + def __init__(self, min_vx, max_vx, min_vy, max_vy, max_samples=1_000_000, **kwargs): """Create a class KBMODV1Search. Parameters @@ -349,7 +486,7 @@ def __init__(self, min_vx, max_vx, min_vy, max_vy, max_samples=1_000_000, *args, The maximum number of samples to generate. Used to avoid infinite loops in KBMOD code. """ - super().__init__(*args, **kwargs) + super().__init__(**kwargs) if max_vx < min_vx or max_vy < min_vy: raise ValueError("Invalid RandomVelocitySearch bounds.") if max_samples <= 0: diff --git a/src/kbmod/work_unit.py b/src/kbmod/work_unit.py index 7746a4663..bf9e932bb 100644 --- a/src/kbmod/work_unit.py +++ b/src/kbmod/work_unit.py @@ -16,6 +16,7 @@ from kbmod.search import ImageStack, LayeredImage, PSF, RawImage, Logging from kbmod.wcs_utils import ( append_wcs_to_hdu_header, + calc_ecliptic_angle, extract_wcs_from_hdu_header, wcs_fits_equal, wcs_from_dict, @@ -204,6 +205,23 @@ def get_wcs(self, img_num): return per_img + def compute_ecliptic_angle(self): + """Return the ecliptic angle (in radians in pixel space) derived from the + images and WCS. + + Returns + ------- + ecliptic_angle : `float` or `None` + The computed ecliptic_angle or ``None`` if data is missing. + """ + + wcs = self.get_wcs(0) + if wcs is None or self.im_stack is None: + logger.warning(f"A valid wcs and ImageStack is needed to compute the ecliptic angle.") + return None + center_pixel = (self.im_stack.get_width() / 2, self.im_stack.get_height() / 2) + return calc_ecliptic_angle(wcs, center_pixel) + def get_all_obstimes(self): """Return a list of the observation times. If the `WorkUnit` was lazily loaded, then the obstimes diff --git a/tests/test_configuration.py b/tests/test_configuration.py index 9a29969a5..e93a49f91 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -52,11 +52,10 @@ def test_from_hdu(self): ["Here3"], ["7"], ["null"], - [ - "[1, 2]", - ], + ["[1, 2]"], + ["{name: test_gen, p1: [1.0, 2.0], p2: 2.0}"], ], - names=("im_filepath", "num_obs", "cluster_type", "ang_arr"), + names=("im_filepath", "num_obs", "cluster_type", "ang_arr", "generator_config"), ) hdu = fits.table_to_hdu(t) @@ -64,6 +63,9 @@ def test_from_hdu(self): self.assertEqual(config["im_filepath"], "Here3") self.assertEqual(config["num_obs"], 7) self.assertEqual(config["ang_arr"], [1, 2]) + self.assertEqual(config["generator_config"]["name"], "test_gen") + self.assertEqual(config["generator_config"]["p1"], [1.0, 2.0]) + self.assertEqual(config["generator_config"]["p2"], 2.0) self.assertIsNone(config["cluster_type"]) def test_to_hdu(self): @@ -73,7 +75,8 @@ def test_to_hdu(self): "cluster_type": None, "do_clustering": False, "res_filepath": "There", - "ang_arr": [1.0, 2.0, 3.0], + "generator_config": {"name": "test_gen", "p1": [1.0, 2.0], "p2": 2.0}, + "basic_array": [1.0, 2.0, 3.0], } config = SearchConfiguration.from_dict(d) hdu = config.to_hdu() @@ -82,7 +85,8 @@ def test_to_hdu(self): self.assertEqual(hdu.data["num_obs"][0], "5\n...") self.assertEqual(hdu.data["cluster_type"][0], "null\n...") self.assertEqual(hdu.data["res_filepath"][0], "There\n...") - self.assertEqual(hdu.data["ang_arr"][0], "[1.0, 2.0, 3.0]") + self.assertEqual(hdu.data["generator_config"][0], "{name: test_gen, p1: [1.0, 2.0], p2: 2.0}") + self.assertEqual(hdu.data["basic_array"][0], "[1.0, 2.0, 3.0]") def test_to_yaml(self): d = { @@ -91,7 +95,7 @@ def test_to_yaml(self): "cluster_type": None, "do_clustering": False, "res_filepath": "There", - "ang_arr": [1.0, 2.0, 3.0], + "generator_config": {"name": "test_gen", "p1": [1.0, 2.0], "p2": 2.0}, } config = SearchConfiguration.from_dict(d) yaml_str = config.to_yaml() @@ -101,9 +105,9 @@ def test_to_yaml(self): self.assertEqual(yaml_dict["num_obs"], 5) self.assertEqual(yaml_dict["cluster_type"], None) self.assertEqual(yaml_dict["res_filepath"], "There") - self.assertEqual(yaml_dict["ang_arr"][0], 1.0) - self.assertEqual(yaml_dict["ang_arr"][1], 2.0) - self.assertEqual(yaml_dict["ang_arr"][2], 3.0) + self.assertEqual(yaml_dict["generator_config"]["name"], "test_gen") + self.assertEqual(yaml_dict["generator_config"]["p1"], [1.0, 2.0]) + self.assertEqual(yaml_dict["generator_config"]["p2"], 2.0) def test_save_and_load_yaml(self): config = SearchConfiguration() diff --git a/tests/test_regression_test.py b/tests/test_regression_test.py index 521a54e8c..7b1159bf1 100644 --- a/tests/test_regression_test.py +++ b/tests/test_regression_test.py @@ -238,32 +238,21 @@ def perform_search(im_stack, res_filename, default_psf): The default PSF value to use when nothing is provided in the PSF file. """ - v_min = 92.0 # Pixels/day - v_max = 550.0 - - # Manually set the average angle that will work with the (manually specified) tracks. - average_angle = 1.1901106654050821 - - # Offset by PI for prograde orbits in lori allen data - ang_below = -np.pi + np.pi / 10.0 # Angle below ecliptic - ang_above = np.pi + np.pi / 10.0 # Angle above ecliptic - v_steps = 51 - ang_steps = 25 - - v_arr = [v_min, v_max, v_steps] - ang_arr = [ang_below, ang_above, ang_steps] - num_obs = 15 - input_parameters = { "im_filepath": "./", "res_filepath": None, "result_filename": res_filename, "psf_val": default_psf, "output_suffix": "", - "v_arr": v_arr, - "average_angle": average_angle, - "ang_arr": ang_arr, - "num_obs": num_obs, + "generator_config": { + "name": "EclipticCenteredSearch", + # Offset by PI for prograde orbits in lori allen data + "angles": [np.pi - np.pi / 10.0, np.pi + np.pi / 10.0, 26], + "velocities": [92.0, 550.0, 52], + "angle_units": "radian", + "given_ecliptic": 1.1901106654050821, + }, + "num_obs": 15, "do_mask": True, "lh_level": 25.0, "sigmaG_lims": [25, 75], diff --git a/tests/test_trajectory_generator.py b/tests/test_trajectory_generator.py index ee8fcb59f..f08fbd9e1 100644 --- a/tests/test_trajectory_generator.py +++ b/tests/test_trajectory_generator.py @@ -1,14 +1,20 @@ +import numpy as np import unittest +from astropy.wcs import WCS + from kbmod.configuration import SearchConfiguration +from kbmod.fake_data.fake_data_creator import FakeDataSet from kbmod.trajectory_generator import ( KBMODV1Search, KBMODV1SearchConfig, + EclipticCenteredSearch, SingleVelocitySearch, RandomVelocitySearch, VelocityGridSearch, create_trajectory_generator, ) +from kbmod.work_unit import WorkUnit class test_trajectory_generator(unittest.TestCase): @@ -44,7 +50,7 @@ def test_VelocityGridSearch(self): self.assertRaises(ValueError, VelocityGridSearch, 3, 2.0, 0.0, 3, -0.25, 0.25) def test_KBMODV1Search(self): - # Note that KBMOD v1's search will never include the upper bound of angle or velocity. + # Note that KBMOD v1's legacy search will never include the upper bound of angle or velocity. gen = KBMODV1Search(3, 0.0, 3.0, 2, -0.25, 0.25) expected_x = [0.0, 0.9689, 1.9378, 0.0, 1.0, 2.0] expected_y = [0.0, -0.247, -0.4948, 0.0, 0.0, 0.0] @@ -62,15 +68,40 @@ def test_KBMODV1Search(self): self.assertAlmostEqual(tbl["vx"][i], expected_x[i], delta=0.001) self.assertAlmostEqual(tbl["vy"][i], expected_y[i], delta=0.001) - # Test invalid number of steps. + # Test invalid number of steps and ranges. self.assertRaises(ValueError, KBMODV1Search, 3, 0.0, 3.0, 0, -0.25, 0.25) self.assertRaises(ValueError, KBMODV1Search, 0, 0.0, 3.0, 2, -0.25, 0.25) self.assertRaises(ValueError, KBMODV1Search, 3, 0.0, 3.0, 2, 0.25, -0.25) self.assertRaises(ValueError, KBMODV1Search, 3, 3.5, 3.0, 2, -0.25, 0.25) + def test_EclipticCenteredSearch(self): + gen = EclipticCenteredSearch( + [0.0, 2.0, 3], [-45.0, 45.0, 3], angle_units="degree", given_ecliptic=0.0 + ) + expected_x = [0.0, 0.707107, 1.41421, 0.0, 1.0, 2.0, 0.0, 0.707107, 1.41421] + expected_y = [0.0, -0.707107, -1.41421, 0.0, 0.0, 0.0, 0.0, 0.707107, 1.41421] + + trjs = [trj for trj in gen] + self.assertEqual(len(trjs), 9) + for i in range(9): + self.assertAlmostEqual(trjs[i].vx, expected_x[i], delta=0.001) + self.assertAlmostEqual(trjs[i].vy, expected_y[i], delta=0.001) + + # Test that we get the correct results if we dump to a table. + tbl = gen.to_table() + self.assertEqual(len(tbl), 9) + for i in range(9): + self.assertAlmostEqual(tbl["vx"][i], expected_x[i], delta=0.001) + self.assertAlmostEqual(tbl["vy"][i], expected_y[i], delta=0.001) + + # Test invalid number of steps and ranges. + self.assertRaises(ValueError, EclipticCenteredSearch, [0.0, 3.0, 3], [-0.25, 0.25, 0]) + self.assertRaises(ValueError, EclipticCenteredSearch, [0.0, 3.0, 0], [-0.25, 0.25, 2]) + self.assertRaises(ValueError, EclipticCenteredSearch, [3.5, 3.0, 3], [-0.25, 0.25, 2]) + def test_KBMODV1SearchConfig(self): # Note that KBMOD v1's search will never include the upper bound of angle or velocity. - gen = KBMODV1SearchConfig([0.0, 3.0, 3], [0.25, 0.25, 2], 0.0) + gen = KBMODV1SearchConfig([0.0, 3.0, 3], [0.25, 0.25, 2], average_angle=0.0) expected_x = [0.0, 0.9689, 1.9378, 0.0, 1.0, 2.0] expected_y = [0.0, -0.247, -0.4948, 0.0, 0.0, 0.0] @@ -127,13 +158,58 @@ def test_create_trajectory_generator(self): self.assertEqual(gen2.vx, 1) self.assertEqual(gen2.vy, 2) + # Create a fake work unit with one image and a WCS with a non-zero ecliptic angle. + fake_wcs = WCS(naxis=2) + fake_wcs.wcs.crpix = [0.0, 0.0] + fake_wcs.wcs.cdelt = np.array([-0.1, 0.1]) + fake_wcs.wcs.crval = [0, -90] + fake_wcs.wcs.ctype = ["RA---TAN-SIP", "DEC--TAN-SIP"] + fake_wcs.wcs.crota = np.array([0.0, 0.0]) + + fake_data = FakeDataSet(10, 10, [0.0]) + base_config = SearchConfiguration() + fake_work = WorkUnit(im_stack=fake_data.stack, config=base_config, wcs=fake_wcs) + fake_ecliptic = fake_work.compute_ecliptic_angle() + self.assertGreater(fake_ecliptic, 1.0) + + # Test we can create a trajectory generator with optional keyword parameters. + config3 = { + "name": "EclipticCenteredSearch", + "angles": [0.0, 45.0, 2], + "velocities": [0.0, 1.0, 2], + "angle_units": "degree", + "given_ecliptic": None, + } + + # Without a given ecliptic, we use the WCS. + gen3 = create_trajectory_generator(config3, work_unit=fake_work) + self.assertTrue(type(gen3) is EclipticCenteredSearch) + self.assertAlmostEqual(gen3.ecliptic_angle, fake_ecliptic) + self.assertAlmostEqual(gen3.min_ang, fake_ecliptic) + self.assertAlmostEqual(gen3.max_ang, fake_ecliptic + np.pi / 4.0) + + # The given_ecliptic has priority over the fake WCS. + config3["given_ecliptic"] = 0.0 + gen4 = create_trajectory_generator(config3, work_unit=fake_work) + self.assertAlmostEqual(gen4.ecliptic_angle, 0.0) + self.assertEqual(gen4.min_ang, 0.0) + self.assertEqual(gen4.max_ang, np.pi / 4.0) + + # Fail with no name or a bad name. + self.assertRaises(KeyError, create_trajectory_generator, {}) + self.assertRaises(KeyError, create_trajectory_generator, {"name": "Invalid_generator"}) + def test_create_trajectory_generator_config(self): config = SearchConfiguration() - config.set("ang_arr", [0.5, 0.5, 30]) - config.set("average_angle", 0.0) - config.set("v_arr", [0.0, 10.0, 100]) + generator_config = { + "name": "KBMODV1SearchConfig", + "ang_arr": [0.5, 0.5, 30], + "average_angle": 0.0, + "v_arr": [0.0, 10.0, 100], + } + config.set("generator_config", generator_config) - # Without a "generator_config" option, we fall back on the legacy generator. + # We process the legacy configuration correctly. gen1 = create_trajectory_generator(config) self.assertTrue(type(gen1) is KBMODV1SearchConfig) self.assertEqual(gen1.vel_steps, 100) @@ -143,12 +219,9 @@ def test_create_trajectory_generator_config(self): self.assertEqual(gen1.min_ang, -0.5) self.assertEqual(gen1.max_ang, 0.5) - # Add a generator configuration. - config.set("generator_config", {"name": "SingleVelocitySearch", "vx": 1, "vy": 2}) - gen2 = create_trajectory_generator(config) - self.assertTrue(type(gen2) is SingleVelocitySearch) - self.assertEqual(gen2.vx, 1) - self.assertEqual(gen2.vy, 2) + # Fail if no generator configuration is provided. + config.set("generator_config", None) + self.assertRaises(ValueError, create_trajectory_generator, config) if __name__ == "__main__": diff --git a/tests/test_work_unit.py b/tests/test_work_unit.py index 14a7d1ba7..dbcf7b8ed 100644 --- a/tests/test_work_unit.py +++ b/tests/test_work_unit.py @@ -317,6 +317,15 @@ def test_save_and_load_fits_global_wcs(self): self.assertIsNotNone(work2.get_wcs(i)) self.assertTrue(wcs_fits_equal(work2.get_wcs(i), self.wcs)) + def test_get_ecliptic_angle(self): + """Check that we can compute an ecliptic angle.""" + work = WorkUnit(self.im_stack, self.config, self.wcs, None) + self.assertAlmostEqual(work.compute_ecliptic_angle(), -0.381541020495931) + + # If we do not have a WCS, we get None for the ecliptic angle. + work2 = WorkUnit(self.im_stack, self.config, None, None) + self.assertIsNone(work2.compute_ecliptic_angle()) + def test_image_positions_to_original_icrs_invalid_format(self): work = WorkUnit( im_stack=self.im_stack,