Split up render_galaxy in base decoder to avoid repeated code in lens…

…ing decoder
prob-ml · Sep 14, 2024 · c02ee4d · c02ee4d
1 parent f4806f5
commit c02ee4d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 25 deletions.
diff --git a/bliss/simulator/decoder.py b/bliss/simulator/decoder.py
@@ -53,17 +53,16 @@ def render_star(self, psf, band, source_params):
         """
         return psf[band].withFlux(source_params["fluxes"][band].item())
 
-    def render_galaxy(self, psf, band, source_params):
-        """Render a galaxy with given params and PSF.
+    def render_bulge_plus_disk(self, band, source_params):
+        """Render a galaxy with given params.
 
         Args:
-            psf (List): a list of PSFs for each band
             band (int): band
             source_params (Tensor): Tensor containing the parameters for a particular source
                 (see prior.py for details about these parameters)
 
         Returns:
-            GSObject: a galsim representation of the rendered galaxy convolved with the PSF
+            GSObject: a galsim representation of the rendered galaxy
         """
         disk_flux = source_params["fluxes"][band] * source_params["galaxy_disk_frac"]
         bulge_frac = 1 - source_params["galaxy_disk_frac"]
@@ -83,7 +82,21 @@ def render_galaxy(self, psf, band, source_params):
             bulge = galsim.DeVaucouleurs(flux=bulge_flux, half_light_radius=bulge_hlr_arcsecs)
             sheared_bulge = bulge.shear(q=source_params["galaxy_bulge_q"].item(), beta=beta)
             components.append(sheared_bulge)
-        galaxy = galsim.Add(components)
+        return galsim.Add(components)
+
+    def render_galaxy(self, psf, band, source_params):
+        """Render a galaxy with given params and PSF.
+
+        Args:
+            psf (List): a list of PSFs for each band
+            band (int): band
+            source_params (Tensor): Tensor containing the parameters for a particular source
+                (see prior.py for details about these parameters)
+
+        Returns:
+            GSObject: a galsim representation of the rendered galaxy convolved with the PSF
+        """
+        galaxy = self.render_bulge_plus_disk(band, source_params)
         return galsim.Convolution(galaxy, psf[band])
 
     @property

diff --git a/case_studies/weak_lensing/lensing_decoder.py b/case_studies/weak_lensing/lensing_decoder.py
@@ -1,5 +1,4 @@
 import galsim
-import numpy as np
 
 from bliss.simulator.decoder import Decoder
 
@@ -17,25 +16,7 @@ def render_galaxy(self, psf, band, source_params):
         Returns:
             GSObject: a galsim representation of the rendered galaxy convolved with the PSF
         """
-        disk_flux = source_params["fluxes"][band] * source_params["galaxy_disk_frac"]
-        bulge_frac = 1 - source_params["galaxy_disk_frac"]
-        bulge_flux = source_params["fluxes"][band] * bulge_frac
-        beta = source_params["galaxy_beta_radians"] * galsim.radians
-
-        components = []
-        if disk_flux > 0:
-            b_d = source_params["galaxy_a_d"] * source_params["galaxy_disk_q"]
-            disk_hlr_arcsecs = np.sqrt(source_params["galaxy_a_d"] * b_d)
-            disk = galsim.Exponential(flux=disk_flux, half_light_radius=disk_hlr_arcsecs)
-            sheared_disk = disk.shear(q=source_params["galaxy_disk_q"].item(), beta=beta)
-            components.append(sheared_disk)
-        if bulge_flux > 0:
-            b_b = source_params["galaxy_a_b"] * source_params["galaxy_bulge_q"]
-            bulge_hlr_arcsecs = np.sqrt(source_params["galaxy_a_b"] * b_b)
-            bulge = galsim.DeVaucouleurs(flux=bulge_flux, half_light_radius=bulge_hlr_arcsecs)
-            sheared_bulge = bulge.shear(q=source_params["galaxy_bulge_q"].item(), beta=beta)
-            components.append(sheared_bulge)
-        galaxy = galsim.Add(components)
+        galaxy = self.render_bulge_plus_disk(band, source_params)
 
         shear = source_params["shear"]
         shear1, shear2 = shear