Remove the demo data file

data/readme.txt

@@ -1,5 +1,4 @@
 The data directory contains a few data sets used in example python files and notebooks:
 - data/demo_image.fits: Contains an image file created using fake_data_creator.py and containing a single fake object.
-- data/demo_data.fits: Contains a WorkUnit generated via the create_fake_data notebook.
 - data/small: Contains 10 small image files created using fake_data_creator.py and containing a single fake object.
 - data/fake_results: Contains the results of running the KBMOD_Demo notebook on the data in data/demo. For a description of the files see the KBMOD documentation.

notebooks/KBMOD_Demo.ipynb

@@ -19,6 +19,7 @@
     "from pathlib import Path\n",
     "\n",
     "from kbmod.configuration import SearchConfiguration\n",
+    "from kbmod.fake_data.demo_helper import make_demo_data\n",
     "from kbmod.run_search import *\n",
     "from kbmod.search import *\n",
     "from kbmod.work_unit import WorkUnit"
@@ -37,7 +38,7 @@
    "source": [
     "In order to run KBMOD you need to have the location of the input data and a `res_filepath` that provides a path to the directory where the output results will be stored. Input data can come from a variety of formats including Rubin’s Bulter, fits files, and `WorkUnit` files. In this demo we use the `WorkUnit` file which is an internal storage format used. For more information on generating a `WorkUnit` from the Butler or fits, see the standardizer notebooks.\n",
     "\n",
-    "If you already have data files, you can use those. Below we use the data in `data/demo_data.fits`. You can also create your own fake data using `fake_data_creator.py`."
+    "If you already have data files, you can use those. Below we create and use data in `data/demo_data.fits`."
    ]
   },
   {
@@ -46,7 +47,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "input_filename = \"../data/demo_data.fits\""
+    "input_filename = \"../data/demo_data.fits\"\n",
+    "\n",
+    "# Create the fake data usering a helper function.\n",
+    "if not Path(input_filename).is_file():\n",
+    "    make_demo_data(filename=input_filename)"
    ]
   },
   {

notebooks/create_fake_data.ipynb

@@ -4,7 +4,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Create Fake Data\n"
+    "# Create Fake Data\n",
+    "\n",
+    "This notebook demonstrates how to create fake data using the functions in `fake_data/fake_data_creator.py`. The data set matches the one created by the `make_demo_data()` in `fake_data/demo_helper.py`."
    ]
   },
   {
@@ -159,7 +161,7 @@
     "config = SearchConfiguration.from_dict(settings)\n",
     "\n",
     "# We comment out this line so as not to overwrite the existing demo data.\n",
-    "ds.save_fake_data_to_work_unit(\"../data/demo_data.fits\", config=config)"
+    "# ds.save_fake_data_to_work_unit(\"../data/demo_data.fits\", config=config)"
    ]
   },
   {

src/kbmod/fake_data/demo_helper.py

@@ -0,0 +1,56 @@
+"""A helper function for creating data used in the demo notebook
+and some of the tests.
+"""
+
+import os
+
+from kbmod.configuration import SearchConfiguration
+from kbmod.fake_data.fake_data_creator import *
+from kbmod.search import *
+
+def make_demo_data(filename=None):
+    """Make the fake demo data.
+
+    Parameters
+    ----------
+    filename : `str`
+        The path and file anem to store the demo data. If ``None`` then
+        does not save the demo data.
+
+    Returns
+    -------
+    work : `WorkUnit`
+        A WorkUnit with the fake data.
+    """
+    # Set the characteristics of the fake data.
+    img_width = 256
+    img_height = 256
+    num_times = 20
+
+    # Create the fake images
+    fake_times = create_fake_times(num_times, t0=57130.2)
+    ds = FakeDataSet(img_width, img_height, fake_times)
+
+    # Insert a fake object with only horizontal velocity.
+    trj = Trajectory(x=50, y=40, vx=10, vy=0, flux=500)
+    ds.insert_object(trj)
+
+    # 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],
+        # Loosen the other filtering parameters.
+        "clip_negative": True,
+        "sigmaG_lims": [15, 60],
+        "mom_lims": [37.5, 37.5, 1.5, 1.0, 1.0],
+        "peak_offset": [3.0, 3.0],
+    }
+    config = SearchConfiguration.from_dict(settings)
+
+    # Create a WorkUnit and save it if needed.
+    work = WorkUnit(im_stack=ds.stack, config=config, wcs=ds.fake_wcs)
+    if filename is not None:
+        work.to_fits(filename)
+    return work

tests/test_end_to_end.py

@@ -2,23 +2,16 @@
 # fail because of a format change, you may need to rerun that notebook to regenerate
 # the data/demo_data.fits file.
 
-import math
-import numpy as np
 import os
 import tempfile
 import unittest
 
 from kbmod.configuration import SearchConfiguration
