Parametric cosmos simulations (#256)

* parametric sims with cosmos cat * updated docstring * tutorial with parametric example * adding exclusion_level cuts * pre-commit changes * resolve pre-commit changes * using catalog directly from galsim * Update 01a-cosmos_tutorial.ipynb fixing kernel metadata * fixing kernel metadata * updating docstring * default exclusion_level as "marginal" * _selection.fits catalog * correcting path * added tests with parametric sims * updated config * removed assert on exclusion_level * updated add_noise Co-authored-by: thuiop <[email protected]> Co-authored-by: Alexandre Boucaud <[email protected]> Co-authored-by: Ismael Mendoza <[email protected]>
LSSTDESC · Feb 20, 2022 · 6b94110 · 6b94110
1 parent 580e052
commit 6b94110
Show file tree

Hide file tree

Showing 8 changed files with 146 additions and 13 deletions.
diff --git a/btk/catalog.py b/btk/catalog.py
@@ -108,19 +108,44 @@ def __init__(self, raw_catalog, galsim_catalog, verbose=False):
         self.galsim_catalog = galsim_catalog
 
     @classmethod
-    def from_file(cls, catalog_files, verbose=False):
-        """Constructs the catalog object from a file.
+    def from_file(cls, catalog_files, exclusion_level="marginal", verbose=False):
+        """Constructs the catalog object from a file. It also places exclusion level cuts.
+
+        For more details: (https://galsim-developers.github.io/GalSim/_build/html/real_gal.html)
 
         Args:
             catalog_files(list): list containing the two paths to the COSMOS data. Please see
                 the tutorial page for more details
                 (https://lsstdesc.org/BlendingToolKit/tutorials.html#using-cosmos-galaxies).
+            exclusion_level(str): Level of additional cuts to make on the galaxies based on the
+                quality of postage stamp definition and/or parametric fit quality [beyond the
+                minimal cuts imposed when making the catalog - see Mandelbaum et
+                al. (2012, MNRAS, 420, 1518) for details].
+                Options:
+                - "none": No cuts.
+                - "bad_stamp": Apply cuts to eliminate galaxies that have failures in
+                    postage stamp definition.  These cuts may also eliminate a small
+                    subset of the good postage stamps as well.
+                - "bad_fits": Apply cuts to eliminate galaxies that have failures in the
+                    parametric fits.  These cuts may also eliminate a small
+                    subset of the good parametric fits as well.
+                - "marginal": Apply the above cuts, plus ones that eliminate some more
+                    marginal cases.
+                Note that the _selection.fits file must be present in the same repo as the real
+                images catalog, Otherwise the "bad_stamp" and "marginal" cuts will fail
+                [default: "marginal"]
             verbose: whether to print verbose info.
         """
+        galsim_catalog = galsim.COSMOSCatalog(catalog_files[0], exclusion_level=exclusion_level)
+
         catalog_coord = astropy.table.Table.read(catalog_files[0])
         catalog_fit = astropy.table.Table.read(catalog_files[1])
+
+        catalog_coord = catalog_coord[galsim_catalog.orig_index]
+        catalog_fit = catalog_fit[galsim_catalog.orig_index]
+
         catalog = astropy.table.hstack([catalog_coord, catalog_fit])
-        galsim_catalog = galsim.COSMOSCatalog(catalog_files[0], exclusion_level="none")
+
         return cls(catalog, galsim_catalog, verbose=verbose)
 
     def _prepare_table(self, raw_catalog):

diff --git a/btk/draw_blends.py b/btk/draw_blends.py
@@ -518,7 +518,72 @@ def render_single(self, entry, filt, psf, survey, extra_data):
 
 
 class CosmosGenerator(DrawBlendsGenerator):
-    """Subclass of DrawBlendsGenerator for drawing real galaxies from the COSMOS catalog."""
+    """Subclass of DrawBlendsGenerator for drawing galaxies from the COSMOS catalog."""
+
+    def __init__(
+        self,
+        catalog,
+        sampling_function,
+        surveys: list,
+        batch_size=8,
+        stamp_size=24,
+        cpus=1,
+        verbose=False,
+        add_noise="all",
+        shifts=None,
+        indexes=None,
+        channels_last=False,
+        save_path=None,
+        seed=DEFAULT_SEED,
+        gal_type="real",
+    ):
+        """Initializes the CosmosGenerator class.
+
+        Args:
+            catalog (btk.catalog.Catalog): BTK catalog object from which galaxies are taken.
+            sampling_function (btk.sampling_function.SamplingFunction): BTK sampling
+                function to use.
+            surveys (list): List of btk Survey objects defining the observing conditions
+            batch_size (int): Number of blends generated per batch
+            stamp_size (float): Size of the stamps, in arcseconds
+            cpus (int): Number of cpus to use; defines the number of minibatches
+            verbose (bool): Indicates whether additionnal information should be printed
+            add_noise (str): Indicates if the blends should be generated with noise.
+                            "all" indicates that all the noise should be applied,
+                            "background" adds only the background noise,
+                            "galaxy" only the galaxy noise, and "none" gives noiseless
+                            images.
+            shifts (list): Contains arbitrary shifts to be applied instead of
+                           random shifts. Must be of length batch_size. Must be used
+                           with indexes. Used mostly for internal testing purposes.
+            indexes (list): Contains the ids of the galaxies to use in the stamp.
+                        Must be of length batch_size. Must be used with shifts.
+                        Used mostly for internal testing purposes.
+            channels_last (bool): Whether to return images as numpy arrays with the channel
+                                (band) dimension as the last dimension or before the pixels
+                                dimensions (default).
+            save_path (str): Path to a directory where results will be saved. If left
+                            as None, results will not be saved.
+            seed (int): Integer seed for reproducible random noise realizations.
+            gal_type (str): string to specify the type of galaxy simulations.
+                            Either "real" (default) or "parametric".
+        """
+        super().__init__(
+            catalog=catalog,
+            sampling_function=sampling_function,
+            surveys=surveys,
+            batch_size=batch_size,
+            stamp_size=stamp_size,
+            cpus=cpus,
+            verbose=verbose,
+            add_noise=add_noise,
+            shifts=shifts,
+            indexes=indexes,
+            channels_last=channels_last,
+            save_path=save_path,
+            seed=seed,
+        )
+        self.gal_type = gal_type
 
     compatible_catalogs = ("CosmosCatalog",)
 
@@ -553,7 +618,7 @@ def render_single(self, entry, filt, psf, survey, extra_data):
             gal_flux = get_flux(entry["ref_mag"], filt, survey)
 
         gal = galsim_catalog.makeGalaxy(
-            entry["btk_index"], gal_type="real", noise_pad_size=0
+            entry["btk_index"], gal_type=self.gal_type, noise_pad_size=0
         ).withFlux(gal_flux)
 
         pix_stamp_size = int(self.stamp_size / survey.pixel_scale)

diff --git a/conf/draw_blends/cosmos.yaml b/conf/draw_blends/cosmos.yaml
@@ -12,3 +12,4 @@ indexes: Null
 channels_last: ${channels_last}
 save_path: ${save_path}
 seed: ${seed}
+gal_type: ${gal_type}
diff --git a/data/cosmos/real_galaxy_catalog_23.5_example_selection.fits b/data/cosmos/real_galaxy_catalog_23.5_example_selection.fits
diff --git a/notebooks/01a-cosmos_tutorial.ipynb b/notebooks/01a-cosmos_tutorial.ipynb
@@ -56,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "id": "0b49f336",
    "metadata": {
     "ExecuteTime": {
@@ -68,7 +68,7 @@
    "source": [
     "stamp_size = 24.0\n",
     "batch_size = 8\n",
-    "catalog = btk.catalog.CosmosCatalog.from_file(COSMOS_CATALOG_PATHS)\n",
+    "catalog = btk.catalog.CosmosCatalog.from_file(COSMOS_CATALOG_PATHS, exclusion_level=\"none\") # To use default galsim exlusion level(\"marginal\"), _selection.fits file must be present \n",
     "sampling_function = btk.sampling_functions.DefaultSampling(stamp_size=stamp_size)\n",
     "survey = btk.survey.get_surveys(\"Rubin\")"
    ]
@@ -85,7 +85,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "id": "de364e04",
    "metadata": {
     "ExecuteTime": {
@@ -95,6 +95,7 @@
    },
    "outputs": [],
    "source": [
+    "gal_type = \"real\" # to choose between real and parametric models\n",
     "draw_generator = btk.draw_blends.CosmosGenerator(\n",
     "        catalog,\n",
     "        sampling_function,\n",
@@ -104,12 +105,13 @@
     "        cpus=1,\n",
     "        add_noise=\"all\",\n",
     "        verbose=False,\n",
+    "        gal_type=gal_type,\n",
     "    )"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "id": "e7bb0a6e",
    "metadata": {
     "ExecuteTime": {

diff --git a/tests/test_catalog.py b/tests/test_catalog.py
@@ -24,15 +24,32 @@ def test_reading_cosmos_catalog():
     return cosmos_catalog
 
 
+def test_reading_cosmos_catalog_without_cuts():
+    """Returns the cosmos catalog"""
+
+    cosmos_catalog = CosmosCatalog.from_file(COSMOS_CATALOG_PATHS, exclusion_level="none")
+    return cosmos_catalog
+
+
 def test_getting_galsim_catalog():
     """Returns the galsim catalog"""
 
     cosmos_catalog = CosmosCatalog.from_file(COSMOS_CATALOG_PATHS)
     galsim_catalog = cosmos_catalog.get_galsim_catalog()
+    # assert that the merged and galsim catalog have same length
+    assert len(galsim_catalog) == len(cosmos_catalog.table), "error reading catalog with cuts"
+    return galsim_catalog
+
+
+def test_getting_galsim_catalog_without_cuts():
+    """Returns the galsim catalog"""
+
+    cosmos_catalog = CosmosCatalog.from_file(COSMOS_CATALOG_PATHS, exclusion_level="none")
+    galsim_catalog = cosmos_catalog.get_galsim_catalog()
     return galsim_catalog
 
 
 def test_verbose():
     """Testing the verbose option"""
     CatsimCatalog.from_file(CATALOG_PATH, verbose=True)
-    CosmosCatalog.from_file(COSMOS_CATALOG_PATHS, verbose=True)
+    CosmosCatalog.from_file(COSMOS_CATALOG_PATHS, verbose=True, exclusion_level="none")
diff --git a/tests/test_cosmos.py b/tests/test_cosmos.py
@@ -16,7 +16,7 @@
 ]
 
 
-def test_cosmos_galaxies():
+def test_cosmos_galaxies_real():
     stamp_size = 24.0
     batch_size = 2
     catalog = CosmosCatalog.from_file(COSMOS_CATALOG_PATHS)
@@ -32,6 +32,29 @@ def test_cosmos_galaxies():
         cpus=1,
         add_noise="all",
         verbose=True,
+        gal_type="real",
+    )
+
+    _ = next(draw_generator)
+
+
+def test_cosmos_galaxies_parametric():
+    stamp_size = 24.0
+    batch_size = 2
+    catalog = CosmosCatalog.from_file(COSMOS_CATALOG_PATHS)
+    sampling_function = DefaultSampling(stamp_size=stamp_size)
+    HST = get_surveys("HST")
+
+    draw_generator = CosmosGenerator(
+        catalog,
+        sampling_function,
+        HST,
+        batch_size=batch_size,
+        stamp_size=stamp_size,
+        cpus=1,
+        add_noise="all",
+        verbose=True,
+        gal_type="parametric",
     )
 
     _ = next(draw_generator)
@@ -40,7 +63,7 @@ def test_cosmos_galaxies():
 def test_cosmos_ext_galaxies():
     stamp_size = 24.0
     batch_size = 2
-    catalog = CosmosCatalog.from_file(COSMOS_EXT_CATALOG_PATHS)
+    catalog = CosmosCatalog.from_file(COSMOS_EXT_CATALOG_PATHS, exclusion_level="none")
     sampling_function = DefaultSampling(stamp_size=stamp_size)
     HST = get_surveys("HST")
 

diff --git a/tests/test_error_cases.py b/tests/test_error_cases.py
@@ -252,7 +252,7 @@ def test_incompatible_catalogs():
             add_noise=add_noise,
         )
 
-    catalog = CosmosCatalog.from_file(COSMOS_CATALOG_PATHS)
+    catalog = CosmosCatalog.from_file(COSMOS_CATALOG_PATHS, exclusion_level="none")
     with pytest.raises(ValueError):
         draw_generator = CatsimGenerator(  # noqa: F841
             catalog,