Merge pull request #466 from dirac-institute/test_cleanups

Consolidate test logic

README.md

@@ -94,9 +94,10 @@ import numpy as np
 # Create a point spread function
 psf = kb.PSF(1.5)
 
-# Create fake data with ten 512x512 pixel images.
+# Create fake data with ten 512x512 pixel images and starting at MJD of 57130.2.
 from kbmod.fake_data_creator import *
-ds = FakeDataSet(512, 512, 10)
+fake_times = create_fake_times(10, t0=57130.2)
+ds = FakeDataSet(512, 512, fake_times)
 imgs = ds.stack.get_images()
 
 # Get the timestamp of the first image.

src/kbmod/fake_data_creator.py

@@ -51,10 +51,49 @@ def add_fake_object(img, x, y, flux, psf=None):
                 sci.interpolated_add(float(initial_x + i), float(initial_y + j), flux * psf.get_value(i, j))
 
 
+def create_fake_times(num_times, t0=0.0, obs_per_day=1, intra_night_gap=0.01, inter_night_gap=1):
+    """Create a list of times based on a cluster of ``obs_per_day`` observations
+    each night spaced out ``intra_night_gap`` within a night and ``inter_night_gap`` data between
+    observing nights.
+
+    Parameters
+    ----------
+    num_times : `int`
+        The number of time steps (number of images).
+    t0 : `float`
+        The time stamp of the first observation [default=0.0]
+    obs_per_day : `int`
+        The number of observations on the same night [default=1]
+    intra_night_gap : `float`
+        The time (in days) between observations in the same night [default=0.01]
+    inter_night_gap : `int`
+        The number of days between observing runs [default=1]
+
+    Returns
+    -------
+    result_times : `list`
+        A list of ``num_times`` floats indicating the different time stamps.
+    """
+    if num_times <= 0:
+        raise ValueError(f"Invalid number of times {num_times}")
+
+    result_times = []
+    seen_on_day = 0  # Number seen so far on the current day
+    day_num = 0
+    for i in range(num_times):
+        result_times.append(t0 + day_num + seen_on_day * intra_night_gap)
+
+        seen_on_day += 1
+        if seen_on_day == obs_per_day:
+            seen_on_day = 0
+            day_num += inter_night_gap
+    return result_times
+
+
 class FakeDataSet:
     """This class creates fake data sets for testing and demo notebooks."""
 
-    def __init__(self, width, height, num_times, noise_level=2.0, psf_val=0.5, obs_per_day=3, use_seed=False):
+    def __init__(self, width, height, times, noise_level=2.0, psf_val=0.5, use_seed=False):
         """The constructor.
 
         Parameters
@@ -63,38 +102,23 @@ def __init__(self, width, height, num_times, noise_level=2.0, psf_val=0.5, obs_p
             The width of the images in pixels.
         height : `int`
             The height of the images in pixels.
-        num_times : `int`
-            The number of time steps (number of images).
+        times : `list`
+            A list of time stamps, such as produced by ``create_fake_times``.
         noise_level : `float`
             The level of the background noise.
         psf_val : `float`
             The value of the default PSF.
-        obs_per_day : `int`
-            The number of observations on the same night.
         use_seed : `bool`
             Use a deterministic seed to avoid flaky tests.
         """
         self.width = width
         self.height = height
         self.psf_val = psf_val
         self.noise_level = noise_level
-        self.num_times = num_times
+        self.num_times = len(times)
         self.use_seed = use_seed
         self.trajectories = []
-
-        # Generate times with multiple observations per night
-        # separated by ~15 minutes.
-        self.times = []
-        seen_on_day = 0
-        day_num = 0
-        for i in range(num_times):
-            t = 57130.2 + day_num + seen_on_day * 0.01
-            self.times.append(t)
-
-            seen_on_day += 1
-            if seen_on_day == obs_per_day:
-                seen_on_day = 0
-                day_num += 1
+        self.times = times
 
         # Make the image stack.
         self.stack = self.make_fake_ImageStack()

tests/test_data_interface.py

@@ -10,7 +10,7 @@
     load_input_from_file,
     load_input_from_individual_files,
 )
-from kbmod.fake_data_creator import FakeDataSet
+from kbmod.fake_data_creator import create_fake_times, FakeDataSet
 from kbmod.search import *
 from kbmod.work_unit import WorkUnit
 from utils.utils_for_tests import get_absolute_data_path
@@ -110,7 +110,8 @@ def test_file_load_workunit(self):
             }
         )
         fake_config = SearchConfiguration()
-        fake_data = FakeDataSet(64, 64, 11, obs_per_day=10, use_seed=True)
+        fake_times = create_fake_times(11, 57130.2, 10, 0.01, 1)
+        fake_data = FakeDataSet(64, 64, fake_times, use_seed=True)
         work = WorkUnit(fake_data.stack, fake_config, fake_wcs, None)
 
         with tempfile.TemporaryDirectory() as dir_name:

