Implemented the resultant image builder. resultant image headers, BFE, optimizations still need to be done

Author: Higgins00

roman_imsim/__init__.py

@@ -16,3 +16,4 @@
 from .skycat import *
 from .stamp import *
 from .wcs import *
+from .resultants import *

roman_imsim/resultants.py

@@ -1,4 +1,4 @@
-import fitsio as fio
+import yaml
 import galsim
 import galsim.config
 from galsim.config import InputLoader, RegisterInputType, RegisterValueType, RegisterObjectType
@@ -7,12 +7,13 @@
 class ResultantDataLoader(object):
     """Read the resultant information from the resultant strategy."""
 
-    _req_params = {"file_name": str, "strategy": str,}
+    _req_params = {"file_name": str, "strategy_name": str,}
 
-    def __init__(self, file_name, strategy, logger=None):
+    def __init__(self, file_name, strategy_name, logger=None):
         self.logger = galsim.config.LoggerWrapper(logger)
         self.file_name = file_name
-        self.strategy = strategy
+        self.strategy_name = strategy_name
+        self.data = {}
 
 
         # try:
@@ -22,31 +23,43 @@ def __init__(self, file_name, strategy, logger=None):
         #     self.logger.warning('Reading visit info from config file.')
 
     def read_resultants(self):
-        """Read resultant info from the resultants file."""
-        if self.file_name is None:
-            raise ValueError("No resultants filename provided, trying to build from config information.")
-        if self.strategy is None:
-            raise ValueError("The strategy must be set when reading resultant strategy info from a resultants file.")
+        """Load the YAML file and get the requested strategy."""
+        self.logger.info("Reading resultants from YAML file: %s", self.file_name)
+        try:
+            with open(self.file_name, "r") as f:
+                all_strategies = yaml.safe_load(f)
+        except Exception as e:
+            raise IOError(f"Could not read YAML file '{self.file_name}': {e}")
 
-        self.logger.warning("Reading info from resultants file %s for strategy %s", self.file_name, self.strategy)
+        if self.strategy_name not in all_strategies:
+            raise ValueError(f"Strategy '{self.strategy_name}' not found in YAML file.")
 
-        data = fio.FITS(self.file_name)[-1][self.strategy]
+        strategy = all_strategies[self.strategy_name]
+        if not isinstance(strategy, list):
+            raise ValueError(f"Invalid strategy format for '{self.strategy_name}': must be a list of lists.")
 
-        self.data = {}
-        self.data["strategy"] = data["strategy"]
-        self.data["dt"] = self.resultants_to_dt()
+        self.data["strategy"] = strategy
+
+    def resultants_to_dt(self, config, base):
+        """Compute dt from list-of-lists."""
+        strategy = self.data["strategy"]
+        if len(strategy) < 2:
+            raise ValueError("Need at least two resultants to compute dt.")
+
+        avg_last = sum(strategy[-1]) / len(strategy[-1])
+        avg_second = sum(strategy[0]) / len(strategy[0])
 
-    def resultants_to_dt(self,config,base):
         if "exptime" in config:
             exptime = galsim.config.ParseValue(config, "exptime", base, float)[0]
         else:
             exptime = roman.exptime
-        dt = exptime/((sum(self.strategy[-1])/len(self.strategy[-1]))-(sum(self.strategy[1])/len(self.strategy[1])))
+
+        dt = exptime / (avg_last - avg_second)
         return dt
-    
+
     def get(self, field, default=None):
         if field not in self.data and default is None:
-            raise KeyError("ResultantData field %s not present in data" % field)
+            raise KeyError(f"Field '{field}' not found in data.")
         return self.data.get(field, default)
 
 def ResultantData(config, base, value_type):
@@ -56,10 +69,14 @@ def ResultantData(config, base, value_type):
     kwargs, safe = galsim.config.GetAllParams(config, base, req=req)
     field = kwargs["field"]
 
-    val = value_type(rdata.get(field))
-    return val, safe
+    if field == "dt":
+        val = rdata.resultants_to_dt(config, base)
+    else:
+        val = rdata.get(field)
+
+    return value_type(val), safe
 
 
 RegisterInputType("resultant_data", InputLoader(ResultantDataLoader, file_scope=True, takes_logger=True))
-RegisterValueType("ResultantData", ResultantData, [float,list], input_type="resultant_data")
+RegisterValueType("ResultantData", ResultantData, [float, list], input_type="resultant_data")
 

