From ca870434802e52ff372da041fdeb10f37066f124 Mon Sep 17 00:00:00 2001
From: Xiangchong Li <mr.superonion@hotmail.com>
Date: Thu, 14 Sep 2023 05:13:55 +0900
Subject: [PATCH] commit

---
 bin/fpfs_process_descsim.py |  5 ++-
 bin/simple_coadd_descsim.py | 64 ++++++++++++++++++++++++++++++++-----
 2 files changed, 58 insertions(+), 11 deletions(-)

diff --git a/bin/fpfs_process_descsim.py b/bin/fpfs_process_descsim.py
index e800931..371e281 100755
--- a/bin/fpfs_process_descsim.py
+++ b/bin/fpfs_process_descsim.py
@@ -111,7 +111,6 @@ def __init__(self, config_name):
         self.sigma_as = cparser.getfloat("FPFS", "sigma_as")
         self.sigma_det = cparser.getfloat("FPFS", "sigma_det")
         self.rcut = cparser.getint("FPFS", "rcut")
-        # order of shear estimator
         self.nnord = cparser.getint("FPFS", "nnord", fallback=4)
         if self.nnord not in [4, 6]:
             raise ValueError(
@@ -223,7 +222,7 @@ def run(self, fname):
 
         det_fname = os.path.join(self.catdir, fname.split("/")[-1])
         det_fname = det_fname.replace("image-", "det-")
-        if os.path.isfile(out_fname):
+        if os.path.isfile(out_fname) and os.path.isfile(det_fname):
             print("Already has measurement for this simulation. ")
             return
         psf_array2, psf_array3, cov_elem = self.prepare_noise_psf(fname)
@@ -236,8 +235,8 @@ def run(self, fname):
         elapsed_time = end_time - start_time
         # Print the elapsed time
         print(f"Elapsed time: {elapsed_time} seconds")
-        fpfs.io.save_catalog(out_fname, cat, dtype="shape", nnord=str(self.nnord))
         fpfs.io.save_catalog(det_fname, det, dtype="position", nnord=str(self.nnord))
+        fpfs.io.save_catalog(out_fname, cat, dtype="shape", nnord=str(self.nnord))
         return
 
 
diff --git a/bin/simple_coadd_descsim.py b/bin/simple_coadd_descsim.py
index 70268ac..fc71913 100755
--- a/bin/simple_coadd_descsim.py
+++ b/bin/simple_coadd_descsim.py
@@ -14,15 +14,38 @@
 # GNU General Public License for more details.
 #
 import os
-import glob
 import schwimmbad
 import numpy as np
 import astropy.io.fits as pyfits
-from fpfs.io import save_image
+
+# from fpfs.io import save_image
 from argparse import ArgumentParser
 from configparser import ConfigParser
 
 
+def get_sim_fname(directory, ftype, min_id, max_id, nshear, nrot, band):
+    """Generate filename for simulations
+    Args:
+        ftype (str):    file type ('src' for source, and 'image' for exposure
+        min_id (int):   minimum id
+        max_id (int):   maximum id
+        nshear (int):   number of shear
+        nrot (int):     number of rotations
+        band (str):     'grizy' or 'a'
+    Returns:
+        out (list):     a list of file name
+    """
+    out = [
+        os.path.join(
+            directory, "%s-%05d_g1-%d_rot%d_%s.fits" % (ftype, fid, gid, rid, band)
+        )
+        for fid in range(min_id, max_id)
+        for gid in range(nshear)
+        for rid in range(nrot)
+    ]
+    return out
+
+
 band_map = {
     "g": 0,
     "r": 1,
@@ -80,10 +103,12 @@ def prepare(self, fname):
     def prepare_image(self, fname):
         blist = ["g", "r", "i", "z"]
         self.prepare(fname)
-        gal_array = np.zeros((self.image_nx, self.image_nx), dtype=np.float32)
+        # dtp = np.float32
+        dtp = np.float64
+        gal_array = np.zeros((self.image_nx, self.image_nx), dtype=dtp)
         weight_all = 0.0
         for band in blist:
-            print("processing %s band" % band)
+            print("%s-band running" % band)
             weight = w_map[band]
             weight_all = weight_all + weight
             fname2 = fname.replace("_g.fits", "_%s.fits" % band)
@@ -93,14 +118,16 @@ def prepare_image(self, fname):
         return gal_array
 
     def run(self, fname):
+        print("running on image: %s" % fname)
         out_dir = self.img_dir
         out_fname = os.path.join(out_dir, fname.split("/")[-1])
         out_fname = out_fname.replace("_g.fits", "_a.fits")
         if os.path.exists(out_fname):
-            print("Already has measurement for this simulation. ")
+            print("Already has image. ")
             return
         exposure = self.prepare_image(fname)
-        save_image(out_fname, exposure)
+        # save_image(out_fname, exposure)
+        pyfits.writeto(out_fname, exposure)
         return
 
 
@@ -112,6 +139,18 @@ def run(self, fname):
         type=str,
         help="configure file name",
     )
+    parser.add_argument(
+        "--min_id",
+        required=True,
+        type=int,
+        help="minimum ID, e.g. 0",
+    )
+    parser.add_argument(
+        "--max_id",
+        required=True,
+        type=int,
+        help="maximum ID, e.g. 4000",
+    )
     group = parser.add_mutually_exclusive_group()
     group.add_argument(
         "--ncores",
@@ -130,8 +169,17 @@ def run(self, fname):
     args = parser.parse_args()
     pool = schwimmbad.choose_pool(mpi=args.mpi, processes=args.n_cores)
     worker = Worker(args.config)
-    fname_list = glob.glob(os.path.join(worker.img_dir, "image-*_g.fits"))
-    fname_list = np.sort(fname_list)  # [0:1]
+
+    band = "g"
+    fname_list = get_sim_fname(
+        worker.img_dir,
+        "image",
+        args.min_id,
+        args.max_id,
+        2,
+        2,
+        band,
+    )
     for r in pool.map(worker.run, fname_list):
         pass
     pool.close()