tests/test_end_to_end.py

@@ -7,6 +7,7 @@
 from kbmod.fake_data_creator import *
 from kbmod.run_search import *
 from kbmod.search import *
+from kbmod.trajectory_utils import make_trajectory
 from kbmod.work_unit import WorkUnit
 
 # from .utils_for_tests import get_absolute_demo_data_path
@@ -99,14 +100,11 @@ def test_demo_stamp_size(self):
     def test_e2e_work_unit(self):
         num_images = 10
 
-        # Create a fake data set with a single bright fake object.
-        ds = FakeDataSet(128, 128, num_images, obs_per_day=10, use_seed=True)
-        trj = Trajectory()
-        trj.x = 50
-        trj.y = 60
-        trj.vx = 5.0
-        trj.vy = 0.0
-        trj.flux = 500.0
+        # Create a fake data set with a single bright fake object and all
+        # the observations on a single day.
+        fake_times = create_fake_times(num_images, 57130.2, 10, 0.01, 1)
+        ds = FakeDataSet(128, 128, fake_times, use_seed=True)
+        trj = make_trajectory(x=50, y=60, vx=5.0, vy=0.0, flux=500.0)
         ds.insert_object(trj)
 
         # Set the configuration to pick up the fake object.

tests/test_fake_data_creator.py

@@ -9,8 +9,22 @@
 
 
 class test_fake_image_creator(unittest.TestCase):
+    def test_create_fake_times(self):
+        times1 = create_fake_times(10, t0=0.0, obs_per_day=3, intra_night_gap=0.01, inter_night_gap=1)
+        expected = [0.0, 0.01, 0.02, 1.0, 1.01, 1.02, 2.0, 2.01, 2.02, 3.0]
+        self.assertEqual(len(times1), 10)
+        for i in range(10):
+            self.assertAlmostEqual(times1[i], expected[i])
+
+        times2 = create_fake_times(7, t0=10.0, obs_per_day=1, intra_night_gap=0.5, inter_night_gap=2)
+        expected = [10.0, 12.0, 14.0, 16.0, 18.0, 20.0, 22.0]
+        self.assertEqual(len(times2), 7)
+        for i in range(7):
+            self.assertAlmostEqual(times2[i], expected[i])
+
     def test_create(self):
-        ds = FakeDataSet(256, 128, 10)
+        times = create_fake_times(10)
+        ds = FakeDataSet(256, 128, times)
         self.assertEqual(ds.stack.img_count(), 10)
 
         last_time = -1.0
@@ -23,8 +37,13 @@ def test_create(self):
             self.assertGreater(t, last_time)
             last_time = t
 
+    def test_create_empty_times(self):
+        ds = FakeDataSet(256, 128, [])
+        self.assertEqual(ds.stack.img_count(), 0)
+
     def test_insert_object(self):
-        ds = FakeDataSet(128, 128, 5, use_seed=True)
+        times = create_fake_times(5, 57130.2, 3, 0.01, 1)
+        ds = FakeDataSet(128, 128, times, use_seed=True)
         self.assertEqual(ds.stack.img_count(), 5)
         self.assertEqual(len(ds.trajectories), 0)
 
@@ -51,7 +70,7 @@ def test_insert_object(self):
 
     def test_save_and_clean(self):
         num_images = 7
-        ds = FakeDataSet(64, 64, num_images)
+        ds = FakeDataSet(64, 64, create_fake_times(num_images))
 
         with tempfile.TemporaryDirectory() as dir_name:
             # Get all the file names.
@@ -75,7 +94,7 @@ def test_save_and_clean(self):
 
     def test_save_times(self):
         num_images = 50
-        ds = FakeDataSet(4, 4, num_images)
+        ds = FakeDataSet(4, 4, create_fake_times(num_images))
 
         with tempfile.TemporaryDirectory() as dir_name:
             file_name = f"{dir_name}/times.dat"
@@ -87,7 +106,7 @@ def test_save_times(self):
 
     def test_save_work_unit(self):
         num_images = 25
-        ds = FakeDataSet(15, 10, num_images)
+        ds = FakeDataSet(15, 10, create_fake_times(num_images))
 
         with tempfile.TemporaryDirectory() as dir_name:
             file_name = os.path.join(dir_name, "fake_work_unit.fits")

tests/test_readme_example.py

@@ -13,7 +13,8 @@ def test_make_and_copy(self):
         psf = kb.PSF(1.5)
 
         # Create fake data with ten 512x512 pixel images.
-        ds = FakeDataSet(512, 512, 10)
+        fake_times = create_fake_times(10, t0=57130.2)
+        ds = FakeDataSet(512, 512, fake_times)
         imgs = ds.stack.get_images()
 
         # Get the timestamp of the first image.