roman_imsim/sca.py

@@ -6,9 +6,10 @@
 from galsim.config import RegisterImageType
 from galsim.config.image_scattered import ScatteredImageBuilder
 from galsim.image import Image
+import gc
 
 
-class RomanSCAImageBuilder(ScatteredImageBuilder):
+class RomanSCAImageBuilderCMOS(ScatteredImageBuilder):
 
     def setup(self, config, base, image_num, obj_num, ignore, logger):
         """Do the initialization and setup for building the image.
@@ -123,8 +124,8 @@ def buildImage(self, config, base, image_num, obj_num, logger):
         full_xsize = base["image_xsize"]
         full_ysize = base["image_ysize"]
         wcs = base["wcs"]
-
-        
+        rdata = galsim.config.GetInputObj("resultant_data", config, base, "ResultantDataLoader")
+        strategy = rdata.get("strategy")    
         full_image = Image(full_xsize, full_ysize, dtype=float)
         full_image.setOrigin(base["image_origin"])
         full_image.wcs = wcs
@@ -155,62 +156,97 @@ def buildImage(self, config, base, image_num, obj_num, logger):
                 "x": {"type": "Random", "min": xmin, "max": xmax},
                 "y": {"type": "Random", "min": ymin, "max": ymax},
             }
-
-        nbatch = self.nobjects // 1000 + 1 
-        full_array = galsim.PhotonArray(0)
-        for batch in range(nbatch):
-            start_obj_num = self.nobjects * batch // nbatch
-            end_obj_num = self.nobjects * (batch + 1) // nbatch
-            nobj_batch = end_obj_num - start_obj_num
-            if nbatch > 1:
-                logger.warning(
-                    "Start batch %d/%d with %d objects [%d, %d)",
-                    batch + 1,
-                    nbatch,
-                    nobj_batch,
-                    start_obj_num,
-                    end_obj_num,
+        #Create index and lists for resultant management
+        max_dt = strategy[-1][-1]
+        resultant_i = 0
+        resultant_buffer= []
+        #Iterate through all dt
+        for dt in np.arange(1,max_dt+1):
+
+            nbatch = self.nobjects // 1000 + 1 
+            full_array = galsim.PhotonArray(0)
+            for batch in range(nbatch):
+                start_obj_num = self.nobjects * batch // nbatch
+                end_obj_num = self.nobjects * (batch + 1) // nbatch
+                nobj_batch = end_obj_num - start_obj_num
+                if nbatch > 1:
+                    logger.warning(
+                        "Start batch %d/%d with %d objects [%d, %d)",
+                        batch + 1,
+                        nbatch,
+                        nobj_batch,
+                        start_obj_num,
+                        end_obj_num,
+                    )
+                stamps, current_vars = galsim.config.BuildStamps(
+                    nobj_batch, base, logger=logger, obj_num=start_obj_num, do_noise=False
                 )
-            stamps, current_vars = galsim.config.BuildStamps(
-                nobj_batch, base, logger=logger, obj_num=start_obj_num, do_noise=False
-            )
-            base["index_key"] = "image_num"
-
-            for k in range(nobj_batch):
-                # This is our signal that the object was skipped.
-                if stamps[k] is None:
-                    continue
-                bounds = full_image.bounds # stamps[k].bounds &
-                if not bounds.isDefined():  # pragma: no cover
-                    # These noramlly show up as stamp==None, but technically it is possible
-                    # to get a stamp that is off the main image, so check for that here to
-                    # avoid an error.  But this isn't covered in the imsim test suite.
-                    continue
-
-                # logger.debug("image %d: full bounds = %s", image_num, str(full_image.bounds))
-                # logger.debug(
-                #     "image %d: stamp %d bounds = %s",
-                #     image_num,
-                #     k + start_obj_num,
-                #     str(stamps[k].bounds),
-                # )
-                # logger.debug("image %d: Overlap = %s", image_num, str(bounds))
-                # full_image[bounds] += stamps[k][bounds]
-                #logger.warning(stamps[k])
-            full_array = galsim.PhotonArray.concatenate([*stamps,full_array])
-
-            stamps = None
-
-        # # Bring the image so far up to a flat noise variance
-        # current_var = FlattenNoiseVariance(
-        #         base, full_image, stamps, current_vars, logger)
-
-        full_array.addTo(full_image)
-        full_array.write("photonarray.fits")
-        full_image.write("phot_image.fits")
+                base["index_key"] = "image_num"
+
+                for k in range(nobj_batch):
+                    # This is our signal that the object was skipped.
+                    if stamps[k] is None:
+                        continue
+                    bounds = full_image.bounds # stamps[k].bounds &
+                    if not bounds.isDefined():  # pragma: no cover
+                        # These noramlly show up as stamp==None, but technically it is possible
+                        # to get a stamp that is off the main image, so check for that here to
+                        # avoid an error.  But this isn't covered in the imsim test suite.
+                        continue
+
+                    # logger.debug("image %d: full bounds = %s", image_num, str(full_image.bounds))
+                    # logger.debug(
+                    #     "image %d: stamp %d bounds = %s",
+                    #     image_num,
+                    #     k + start_obj_num,
+                    #     str(stamps[k].bounds),
+                    # )
+                    # logger.debug("image %d: Overlap = %s", image_num, str(bounds))
+                    # full_image[bounds] += stamps[k][bounds]
+                    #logger.warning(stamps[k])
+                full_array = galsim.PhotonArray.concatenate([*stamps,full_array])
+                
+                stamps = None
+
+            # # Bring the image so far up to a flat noise variance
+            # current_var = FlattenNoiseVariance(
+            #         base, full_image, stamps, current_vars, logger)
+            #TODO : Apply BFE photon operation (uses current pre-read image and next photon array)
+
+            # Turn full_image into running pre-read image
+            full_array.addTo(full_image)
+            del full_array
+            gc.collect()
+            # Decide what to do with readout based on resultant strategy
+            if (np.array([item for sub in strategy for item in sub]) == dt).any():
+                if (np.array(strategy[resultant_i]) == dt).any():
+                    readout_im = full_image.copy()
+                    readout_im = self.addNoiseToImage(readout_im, config, base, logger) 
+                    resultant_buffer.extend([readout_im])
+                    if len(resultant_buffer)>1:
+                        #combine readout images
+                        resultant_buffer[0].array = resultant_buffer[0].array + resultant_buffer[1].array
+                        del resultant_buffer[-1]
+                        gc.collect()
+   
+                if np.array(strategy[resultant_i][-1]) == dt:
+                    divisor = len(np.array(strategy[resultant_i]))
+                    #divide summed images by the length of resultant to get the average and apply headers to the image array
+                    #TODO:apply header to the image array
+                    resultant_buffer[0].array = resultant_buffer[0].array/divisor
+                    resultant_buffer[0].write('resultant_{0}.fits'.format(resultant_i))
+                    resultant_i+=1
+                    resultant_buffer = []
+                    logger.warning('resultant{0} done'.format(resultant_i))
+
+                    
+
+        # full_array.write("photonarray.fits")
+        # full_image.write("phot_image.fits")
+
         return full_image, None
 
-    def addNoise(self, image, config, base, image_num, obj_num, current_var, logger):
+    def addNoiseToImage(self, image, config, base, logger):
         """Add the final noise to a Scattered image
 
         Parameters:
