Merge pull request #222 from mr-superonion/master

Allow user to set separation for grid and hex layouts

descwl_shear_sims/galaxies.py

@@ -48,7 +48,7 @@ def make_galaxy_catalog(
         Can be sent for fixed galaxy catalog.  See DEFAULT_FIXED_GAL_CONFIG
         for defaults mag, hlr and morph
     sep: float, optional
-        Separation of pair in arcsec for layout='pair'
+        Separation of pair in arcsec for layout='pair', 'grid' or 'hex'
     """
 
     if layout is None and gal_type == 'wldeblend':
@@ -89,6 +89,7 @@ def make_galaxy_catalog(
             galaxy_catalog = WLDeblendGalaxyCatalog(
                 rng=rng,
                 layout=layout,
+                sep=sep,
             )
         elif gal_type in ['fixed', 'varying', 'exp']:  # TODO remove exp
             gal_config = get_fixed_gal_config(config=gal_config)
@@ -104,6 +105,7 @@ def make_galaxy_catalog(
                 hlr=gal_config['hlr'],
                 morph=gal_config['morph'],
                 layout=layout,
+                sep=sep,
             )
 
         else:
@@ -164,6 +166,8 @@ class FixedGalaxyCatalog(object):
         for layout 'grid'.  Default 0.
     pixel_scale: float, optional
         pixel scale in arcsec
+    sep: float | None
+        Separation of galaxies in arcsec
     """
     def __init__(
         self, *,
@@ -175,6 +179,7 @@ def __init__(
         coadd_dim=None,
         buff=0,
         pixel_scale=SCALE,
+        sep=None,
     ):
         self.gal_type = 'fixed'
         self.morph = morph
@@ -188,6 +193,7 @@ def __init__(
             self.layout = layout
         self.shifts_array = self.layout.get_shifts(
             rng=rng,
+            sep=sep,
         )
 
     def __len__(self):
@@ -284,6 +290,8 @@ class GalaxyCatalog(FixedGalaxyCatalog):
         for layout 'grid'.  Default 0.
     pixel_scale: float
         pixel scale in arcsec
+    sep: float | None
+        Separation of galaxies in arcsec
     """
     def __init__(
         self, *,
@@ -295,6 +303,7 @@ def __init__(
         coadd_dim=None,
         buff=0,
         pixel_scale=SCALE,
+        sep=None,
     ):
         super().__init__(
             rng=rng,
@@ -305,6 +314,7 @@ def __init__(
             mag=mag,
             hlr=hlr,
             morph=morph,
+            sep=sep,
         )
         self.gal_type = 'varying'
 
@@ -637,6 +647,11 @@ class WLDeblendGalaxyCatalog(object):
         lower limits of the slection cuts
     select_upper_limit: list | ndarray
         upper limits of the slection cuts
+    sep: float
+        Separation of galaxies in arcsec
+    indice_id: None | int
+        galaxy index to use, use galaxies in the range between indice_id * num
+        and (indice_id + 1) * num
     """
     def __init__(
         self,
@@ -649,6 +664,8 @@ def __init__(
         select_observable=None,
         select_lower_limit=None,
         select_upper_limit=None,
+        sep=None,
+        indice_id=None,
     ):
         self.gal_type = 'wldeblend'
         self.rng = rng
@@ -661,6 +678,8 @@ def __init__(
 
         # one square degree catalog, convert to arcmin
         density = self._wldeblend_cat.size / (60 * 60)
+        if buff is None:
+            buff = 0
         if isinstance(layout, str):
             self.layout = Layout(layout, coadd_dim, buff, pixel_scale)
         else:
@@ -669,15 +688,27 @@ def __init__(
         self.shifts_array = self.layout.get_shifts(
             rng=rng,
             density=density,
+            sep=sep,
         )
 
         # randomly sample from the catalog
         num = len(self)
-        self.indices = self.rng.randint(
-            0,
-            self._wldeblend_cat.size,
-            size=num,
-        )
+        if indice_id is None:
+            self.indices = self.rng.randint(
+                0,
+                self._wldeblend_cat.size,
+                size=num,
+            )
+        else:
+            indice_min = indice_id * num
+            indice_max = min(indice_min + num, self._wldeblend_cat.size)
+            if indice_min >= self._wldeblend_cat.size:
+                raise ValueError("indice_min too large")
+            self.indices = np.arange(
+                indice_min,
+                indice_max,
+                dtype=int,
+            )
         # do a random rotation for each galaxy
         self.angles = self.rng.uniform(low=0, high=360, size=num)
 

descwl_shear_sims/layout/layout.py

@@ -37,7 +37,7 @@ def __init__(
         ----------
         layout_name: string
             'grid', 'pair', 'hex', or 'random'
-        coadd_dim: int
+        coadd_dim: int | None
             Dimensions of final coadd
         buff: int, optional
             Buffer region where no objects will be drawn.  Default 0.
@@ -117,7 +117,7 @@ def get_shifts(
         density: float, optional
             galaxy number density [/arcmin^2] ,default set to RANDOM_DENSITY
         sep: float, optional
-            The separation in arcseconds for layout='pair'
+            The separation in arcseconds for layout='pair', 'grid' or 'hex'
         """
 
         if self.layout_name == 'pair':
@@ -130,20 +130,24 @@ def get_shifts(
             )
         else:
             if self.layout_name == 'grid':
+                if sep is None:
+                    sep = GRID_SPACING
                 shifts = get_grid_shifts(
                     rng=rng,
                     dim=self.coadd_dim,
                     buff=self.buff,
                     pixel_scale=self.pixel_scale,
-                    spacing=GRID_SPACING,
+                    spacing=sep,
                 )
             elif self.layout_name == 'hex':
+                if sep is None:
+                    sep = HEX_SPACING
                 shifts = get_hex_shifts(
                     rng=rng,
                     dim=self.coadd_dim,
                     buff=self.buff,
                     pixel_scale=self.pixel_scale,
-                    spacing=HEX_SPACING,
+                    spacing=sep,
                 )
             elif self.layout_name == 'random':
                 # area covered by objects

descwl_shear_sims/sim.py

@@ -35,7 +35,7 @@
     "se_dim": None,
     "buff": 0,
     "layout": "grid",
-    "sep": None,  # sep in arcsec for layout=pair
+    "sep": None,
     "dither": False,
     "rotate": False,
     "bands": ["i"],