Fix plotting of composite source models in source plane

Herculens · Jun 3, 2024 · e61183c · e61183c
1 parent bed8536
commit e61183c
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 17 deletions.
diff --git a/herculens/Analysis/plot.py b/herculens/Analysis/plot.py
@@ -138,15 +138,14 @@ def model_summary(self, lens_image, kwargs_result,
 
         if show_source:
             kwargs_source = copy.deepcopy(kwargs_result['kwargs_source'])
-            if lens_image.SourceModel.has_pixels:
-                src_idx = lens_image.SourceModel.pixelated_index
-                source_model = kwargs_source[src_idx]['pixels']
-                _, _, src_extent = lens_image.get_source_coordinates(
-                    kwargs_result['kwargs_lens'], return_plt_extent=True)
-            elif kwargs_grid_source is not None:
+            if kwargs_grid_source is not None:
                 grid_src = lens_image.Grid.create_model_grid(**kwargs_grid_source)
                 x_grid_src, y_grid_src = grid_src.pixel_coordinates
-                source_model = lens_image.SourceModel.surface_brightness(x_grid_src, y_grid_src, kwargs_source)
+                source_model = lens_image.eval_source_surface_brightness(
+                    x_grid_src, y_grid_src, 
+                    kwargs_source, kwargs_lens=kwargs_result['kwargs_lens'],
+                    k=None, k_lens=k_lens, de_lensed=True,
+                )
                 source_model *= lens_image.Grid.pixel_area
                 src_extent = grid_src.plt_extent
             else:
@@ -365,7 +364,7 @@ def model_summary(self, lens_image, kwargs_result,
                     ax.set_xlim(src_extent[0], src_extent[1])
                     ax.set_ylim(src_extent[2], src_extent[3])
             if ps_src_pos is not None:
-                ax.scatter(*ps_src_pos, s=30, c='tab:blue', marker='x', linewidths=0.5, 
+                ax.scatter(*ps_src_pos, s=100, c='tab:green', marker='*', linewidths=0.5, 
                            label="point source")
                 ax.legend()
             ax = axes[i_row, 2]

diff --git a/herculens/LensImage/lens_image.py b/herculens/LensImage/lens_image.py
@@ -128,24 +128,29 @@ def source_surface_brightness(self, kwargs_source, kwargs_lens=None,
         if len(self.SourceModel.profile_type_list) == 0:
             return jnp.zeros(self.Grid.num_pixel_axes)
         x_grid_img, y_grid_img = self.ImageNumerics.coordinates_evaluate
+        source_light = self.eval_source_surface_brightness(
+            x_grid_img, y_grid_img,
+            kwargs_source, kwargs_lens=kwargs_lens,
+            k=k, k_lens=k_lens, de_lensed=de_lensed,
+        )
+        if not supersampled:
+            source_light = self.ImageNumerics.re_size_convolve(
+                source_light, unconvolved=unconvolved)
+        return source_light
+
+    def eval_source_surface_brightness(self, x, y, kwargs_source, kwargs_lens=None, 
+                                       k=None, k_lens=None, de_lensed=False):
         if self._src_adaptive_grid:
             pixels_x_coord, pixels_y_coord, _ = self.adapt_source_coordinates(kwargs_lens, k_lens=k_lens)
         else:
             pixels_x_coord, pixels_y_coord = None, None  # fall back on fixed, user-defined coordinates
         if de_lensed is True:
-            if self._src_adaptive_grid:
-                offset_x, offset_y = pixels_x_coord.mean(), pixels_y_coord.mean()
-            else:
-                offset_x, offset_y = 0., 0.
-            source_light = self.SourceModel.surface_brightness(x_grid_img+offset_x, y_grid_img+offset_y, kwargs_source, k=k,
+            source_light = self.SourceModel.surface_brightness(x, y, kwargs_source, k=k,
                                                                pixels_x_coord=pixels_x_coord, pixels_y_coord=pixels_y_coord)
         else:
-            x_grid_src, y_grid_src = self.MassModel.ray_shooting(x_grid_img, y_grid_img, kwargs_lens, k=k_lens)
+            x_grid_src, y_grid_src = self.MassModel.ray_shooting(x, y, kwargs_lens, k=k_lens)
             source_light = self.SourceModel.surface_brightness(x_grid_src, y_grid_src, kwargs_source, k=k,
                                                                pixels_x_coord=pixels_x_coord, pixels_y_coord=pixels_y_coord)
-        if not supersampled:
-            source_light = self.ImageNumerics.re_size_convolve(
-                source_light, unconvolved=unconvolved)
         return source_light
 
     def lens_surface_brightness(self, kwargs_lens_light, unconvolved=False,