From fd87556ef6ec7cd91f841c5b53e3687ce37c508e Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 4 Sep 2024 13:48:35 -0400 Subject: [PATCH 1/3] Remove some of the tests that read from files --- tests/test_layered_image.py | 74 -------------- tests/test_regression_test.py | 176 ++++------------------------------ 2 files changed, 19 insertions(+), 231 deletions(-) diff --git a/tests/test_layered_image.py b/tests/test_layered_image.py index a9e45698b..cf089273e 100644 --- a/tests/test_layered_image.py +++ b/tests/test_layered_image.py @@ -6,7 +6,6 @@ from astropy.io import fits -from kbmod.data_interface import load_deccam_layered_image, save_deccam_layered_image from kbmod.fake_data.fake_data_creator import add_fake_object, make_fake_layered_image from kbmod.search import * @@ -404,79 +403,6 @@ def test_subtract_template(self): template2.set_all(0.0) self.assertRaises(RuntimeError, self.image.sub_template, template2) - def test_read_write_files(self): - with tempfile.TemporaryDirectory() as dir_name: - im1 = make_fake_layered_image( - 15, # dim_x = 15 pixels, - 20, # dim_y = 20 pixels, - 2.0, # noise_level - 4.0, # variance - 10.0, # time = 10.0 - self.p, - ) - - # Make some changes to the mask to ensure that - # layer has something to compare. - mask1 = im1.get_mask() - mask1.set_pixel(3, 5, 1.0) - mask1.set_pixel(5, 3, 1.0) - - # Save the test data. - full_path = os.path.join(dir_name, "tmp_layered_test_data.fits") - save_deccam_layered_image(im1, full_path) - - # Reload the test data and check that it matches. - im2 = load_deccam_layered_image(full_path, self.p) - self.assertEqual(im1.get_height(), im2.get_height()) - self.assertEqual(im1.get_width(), im2.get_width()) - self.assertEqual(im1.get_npixels(), im2.get_npixels()) - self.assertEqual(im1.get_obstime(), im2.get_obstime()) - - sci1 = im1.get_science() - sci2 = im2.get_science() - self.assertEqual(sci1.obstime, sci2.obstime) - - var1 = im1.get_variance() - mask1 = im1.get_mask() - var2 = im2.get_variance() - mask2 = im2.get_mask() - for x in range(im1.get_width()): - for y in range(im2.get_height()): - self.assertEqual(sci1.get_pixel(y, x), sci2.get_pixel(y, x)) - self.assertEqual(var1.get_pixel(y, x), var2.get_pixel(y, x)) - self.assertEqual(mask1.get_pixel(y, x), mask2.get_pixel(y, x)) - - def test_overwrite_files(self): - with tempfile.TemporaryDirectory() as dir_name: - full_path = os.path.join(dir_name, "tmp_layered_test_data2.fits") - - # Save the test image. - img1 = make_fake_layered_image(15, 20, 2.0, 4.0, 10.0, self.p) - save_deccam_layered_image(img1, full_path) - with fits.open(full_path) as hdulist: - self.assertEqual(len(hdulist), 4) - self.assertEqual(hdulist[1].header["NAXIS1"], 15) - self.assertEqual(hdulist[1].header["NAXIS2"], 20) - - # Save a new test image over the first and check - # that it replaces it. - img2 = make_fake_layered_image(25, 40, 2.0, 4.0, 10.0, self.p) - save_deccam_layered_image(img2, full_path) - with fits.open(full_path) as hdulist2: - self.assertEqual(len(hdulist2), 4) - self.assertEqual(hdulist2[1].header["NAXIS1"], 25) - self.assertEqual(hdulist2[1].header["NAXIS2"], 40) - - # Check that we get an error if we set overwrite = False - self.assertRaises( - ValueError, - save_deccam_layered_image, - img1, - full_path, - None, - False, - ) - def test_stats_string(self): result = self.image.stats_string() self.assertGreater(len(result), 0) diff --git a/tests/test_regression_test.py b/tests/test_regression_test.py index 2b7d0ba38..c5805efee 100644 --- a/tests/test_regression_test.py +++ b/tests/test_regression_test.py @@ -13,14 +13,16 @@ import numpy as np from astropy.io import fits +import astropy.wcs -from kbmod.data_interface import save_deccam_layered_image +from kbmod.configuration import SearchConfiguration from kbmod.fake_data.fake_data_creator import add_fake_object, make_fake_layered_image from kbmod.file_utils import * from kbmod.results import Results from kbmod.run_search import SearchRunner from kbmod.search import * -from kbmod.wcs_utils import append_wcs_to_hdu_header, make_fake_wcs_info +from kbmod.wcs_utils import make_fake_wcs_info +from kbmod.work_unit import WorkUnit logger = logging.getLogger(__name__) @@ -185,12 +187,7 @@ def make_fake_ImageStack(times, trjs, psf_vals): p = PSF(psf_vals[i]) time = times[i] - t0 - # For each odd case, don't save the time. These will be provided by the time file. - saved_time = times[i] - if i % 2 == 1: - saved_time = 0.0 - - img = make_fake_layered_image(dim_x, dim_y, noise_level, variance, saved_time, p, seed=i) + img = make_fake_layered_image(dim_x, dim_y, noise_level, variance, times[i], p, seed=i) for trj in trjs: px = trj.x + time * trj.vx + 0.5 @@ -202,96 +199,6 @@ def make_fake_ImageStack(times, trjs, psf_vals): return stack -def add_wcs_header_data(full_file_name): - """Add (fixed) WCS data to a fits file. - - Parameters - ---------- - full_file_name : `str` - The path and filename of the FITS file to modify. - """ - hdul = fits.open(full_file_name) - - # Create a fake WCS with some minimal information. - wcs_dict = make_fake_wcs_info(200.615, -7.789, 2068, 4088) - append_wcs_to_hdu_header(wcs_dict, hdul[1].header) - - # Add rotational and scale parameters that will allow us to use a matching - # set of search angles (using KBMOD angle suggestion function). - hdul[1].header["CD1_1"] = -1.13926485986789e-07 - hdul[1].header["CD1_2"] = 7.31839748843125e-05 - hdul[1].header["CD2_1"] = -7.30064978350695e-05 - hdul[1].header["CD2_2"] = -1.27520156332774e-07 - - # Save the augmented header. - hdul.writeto(full_file_name, overwrite=True) - hdul.close() - - -def save_fake_data(data_dir, stack, times, psf_vals, default_psf_val=1.0): - """Save the fake data to files. - - Parameters - ---------- - data_dir : `str` - The directory to place the fake data. - stack : `kbmod.search.ImageStack` - The image data. - times : `list` - A list of all times. - psf_vals : `list` - A list of PSF values for each time. - default_psf_val : `float` - The default PSF time if there is not a corresponding value in psf_vals. - """ - # Make the subdirectory if needed. - dir_path = Path(data_dir) - if not dir_path.is_dir(): - logger.debug("Directory '%s' does not exist. Creating." % data_dir) - os.mkdir(data_dir) - - # Make the subdirectory if needed. - img_dir = data_dir + "/imgs" - dir_path = Path(img_dir) - if not dir_path.is_dir(): - logger.debug("Directory '%s' does not exist. Creating." % img_dir) - os.mkdir(img_dir) - - # Save each of the image files. - for i in range(stack.img_count()): - img = stack.get_single_image(i) - filename = os.path.join(img_dir, ("%06i.fits" % i)) - logger.debug("Saving file: %s" % filename) - - # If the file already exists, delete it. - if Path(filename).exists(): - os.remove(filename) - - # Save the file. - save_deccam_layered_image(img, filename) - - # Open the file and insert fake WCS data. - add_wcs_header_data(filename) - - # Save the psf file. - psf_file_name = data_dir + "/psf_vals.dat" - logger.debug("Creating psf file: %s" % psf_file_name) - with open(psf_file_name, "w") as file: - file.write("# visit_id psf_val\n") - for i in range(len(times)): - if psf_vals[i] != default_psf_val: - file.write("%06i %f\n" % (i, psf_vals[i])) - - # Save the time file, but only include half the file times (odd indices). - time_file_name = data_dir + "/times.dat" - time_mapping = {} - for i in range(len(times)): - if i % 2 == 1: - id_str = "%06i" % i - time_mapping[id_str] = times[i] - FileUtils.save_time_dictionary(time_file_name, time_mapping) - - def load_trajectories_from_file(filename): """Load in the result trajectories from their file. @@ -317,18 +224,14 @@ def load_trajectories_from_file(filename): return trjs -def perform_search(im_filepath, time_file, psf_file, res_filename, default_psf): +def perform_search(im_stack, res_filename, default_psf): """ Run the core search algorithm. Parameters ---------- - im_filepath : `str` - The file path (directory) for the image files. - time_file : `str` - The path and file name of the file of timestamps. - psf_file : `str` - The path and file name of the psf values. + im_stack : `ImageStack` + The images to search. res_filename : `str` The path (directory) for the new result files. default_psf : `float` @@ -338,6 +241,9 @@ def perform_search(im_filepath, time_file, psf_file, res_filename, default_psf): 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 @@ -348,45 +254,16 @@ def perform_search(im_filepath, time_file, psf_file, res_filename, default_psf): 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, + "im_filepath": "./", "res_filepath": None, "result_filename": res_filename, - "time_file": time_file, - "psf_file": psf_file, + "time_file": None, + "psf_file": None, "psf_val": default_psf, "output_suffix": "", "v_arr": v_arr, + "average_angle": average_angle, "ang_arr": ang_arr, "num_obs": num_obs, "do_mask": True, @@ -400,17 +277,15 @@ def perform_search(im_filepath, time_file, psf_file, res_filename, default_psf): "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, "x_pixel_buffer": 10, "y_pixel_buffer": 10, "debug": True, } + config = SearchConfiguration.from_dict(input_parameters) + wu = WorkUnit(im_stack=im_stack, config=config) # , wcs=fake_wcs) rs = SearchRunner() - rs.run_search_from_config(input_parameters) + rs.run_search_from_work_unit(wu) def run_full_test(): @@ -467,24 +342,11 @@ def run_full_test(): # Set PSF values between +/- 0.1 around the default value. psf_vals.append(default_psf - 0.1 + 0.1 * (i % 3)) - # Add several instances to the end that will be filtered by the time bounds. - for i in range(3): - times.append(67130.2 + i) - psf_vals.append(default_psf + 0.01) - stack = make_fake_ImageStack(times, trjs, psf_vals) - save_fake_data(dir_name, stack, times, psf_vals, default_psf) # Do the search. - logger.debug("Running search with data in %s/" % dir_name) result_filename = os.path.join(dir_name, "results.ecsv") - perform_search( - os.path.join(dir_name, "imgs"), - os.path.join(dir_name, "times.dat"), - os.path.join(dir_name, "psf_vals.dat"), - result_filename, - default_psf, - ) + perform_search(stack, result_filename, default_psf) # Load the results from the results file and extract a list of trajectories. loaded_data = Results.read_table(result_filename) From 75877c0fb2779247c89d459ab8edbd80aba863ca Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 4 Sep 2024 15:09:05 -0400 Subject: [PATCH 2/3] Remove more file-based tests --- notebooks/create_fake_data.ipynb | 48 +++----------------------------- src/kbmod/data_interface.py | 48 -------------------------------- tests/test_fake_data_creator.py | 36 ------------------------ 3 files changed, 4 insertions(+), 128 deletions(-) diff --git a/notebooks/create_fake_data.ipynb b/notebooks/create_fake_data.ipynb index 71dbd830e..7f81c921c 100644 --- a/notebooks/create_fake_data.ipynb +++ b/notebooks/create_fake_data.ipynb @@ -4,8 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Create Fake Data\n", - "**Warning**: For the purposes this example we create (and possibly delete) a data directory at base_dir/fake_data. If that directory already exists, this notebook may overwrite the contents." + "# Create Fake Data\n" ] }, { @@ -86,45 +85,6 @@ " print(f\"{i}: t={ti:.3f} at ({px}, {py})\")" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Save the fake image files\n", - "\n", - "We save the fake images to a given base directory." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dir_path = \"./fake_data\"\n", - "ds.save_fake_data_to_dir(dir_path)\n", - "print(os.listdir(dir_path))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Delete the fake data files.\n", - "\n", - "We can (optionally) delete the fake data files using the delete_fake_data() function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ds.delete_fake_data_dir(dir_path)\n", - "print(os.listdir(dir_path))" - ] - }, { "cell_type": "code", "execution_count": null, @@ -135,9 +95,9 @@ ], "metadata": { "kernelspec": { - "display_name": "Python (.conda-kbmod)", + "display_name": "Jeremy's KBMOD", "language": "python", - "name": "conda-env-.conda-kbmod-py" + "name": "kbmod_jk" }, "language_info": { "codemirror_mode": { @@ -149,7 +109,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.1" + "version": "3.12.2" } }, "nbformat": 4, diff --git a/src/kbmod/data_interface.py b/src/kbmod/data_interface.py index daef33fd6..82c634832 100644 --- a/src/kbmod/data_interface.py +++ b/src/kbmod/data_interface.py @@ -92,54 +92,6 @@ def load_deccam_layered_image(filename, psf): return img -def save_deccam_layered_image(img, filename, wcs=None, overwrite=True): - """Save a layered image to the legacy deccam format. - - Parameters - ---------- - img : `LayeredImage` - The image to save. - filename : `str` - The name of the file to save. - wcs : `astropy.wcs.WCS`, optional - The WCS of the image. If provided appends this information to the header. - overwrite : `bool` - Indicates whether to overwrite the current file if it exists. - - Raises - ------ - Raises a ``ValueError`` if the file exists and ``overwrite`` is ``False``. - """ - if Path(filename).is_file() and not overwrite: - raise ValueError(f"{filename} already exists") - - hdul = fits.HDUList() - - # Create the primary header. - pri = fits.PrimaryHDU() - pri.header["MJD"] = img.get_obstime() - if wcs is not None: - append_wcs_to_hdu_header(wcs, pri.header) - hdul.append(pri) - - # Append the science layer. - sci_hdu = raw_image_to_hdu(img.get_science(), wcs) - sci_hdu.name = f"science" - hdul.append(sci_hdu) - - # Append the mask layer. - msk_hdu = raw_image_to_hdu(img.get_mask(), wcs) - msk_hdu.name = f"mask" - hdul.append(msk_hdu) - - # Append the variance layer. - var_hdu = raw_image_to_hdu(img.get_variance(), wcs) - var_hdu.name = f"variance" - hdul.append(var_hdu) - - hdul.writeto(filename, overwrite=overwrite) - - def load_input_from_individual_files( im_filepath, time_file, diff --git a/tests/test_fake_data_creator.py b/tests/test_fake_data_creator.py index 3428c1a5a..ea92205b8 100644 --- a/tests/test_fake_data_creator.py +++ b/tests/test_fake_data_creator.py @@ -83,42 +83,6 @@ def test_insert_object(self): pix_val = ds.stack.get_single_image(i).get_science().get_pixel(py, px) self.assertGreaterEqual(pix_val, 50.0) - def test_save_and_clean(self): - num_images = 7 - ds = FakeDataSet(64, 64, create_fake_times(num_images)) - - with tempfile.TemporaryDirectory() as dir_name: - # Get all the file names. - filenames = [] - for i in range(num_images): - filenames.append(os.path.join(dir_name, "%06i.fits" % i)) - - # Check no data exists yet. - for name in filenames: - self.assertFalse(Path(name).exists()) - - # Save the data and check the data now exists. - ds.save_fake_data_to_dir(dir_name) - for name in filenames: - self.assertTrue(Path(name).exists()) - - # Clean the data and check the data no longer exists. - ds.delete_fake_data_dir(dir_name) - for name in filenames: - self.assertFalse(Path(name).exists()) - - def test_save_times(self): - num_images = 50 - ds = FakeDataSet(4, 4, create_fake_times(num_images)) - - with tempfile.TemporaryDirectory() as dir_name: - file_name = f"{dir_name}/times.dat" - ds.save_time_file(file_name) - self.assertTrue(Path(file_name).exists()) - - time_load = FileUtils.load_time_dictionary(file_name) - self.assertEqual(len(time_load), 50) - def test_save_work_unit(self): num_images = 25 ds = FakeDataSet(15, 10, create_fake_times(num_images)) From 9e19f01b910aa76027fcf18501d16a55221ba960 Mon Sep 17 00:00:00 2001 From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com> Date: Wed, 4 Sep 2024 15:17:00 -0400 Subject: [PATCH 3/3] Update fake_data_creator.py --- src/kbmod/fake_data/fake_data_creator.py | 55 ------------------------ 1 file changed, 55 deletions(-) diff --git a/src/kbmod/fake_data/fake_data_creator.py b/src/kbmod/fake_data/fake_data_creator.py index f9e06c3b4..9306f4e01 100644 --- a/src/kbmod/fake_data/fake_data_creator.py +++ b/src/kbmod/fake_data/fake_data_creator.py @@ -14,7 +14,6 @@ from astropy.io import fits from kbmod.configuration import SearchConfiguration -from kbmod.data_interface import save_deccam_layered_image from kbmod.file_utils import * from kbmod.search import * from kbmod.search import Logging @@ -270,60 +269,6 @@ def insert_random_object(self, flux): return t - def save_fake_data_to_dir(self, data_dir): - """Create the fake data in a given directory. - - Parameters - ---------- - data_dir : `str` - The path of the directory for the fake data. - """ - # Make the subdirectory if needed. - dir_path = Path(data_dir) - if not dir_path.is_dir(): - logger.info(f"Directory {data_dir} does not exist. Creating.") - os.mkdir(data_dir) - - # Save each of the image files. - for i in range(self.stack.img_count()): - img = self.stack.get_single_image(i) - filename = os.path.join(data_dir, ("%06i.fits" % i)) - - # If the file already exists, delete it. - if Path(filename).exists(): - os.remove(filename) - - # Save the file. - save_deccam_layered_image(img, filename, wcs=self.fake_wcs) - - def save_time_file(self, file_name): - """Save the mapping of visit ID -> timestamp to a file. - - Parameters - ---------- - file_name : `str` - The file name for the timestamp file. - """ - mapping = {} - for i in range(self.num_times): - id_str = "%06i" % i - mapping[id_str] = self.times[i] - FileUtils.save_time_dictionary(file_name, mapping) - - def delete_fake_data_dir(self, data_dir): - """Remove the fake data in a given directory. - - Parameters - ---------- - data_dir : `str` - The path of the directory for the fake data. - """ - for i in range(self.stack.img_count()): - img = self.stack.get_single_image(i) - filename = os.path.join(data_dir, ("%06i.fits" % i)) - if Path(filename).exists(): - os.remove(filename) - def save_fake_data_to_work_unit(self, filename, config=None): """Create the fake data in a WorkUnit file.