fixed issue with codeconv test coverage

LSST-strong-lensing · Feb 24, 2025 · b6aa384 · b6aa384
1 parent 1802e37
commit b6aa384
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/slsim/Sources/source.py b/slsim/Sources/source.py
@@ -587,6 +587,12 @@ def kwargs_extended_source_light(self, center_lens, draw_area, band=None):
             z_image = self.source_dict["z_data"]
             pixel_width = self.source_dict["pixel_width_data"]
             pixel_width *= z_scale_factor(z_old=z_image, z_new=self.redshift)
+
+            image = self.source_dict["image"]
+
+            if isinstance(image, (list, np.ndarray)):
+                image = image[0]
+
             kwargs_extended_source = [
                 {
                     "magnitude": mag_source,

diff --git a/tests/test_Source/test_source.py b/tests/test_Source/test_source.py
@@ -661,34 +661,36 @@ def test_extended_source_position_interpolated(self):
         center_lens = np.array([0, 0])
         draw_area = 4 * np.pi
 
-        result_with_center = self.source_interp1.extended_source_position(
-            center_lens, draw_area
-        )
+        # Ensure the method recomputes the position
+        if hasattr(self.source_interp1, "_center_source"):
+            del self.source_interp1._center_source
+
+        result_with_center = self.source_interp1.extended_source_position(center_lens, draw_area)
         expected_position_with_center = np.array(
             [
                 self.source_interp1.source_dict["center_x"][0],
                 self.source_interp1.source_dict["center_y"][0],
             ]
         )
-
         np.testing.assert_array_almost_equal(
             result_with_center.flatten(),
             expected_position_with_center,
             decimal=5,
             err_msg="Center position from source_dict does not match expected when center_x and center_y exist.",
         )
 
-        result_no_center = self.source_interp2.extended_source_position(
-            center_lens, draw_area
-        )
+        # Do the same for the no-center branch:
+        if hasattr(self.source_interp2, "_center_source"):
+            del self.source_interp2._center_source
 
+        result_no_center = self.source_interp2.extended_source_position(center_lens, draw_area)
         test_area_radius = np.sqrt(draw_area / np.pi)
         r = np.linalg.norm(result_no_center - center_lens)
-
         assert (
             0 <= r <= test_area_radius
         ), f"Generated position {result_no_center} is out of expected test area range."
 
+
     def test_redshift(self):
         assert self.source.redshift == 0.5
         assert self.source11.redshift == 3.123