@@ -224,7 +260,7 @@ def addNoise(self, image, config, base, image_num, obj_num, current_var, logger)
         """
         # check ignore noise
         if self.ignore_noise:
-            return
+            return image
 
         base["current_noise_image"] = base["current_image"]
         wcs = base["wcs"]
@@ -314,7 +350,11 @@ def addNoise(self, image, config, base, image_num, obj_num, current_var, logger)
             sky_image /= roman.gain
             sky_image.quantize()
             image -= sky_image
+        return image
+    
+    def addNoise(self, image, config, base, image_num, obj_num, current_var, logger):
+        pass
 
 
 # Register this as a valid type
-RegisterImageType("roman_sca", RomanSCAImageBuilder())
+RegisterImageType("roman_sca_cmos", RomanSCAImageBuilderCMOS())

roman_imsim/stamp.py

@@ -503,9 +503,9 @@ def add_poisson_noise(self, fft_image):
 
 # Pick the right function to be _fix_seds.
 if galsim.__version_info__ < (2, 5):
-    Roman_stamp.fix_seds = Roman_stamp._fix_seds_24
+    Roman_stamp_CMOS.fix_seds = Roman_stamp_CMOS._fix_seds_24
 else:
-    Roman_stamp.fix_seds = Roman_stamp._fix_seds_25
+    Roman_stamp_CMOS.fix_seds = Roman_stamp_CMOS._fix_seds_25
 
 
 # Register this as a valid type