From 845725ead423418310688f878dc7df067f08ee82 Mon Sep 17 00:00:00 2001 From: DinoBektesevic Date: Tue, 26 Mar 2024 09:39:10 -0700 Subject: [PATCH] Add vectorized image creation. --- src/kbmod/mocking/catalogs.py | 34 ++++---------- src/kbmod/mocking/fits.py | 55 +++++++---------------- src/kbmod/mocking/fits_data.py | 74 ++++++++++--------------------- src/kbmod/mocking/test_mocking.py | 20 ++++----- 4 files changed, 55 insertions(+), 128 deletions(-) diff --git a/src/kbmod/mocking/catalogs.py b/src/kbmod/mocking/catalogs.py index 46085314..06fa5dad 100644 --- a/src/kbmod/mocking/catalogs.py +++ b/src/kbmod/mocking/catalogs.py @@ -34,8 +34,6 @@ def gen_catalog(n, param_ranges, seed=None): return cat -def gen_source_catalog(n, param_ranges, seed=None): - class CatalogFactory(abc.ABC): @abc.abstractmethod @@ -72,41 +70,25 @@ def gen_realization(self, *args, t=None, dt=None, **kwargs): return self.table.copy() return self.table - def mock(self, n=1): - if n == 1: - return self.table - return [self.table for i in range(n)] - class SimpleObjectCatalog(CatalogFactory): base_param_ranges = { - "amplitude": [50, 500], + "amplitude": [1, 100], "x_mean": [0, 4096], "y_mean": [0, 2048], - "vx": [500, 5000], - "vy": [500, 5000], - "stddev": [1, 2], + "vx": [500, 1000], + "vy": [500, 1000], + "stddev": [1, 1.8], "theta": [0, np.pi], } - dtype = np.dtype([ - ("amplitude", np.float32), - ("x_mean", np.float32), - ("y_mean", np.float32), - ("vx", np.float32), - ("vy", np.float32), - ("x_stddev", np.float32), - ("y_stddev", np.float32), - ("thetae", np.float32) - ] - def __init__(self, table, obstime=None): self.table = table self._realization = table.copy() self.obstime = 0 if obstime is None else obstime @classmethod - def from_params(cls, n=10, param_ranges=None): + def from_params(cls, n=100, param_ranges=None): param_ranges = {} if param_ranges is None else param_ranges tmp = cls.base_param_ranges.copy() tmp.update(param_ranges) @@ -121,7 +103,7 @@ def gen_realization(self, t=None, dt=None, **kwargs): self._realization["y_mean"] += self._realization["vy"] * dt return self._realization - def mock(self, n=1, dt=0.001): + def mock(self, n=1, **kwargs): if n == 1: - return self.gen_realization() - + return self.gen_realization(**kwargs) + return [self.gen_realization(**kwargs).copy() for i in range(n)] diff --git a/src/kbmod/mocking/fits.py b/src/kbmod/mocking/fits.py index 043a44d9..9a2dbc37 100644 --- a/src/kbmod/mocking/fits.py +++ b/src/kbmod/mocking/fits.py @@ -171,13 +171,13 @@ def __init__(self, shape, noise=0, noise_std=1, dt=0.001, source_catalog=None, o self.variance = HeaderFactory.from_base_ext({"EXTNAME": "VARIANCE", **dims}) self.mask = HeaderFactory.from_base_ext({"EXTNAME": "MASK", **dims}) - self.img_data = SimpleImage.simulate( + self.img_data = SimpleImage( shape=shape, noise=noise, - catalog=self.src_cat, - return_copy=False if self.obj_cat is None else True, + noise_std=noise_std, + catalog=self.src_cat ) - self.var_data = SimpleVariance(self.img_data.base_data, read_noise=noise, gain=1.0) + self.var_data = SimpleVariance(self.img_data.base, read_noise=noise, gain=1.0) self.mask_data = ZeroedData(np.zeros(shape)) # Now we can build the HDU map and the HDUList layout @@ -201,52 +201,27 @@ def from_defaults(cls, shape=(2048, 4096), add_static_sources=False, add_moving_ return cls(shape=shape, source_catalog=source_catalog, object_catalog=object_catalog, **kwargs) - def mock(self, n=1, catalog=None): - if n == 1: - # a valiant effort, but too slow... - return self.generate() - - shape = (n*2, *self.shape) - images = np.zeros(shape, dtype=np.float32) - mask_data = self.mask_data.mock() - - if self.noise != 0: - rng = np.random.default_rng() - rng.standard_normal(size=shape, dtype=np.float32, out=images) - images *= self.noise_std - images += self.noise - - # update the base image (static objects, bad columns etc.) - if self.src_cat is not None: - images += self.img_data.base_data + def mock(self, n=1): + obj_cats = None - # update the variance images if needed only - if self.var_data.gain != 1.0: - images[n:] /= self.var_data.gain - if self.var_data.read_noise != 0.0: - images[n:] += self.var_data.read_noise**2 + if self.obj_cat is not None: + obj_cats = self.obj_cat.mock(n, dt=self.dt) - # shared headers + images = self.img_data.mock(n, obj_cats) + variances = self.var_data.mock(images=images) + mask = self.mask_data.mock() imghdr, varhdr, maskhdr = self.image.mock(), self.variance.mock(), self.mask.mock() hduls = [] for i in range(n): - if self.obj_cat is not None: - new_cat = self.obj_cat.gen_realization(dt=self.dt) - new_img = add_model_objects(np.zeros(images[0].shape), new_cat, models.Gaussian2D(x_stddev=1, y_stddev=1)) - images[i] += new_img - images[n+1] += new_img - hduls.append(HDUList( hdus=[ PrimaryHDU(header=self.primary.mock()), CompImageHDU(header=imghdr, data=images[i]), - CompImageHDU(header=varhdr, data=images[n+i]), - CompImageHDU(header=maskhdr, data=mask_data), + CompImageHDU(header=varhdr, data=variances[i]), + CompImageHDU(header=maskhdr, data=mask), ] )) - - self._idx += n return hduls @@ -276,6 +251,6 @@ def __init__(self): def mock(self, n=1): if n==1: - return self.step() - hduls = [self.step() for i in range(n)] + return self.generate() + hduls = [self.generate() for i in range(n)] return hduls diff --git a/src/kbmod/mocking/fits_data.py b/src/kbmod/mocking/fits_data.py index d42f3e42..5519ed93 100644 --- a/src/kbmod/mocking/fits_data.py +++ b/src/kbmod/mocking/fits_data.py @@ -51,7 +51,7 @@ class DataFactory: 32: np.float32, 64: np.float64, # classic IEEE float and double - -32: np.float32 + -32: np.float32, -64: np.float64, } @@ -78,6 +78,11 @@ def __init__(self, image, read_noise, gain, return_copy=False, mutable=False): self.gain = gain super().__init__(base=image / gain + read_noise**2, return_copy=return_copy, mutable=mutable) + def mock(self, images=None): + if images is None: + return self.base + return images/self.gain + self.read_noise**2 + class SimpleMask(DataFactory): def __init__(self, mask, return_copy=False, mutable=False): @@ -154,31 +159,27 @@ class SimpleImage(DataFactory): noise_gen = rng.standard_normal model = models.Gaussian2D - @classmethod - def simulate(cls, shape, noise=0, catalog=None, **kwargs): - # make a blank image - img = np.zeros(shape, dtype=np.float32) + def __init__(self, image=None, shape=(1000, 1000), noise=0, noise_std=1.0, src_cat=None, **kwargs): + if image is None: + image = np.zeros(shape, dtype=np.float32) + self.shape = shape + else: + image = image + self.shape = image.shape - # add static sources to it - if catalog is not None: - img = add_model_objects(img, catalog.table, src_model) + if src_cat is not None: + add_model_objects(image, src_cat.table, self.model) - return cls(img, noise=noise, **kwargs) + super().__init__(image, False, False) - def __init__(self, image=None, shape=(1000, 1000), noise=0, noise_std=1.0, dtype=np.float32, return_copy=False, mutable=False): - if image is None: - image = np.zeros(self.shape, dtype=dtype) - super().__init__(image, return_copy, mutable) self._base_contains_data = image.sum() != 0 self.noise = noise + self.noise_std = noise_std - def mock(self, n=1, catalog=None, hdu=None, **kwargs): + def mock(self, n=1, obj_cats=None, **kwargs): shape = (n, *self.shape) images = np.zeros(shape, dtype=np.float32) - if self._base_contains_data: - images += self.base - if self.noise != 0: rng = np.random.default_rng() rng.standard_normal(size=shape, dtype=np.float32, out=images) @@ -186,41 +187,14 @@ def mock(self, n=1, catalog=None, hdu=None, **kwargs): images *= self.noise_std images += self.noise - if catalog is not None: - for img in images: - add_model_objects(img, catalog, self.model(x_stddev=1, y_stddev=1)) + if self._base_contains_data: + images += self.base + + if obj_cats is not None: + for i, (img, cat) in enumerate(zip(images, obj_cats)): + add_model_objects(img, cat, self.model(x_stddev=1, y_stddev=1)) return images -# # if no catalog of moving obj were given just get -# # the base image, otherwise add the new sources -# if catalog is None: -# base = self.get_base() -# else: -# base = self.draw_new(catalog) -# -# # finally, create poisson background every time and add it to -# # the base of the image -# bckg = self.noise_gen(size=base.shape) -# return base + bckg -# -# def get_base(self): -# if self.return_copy: -# return self.base.copy() -# return self.base -# -# def draw_new(self, catalog): -# if not self.return_copy and not self.mutable: -# raise RuntimeError( -# "Can not get_realization an image that is neither mutable nor " -# "copyable. Setting the array to mutable will update it in-place, " -# "while setting copyable returns a copy of the array." -# ) -# -# model = models.Gaussian2D(x_stddev=1, y_stddev=1) -# if self.mutable: -# return add_model_objects(self.base, catalog, model) -# else: -# return add_model_objects(self.base.copy(), catalog, model) class SimulatedImage(DataFactory): diff --git a/src/kbmod/mocking/test_mocking.py b/src/kbmod/mocking/test_mocking.py index d5c8882f..43e13557 100644 --- a/src/kbmod/mocking/test_mocking.py +++ b/src/kbmod/mocking/test_mocking.py @@ -5,36 +5,32 @@ import matplotlib.pyplot as plt start = timeit.default_timer() -#fits_files = DECamImdiffs().mock(50) +fits_files = DECamImdiffs().mock(50) print(f"DECam Imdiff Fits: {timeit.default_timer() - start} seconds") -#del fits_files +del fits_files start = timeit.default_timer() -#fits_files = EmptyFits().mock(50) +fits_files = EmptyFits().mock(50) print(f"EmptyFits: {timeit.default_timer() - start} seconds") -#del fits_files +del fits_files start = timeit.default_timer() fits_files = SimpleFits.from_defaults( - shape=(500, 500), + shape=(1000, 1000), add_static_sources=False, add_moving_objects=True, - noise=1000, + noise=100, + noise_std=20 ).mock(50) print(f"SimpleFits: {timeit.default_timer() - start} seconds") -#fits_factory = SimpleFits.from_defaults(add_moving_objects=False) -#print(f"Fits factory took: {timeit.default_timer() - start} seconds") - -#fits_files = [fits_factory.mock() for i in range(50)] -#print(f"Generated 100 images in: {timeit.default_timer() - start} seconds") start = timeit.default_timer() ic = ImageCollection.fromTargets(fits_files, force="TestDataStd") print(f"Image collection took: {timeit.default_timer() - start} seconds") print(ic) -_ = [plt.imsave(f"/home/dino/Downloads/tmp_fits/{i}.png", h[1].data) for i, h in enumerate(fits_files)] +_ = [plt.imsave(f"/home/dino/Downloads/tmp_fits/{i}.png", h[1].data, vmin=50, vmax=500) for i, h in enumerate(fits_files)] breakpoint()