Add vectorized image creation.

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)]

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

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,73 +159,42 @@ 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)
             if self.noise_std != 1.0:
                 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):

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()