tests/test_search_encode.py

@@ -2,8 +2,9 @@
 
 import numpy as np
 
-from kbmod.fake_data_creator import add_fake_object
+from kbmod.fake_data_creator import FakeDataSet
 from kbmod.search import *
+from kbmod.trajectory_utils import make_trajectory
 
 
 class test_search_filter(unittest.TestCase):
@@ -14,12 +15,9 @@ def setUp(self):
         self.flux_error = 0.25
 
         # image properties
-        self.imCount = 20
+        self.img_count = 20
         self.dim_x = 100
         self.dim_y = 110
-        self.noise_level = 4.0
-        self.variance = self.noise_level**2
-        self.p = PSF(1.0)
 
         # object properties
         self.object_flux = 250.0
@@ -29,11 +27,7 @@ def setUp(self):
         self.vyel = 19.0
 
         # create a Trajectory for the object
-        self.trj = Trajectory()
-        self.trj.x = self.start_x
-        self.trj.y = self.start_y
-        self.trj.vx = self.vxel
-        self.trj.vy = self.vyel
+        self.trj = make_trajectory(self.start_x, self.start_y, self.vxel, self.vyel, flux=self.object_flux)
 
         # search parameters
         self.angle_steps = 150
@@ -49,19 +43,17 @@ def setUp(self):
         self.lh_level = 10.0
 
         # create image set with single moving object
-        self.imlist = []
-        for i in range(self.imCount):
-            time = i / self.imCount
-            im = LayeredImage(self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i)
-            add_fake_object(
-                im,
-                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 = ImageStack(self.imlist)
+        fake_times = [i / self.img_count for i in range(self.img_count)]
+        fake_ds = FakeDataSet(
+            self.dim_x,
+            self.dim_y,
+            fake_times,
+            noise_level=2.0,
+            psf_val=1.0,
+            use_seed=True,
+        )
+        fake_ds.insert_object(self.trj)
+        self.stack = fake_ds.stack
 
     @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
     def test_different_encodings(self):
@@ -76,7 +68,7 @@ def test_different_encodings(self):
                     self.max_angle,
                     self.min_vel,
                     self.max_vel,
-                    int(self.imCount / 2),
+                    int(self.img_count / 2),
                 )
 
                 results = search.get_results(0, 10)

tests/test_search_filter.py

@@ -2,8 +2,9 @@
 
 import numpy as np
 
-from kbmod.fake_data_creator import add_fake_object
+from kbmod.fake_data_creator import FakeDataSet
 from kbmod.search import *
+from kbmod.trajectory_utils import make_trajectory
 
 
 class test_search_filter(unittest.TestCase):
@@ -14,12 +15,9 @@ def setUp(self):
         self.flux_error = 0.15
 
         # image properties
-        self.imCount = 20
+        self.img_count = 20
         self.dim_x = 80
         self.dim_y = 60
-        self.noise_level = 4.0
-        self.variance = self.noise_level**2
-        self.p = PSF(1.0)
 
         # object properties
         self.object_flux = 250.0
@@ -29,11 +27,7 @@ def setUp(self):
         self.vyel = 16.0
 
         # create a Trajectory for the object
-        self.trj = Trajectory()
-        self.trj.x = self.start_x
-        self.trj.y = self.start_y
-        self.trj.vx = self.vxel
-        self.trj.vy = self.vyel
+        self.trj = make_trajectory(self.start_x, self.start_y, self.vxel, self.vyel, flux=self.object_flux)
 
         # search parameters
         self.angle_steps = 150
@@ -49,21 +43,18 @@ def setUp(self):
         self.lh_level = 10.0
 
         # create image set with single moving object
-        self.imlist = []
-        for i in range(self.imCount):
-            time = i / self.imCount
-            im = LayeredImage(self.dim_x, self.dim_y, self.noise_level, self.variance, time, self.p, i)
-            add_fake_object(
-                im,
-                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 = ImageStack(self.imlist)
+        fake_times = [i / self.img_count for i in range(self.img_count)]
+        fake_ds = FakeDataSet(
+            self.dim_x,
+            self.dim_y,
+            fake_times,
+            noise_level=2.0,
+            psf_val=1.0,
+            use_seed=True,
+        )
+        fake_ds.insert_object(self.trj)
 
-        self.search = StackSearch(self.stack)
+        self.search = StackSearch(fake_ds.stack)
         self.search.enable_gpu_sigmag_filter(self.sigmaG_lims, self.sigmaG_coeff, self.lh_level)
         self.search.search(
             self.angle_steps,
@@ -72,7 +63,7 @@ def setUp(self):
             self.max_angle,
             self.min_vel,
             self.max_vel,
-            int(self.imCount / 2),
+            int(self.img_count / 2),
         )
 
     @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")