flat field mode hack tested

autocti/charge_injection/imaging/simulator.py

@@ -27,6 +27,7 @@ def __init__(
         read_noise: Optional[float] = None,
         charge_noise: Optional[float] = None,
         stray_light : Optional[Tuple[float, float]] = None,
+        flat_field_mode : bool = False,
         noise_if_add_noise_false: float = 0.1,
         noise_seed: int = -1,
         ci_seed: int = -1,
@@ -56,6 +57,7 @@ def __init__(
         self.read_noise = read_noise
         self.charge_noise = charge_noise
         self.stray_light = stray_light
+        self.flat_field_mode = flat_field_mode
 
         self.ci_seed = ci_seed
 
@@ -178,6 +180,15 @@ def via_layout_from(
             Seed for the read-noises added to the image.
         """
 
+        if self.flat_field_mode:
+
+            return self.via_flat_field_mode(
+                layout=layout,
+                clocker=clocker,
+                cti=cti,
+                cosmic_ray_map=cosmic_ray_map,
+            )
+
         if self.column_sigma is not None:
             pre_cti_data = self.pre_cti_data_non_uniform_from(layout=layout)
         else:
@@ -214,8 +225,6 @@ def via_pre_cti_data_from(
                 [m * i + c for i in range(0, region.y1)]
             )
 
-            print(stray_light)
-
             pre_cti_data[0: region.y1, region.x0: region.x1] += stray_light[:, None]
 
 
@@ -287,3 +296,44 @@ def via_post_cti_data_from(
             cosmic_ray_map=cosmic_ray_map,
             layout=layout,
         )
+
+    def via_flat_field_mode(
+            self,
+            layout: Layout2DCI,
+            clocker: Optional[Clocker2D],
+            cti: Optional[CTI2D],
+            cosmic_ray_map: Optional[aa.Array2D] = None,
+    ):
+
+        pre_cti_data = np.zeros(layout.shape_2d)
+        pre_cti_data[
+            layout.parallel_overscan.x0: layout.parallel_overscan.x1,
+            0:layout.parallel_overscan.y0
+        ] = self.norm
+
+        pre_cti_data_poisson = np.random.poisson(pre_cti_data, pre_cti_data.shape)
+
+        pre_cti_data_poisson = copy.copy(pre_cti_data_poisson)
+
+        pre_cti_data_poisson = aa.Array2D.no_mask(
+            values=pre_cti_data_poisson, pixel_scales=self.pixel_scales
+        )
+
+        if cosmic_ray_map is not None:
+            pre_cti_data_poisson += cosmic_ray_map.native
+
+        if cti is not None:
+            post_cti_data = clocker.add_cti(data=pre_cti_data_poisson, cti=cti)
+        else:
+            post_cti_data = copy.copy(pre_cti_data)
+
+        pre_cti_data = aa.Array2D.no_mask(
+            values=pre_cti_data, pixel_scales=self.pixel_scales
+        )
+
+        return self.via_post_cti_data_from(
+            post_cti_data=post_cti_data,
+            pre_cti_data=pre_cti_data,
+            layout=layout,
+            cosmic_ray_map=cosmic_ray_map,
+        )

test_autocti/charge_injection/imaging/test_simulator_ci.py

@@ -398,3 +398,40 @@ def test__from_post_cti_data(parallel_clocker_2d, traps_x2, ccd):
     assert (dataset.noise_map == dataset_via.noise_map).all()
     assert (dataset.pre_cti_data == dataset_via.pre_cti_data).all()
     assert (dataset.cosmic_ray_map == dataset_via.cosmic_ray_map).all()
+
+
+def test__flat_field_mode(parallel_clocker_2d, traps_x2, ccd):
+    layout = ac.Layout2DCI(
+        shape_2d=(5, 5),
+        region_list=[(0, 3, 0, 4)],
+        parallel_overscan=ac.Region2D((4, 5, 0, 4)),
+    )
+
+    simulator = ac.SimulatorImagingCI(
+        pixel_scales=1.0, norm=100.0, read_noise=None, noise_seed=1, flat_field_mode=True
+    )
+
+    cti = ac.CTI2D(parallel_trap_list=traps_x2, parallel_ccd=ccd)
+
+    dataset = simulator.via_layout_from(
+        layout=layout,
+        clocker=parallel_clocker_2d,
+        cti=cti,
+    )
+
+    assert dataset.data[0,0] > 20.0
+    assert dataset.pre_cti_data[0,0] == pytest.approx(100.0, 1.0e-4)
+
+    simulator = ac.SimulatorImagingCI(
+        pixel_scales=1.0, norm=100.0, read_noise=4.0, noise_seed=1, flat_field_mode=True
+    )
+
+    cti = ac.CTI2D(parallel_trap_list=traps_x2, parallel_ccd=ccd)
+
+    dataset = simulator.via_layout_from(
+        layout=layout,
+        clocker=parallel_clocker_2d,
+        cti=cti,
+    )
+
+    assert dataset.noise_map[0,0] == pytest.approx(4.0, 1.0e-4)