From c06e497dae44fa43e284bb9a7629b702348d8233 Mon Sep 17 00:00:00 2001
From: Axel <axel.guinot.astro@gmail.com>
Date: Mon, 12 Jul 2021 16:51:26 +0200
Subject: [PATCH 1/7] Implementation of correlated noise

---
 descwl_shear_sims/sim.py                      |  41 ++-
 .../test_shear_meas_corr_noise.py             | 260 ++++++++++++++++++
 2 files changed, 300 insertions(+), 1 deletion(-)
 create mode 100644 shear_meas_tests/test_shear_meas_corr_noise.py

diff --git a/descwl_shear_sims/sim.py b/descwl_shear_sims/sim.py
index 5a5cdc1..bdf03cd 100644
--- a/descwl_shear_sims/sim.py
+++ b/descwl_shear_sims/sim.py
@@ -38,6 +38,7 @@
 def make_sim(
     *,
     rng,
+    galsim_rng,
     galaxy_catalog,
     coadd_dim,
     g1,
@@ -54,6 +55,9 @@ def make_sim(
     cosmic_rays=False,
     bad_columns=False,
     star_bleeds=False,
+    corr_noise=False,
+    g1_noise=0.01,
+    g2_noise=0.01,
 ):
     """
     Make simulation data
@@ -62,6 +66,8 @@ def make_sim(
     ----------
     rng: numpy.random.RandomState
         Numpy random state
+    galsim_rng: galsim.BaseDeviate
+        Galsim random state (used for noise realisation)
     galaxy_catalog: catalog
         E.g. WLDeblendGalaxyCatalog or FixedGalaxyCatalog
     coadd_dim: int
@@ -94,6 +100,12 @@ def make_sim(
         If True, add cosmic rays
     bad_columns: bool
         If True, add bad columns
+    corr_noise: bool
+        If True, make correlated noise
+    g1_noise: float
+        g1 for shear correlated noise
+    g2_noise: float
+        g2 for shear correlated noise
     """
 
     coadd_wcs, coadd_bbox = make_coadd_dm_wcs(coadd_dim)
@@ -132,6 +144,7 @@ def make_sim(
         for epoch in range(epochs_per_band):
             exp = make_exp(
                 rng=rng,
+                galsim_rng=galsim_rng,
                 band=band,
                 noise=noise_per_epoch,
                 objlist=objlist,
@@ -149,6 +162,9 @@ def make_sim(
                 cosmic_rays=cosmic_rays,
                 bad_columns=bad_columns,
                 star_bleeds=star_bleeds,
+                corr_noise=corr_noise,
+                g1_noise=g1_noise,
+                g2_noise=g2_noise,
             )
             if galaxy_catalog.gal_type == 'wldeblend':
                 rescale_wldeblend_exp(
@@ -181,6 +197,7 @@ def make_sim(
 def make_exp(
     *,
     rng,
+    galsim_rng,
     band,
     noise,
     objlist,
@@ -198,6 +215,9 @@ def make_exp(
     cosmic_rays=False,
     bad_columns=False,
     star_bleeds=False,
+    corr_noise=False,
+    g1_noise=0.01,
+    g2_noise=0.01,
 ):
     """
     Make an SEObs
@@ -206,6 +226,8 @@ def make_exp(
     ----------
     rng: numpy.random.RandomState
         The random number generator
+    galsim_rng: galsim.BaseDeviate
+        Galsim random number generator
     band: str
         Band as a string, e.g. 'i'
     noise: float
@@ -243,6 +265,12 @@ def make_exp(
         If True, put in bad columns
     star_bleeds: bool
         If True, add bleed trails to stars
+    corr_noise: bool
+        If True, make correlated noise
+    g1_noise: float
+        g1 for shear correlated noise
+    g2_noise: float
+        g2 for shear correlated noise
     """
 
     dims = [dim]*2
@@ -289,7 +317,18 @@ def make_exp(
         offset=offset,
     )
 
-    image.array[:, :] += rng.normal(scale=noise, size=dims)
+    # Make correlated noise
+    if corr_noise:
+        corr_noise = galsim.UncorrelatedNoise(variance=noise**2.,
+                                          rng=galsim_rng,
+                                          scale=SCALE)
+        # Shear correlation in the noise
+        corr_noise.shear(g1=g1_noise, g2=g2_noise)
+        image.addNoise(corr_noise)
+    else:
+        # Maybe change this to use galsim noise generator?
+        # To fit with the correlated noise option
+        image.array[:, :] += rng.normal(scale=noise, size=dims)
 
     bmask = get_bmask(
         image=image,
diff --git a/shear_meas_tests/test_shear_meas_corr_noise.py b/shear_meas_tests/test_shear_meas_corr_noise.py
new file mode 100644
index 0000000..7d3cee7
--- /dev/null
+++ b/shear_meas_tests/test_shear_meas_corr_noise.py
@@ -0,0 +1,260 @@
+import time
+import copy
+import numpy as np
+import ngmix
+import tqdm
+import joblib
+import galsim
+
+import pytest
+
+import metadetect.lsst_metadetect as lsst_metadetect
+import descwl_shear_sims as sim
+from descwl_coadd.coadd import make_coadd_obs
+
+CONFIG = {
+    "model": "wmom",
+    "bmask_flags": 0,
+    "metacal": {
+        "use_noise_image": True,
+        "psf": "fitgauss",
+    },
+    "psf": {
+        "model": "gauss",
+        "lm_pars": {},
+        "ntry": 2,
+    },
+    "weight": {
+        "fwhm": 1.2,
+    },
+    "detect": {
+        "thresh": 10.0,
+    },
+    'meds': {},
+}
+
+
+def _make_lsst_sim(*, seed, g1, g2, g1_noise, g2_noise, layout):
+    rng = np.random.RandomState(seed=seed)
+    galsim_rng = galsim.BaseDeviate(seed=seed)
+
+    galaxy_catalog = sim.galaxies.make_galaxy_catalog(
+        rng=rng,
+        coadd_dim=sim.sim.DEFAULT_SIM_CONFIG["coadd_dim"],
+        buff=sim.sim.DEFAULT_SIM_CONFIG["buff"],
+        layout=layout,
+        gal_type='exp',
+    )
+
+    psf = sim.psfs.make_fixed_psf(psf_type='gauss')
+
+    sim_data = sim.make_sim(
+        rng=rng,
+        galsim_rng=galsim_rng,
+        galaxy_catalog=galaxy_catalog,
+        coadd_dim=sim.sim.DEFAULT_SIM_CONFIG["coadd_dim"],
+        g1=g1,
+        g2=g2,
+        psf=psf,
+        corr_noise=True,
+        g1_noise=g1_noise,
+        g2_noise=g2_noise,
+    )
+    return sim_data
+
+
+def _shear_cuts(arr):
+    msk = (
+        (arr['flags'] == 0)
+        & (arr['wmom_s2n'] > 10)
+        & (arr['wmom_T_ratio'] > 1.2)
+    )
+    return msk
+
+
+def _meas_shear_data(res):
+    msk = _shear_cuts(res['noshear'])
+    g1 = np.mean(res['noshear']['wmom_g'][msk, 0])
+    g2 = np.mean(res['noshear']['wmom_g'][msk, 1])
+
+    msk = _shear_cuts(res['1p'])
+    g1_1p = np.mean(res['1p']['wmom_g'][msk, 0])
+    msk = _shear_cuts(res['1m'])
+    g1_1m = np.mean(res['1m']['wmom_g'][msk, 0])
+    R11 = (g1_1p - g1_1m) / 0.02
+
+    msk = _shear_cuts(res['2p'])
+    g2_2p = np.mean(res['2p']['wmom_g'][msk, 1])
+    msk = _shear_cuts(res['2m'])
+    g2_2m = np.mean(res['2m']['wmom_g'][msk, 1])
+    R22 = (g2_2p - g2_2m) / 0.02
+
+    dt = [
+        ('g1', 'f8'),
+        ('g2', 'f8'),
+        ('R11', 'f8'),
+        ('R22', 'f8')]
+    return np.array([(g1, g2, R11, R22)], dtype=dt)
+
+
+def _bootstrap_stat(d1, d2, func, seed, nboot=500):
+    dim = d1.shape[0]
+    rng = np.random.RandomState(seed=seed)
+    stats = []
+    for _ in tqdm.trange(nboot, leave=False):
+        ind = rng.choice(dim, size=dim, replace=True)
+        stats.append(func(d1[ind], d2[ind]))
+    return stats
+
+
+def _meas_m_c_cancel(pres, mres):
+    x = np.mean(pres['g1'] - mres['g1'])/2
+    y = np.mean(pres['R11'] + mres['R11'])/2
+    m = x/y/0.02 - 1
+
+    x = np.mean(pres['g2'] + mres['g2'])/2
+    y = np.mean(pres['R22'] + mres['R22'])/2
+    c = x/y
+
+    return m, c
+
+
+def _boostrap_m_c(pres, mres):
+    m, c = _meas_m_c_cancel(pres, mres)
+    bdata = _bootstrap_stat(pres, mres, _meas_m_c_cancel, 14324, nboot=500)
+    merr, cerr = np.std(bdata, axis=0)
+    return m, merr, c, cerr
+
+
+def _run_sim_one(*, seed, mdet_seed, g1, g2, g1_noise, g2_noise, **kwargs):
+    sim_data = _make_lsst_sim(
+        seed=seed, 
+        g1=g1, g2=g2,
+        g1_noise=g1_noise, g2_noise=g2_noise, 
+        **kwargs,
+        )
+    coadd_obs = make_coadd_obs(
+        exps=sim_data['band_data']['i'],
+        coadd_wcs=sim_data['coadd_wcs'],
+        coadd_bbox=sim_data['coadd_bbox'],
+        psf_dims=sim_data['psf_dims'],
+        rng=np.random.RandomState(seed),
+        remove_poisson=False,  # no object poisson noise in sims
+    )
+
+    coadd_mbobs = ngmix.MultiBandObsList()
+    obslist = ngmix.ObsList()
+    obslist.append(coadd_obs)
+    coadd_mbobs.append(obslist)
+
+    md = lsst_metadetect.LSSTMetadetect(
+        copy.deepcopy(CONFIG),
+        coadd_mbobs,
+        np.random.RandomState(seed=mdet_seed),
+    )
+    md.go()
+    return md.result
+
+
+def run_sim(seed, mdet_seed, **kwargs):
+    # positive shear
+    _pres = _run_sim_one(
+        seed=seed, 
+        mdet_seed=mdet_seed, 
+        g1=0.02, g2=0, 
+        g1_noise=0., g2_noise=0.1,
+        **kwargs,
+        )
+    if _pres is None:
+        return None
+
+    # negative shear
+    _mres = _run_sim_one(
+        seed=seed,
+        mdet_seed=mdet_seed,
+        g1=-0.02, g2=0,
+        g1_noise=0., g2_noise=0.1,
+        **kwargs,
+        )
+    if _mres is None:
+        return None
+
+    return _meas_shear_data(_pres), _meas_shear_data(_mres)
+
+
+@pytest.mark.parametrize(
+    'layout,ntrial', [('grid', 500), ('random', 2500)]
+)
+def test_shear_meas(layout, ntrial):
+    nsub = max(ntrial // 100, 10)
+    nitr = ntrial // nsub
+    rng = np.random.RandomState(seed=116)
+    seeds = rng.randint(low=1, high=2**29, size=ntrial)
+    mdet_seeds = rng.randint(low=1, high=2**29, size=ntrial)
+
+    tm0 = time.time()
+
+    print("")
+
+    pres = []
+    mres = []
+    loc = 0
+    for itr in tqdm.trange(nitr):
+        jobs = [
+            joblib.delayed(run_sim)(
+                seeds[loc+i], mdet_seeds[loc+i], layout=layout,
+            )
+            for i in range(nsub)
+        ]
+        outputs = joblib.Parallel(n_jobs=-2, verbose=0, backend='loky')(jobs)
+
+        for out in outputs:
+            if out is None:
+                continue
+            pres.append(out[0])
+            mres.append(out[1])
+        loc += nsub
+
+        m, merr, c, cerr = _boostrap_m_c(
+            np.concatenate(pres),
+            np.concatenate(mres),
+        )
+        print(
+            (
+                "\n"
+                "nsims: %d\n"
+                "m [1e-3, 3sigma]: %s +/- %s\n"
+                "c [1e-5, 3sigma]: %s +/- %s\n"
+                "\n"
+            ) % (
+                len(pres),
+                m/1e-3,
+                3*merr/1e-3,
+                c/1e-5,
+                3*cerr/1e-5,
+            ),
+            flush=True,
+        )
+
+    total_time = time.time()-tm0
+    print("time per:", total_time/ntrial, flush=True)
+
+    pres = np.concatenate(pres)
+    mres = np.concatenate(mres)
+    m, merr, c, cerr = _boostrap_m_c(pres, mres)
+
+    print(
+        (
+            "\n\nm [1e-3, 3sigma]: %s +/- %s"
+            "\nc [1e-5, 3sigma]: %s +/- %s"
+        ) % (
+            m/1e-3,
+            3*merr/1e-3,
+            c/1e-5,
+            3*cerr/1e-5,
+        ),
+        flush=True,
+    )
+
+    assert np.abs(m) < max(1e-3, 3*merr)
+    assert np.abs(c) < 3*cerr

From 02fd39dbe4c982c575193ea861bbaa51b99466f0 Mon Sep 17 00:00:00 2001
From: Axel <axel.guinot.astro@gmail.com>
Date: Mon, 12 Jul 2021 17:19:16 +0200
Subject: [PATCH 2/7] Fix bugs

Pylint, compatibility with original scripts (gal_rng keyword)
---
 descwl_shear_sims/sim.py                      | 20 ++++++++++---------
 .../test_shear_meas_corr_noise.py             |  4 ++--
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/descwl_shear_sims/sim.py b/descwl_shear_sims/sim.py
index bdf03cd..793fc75 100644
--- a/descwl_shear_sims/sim.py
+++ b/descwl_shear_sims/sim.py
@@ -38,7 +38,6 @@
 def make_sim(
     *,
     rng,
-    galsim_rng,
     galaxy_catalog,
     coadd_dim,
     g1,
@@ -55,6 +54,7 @@ def make_sim(
     cosmic_rays=False,
     bad_columns=False,
     star_bleeds=False,
+    galsim_rng=None,
     corr_noise=False,
     g1_noise=0.01,
     g2_noise=0.01,
@@ -66,8 +66,6 @@ def make_sim(
     ----------
     rng: numpy.random.RandomState
         Numpy random state
-    galsim_rng: galsim.BaseDeviate
-        Galsim random state (used for noise realisation)
     galaxy_catalog: catalog
         E.g. WLDeblendGalaxyCatalog or FixedGalaxyCatalog
     coadd_dim: int
@@ -100,6 +98,8 @@ def make_sim(
         If True, add cosmic rays
     bad_columns: bool
         If True, add bad columns
+    galsim_rng: galsim.BaseDeviate
+        Galsim random state (used for noise realisation)
     corr_noise: bool
         If True, make correlated noise
     g1_noise: float
@@ -197,7 +197,6 @@ def make_sim(
 def make_exp(
     *,
     rng,
-    galsim_rng,
     band,
     noise,
     objlist,
@@ -215,6 +214,7 @@ def make_exp(
     cosmic_rays=False,
     bad_columns=False,
     star_bleeds=False,
+    galsim_rng=None,
     corr_noise=False,
     g1_noise=0.01,
     g2_noise=0.01,
@@ -226,8 +226,6 @@ def make_exp(
     ----------
     rng: numpy.random.RandomState
         The random number generator
-    galsim_rng: galsim.BaseDeviate
-        Galsim random number generator
     band: str
         Band as a string, e.g. 'i'
     noise: float
@@ -265,6 +263,8 @@ def make_exp(
         If True, put in bad columns
     star_bleeds: bool
         If True, add bleed trails to stars
+    galsim_rng: galsim.BaseDeviate
+        Galsim random number generator
     corr_noise: bool
         If True, make correlated noise
     g1_noise: float
@@ -319,9 +319,11 @@ def make_exp(
 
     # Make correlated noise
     if corr_noise:
-        corr_noise = galsim.UncorrelatedNoise(variance=noise**2.,
-                                          rng=galsim_rng,
-                                          scale=SCALE)
+        corr_noise = galsim.UncorrelatedNoise(
+            variance=noise**2., 
+            rng=galsim_rng,
+            scale=SCALE,
+            )
         # Shear correlation in the noise
         corr_noise.shear(g1=g1_noise, g2=g2_noise)
         image.addNoise(corr_noise)
diff --git a/shear_meas_tests/test_shear_meas_corr_noise.py b/shear_meas_tests/test_shear_meas_corr_noise.py
index 7d3cee7..dbb7a4f 100644
--- a/shear_meas_tests/test_shear_meas_corr_noise.py
+++ b/shear_meas_tests/test_shear_meas_corr_noise.py
@@ -162,7 +162,7 @@ def run_sim(seed, mdet_seed, **kwargs):
         seed=seed, 
         mdet_seed=mdet_seed, 
         g1=0.02, g2=0, 
-        g1_noise=0., g2_noise=0.1,
+        g1_noise=0., g2_noise=0.1,      # High value for the purpose of testing
         **kwargs,
         )
     if _pres is None:
@@ -173,7 +173,7 @@ def run_sim(seed, mdet_seed, **kwargs):
         seed=seed,
         mdet_seed=mdet_seed,
         g1=-0.02, g2=0,
-        g1_noise=0., g2_noise=0.1,
+        g1_noise=0., g2_noise=0.1,      # High value for the purpose of testing
         **kwargs,
         )
     if _mres is None:

From d28658205f823fec132b38e52fff51760b0b0766 Mon Sep 17 00:00:00 2001
From: Axel <axel.guinot.astro@gmail.com>
Date: Tue, 13 Jul 2021 10:50:05 +0200
Subject: [PATCH 3/7] Update the handling of galsim rng

---
 descwl_shear_sims/sim.py                       | 12 ++++--------
 shear_meas_tests/test_shear_meas_corr_noise.py |  8 +++-----
 2 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/descwl_shear_sims/sim.py b/descwl_shear_sims/sim.py
index 793fc75..7a10000 100644
--- a/descwl_shear_sims/sim.py
+++ b/descwl_shear_sims/sim.py
@@ -54,7 +54,6 @@ def make_sim(
     cosmic_rays=False,
     bad_columns=False,
     star_bleeds=False,
-    galsim_rng=None,
     corr_noise=False,
     g1_noise=0.01,
     g2_noise=0.01,
@@ -98,8 +97,6 @@ def make_sim(
         If True, add cosmic rays
     bad_columns: bool
         If True, add bad columns
-    galsim_rng: galsim.BaseDeviate
-        Galsim random state (used for noise realisation)
     corr_noise: bool
         If True, make correlated noise
     g1_noise: float
@@ -144,7 +141,6 @@ def make_sim(
         for epoch in range(epochs_per_band):
             exp = make_exp(
                 rng=rng,
-                galsim_rng=galsim_rng,
                 band=band,
                 noise=noise_per_epoch,
                 objlist=objlist,
@@ -263,8 +259,6 @@ def make_exp(
         If True, put in bad columns
     star_bleeds: bool
         If True, add bleed trails to stars
-    galsim_rng: galsim.BaseDeviate
-        Galsim random number generator
     corr_noise: bool
         If True, make correlated noise
     g1_noise: float
@@ -319,11 +313,13 @@ def make_exp(
 
     # Make correlated noise
     if corr_noise:
+        galsim_seed = rng.randint(low=1, high=2**29, size=1)
+        galsim_rng = galsim.BaseDeviate(seed=galsim_seed)
         corr_noise = galsim.UncorrelatedNoise(
-            variance=noise**2., 
+            variance=noise**2.,
             rng=galsim_rng,
             scale=SCALE,
-            )
+        )
         # Shear correlation in the noise
         corr_noise.shear(g1=g1_noise, g2=g2_noise)
         image.addNoise(corr_noise)
diff --git a/shear_meas_tests/test_shear_meas_corr_noise.py b/shear_meas_tests/test_shear_meas_corr_noise.py
index dbb7a4f..c54d6f3 100644
--- a/shear_meas_tests/test_shear_meas_corr_noise.py
+++ b/shear_meas_tests/test_shear_meas_corr_noise.py
@@ -36,7 +36,6 @@
 
 def _make_lsst_sim(*, seed, g1, g2, g1_noise, g2_noise, layout):
     rng = np.random.RandomState(seed=seed)
-    galsim_rng = galsim.BaseDeviate(seed=seed)
 
     galaxy_catalog = sim.galaxies.make_galaxy_catalog(
         rng=rng,
@@ -50,7 +49,6 @@ def _make_lsst_sim(*, seed, g1, g2, g1_noise, g2_noise, layout):
 
     sim_data = sim.make_sim(
         rng=rng,
-        galsim_rng=galsim_rng,
         galaxy_catalog=galaxy_catalog,
         coadd_dim=sim.sim.DEFAULT_SIM_CONFIG["coadd_dim"],
         g1=g1,
@@ -132,7 +130,7 @@ def _run_sim_one(*, seed, mdet_seed, g1, g2, g1_noise, g2_noise, **kwargs):
         g1=g1, g2=g2,
         g1_noise=g1_noise, g2_noise=g2_noise, 
         **kwargs,
-        )
+    )
     coadd_obs = make_coadd_obs(
         exps=sim_data['band_data']['i'],
         coadd_wcs=sim_data['coadd_wcs'],
@@ -164,7 +162,7 @@ def run_sim(seed, mdet_seed, **kwargs):
         g1=0.02, g2=0, 
         g1_noise=0., g2_noise=0.1,      # High value for the purpose of testing
         **kwargs,
-        )
+    )
     if _pres is None:
         return None
 
@@ -175,7 +173,7 @@ def run_sim(seed, mdet_seed, **kwargs):
         g1=-0.02, g2=0,
         g1_noise=0., g2_noise=0.1,      # High value for the purpose of testing
         **kwargs,
-        )
+    )
     if _mres is None:
         return None
 

From 2dfb20d6c75d6667bddd5261d4ef3c58821b3d30 Mon Sep 17 00:00:00 2001
From: "Matthew R. Becker" <beckermr@users.noreply.github.com>
Date: Tue, 13 Jul 2021 05:29:24 -0500
Subject: [PATCH 4/7] Update test_shear_meas.py

---
 shear_meas_tests/test_shear_meas.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/shear_meas_tests/test_shear_meas.py b/shear_meas_tests/test_shear_meas.py
index 23bfe37..2bf0de6 100644
--- a/shear_meas_tests/test_shear_meas.py
+++ b/shear_meas_tests/test_shear_meas.py
@@ -30,6 +30,7 @@
         "thresh": 10.0,
     },
     'meds': {},
+    'subtract_sky': False,
 }
 
 

From 3b7d4ae1d926ae8698f7c03a4dd668dbb860eb7a Mon Sep 17 00:00:00 2001
From: "Matthew R. Becker" <beckermr@users.noreply.github.com>
Date: Tue, 13 Jul 2021 07:29:49 -0500
Subject: [PATCH 5/7] Update shear_meas_tests/test_shear_meas_corr_noise.py

---
 shear_meas_tests/test_shear_meas_corr_noise.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/shear_meas_tests/test_shear_meas_corr_noise.py b/shear_meas_tests/test_shear_meas_corr_noise.py
index c54d6f3..bcf8f42 100644
--- a/shear_meas_tests/test_shear_meas_corr_noise.py
+++ b/shear_meas_tests/test_shear_meas_corr_noise.py
@@ -31,6 +31,7 @@
         "thresh": 10.0,
     },
     'meds': {},
+    'subtract_sky': False,
 }
 
 

From ebb9336dccd2fb7931162e44d57266b018fbcb90 Mon Sep 17 00:00:00 2001
From: Axel <axel.guinot.astro@gmail.com>
Date: Mon, 20 Sep 2021 11:13:38 +0200
Subject: [PATCH 6/7] Bug fixes, enhancements

Add command line options
---
 descwl_shear_sims/sim.py                      |  2 +-
 shear_meas_tests/conftest.py                  | 63 +++++++++++++++++++
 .../test_shear_meas_corr_noise.py             | 62 ++++++++++++++----
 3 files changed, 114 insertions(+), 13 deletions(-)
 create mode 100644 shear_meas_tests/conftest.py

diff --git a/descwl_shear_sims/sim.py b/descwl_shear_sims/sim.py
index 7a10000..e6d9e96 100644
--- a/descwl_shear_sims/sim.py
+++ b/descwl_shear_sims/sim.py
@@ -321,7 +321,7 @@ def make_exp(
             scale=SCALE,
         )
         # Shear correlation in the noise
-        corr_noise.shear(g1=g1_noise, g2=g2_noise)
+        corr_noise = corr_noise.shear(g1=g1_noise, g2=g2_noise)
         image.addNoise(corr_noise)
     else:
         # Maybe change this to use galsim noise generator?
diff --git a/shear_meas_tests/conftest.py b/shear_meas_tests/conftest.py
new file mode 100644
index 0000000..1864fa9
--- /dev/null
+++ b/shear_meas_tests/conftest.py
@@ -0,0 +1,63 @@
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--save",
+        action="store_true",
+        default=False,
+        help="--save: Wether to save the results in a separate file"
+    )
+    parser.addoption(
+        "--save_dir",
+        action="store",
+        default='.',
+        type=str,
+        help="--save_dir: Directory where to store the output"
+    )
+    parser.addoption(
+        "--g1_noise",
+        action="store",
+        default=0.,
+        type=float,
+        help="--g1_noise: g1 shear to apply on the noise"
+    )
+    parser.addoption(
+        "--g2_noise",
+        action="store",
+        default=0.,
+        type=float,
+        help="--g2_noise: g2 shear to apply on the noise"
+    )
+    parser.addoption(
+        "--n_jobs",
+        action="store",
+        default=2,
+        type=int,
+        help="--n_jobs: Number of cores to use"
+    )
+
+
+@pytest.fixture()
+def save(request):
+    return request.config.getoption("--save")
+
+
+@pytest.fixture()
+def save_dir(request):
+    return request.config.getoption("--save_dir")
+
+
+@pytest.fixture()
+def g1_noise(request):
+    return request.config.getoption("--g1_noise")
+
+
+@pytest.fixture()
+def g2_noise(request):
+    return request.config.getoption("--g2_noise")
+
+
+@pytest.fixture()
+def n_jobs(request):
+    return request.config.getoption("--n_jobs")
diff --git a/shear_meas_tests/test_shear_meas_corr_noise.py b/shear_meas_tests/test_shear_meas_corr_noise.py
index c54d6f3..e009ac6 100644
--- a/shear_meas_tests/test_shear_meas_corr_noise.py
+++ b/shear_meas_tests/test_shear_meas_corr_noise.py
@@ -4,7 +4,6 @@
 import ngmix
 import tqdm
 import joblib
-import galsim
 
 import pytest
 
@@ -126,9 +125,9 @@ def _boostrap_m_c(pres, mres):
 
 def _run_sim_one(*, seed, mdet_seed, g1, g2, g1_noise, g2_noise, **kwargs):
     sim_data = _make_lsst_sim(
-        seed=seed, 
+        seed=seed,
         g1=g1, g2=g2,
-        g1_noise=g1_noise, g2_noise=g2_noise, 
+        g1_noise=g1_noise, g2_noise=g2_noise,
         **kwargs,
     )
     coadd_obs = make_coadd_obs(
@@ -157,10 +156,9 @@ def _run_sim_one(*, seed, mdet_seed, g1, g2, g1_noise, g2_noise, **kwargs):
 def run_sim(seed, mdet_seed, **kwargs):
     # positive shear
     _pres = _run_sim_one(
-        seed=seed, 
-        mdet_seed=mdet_seed, 
-        g1=0.02, g2=0, 
-        g1_noise=0., g2_noise=0.1,      # High value for the purpose of testing
+        seed=seed,
+        mdet_seed=mdet_seed,
+        g1=0.02, g2=0,
         **kwargs,
     )
     if _pres is None:
@@ -171,7 +169,6 @@ def run_sim(seed, mdet_seed, **kwargs):
         seed=seed,
         mdet_seed=mdet_seed,
         g1=-0.02, g2=0,
-        g1_noise=0., g2_noise=0.1,      # High value for the purpose of testing
         **kwargs,
     )
     if _mres is None:
@@ -181,9 +178,13 @@ def run_sim(seed, mdet_seed, **kwargs):
 
 
 @pytest.mark.parametrize(
-    'layout,ntrial', [('grid', 500), ('random', 2500)]
+    'layout,ntrial',
+    [('grid', 100),
+     ('random', 2500)]
 )
-def test_shear_meas(layout, ntrial):
+def test_shear_meas(layout, ntrial, g1_noise, g2_noise, save, save_dir, n_jobs):
+    if save:
+        date = time.strftime('_%d-%m-%Y_%H-%M-%S')
     nsub = max(ntrial // 100, 10)
     nitr = ntrial // nsub
     rng = np.random.RandomState(seed=116)
@@ -200,11 +201,17 @@ def test_shear_meas(layout, ntrial):
     for itr in tqdm.trange(nitr):
         jobs = [
             joblib.delayed(run_sim)(
-                seeds[loc+i], mdet_seeds[loc+i], layout=layout,
+                seeds[loc+i],
+                mdet_seeds[loc+i],
+                layout=layout,
+                g1_noise=g1_noise,
+                g2_noise=g2_noise,
             )
             for i in range(nsub)
         ]
-        outputs = joblib.Parallel(n_jobs=-2, verbose=0, backend='loky')(jobs)
+        outputs = joblib.Parallel(n_jobs=n_jobs,
+                                  verbose=0,
+                                  backend='loky')(jobs)
 
         for out in outputs:
             if out is None:
@@ -254,5 +261,36 @@ def test_shear_meas(layout, ntrial):
         flush=True,
     )
 
+    if save:
+        base_name = '/results{}.txt'
+        add_name = ''
+        if g1_noise != 0:
+            add_name += '_corr_g1'
+        if g2_noise != 0:
+            add_name += '_corr_g2'
+        if (g1_noise == 0) and (g2_noise == 0):
+            add_name += '_no_corr'
+        add_name += date
+        res_name = base_name.format(add_name)
+        res_file = open(save_dir + res_name, 'w+')
+        res_file.write((
+            "g1 noise: %s"
+            "g2 noise: %s"
+            "nsims: %d"
+            "\nm [1e-3, 3sigma]: %s +/- %s"
+            "\nc [1e-5, 3sigma]: %s +/- %s"
+            "\ntotal time: %s s"
+        ) % (
+            g1_noise,
+            g2_noise,
+            ntrial,
+            m/1e-3,
+            3*merr/1e-3,
+            c/1e-5,
+            3*cerr/1e-5,
+            total_time,
+        ))
+        res_file.close()
+
     assert np.abs(m) < max(1e-3, 3*merr)
     assert np.abs(c) < 3*cerr

From 5b166250857440643d1b44b00518237fef1ac4ae Mon Sep 17 00:00:00 2001
From: Axel <axel.guinot.astro@gmail.com>
Date: Thu, 23 Sep 2021 13:53:54 +0200
Subject: [PATCH 7/7] Add support for gal_mag in command line

---
 descwl_shear_sims/sim.py                      |  1 -
 shear_meas_tests/conftest.py                  | 12 ++++++++
 .../test_shear_meas_corr_noise.py             | 30 +++++++++++++++++--
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/descwl_shear_sims/sim.py b/descwl_shear_sims/sim.py
index cc7ad8e..0078dd7 100644
--- a/descwl_shear_sims/sim.py
+++ b/descwl_shear_sims/sim.py
@@ -233,7 +233,6 @@ def make_exp(
     star_bleeds=False,
     sky_n_sigma=None,
     draw_method='auto',
-    galsim_rng=None,
     corr_noise=False,
     g1_noise=0.01,
     g2_noise=0.01,
diff --git a/shear_meas_tests/conftest.py b/shear_meas_tests/conftest.py
index 1864fa9..a755855 100644
--- a/shear_meas_tests/conftest.py
+++ b/shear_meas_tests/conftest.py
@@ -29,6 +29,13 @@ def pytest_addoption(parser):
         type=float,
         help="--g2_noise: g2 shear to apply on the noise"
     )
+    parser.addoption(
+        "--gal_mag",
+        action="store",
+        default=None,
+        type=float,
+        help="--gal_mag: magnitude of the galaxies drawn"
+    )
     parser.addoption(
         "--n_jobs",
         action="store",
@@ -58,6 +65,11 @@ def g2_noise(request):
     return request.config.getoption("--g2_noise")
 
 
+@pytest.fixture()
+def gal_mag(request):
+    return request.config.getoption("--gal_mag")
+
+
 @pytest.fixture()
 def n_jobs(request):
     return request.config.getoption("--n_jobs")
diff --git a/shear_meas_tests/test_shear_meas_corr_noise.py b/shear_meas_tests/test_shear_meas_corr_noise.py
index ec7133a..67a03f6 100644
--- a/shear_meas_tests/test_shear_meas_corr_noise.py
+++ b/shear_meas_tests/test_shear_meas_corr_noise.py
@@ -34,15 +34,20 @@
 }
 
 
-def _make_lsst_sim(*, seed, g1, g2, g1_noise, g2_noise, layout):
+def _make_lsst_sim(*, seed, g1, g2, g1_noise, g2_noise, layout, galaxy_mag=None):
     rng = np.random.RandomState(seed=seed)
 
+    gal_config = None
+    if not isinstance(galaxy_mag, type(None)):
+        gal_config = dict(mag=galaxy_mag)
+
     galaxy_catalog = sim.galaxies.make_galaxy_catalog(
         rng=rng,
         coadd_dim=sim.sim.DEFAULT_SIM_CONFIG["coadd_dim"],
         buff=sim.sim.DEFAULT_SIM_CONFIG["buff"],
         layout=layout,
         gal_type='exp',
+        gal_config=gal_config,
     )
 
     psf = sim.psfs.make_fixed_psf(psf_type='gauss')
@@ -124,11 +129,20 @@ def _boostrap_m_c(pres, mres):
     return m, merr, c, cerr
 
 
-def _run_sim_one(*, seed, mdet_seed, g1, g2, g1_noise, g2_noise, **kwargs):
+def _run_sim_one(
+    *,
+    seed,
+    mdet_seed,
+    g1, g2,
+    g1_noise, g2_noise,
+    gal_mag=None,
+    **kwargs
+):
     sim_data = _make_lsst_sim(
         seed=seed,
         g1=g1, g2=g2,
         g1_noise=g1_noise, g2_noise=g2_noise,
+        galaxy_mag=gal_mag,
         **kwargs,
     )
     coadd_obs = make_coadd_obs(
@@ -183,7 +197,16 @@ def run_sim(seed, mdet_seed, **kwargs):
     [('grid', 100),
      ('random', 2500)]
 )
-def test_shear_meas(layout, ntrial, g1_noise, g2_noise, save, save_dir, n_jobs):
+def test_shear_meas(
+    layout,
+    ntrial,
+    g1_noise,
+    g2_noise,
+    gal_mag,
+    save,
+    save_dir,
+    n_jobs
+):
     if save:
         date = time.strftime('_%d-%m-%Y_%H-%M-%S')
     nsub = max(ntrial // 100, 10)
@@ -207,6 +230,7 @@ def test_shear_meas(layout, ntrial, g1_noise, g2_noise, save, save_dir, n_jobs):
                 layout=layout,
                 g1_noise=g1_noise,
                 g2_noise=g2_noise,
+                gal_mag=gal_mag,
             )
             for i in range(nsub)
         ]