Repository: LSSTDESC/BlendingToolKit

Issue:
add a unit test on drawing real galaxies

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC) Report
Attention: Patch coverage is `85.71429%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 63.90%. Comparing base [(`3eebcbd`)](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/commit/3eebcbdd43a09a562fbf9c42d5f59c5be51dbc22?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC) to head [(`cd4605d`)](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC).

| [Files](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC) | Patch % | Lines |
|---|---|---|
| [btk/draw\_blends.py](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485?src=pr&el=tree&filepath=btk%2Fdraw_blends.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC#diff-YnRrL2RyYXdfYmxlbmRzLnB5) | 85.71% | [1 Missing :warning: ](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##              dev     #485      +/-   ##
==========================================
+ Coverage   61.69%   63.90%   +2.20%     
==========================================
  Files          19       19              
  Lines        1569     1568       -1     
==========================================
+ Hits          968     1002      +34     
+ Misses        601      566      -35     
```

| [Flag](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC) | `63.90% <85.71%> (+2.20%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC#carryforward-flags-in-the-pull-request-comment) to find out more.


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/LSSTDESC/BlendingToolKit/pull/485?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=LSSTDESC).


Patch:
diff --git a/btk/draw_blends.py b/btk/draw_blends.py
index 2486a5e..db14b70 100644
--- a/btk/draw_blends.py
+++ b/btk/draw_blends.py
@@ -14,7 +14,7 @@ from tqdm.auto import tqdm
 
 from btk.blend_batch import BlendBatch, MultiResolutionBlendBatch
 from btk.blend_generator import BlendGenerator
-from btk.catalog import Catalog
+from btk.catalog import Catalog, CosmosCatalog
 from btk.multiprocess import get_current_process, multiprocess
 from btk.sampling_functions import SamplingFunction
 from btk.survey import Filter, Survey, make_wcs
@@ -518,7 +518,7 @@ class CosmosGenerator(DrawBlendsGenerator):
 
     def __init__(
         self,
-        catalog: Catalog,
+        catalog: CosmosCatalog,
         sampling_function: SamplingFunction,
         surveys: List[Survey],
         batch_size: int = 8,
@@ -531,6 +531,7 @@ class CosmosGenerator(DrawBlendsGenerator):
         apply_shear: bool = False,
         augment_data: bool = False,
         gal_type: str = "real",
+        noise_pad_size: float = 0,
     ):
         """Initializes the CosmosGenerator class. See parent class for most attributes.
 
@@ -549,6 +550,12 @@ class CosmosGenerator(DrawBlendsGenerator):
             augment_data: See parent class.
             gal_type: string to specify the type of galaxy simulations.
                             Either "real" (default) or "parametric".
+            noise_pad_size: For realistic galaxies, the size of region to pad with noise for
+                            each individual galaxy before forming the blend. This is the argument
+                            to the `COSMOSCatalog.makeGalaxy` function inside `galsim`. In BTK, we
+                            add noise after the blend is produced according to the `sky_level`,
+                            so the default for this argument is 0, as otherwise, noise would
+                            be added to the image twice.
         """
         super().__init__(
             catalog,
@@ -570,6 +577,7 @@ class CosmosGenerator(DrawBlendsGenerator):
                 f"gal_type must be either 'real' or 'parametric', but you provided {gal_type}"
             )
         self.gal_type = gal_type
+        self.noise_pad_size = noise_pad_size
 
     def _check_compatibility(self, survey: Survey) -> None:
         if type(self.catalog).__name__ not in self.compatible_catalogs:
@@ -577,24 +585,26 @@ class CosmosGenerator(DrawBlendsGenerator):
                 f"The catalog provided is of the wrong type. The types of "
                 f"catalogs available for the {type(self).__name__} are {self.compatible_catalogs}"
             )
-        for band in survey.available_filters:
-            if f"{survey.name}_{band}" not in self.catalog.table.keys():
-                raise ValueError(
-                    f"The {band} filter of the survey {survey.name} "
-                    f"has no associated magnitude in the given catalog."
-                )
 
     def render_single(self, entry: Table, filt: Filter, psf: galsim.GSObject, survey: Survey):
         """Returns the Galsim Image of an isolated galaxy."""
-        galsim_catalog = self.catalog.get_galsim_catalog()
+        galsim_catalog: galsim.COSMOSCatalog = self.catalog.get_galsim_catalog()
+
+        # we optionally check if additional entres if of the form `f'{survey_name}_{filter_name}`
+        # were added to the catalog in which case we use that (assuming it's an AB magnitude).
+        # otherwise we simply use the COSMOS magnitudes which AB magnitudes so we can apply
+        # the standard approach to converting to fluxes.
+        if f"{survey.name}_{filt.name}" in entry.colnames:
+            gal_mag = entry[f"{survey.name}_{filt.name}"]
+        else:
+            gal_mag = entry["MAG"]
 
