WIP: implement a simple z-dependent shear

descwl_shear_sims/sim/__init__.py

@@ -11,3 +11,4 @@
 )
 from .star_catalogs import StarCatalog
 from .psfs import make_psf, make_ps_psf
+from .shear import ZSliceShear

descwl_shear_sims/sim/galaxy_catalogs.py

@@ -7,7 +7,6 @@
 from .constants import SCALE
 from ..cache_tools import cached_catalog_read
 
-
 DEFAULT_FIXED_GAL_CONFIG = {
     "mag": 17.0,
     "hlr": 0.5,
@@ -244,14 +243,21 @@ def get_objlist(self, *, survey, g1, g2):
 
         # object is already shifted, so this results in the scene
         # being sheared
-        objlist = [
-            self._get_galaxy(builder, band, i).shear(g1=g1, g2=g2)
-            for i in range(num)
-        ]
+        objlist = []
+        for i in range(num):
+            gal = self._get_galaxy(builder, band, i)
+            z = self._get_redshift(i)
+            objlist.append(
+                self._shear(gal, g1, g2, z)
+            )
 
         shifts = self.shifts.copy()
         return objlist, shifts
 
+    def _get_redshift(self, i):
+        index = self.indices[i]
+        return self._wldeblend_cat[index]['redshift']
+
     def _get_galaxy(self, builder, band, i):
         """
         Get a galaxy
@@ -289,6 +295,17 @@ def _get_galaxy(self, builder, band, i):
 
         return galaxy
 
+    def _shear(self, gal, g1, g2, redshift):
+        if isinstance(g1, float):
+            g1_use = g1
+        else:
+            g1_use = g1(redshift)
+        if isinstance(g2, float):
+            g2_use = g2
+        else:
+            g2_use = g2(redshift)
+        return gal.shear(g1=g1_use, g2=g2_use)
+
 
 def read_wldeblend_cat(rng):
     """

descwl_shear_sims/sim/shear.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+
+class ZSliceShear(object):
+    """
+    Class for storing a "z-slice"
+    shear. This has one shear value
+    for zmin<z<=zmax, and another
+    otherwise
+
+    Parameters
+    ----------
+    g: float
+        shear value for objects with
+    redshift zmin<z<=zmax
+    gother: float
+        shear value outside the interval
+    zmin: float
+        minimum redshift of interval
+    zmax: float
+        maximum redshift of interval
+    """
+    def __init__(self, g,
+                 gother, zmin,
+                 zmax):
+        self.g = g
+        self.zmin = zmin
+        self.zmax = zmax
+        self.gother = gother
+
+    def __call__(self, z):
+        """
+        Return the shear for
+        redshift z
+        """
+        return np.where((z > self.zmin) * (z <= self.zmax),
+                        self.g, self.gother)

descwl_shear_sims/tests/test_zslice_shear.py

@@ -0,0 +1,24 @@
+import numpy as np
+from ..sim import ZSliceShear
+
+
+def test_zslice_shear():
+    # input parameters
+    # for ZSliceShear
+    g = 0.2
+    gother = -0.2
+    zmin = 0.3
+    zmax = 0.7
+    s = ZSliceShear(g, gother,
+                    zmin, zmax)
+    # Generate some redshifts and
+    # assert that they get assigned
+    # the correct shears
+    zs = np.arange(0., 3., 300)
+    shears = s(zs)
+
+    # check the objects inside and outside
+    # the slice got the right shear values
+    in_zslice = (zs > zmin) * (zs <= zmax)
+    assert np.allclose(shears[in_zslice], g)
+    assert np.allclose(shears[~in_zslice], gother)