-from kbmod.fake_data.fake_data_creator import *
-from kbmod.run_search import *
-from kbmod.search import *
-from kbmod.wcs_utils import make_fake_wcs
+from kbmod.fake_data.demo_helper import make_demo_data
+from kbmod.run_search import SearchRunner
+from kbmod.search import HAS_GPU
 from kbmod.work_unit import WorkUnit
 
-# from .utils_for_tests import get_absolute_demo_data_path
-# import utils_for_tests
-from utils.utils_for_tests import get_absolute_demo_data_path
-
 
 # this is the first test to actually test things like get_all_stamps from
 # analysis utils. For now stamps have to be RawImages (because methods like
@@ -28,67 +21,47 @@
 # these operations into functions and calling on the .image attribute
 # apply_stamp_filter for example is literal copy of the C++ code in RawImage?
 class test_end_to_end(unittest.TestCase):
-    def setUp(self):
-        # Define the path for the data.
-        self.im_filename = get_absolute_demo_data_path("demo_data.fits")
-
     @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
     def test_demo_defaults(self):
-        # Load the WorkUnit.
-        input_data = WorkUnit.from_fits(self.im_filename)
-
-        rs = SearchRunner()
-        keep = rs.run_search_from_work_unit(input_data)
-        self.assertGreaterEqual(len(keep), 1)
-        self.assertEqual(keep["stamp"][0].shape, (21, 21))
-
-    @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
-    def test_demo_stamp_size(self):
-        input_data = WorkUnit.from_fits(self.im_filename)
-
-        # Override the stamp settings of the configuration
-        input_data.config.set("stamp_radius", 15)
-        input_data.config.set("mom_lims", [80.0, 80.0, 50.0, 20.0, 20.0])
-        input_data.config.set("save_all_stamps", True)
-
-        rs = SearchRunner()
-        keep = rs.run_search_from_work_unit(input_data)
-        self.assertGreaterEqual(len(keep), 1)
+        with tempfile.TemporaryDirectory() as dir_name:
+            # Create a fake data file.
+            filename = os.path.join(dir_name, "test_workunit1.fits")
+            make_demo_data(filename)
 
-        self.assertIsNotNone(keep["stamp"][0])
-        self.assertEqual(keep["stamp"][0].shape, (31, 31))
+            # Load the WorkUnit.
+            input_data = WorkUnit.from_fits(filename)
 
-        self.assertIsNotNone(keep["all_stamps"][0])
-        for s in keep["all_stamps"][0]:
-            self.assertEqual(s.shape, (31, 31))
+            rs = SearchRunner()
+            keep = rs.run_search_from_work_unit(input_data)
+            self.assertGreaterEqual(len(keep), 1)
+            self.assertEqual(keep["stamp"][0].shape, (21, 21))
 
     @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)")
-    def test_e2e_work_unit(self):
-        num_images = 10
-
-        # 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 = 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.
-        config = SearchConfiguration()
-        config.set("ang_arr", [math.pi, math.pi, 16])
-        config.set("v_arr", [0, 10.0, 20])
+    def test_demo_stamp_size(self):
+        with tempfile.TemporaryDirectory() as dir_name:
+            # Create a fake data file.
+            filename = os.path.join(dir_name, "test_workunit1.fits")
+            make_demo_data(filename)
 
-        fake_wcs = make_fake_wcs(10.0, 10.0, 128, 128)
-        work = WorkUnit(im_stack=ds.stack, config=config, wcs=fake_wcs)
+            # Load the WorkUnit.
+            input_data = WorkUnit.from_fits(filename)
 
-        with tempfile.TemporaryDirectory() as dir_name:
-            file_path = os.path.join(dir_name, "test_workunit.fits")
-            work.to_fits(file_path)
+            # Override the stamp settings of the configuration
+            input_data.config.set("stamp_radius", 15)
+            input_data.config.set("mom_lims", [80.0, 80.0, 50.0, 20.0, 20.0])
+            input_data.config.set("save_all_stamps", True)
 
             rs = SearchRunner()
-            keep = rs.run_search_from_file(file_path)
+            keep = rs.run_search_from_work_unit(input_data)
             self.assertGreaterEqual(len(keep), 1)
 
+            self.assertIsNotNone(keep["stamp"][0])
+            self.assertEqual(keep["stamp"][0].shape, (31, 31))
+
+            self.assertIsNotNone(keep["all_stamps"][0])
+            for s in keep["all_stamps"][0]:
+                self.assertEqual(s.shape, (31, 31))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/utils/utils_for_tests.py

@@ -5,9 +5,3 @@ def get_absolute_data_path(file_or_directory):
     test_dir = path.abspath(path.dirname(path.dirname(__file__)))
     data_dir = path.join(test_dir, "data")
     return path.join(data_dir, file_or_directory)
-
-
-def get_absolute_demo_data_path(file_or_directory):
-    project_root_dir = path.abspath(path.dirname(path.dirname(path.dirname(__file__))))
-    data_dir = path.join(project_root_dir, "data")
-    return path.join(data_dir, file_or_directory)