-        # get galaxy flux
-        mag_name = f"{survey.name}_{filt.name}"
-        gal_mag = entry[mag_name]
         gal_flux = mag2counts(gal_mag, survey, filt).to_value("electron")
 
         index = entry["btk_index"]
-        gal = galsim_catalog.makeGalaxy(index, gal_type=self.gal_type, noise_pad_size=0)
+        gal = galsim_catalog.makeGalaxy(
+            index, gal_type=self.gal_type, noise_pad_size=self.noise_pad_size
+        )
         gal = gal.withFlux(gal_flux)
         gal = gal.rotate(galsim.Angle(entry["btk_rotation"], unit=galsim.degrees))
         if self.apply_shear:


Test patch:
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index e61c905..b7df278 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -21,7 +21,6 @@ def test_pipeline(data_dir):
     fltr = survey.get_filter("r")
     assert hasattr(fltr, "psf")
 
-    # single bright galaxy
     stamp_size = 24.0
     max_shift = 1.0
     max_n_sources = 2
@@ -115,6 +114,57 @@ def test_pipeline(data_dir):
         assert deblend_batch.image_size == deblend_batch2.image_size
 
 
+def test_cosmos_generator(data_dir):
+    """Test the pipeline as a whole for a single deblender."""
+    cosmos_catalog_paths = [
+        data_dir / "cosmos" / "real_galaxy_catalog_23.5_example.fits",
+        data_dir / "cosmos" / "real_galaxy_catalog_23.5_example_fits.fits",
+    ]
+    cosmos_catalog_files = [p.as_posix() for p in cosmos_catalog_paths]
+    catalog = btk.catalog.CosmosCatalog.from_file(cosmos_catalog_files)
+
+    _ = catalog.get_raw_catalog()
+
+    survey: Survey = btk.survey.get_surveys("LSST")
+    fltr = survey.get_filter("r")
+    assert hasattr(fltr, "psf")
+
+    stamp_size = 24.0
+    max_shift = 1.0
+    max_n_sources = 2
+    sampling_function = btk.sampling_functions.DefaultSampling(
+        max_number=max_n_sources,
+        min_number=1,
+        stamp_size=stamp_size,
+        max_shift=max_shift,
+        min_mag=20,
+        max_mag=21,
+        seed=SEED,
+        mag_name="MAG",
+    )
+
+    batch_size = 10
+
+    draw_generator = btk.draw_blends.CosmosGenerator(
+        catalog,
+        sampling_function,
+        survey,
+        batch_size=batch_size,
+        stamp_size=stamp_size,
+        njobs=1,
+        add_noise="all",
+        seed=SEED,
+        gal_type="real",
+    )
+
+    # generate batch 100 blend catalogs and images.
+    blend_batch = next(draw_generator)
+    assert len(blend_batch.catalog_list) == batch_size
+    assert blend_batch.blend_images.shape == (batch_size, 6, stamp_size / 0.2, stamp_size / 0.2)
+    iso_shape = (batch_size, max_n_sources, 6, stamp_size / 0.2, stamp_size / 0.2)
+    assert blend_batch.isolated_images.shape == iso_shape
+
+
 def test_sep(data_dir):
     """Check we always detect single bright objects."""