diff --git a/.gitignore b/.gitignore index 59a41bded..b60793bbc 100644 --- a/.gitignore +++ b/.gitignore @@ -31,4 +31,5 @@ docs/source/examples/_notebooks .#* \#*# # Exclude files starting with exclude (local test and development) -exclude* \ No newline at end of file +exclude* +data/demo_data.fits diff --git a/data/demo/000001.fits b/data/demo/000001.fits deleted file mode 100644 index 99c87f414..000000000 Binary files a/data/demo/000001.fits and /dev/null differ diff --git a/data/demo/000002.fits b/data/demo/000002.fits deleted file mode 100644 index 003cbd863..000000000 Binary files a/data/demo/000002.fits and /dev/null differ diff --git a/data/demo/000003.fits b/data/demo/000003.fits deleted file mode 100644 index 2f00c96a2..000000000 Binary files a/data/demo/000003.fits and /dev/null differ diff --git a/data/demo/000004.fits b/data/demo/000004.fits deleted file mode 100644 index 8a5e3242f..000000000 Binary files a/data/demo/000004.fits and /dev/null differ diff --git a/data/demo/000005.fits b/data/demo/000005.fits deleted file mode 100644 index 97795129f..000000000 Binary files a/data/demo/000005.fits and /dev/null differ diff --git a/data/demo/000006.fits b/data/demo/000006.fits deleted file mode 100644 index b37aa7d7c..000000000 Binary files a/data/demo/000006.fits and /dev/null differ diff --git a/data/demo/000007.fits b/data/demo/000007.fits deleted file mode 100644 index 4a0404d21..000000000 Binary files a/data/demo/000007.fits and /dev/null differ diff --git a/data/demo/000008.fits b/data/demo/000008.fits deleted file mode 100644 index c31a32f2c..000000000 Binary files a/data/demo/000008.fits and /dev/null differ diff --git a/data/demo/000009.fits b/data/demo/000009.fits deleted file mode 100644 index 445eba737..000000000 Binary files a/data/demo/000009.fits and /dev/null differ diff --git a/data/demo_config.yml b/data/demo_config.yml deleted file mode 100644 index 775e2db53..000000000 --- a/data/demo_config.yml +++ /dev/null @@ -1,82 +0,0 @@ -ang_arr: -- 0.5 -- 0.5 -- 11 -average_angle: 0.0 -center_thresh: 0.0 -chunk_size: 1000000 -clip_negative: true -cluster_type: position -debug: true -do_clustering: true -do_mask: true -do_stamp_filter: true -encode_num_bytes: -1 -eps: 0.03 -flag_keys: -- BAD -- CR -- INTRP -- NO_DATA -- SENSOR_EDGE -- SAT -- SUSPECT -- CLIPPED -- REJECTED -- DETECTED_NEGATIVE -gpu_filter: true -im_filepath: data/pg300_ccd10 -known_obj_jpl: false -known_obj_obs: 3 -known_obj_thresh: null -lh_level: 10.0 -mask_bits_dict: - BAD: 0 - CLIPPED: 9 - CR: 3 - DETECTED: 5 - DETECTED_NEGATIVE: 6 - EDGE: 4 - INEXACT_PSF: 10 - INTRP: 2 - NOT_DEBLENDED: 11 - NO_DATA: 8 - REJECTED: 12 - SAT: 1 - SENSOR_EDGE: 13 - SUSPECT: 7 -mask_grow: 10 -mask_num_images: 10 -mask_threshold: null -max_lh: 1000.0 -mjd_lims: null -mom_lims: -- 37.5 -- 37.5 -- 1.5 -- 1.0 -- 1.0 -num_obs: 7 -output_suffix: DEMO -peak_offset: -- 3.0 -- 3.0 -psf_file: null -psf_val: 1.4 -repeated_flag_keys: -- DETECTED -res_filepath: null -sigmaG_lims: -- 15 -- 60 -stamp_radius: 10 -stamp_type: median -time_file: null -v_arr: -- 0.0 -- 20.0 -- 21 -x_pixel_bounds: null -x_pixel_buffer: null -y_pixel_bounds: null -y_pixel_buffer: null diff --git a/data/demo/000000.fits b/data/demo_image.fits similarity index 100% rename from data/demo/000000.fits rename to data/demo_image.fits diff --git a/data/demo_times.dat b/data/demo_times.dat deleted file mode 100644 index 138dafa71..000000000 --- a/data/demo_times.dat +++ /dev/null @@ -1,21 +0,0 @@ -# visit_id timestamp -000000 57130.2 -000001 57130.21 -000002 57130.22 -000003 57131.2 -000004 57131.21 -000005 57131.22 -000006 57132.2 -000007 57132.21 -000008 57132.22 -000009 57133.2 -000010 57133.21 -000011 57133.22 -000012 57134.2 -000013 57134.21 -000014 57134.22 -000015 57135.2 -000016 57135.21 -000017 57135.22 -000018 57136.2 -000019 57136.21 diff --git a/data/readme.txt b/data/readme.txt index a2297ee80..8036ec2b9 100644 --- a/data/readme.txt +++ b/data/readme.txt @@ -1,5 +1,4 @@ The data directory contains a few data sets used in example python files and notebooks: -- data/demo: Contains 10 image files created using fake_data_creator.py and containing a single fake object. -- data/demo_times.dat: The external time file for the images in data/demo. +- data/demo_image.fits: Contains an image file created using fake_data_creator.py and containing a single fake object. - 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. diff --git a/notebooks/KBMOD_Demo.ipynb b/notebooks/KBMOD_Demo.ipynb index 68fb0cc25..c21cdea39 100644 --- a/notebooks/KBMOD_Demo.ipynb +++ b/notebooks/KBMOD_Demo.ipynb @@ -18,8 +18,11 @@ "\n", "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 *" + "from kbmod.search import *\n", + "from kbmod.work_unit import WorkUnit" ] }, { @@ -33,13 +36,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "There are at least two file paths you need to setup in order to run kbmod:\n", - "1. The im_filepath provides a path to the input images.\n", - "1. The res_filepath provides a path to the directory where the output results will be stored.\n", + "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", - "A time and psf file can optionally be specified.\n", - "\n", - "If you already have data files, you can use those. Below we use the data in `data/demo`. 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`." ] }, { @@ -48,9 +47,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Define the path for the data.\n", - "im_filepath = \"../data/demo\"\n", - "print(os.listdir(im_filepath))" + "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)" ] }, { @@ -108,21 +109,19 @@ "results_suffix = \"DEMO\"\n", "\n", "# The demo data has an object moving at x_v=10 px/day\n", - "# and y_v = 0 px/day. So we search velocities [0, 20].\n", + "# and y_v = 0 px/day. So we search velocities [0, 19].\n", "v_min = 0\n", "v_max = 20\n", - "v_steps = 21\n", + "v_steps = 20\n", "v_arr = [v_min, v_max, v_steps]\n", "\n", - "# and angles [-0.5, 0.5]\n", + "# and angles [-0.5, 0.5)\n", "ang_below = 0.5\n", "ang_above = 0.5\n", - "ang_steps = 11\n", + "ang_steps = 10\n", "ang_arr = [ang_below, ang_above, ang_steps]\n", "\n", "input_parameters = {\n", - " # Input parameters\n", - " \"im_filepath\": im_filepath,\n", " # Grid search parameters\n", " \"v_arr\": v_arr,\n", " \"ang_arr\": ang_arr,\n", @@ -134,12 +133,15 @@ " \"do_mask\": True, # <-- Apply the masks specified in the FITS files.\n", " \"mask_num_images\": 10,\n", " # Basic filtering (always applied)\n", - " \"num_obs\": 7, # <-- Filter anything with fewer than 7 observations\n", + " \"num_obs\": 15, # <-- Filter anything with fewer than 15 observations\n", " \"lh_level\": 10.0, # <-- Filter anything with a likelihood < 10.0\n", " # SigmaG clipping parameters\n", " \"sigmaG_lims\": [15, 60], # <-- Clipping parameters (lower and upper percentile)\n", " \"gpu_filter\": True, # <-- Apply clipping and filtering on the GPU\n", " \"clip_negative\": True,\n", + " # 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", @@ -160,7 +162,24 @@ "metadata": {}, "outputs": [], "source": [ - "# print(config)" + "print(config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We load the data as a `WorkUnit` object. In general `WorkUnit`'s include a copy of their own configuration so they have all the information they need for a full run. We overwrite the stored configuration with the one we defined above." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "input_data = WorkUnit.from_fits(input_filename)\n", + "input_data.config = config" ] }, { @@ -188,7 +207,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Once we have defined the search parameters, we can create a ``SearchRunner`` and use one of the run_search functions. In this case we use ``run_search_from_config`` which uses the config to search for the input files." + "Once we have defined the search parameters, we can create a `SearchRunner` and use one of the run_search functions. In this case we use `run_search_from_work_unit` which uses the `WorkUnit` to define both the image data and the configuration information." ] }, { @@ -198,7 +217,7 @@ "outputs": [], "source": [ "rs = SearchRunner()\n", - "results = rs.run_search_from_config(config)" + "results = rs.run_search_from_work_unit(input_data)" ] }, { @@ -289,13 +308,11 @@ "outputs": [], "source": [ "# Turn on filtered tracking\n", - "input_parameters[\"track_filtered\"] = True\n", + "input_data.config.set(\"track_filtered\", True)\n", "\n", "# Turn up filtering of stamp filtering. This will require 100% of the stamp's flux\n", "# to be at the center pixel and effectively filter every candidate trajectory.\n", - "input_parameters[\"center_thresh\"] = 1.0\n", - "\n", - "config = SearchConfiguration.from_dict(input_parameters)" + "input_data.config.set(\"center_thresh\", 1.0)" ] }, { @@ -312,7 +329,7 @@ "outputs": [], "source": [ "rs = SearchRunner()\n", - "results = rs.run_search_from_config(config)\n", + "results = rs.run_search_from_work_unit(input_data)\n", "print(f\"Search found {len(results)} results.\")" ] }, @@ -355,7 +372,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Now we search for one of the expected trajectories (starting at pixel (106, 44) at the first time step) by using the table's search functionality." + "Now we search for one of the expected trajectories (starting at pixel (50, 40) at the first time step) by using the table's search functionality." ] }, { @@ -364,7 +381,7 @@ "metadata": {}, "outputs": [], "source": [ - "subset = results.table[(results.table[\"x\"] == 106) & (results.table[\"y\"] == 44)]\n", + "subset = results.table[(results.table[\"x\"] == 50) & (results.table[\"y\"] == 40)]\n", "print(subset)" ] }, @@ -374,6 +391,13 @@ "source": [ "As we can see all of the potential trajectories were filtered by the stamp filter. We can use this information to help tune different filtering stages." ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/notebooks/Kbmod_Reference.ipynb b/notebooks/Kbmod_Reference.ipynb index e34b6ebf6..40193c8dd 100644 --- a/notebooks/Kbmod_Reference.ipynb +++ b/notebooks/Kbmod_Reference.ipynb @@ -40,7 +40,7 @@ "import matplotlib.pyplot as plt\n", "import os\n", "\n", - "im_path = \"../data/demo/\"\n", + "im_file = \"../data/demo_image.fits\"\n", "res_path = \"./results\"" ] }, @@ -183,7 +183,7 @@ "source": [ "from kbmod.data_interface import load_deccam_layered_image\n", "\n", - "im = load_deccam_layered_image(im_path + \"000000.fits\", p)\n", + "im = load_deccam_layered_image(im_file, p)\n", "print(f\"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_obstime()}\")" ] }, diff --git a/notebooks/create_fake_data.ipynb b/notebooks/create_fake_data.ipynb index 7f81c921c..489789c42 100644 --- a/notebooks/create_fake_data.ipynb +++ b/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`." ] }, { @@ -15,6 +17,7 @@ "source": [ "import os\n", "\n", + "import matplotlib.pyplot as plt\n", "from kbmod.fake_data.fake_data_creator import *\n", "from kbmod.search import *" ] @@ -46,7 +49,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### Insert a fake moving object\n", + "### Insert fake moving objects\n", "\n", "This function creates a random moving object with a given flux that stays within the image for the entire time. The trajectory is defined by starting pixels (x, y) and velocities (x_v, y_v) of pixels per day." ] @@ -85,6 +88,82 @@ " print(f\"{i}: t={ti:.3f} at ({px}, {py})\")" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also create a known trajectory and insert an object using that trajectory. Here we create an object starting at x=50, y=10 and moving at 10 pixels per day in x and 0 pixels per day in y." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "trj2 = Trajectory(x=50, y=40, vx=10, vy=0, flux=500)\n", + "ds.insert_object(trj2)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can plot the first and last images and we clearly see our two objects." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plt.imshow(ds.stack.get_single_image(0).get_science().image)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plt.imshow(ds.stack.get_single_image(num_times - 1).get_science().image)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Save the fake image files\n", + "\n", + "We can store the images in a `WorkUnit` and save that to disk (to be used in other notebooks). In the `WorkUnit` we store a configuration with the minimal search parameters needed to retrieve the object we inserted. Specifically, the data has a known object moving at x_v=10 px/day and y_v = 0 px/day, so we set the search search velocities to [0, 20) and the search angles to [-0.5, 0.5). Note that these bounds may not find the other (random) trajectory we inserted." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from kbmod.configuration import SearchConfiguration\n", + "\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", + " # Loosen the other filtering parameters.\n", + " \"clip_negative\": True,\n", + " \"sigmaG_lims\": [15, 60],\n", + " \"mom_lims\": [37.5, 37.5, 1.5, 1.0, 1.0],\n", + " \"peak_offset\": [3.0, 3.0],\n", + "}\n", + "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)" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/src/kbmod/fake_data/demo_helper.py b/src/kbmod/fake_data/demo_helper.py new file mode 100644 index 000000000..62c7f9fc5 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/diff_test.py b/tests/diff_test.py deleted file mode 100644 index 675f0f9a9..000000000 --- a/tests/diff_test.py +++ /dev/null @@ -1,311 +0,0 @@ -import argparse -import os -import sys -import tempfile -from pathlib import Path - -import numpy as np - -from kbmod.run_search import SearchRunner - - -def check_and_create_goldens_dir(): - """ - Test whether the goldens directory exists and create it if not. - """ - dir_path = Path("goldens") - if not dir_path.is_dir(): - os.mkdir("goldens") - - -def check_goldens_exist(results_suffix): - """ - Test whether the needed goldens files exist. - - Arguments: - results_suffix - The suffix of the goldens file. - """ - file_path = Path("goldens/results_%s.txt" % results_suffix) - if not file_path.is_file(): - return False - - file_path = Path("goldens/ps_%s.txt" % results_suffix) - if not file_path.is_file(): - return False - return True - - -def compare_ps_files(goldens_file, new_results_file, delta=0.00001): - """ - Compare two PS result files. - - Arguments: - goldens_file - The path and filename of the golden results. - new_results_file - The path and filename of the new results. - delta - The maximum difference in any entry. - - Returns: - A Boolean indicating whether the files are the same. - """ - files_equal = True - - res_new = np.loadtxt(new_results_file, dtype=str) - print("Loaded %i new results from %s." % (len(res_new), new_results_file)) - res_old = np.loadtxt(goldens_file, dtype=str) - print("Loaded %i old results from %s." % (len(res_old), goldens_file)) - - # Check that the number of results matches up. - if len(res_new) != len(res_old): - print("Mismatched number of results (%i vs %i)." % (len(res_old), len(res_new))) - files_equal = False - else: - # Check each line to see if it matches. - for i in range(len(res_new)): - old_line = res_old[i] - new_line = res_new[i] - - found_diff = False - if len(new_line) != len(old_line): - found_diff = True - else: - for d in range(len(new_line)): - if (abs(float(old_line[d]) - float(new_line[d]))) > delta: - found_diff = True - - if found_diff: - files_equal = False - print("Found a difference in line %i:" % i) - print(" [OLD] %s" % old_line[i]) - print(" [NEW] %s" % new_line[i]) - return files_equal - - -def compare_result_files(goldens_file, new_results_file, delta=0.001): - """ - Compare two result files. - - Arguments: - goldens_file - The path and filename of the golden results. - new_results_file - The path and filename of the new results. - delta - The maximum difference in numerical values. - - Returns: - A Boolean indicating whether the files are the same. - """ - files_equal = True - - res_new = np.loadtxt(new_results_file, dtype=str) - print("Loaded %i new results from %s." % (len(res_new), new_results_file)) - res_old = np.loadtxt(goldens_file, dtype=str) - print("Loaded %i old results from %s." % (len(res_old), goldens_file)) - - # Check that the number of results matches up. - if len(res_new) != len(res_old): - print("Mismatched number of results (%i vs %i)." % (len(res_old), len(res_new))) - files_equal = False - else: - # Check each line to see if it matches. - for i in range(len(res_new)): - old_line = res_old[i] - new_line = res_new[i] - - found_diff = False - if len(new_line) != len(old_line): - found_diff = True - else: - for d in range(len(new_line)): - if d % 2 == 1: - # Allow a small difference in estimated velocities because - # the search might find multiple results whose scores tie. - if abs(float(old_line[d]) - float(new_line[d])) > delta: - found_diff = True - else: - if old_line[d] != new_line[d]: - found_diff = True - - if found_diff: - files_equal = False - print("Found a difference in line %i:" % i) - print(" [OLD] %s" % old_line) - print(" [NEW] %s" % new_line) - return files_equal - - -def perform_search(im_filepath, time_file, psf_file, res_filepath, res_suffix, shorten_search, stamp_type): - """Run the core search algorithm. - - Parameters - ---------- - im_filepath : string - The file path (directory) for the image files. - time_file : string - The path and file name of the file of timestamps. - res_filepath : string - The path (directory) for the new result files. - res_suffix : string - The file suffix to use for the new results. - shorten_search : bool - A Boolean indicating whether to use a coarser grid for a fast but less thorough search. - stamp_type : string - The stamp type to use for coadds. - """ - v_min = 92.0 # Pixels/day - v_max = 550.0 - - # 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 = 512 - ang_steps = 256 - if shorten_search: - 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 - - mask_bits_dict = { - "BAD": 0, - "CLIPPED": 9, - "CR": 3, - "DETECTED": 5, - "DETECTED_NEGATIVE": 6, - "EDGE": 4, - "INEXACT_PSF": 10, - "INTRP": 2, - "NOT_DEBLENDED": 11, - "NO_DATA": 8, - "REJECTED": 12, - "SAT": 1, - "SENSOR_EDGE": 13, - "SUSPECT": 7, - } - flag_keys = [ - "BAD", - "CR", - "INTRP", - "NO_DATA", - "SENSOR_EDGE", - "SAT", - "SUSPECT", - "CLIPPED", - "REJECTED", - "DETECTED_NEGATIVE", - ] - repeated_flag_keys = ["DETECTED"] - - input_parameters = { - "im_filepath": im_filepath, - "res_filepath": res_filepath, - "time_file": time_file, - "psf_file": psf_file, - "output_suffix": results_suffix, - "v_arr": v_arr, - "ang_arr": ang_arr, - "num_obs": num_obs, - "do_mask": True, - "lh_level": 10.0, - "sigmaG_lims": [25, 75], - "mom_lims": [37.5, 37.5, 1.5, 1.0, 1.0], - "peak_offset": [3.0, 3.0], - "chunk_size": 1000000, - "stamp_type": stamp_type, - "eps": 0.03, - "gpu_filter": True, - "clip_negative": True, - "mask_num_images": 10, - "mask_bits_dict": mask_bits_dict, - "flag_keys": flag_keys, - "repeated_flag_keys": repeated_flag_keys, - "known_obj_thresh": None, - "known_obj_jpl": False, - "encode_num_bytes": -1, - } - - rs = SearchRunner() - rs.run_search_from_config(input_parameters) - - -if __name__ == "__main__": - # Parse the command line arguments. - parser = argparse.ArgumentParser() - parser.add_argument( - "--generate_goldens", - default=False, - action="store_true", - help="Generate the golden files (WARNING overwrites files).", - ) - parser.add_argument( - "--data_filepath", default="data/pg300_ccd10", help="The filepath for the images files." - ) - parser.add_argument( - "--time_filepath", default="./loriallen_times.dat", help="The filepath for the time stamps file." - ) - parser.add_argument("--psf_filepath", default=None, help="The filepath for the psf values file.") - parser.add_argument( - "--short", default=False, action="store_true", help="Use a coarse grid for a fast search." - ) - parser.add_argument("--stamp_type", default="sum", help="What stamp type to use.") - args = parser.parse_args() - - # Set up the file path information. - im_filepath = args.data_filepath - time_file = args.time_filepath - psf_file = args.psf_filepath - results_suffix = args.stamp_type - if args.short: - results_suffix = f"{args.stamp_type}_short" - - # Test whether we are regenerating the goldens. - if args.generate_goldens: - # Check whether the goldens directory exists (and create it - # if not) and check whether the golden files already exist - # (and prompt the user if so). - check_and_create_goldens_dir() - if check_goldens_exist(results_suffix): - print("*** WARNING (re)generating golden files.") - print("This will overwrite previous golden files.") - answer = input("Proceed [y/N]: ") - if answer != "y" and answer != "Y": - print("Canceling and exiting.") - sys.exit() - - # Run the search code and save the results to goldens/ - perform_search( - im_filepath, time_file, psf_file, "goldens", results_suffix, args.short, args.stamp_type - ) - else: - if not check_goldens_exist(results_suffix): - print("ERROR: Golden files do not exist. Generate new goldens using " "'--generate_goldens'") - else: - with tempfile.TemporaryDirectory() as dir_name: - dir_name = "tmp" - print("Running diff test with data in %s/" % im_filepath) - print("Time file: %s" % time_file) - print("PSF file: %s" % psf_file) - if args.short: - print("Using a reduced parameter search (approximate).") - - # Do the search. - perform_search( - im_filepath, time_file, psf_file, dir_name, results_suffix, args.short, args.stamp_type - ) - - # Compare the result files. - goldens_file = "goldens/results_%s.txt" % results_suffix - new_results_file = "%s/results_%s.txt" % (dir_name, results_suffix) - print("Comparing %s and %s" % (goldens_file, new_results_file)) - success = compare_result_files(goldens_file, new_results_file) - - # Compare the PS files. - if success: - goldens_file = "goldens/ps_%s.txt" % results_suffix - new_results_file = "%s/ps_%s.txt" % (dir_name, results_suffix) - print("Comparing %s and %s" % (goldens_file, new_results_file)) - success = compare_ps_files(goldens_file, new_results_file) - - if success: - print("\n*****\nDiff test PASSED.") - else: - print("\n*****\nDiff test FAILED.") diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py index 799d46158..ec6aa20ba 100644 --- a/tests/test_end_to_end.py +++ b/tests/test_end_to_end.py @@ -1,19 +1,17 @@ -import math -import numpy as np +# The tests in this file use data generated by the create_fake_data.py notebook. If the tests +# fail because of a format change, you may need to rerun that notebook to regenerate +# the data/demo_data.fits file. + import os import tempfile import unittest -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.configuration import SearchConfiguration +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 @@ -23,107 +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. - im_filepath = get_absolute_demo_data_path("demo") - - # The demo data has an object moving at x_v=10 px/day - # and y_v = 0 px/day. So we search velocities [0, 20] - # and angles [-0.5, 0.5]. - v_arr = [0, 20, 21] - ang_arr = [0.5, 0.5, 11] - - self.input_parameters = { - # Required - "im_filepath": im_filepath, - "res_filepath": None, - "time_file": None, - "output_suffix": "DEMO", - "v_arr": v_arr, - "ang_arr": ang_arr, - # Important - "num_obs": 7, - "do_mask": True, - "lh_level": 10.0, - "gpu_filter": True, - # Fine tuning - "sigmaG_lims": [15, 60], - "mom_lims": [37.5, 37.5, 1.5, 1.0, 1.0], - "peak_offset": [3.0, 3.0], - "chunk_size": 1000000, - "stamp_type": "cpp_median", - "eps": 0.03, - "clip_negative": True, - "mask_num_images": 10, - "cluster_type": "position", - # Override the ecliptic angle for the demo data since we - # know the true angle in pixel space. - "average_angle": 0.0, - "save_all_stamps": True, - } - @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_demo_defaults(self): - rs = SearchRunner() - keep = rs.run_search_from_config(self.input_parameters) - 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_config_file(self): - im_filepath = get_absolute_demo_data_path("demo") - config_file = get_absolute_demo_data_path("demo_config.yml") - rs = SearchRunner() - keep = rs.run_search_from_file( - config_file, - overrides={"im_filepath": im_filepath}, - ) - 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): - self.input_parameters["stamp_radius"] = 15 - self.input_parameters["mom_lims"] = [80.0, 80.0, 50.0, 20.0, 20.0] - - rs = SearchRunner() - keep = rs.run_search_from_config(self.input_parameters) - 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_workunit2.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() diff --git a/tests/utils/utils_for_tests.py b/tests/utils/utils_for_tests.py index 3fbadec89..975eb0f74 100644 --- a/tests/utils/utils_for_tests.py +++ b/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)