Restrict model fitting window to a disk (#596)

* test: restrict fitting area to a disk

* revert change in residual block

* you can't fix stupid

* fitting window with checkimage WIP

* Fix python configuration problem

* Implement all window types

* fix code factor problem

* wip: add extra parameter to store x and y source size

* add configuration for width and height of sources in no detect mode

* minor style fixes

* Added rotated ellipse

* added ellipse_scale config option, fixed windows on the borders

Author: marcschefer

SEImplementation/SEImplementation/CheckImages/CheckImages.h

@@ -117,6 +117,8 @@ class CheckImages : public Configurable {
 
   std::shared_ptr<WriteableImage<MeasurementImage::PixelType>> getModelFittingImage(unsigned int frame_number);
 
+  std::shared_ptr<WriteableImage<int>> getFittingWindowImage(unsigned int frame_number);
+
   std::shared_ptr<WriteableImage<MeasurementImage::PixelType>> getPsfImage(unsigned int frame_number);
 
   std::shared_ptr<WriteableImage<float>> getMLDetectionImage(unsigned int plane_number, size_t index);
@@ -186,6 +188,7 @@ class CheckImages : public Configurable {
   std::map<unsigned int, std::shared_ptr<WriteableImage<int>>> m_measurement_aperture_images;
   std::map<unsigned int, std::shared_ptr<WriteableImage<int>>> m_measurement_auto_aperture_images;
   std::map<unsigned int, std::shared_ptr<WriteableImage<MeasurementImage::PixelType>>> m_check_image_model_fitting, m_check_image_psf;
+  std::map<unsigned int, std::shared_ptr<WriteableImage<int>>> m_check_image_fitting_window;
   std::vector<std::map<unsigned int, std::shared_ptr<WriteableImage<float>>>> m_check_image_ml_detection;
 
   std::vector<std::shared_ptr<DetectionImage>> m_detection_images;
@@ -202,6 +205,7 @@ class CheckImages : public Configurable {
   std::vector<std::shared_ptr<CoordinateSystem>> m_coordinate_systems;
 
   boost::filesystem::path m_model_fitting_image_filename;
+  boost::filesystem::path m_fitting_window_image_filename;
   boost::filesystem::path m_residual_filename;
   boost::filesystem::path m_background_filename;
   boost::filesystem::path m_variance_filename;

SEImplementation/SEImplementation/Configuration/CheckImagesConfig.h

@@ -45,6 +45,10 @@ class CheckImagesConfig : public Euclid::Configuration::Configuration {
     return m_model_fitting_filename;
   }
 
+  const std::string& getFittingWindowImageFilename() const {
+    return m_fitting_window_filename;
+  }
+
   const std::string& getModelFittingResidualFilename() const {
     return m_model_fitting_residual_filename;
   }
@@ -112,6 +116,7 @@ class CheckImagesConfig : public Euclid::Configuration::Configuration {
 private:
 
   std::string m_model_fitting_filename;
+  std::string m_fitting_window_filename;
   std::string m_model_fitting_residual_filename;
   std::string m_background_filename;
   std::string m_variance_filename;

SEImplementation/SEImplementation/Configuration/ModelFittingConfig.h

@@ -23,10 +23,14 @@
 #define _SEIMPLEMENTATION_CONFIGURATION_MODELFITTINGCONFIG_H
 
 #include <vector>
+
 #include <SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingParameter.h>
 #include <SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingModel.h>
 #include <SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingFrame.h>
 #include <SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingPrior.h>
+
+#include <SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingIterativeTask.h>
+
 #include <Configuration/Configuration.h>
 
 namespace SourceXtractor {
@@ -60,6 +64,8 @@ class ModelFittingConfig : public Euclid::Configuration::Configuration {
   int getMetaIterations() const { return m_meta_iterations; }
   double getDeblendFactor() const { return m_deblend_factor; }
   double getMetaIterationStop() const { return m_meta_iteration_stop; }
+  FlexibleModelFittingIterativeTask::WindowType getWindowType() const { return m_window_type; }
+  double getEllipseScale() const { return m_ellipse_scale; }
 
 private:
   std::string m_least_squares_engine;
@@ -69,6 +75,9 @@ class ModelFittingConfig : public Euclid::Configuration::Configuration {
   int m_meta_iterations { 3 };
   double m_deblend_factor { 1.0 };
   double m_meta_iteration_stop { 0.0001 };
+  FlexibleModelFittingIterativeTask::WindowType m_window_type
+      { FlexibleModelFittingIterativeTask::WindowType::RECTANGLE };
+  double m_ellipse_scale { 3.0 };
 
   std::map<int, std::shared_ptr<FlexibleModelFittingParameter>> m_parameters;
   std::map<int, std::shared_ptr<FlexibleModelFittingModel>> m_models;

SEImplementation/SEImplementation/Plugin/AssocMode/AssocMode.h

@@ -40,13 +40,15 @@ class AssocMode : public Property {
    */
   virtual ~AssocMode() = default;
 
-  AssocMode() : m_has_assoc(false), m_assoc_data(std::vector<size_t>({0})), m_ref_frame_pixel_radius(1), m_group_id(0) {
+  AssocMode() : m_has_assoc(false), m_assoc_data(std::vector<size_t>({0})),
+      m_ref_frame_pixel_width(1), m_ref_frame_pixel_height(1), m_group_id(0) {
   }
 
   AssocMode(bool has_assoc, const std::vector<double>& assoc_data,
-      double ref_frame_pixel_radius=0.0, unsigned int group_id=0) :
+      double ref_frame_pixel_width=0.0, double ref_frame_pixel_height=0.0, unsigned int group_id=0) :
       m_has_assoc(has_assoc), m_assoc_data(std::vector<size_t>({assoc_data.size()})),
-      m_ref_frame_pixel_radius(ref_frame_pixel_radius), m_group_id(group_id) {
+      m_ref_frame_pixel_width(ref_frame_pixel_width), m_ref_frame_pixel_height(ref_frame_pixel_height),
+      m_group_id(group_id) {
     for (size_t i=0; i<assoc_data.size(); i++) {
       m_assoc_data.at(i) = assoc_data[i];
     }
@@ -60,8 +62,12 @@ class AssocMode : public Property {
     return m_assoc_data;
   }
 
-  double getRefFramePixelRadius() const {
-    return m_ref_frame_pixel_radius;
+  double getRefFramePixelWidth() const {
+    return m_ref_frame_pixel_width;
+  }
+
+  double getRefFramePixelHeight() const {
+    return m_ref_frame_pixel_height;
   }
 
   unsigned int getGroupId() const {
@@ -72,7 +78,9 @@ class AssocMode : public Property {
   bool m_has_assoc;
 
   NdArray<SeFloat> m_assoc_data;
-  double m_ref_frame_pixel_radius;
+  //double m_ref_frame_pixel_radius;
+  double m_ref_frame_pixel_width;
+  double m_ref_frame_pixel_height;
   unsigned int m_group_id;
 };
 

SEImplementation/SEImplementation/Plugin/AssocMode/AssocModeConfig.h

@@ -60,7 +60,10 @@ class AssocModeConfig : public Euclid::Configuration::Configuration {
     WorldCoordinate world_coord;
     double weight;
     std::vector<double> assoc_columns;
-    double source_radius_pixels;
+
+    double source_pixel_width;
+    double source_pixel_height;
+
     unsigned int group_id;
   };
 
@@ -107,7 +110,8 @@ class AssocModeConfig : public Euclid::Configuration::Configuration {
   AssocMode m_assoc_mode;
   double m_assoc_radius;
   double m_default_pixel_size;
-  int m_pixel_size_column;
+  int m_pixel_width_column;
+  int m_pixel_height_column;
   int m_group_id_column;
 
   std::vector<std::vector<CatalogEntry>> m_catalogs;

SEImplementation/SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingIterativeTask.h

@@ -26,8 +26,7 @@
 
 #include "SEFramework/Image/VectorImage.h"
 #include "SEFramework/Task/GroupTask.h"
-
-//#include "SEImplementation/Configuration/SamplingConfig.h"
+#include "SEFramework/Source/SourceFlags.h"
 
 #include "SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingParameter.h"
 #include "SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingFrame.h"
@@ -38,6 +37,18 @@ namespace SourceXtractor {
 class FlexibleModelFittingIterativeTask : public GroupTask {
 
 public:
+  enum class WindowType {
+    RECTANGLE = 1,
+    SQUARE_MIN = 2,
+    SQUARE_MAX = 3,
+    SQUARE_AREA = 4,
+    DISK_MIN = 5,
+    DISK_MAX = 6,
+    DISK_AREA = 7,
+    ALIGNED_ELLIPSE = 8,
+    ROTATED_ELLIPSE = 9
+  };
+
   FlexibleModelFittingIterativeTask(const std::string &least_squares_engine,
       unsigned int max_iterations, double modified_chi_squared_scale,
       std::vector<std::shared_ptr<FlexibleModelFittingParameter>> parameters,
@@ -47,7 +58,9 @@ class FlexibleModelFittingIterativeTask : public GroupTask {
       int meta_iterations=3,
       double deblend_factor=1.0,
       double meta_iteration_stop=0.0001,
-      size_t max_fit_size=100
+      size_t max_fit_size=100,
+      WindowType window_type = WindowType::RECTANGLE,
+      double ellipse_scale=3.0
       );
 
   virtual ~FlexibleModelFittingIterativeTask();
@@ -75,9 +88,21 @@ class FlexibleModelFittingIterativeTask : public GroupTask {
     std::vector<SourceState> source_states;
   };
 
+  struct FittingEllipse {
+    double m_x;
+    double m_y;
+    double m_a;
+    double m_b;
+    double m_theta;
+  };
+
+  PixelRectangle getFittingRect(SourceInterface& source, int frame_index) const;
   std::shared_ptr<VectorImage<SeFloat>> createDeblendImage(
       SourceGroupInterface& group, SourceInterface& source, int source_index,
       std::shared_ptr<FlexibleModelFittingFrame> frame, FittingState& state) const;
+  std::shared_ptr<VectorImage<SeFloat>> createWeightImage(SourceInterface& source, int frame_index) const;
+  bool isFrameValid(SourceInterface& source, int frame_index) const;
+  std::shared_ptr<VectorImage<SeFloat>> createImageCopy(SourceInterface& source, int frame_index) const;
 
   void fitSource(SourceGroupInterface& group, SourceInterface& source, int index, FittingState& state) const;
   void updateCheckImages(SourceGroupInterface& group, double pixel_scale, FittingState& state) const;
@@ -97,6 +122,12 @@ class FlexibleModelFittingIterativeTask : public GroupTask {
       SeFloat avg_reduced_chi_squared, SeFloat duration, unsigned int iterations, unsigned int stop_reason, Flags flags,
       ModelFitting::LeastSquareSummary solution,
       int index, FittingState& state) const;
+  FlexibleModelFittingIterativeTask::FittingEllipse getFittingEllipse(SourceInterface& source, int frame_index) const;
+  PixelRectangle getEllipseRect(FittingEllipse ellipse) const;
+  FlexibleModelFittingIterativeTask::FittingEllipse transformEllipse(
+      FittingEllipse ellipse, SourceInterface& source, int frame_index) const;
+  PixelRectangle clipFittingRect(PixelRectangle fitting_rect, SourceInterface& source, int frame_index) const;
+  PixelRectangle getUnclippedFittingRect(SourceInterface& source, int frame_index) const;
 
   // Task configuration
   std::string m_least_squares_engine;
@@ -111,6 +142,9 @@ class FlexibleModelFittingIterativeTask : public GroupTask {
   std::vector<std::shared_ptr<FlexibleModelFittingParameter>> m_parameters;
   std::vector<std::shared_ptr<FlexibleModelFittingFrame>> m_frames;
   std::vector<std::shared_ptr<FlexibleModelFittingPrior>> m_priors;
+
+  WindowType m_window_type { WindowType::RECTANGLE };
+  double m_ellipse_scale = 3.0;
 };
 
 }

SEImplementation/SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingTaskFactory.h

@@ -26,7 +26,7 @@
 
 #include "SEFramework/Task/TaskFactory.h"
 
-#include "SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingTask.h"
+#include "SEImplementation/Plugin/FlexibleModelFitting/FlexibleModelFittingIterativeTask.h"
 
 namespace SourceXtractor {
 
@@ -52,6 +52,8 @@ class FlexibleModelFittingTaskFactory : public TaskFactory {
   int m_meta_iterations { 3 };
   double m_deblend_factor { 1.0 };
   double m_meta_iteration_stop { 0.0001 };
+  FlexibleModelFittingIterativeTask::WindowType m_window_type { FlexibleModelFittingIterativeTask::WindowType::RECTANGLE };
+  double m_ellipse_scale { 3.0 };
 
   std::vector<std::shared_ptr<FlexibleModelFittingParameter>> m_parameters;
   std::vector<std::shared_ptr<FlexibleModelFittingFrame>> m_frames;

SEImplementation/python/sourcextractor/config/__init__.py

@@ -26,4 +26,4 @@
                             get_flux_parameter, get_pos_parameters, get_position_angle,
                             get_separation_angle, get_world_parameters,
                             get_world_position_parameters, pixel_to_world_coordinate,
-                            print_model_fitting_info, radius_to_wc_angle, ModelFitting)
+                            print_model_fitting_info, radius_to_wc_angle, ModelFitting, WindowType)

SEImplementation/python/sourcextractor/config/compat.py

@@ -119,3 +119,14 @@ def set_modified_chi_squared_scale(scale):
 @_compat_doc_helper(copy_doc_from=ModelFitting.use_iterative_fitting)
 def use_iterative_fitting(use):
     global_measurement_config.model_fitting.use_iterative_fitting(use)
+
+
+@_compat_doc_helper(copy_doc_from=ModelFitting.set_window_type)
+def set_window_type(window_type):
+    global_measurement_config.model_fitting.set_window_type(window_type)
+
+@_compat_doc_helper(copy_doc_from=ModelFitting.set_ellipse_scale)
+def set_ellipse_scale(ellipse_scale):
+    global_measurement_config.model_fitting.set_ellipse_scale(ellipse_scale)
+
+

SEImplementation/python/sourcextractor/config/model_fitting.py

@@ -1034,6 +1034,16 @@ def wc_ratio_func(x, y, radius, angle, ratio):
 
     return ra, dec, wc_rad, wc_angle, wc_ratio
 
+class WindowType(Enum):
+    RECTANGLE = 1
+    SQUARE_MIN = 2
+    SQUARE_MAX = 3
+    SQUARE_AREA = 4
+    DISK_MIN = 5
+    DISK_MAX = 6
+    DISK_AREA = 7
+    ALIGNED_ELLIPSE = 8
+    ROTATED_ELLIPSE = 9
 
 class ModelFitting:
     def __init__(self):
@@ -1050,7 +1060,9 @@ def __init__(self):
         self.onnx_model_dict = {}
         self.params_dict = {"max_iterations": 200, "modified_chi_squared_scale": 10, "engine": "",
                             "use_iterative_fitting": True, "meta_iterations": 5,
-                            "deblend_factor": 0.95, "meta_iteration_stop": 0.0001}
+                            "deblend_factor": 0.95, "meta_iteration_stop": 0.0001,
+                            "window_type": WindowType.RECTANGLE, "ellipse_scale": 3.0
+                            }
 
     def _set_model_to_frames(self, group, model):
         for x in group:
@@ -1217,6 +1229,31 @@ def set_meta_iteration_stop(self, meta_iteration_stop):
 
         """
         self.params_dict["meta_iteration_stop"] = meta_iteration_stop
+        
+    def set_window_type(self, window_type):
+        """
+        Parameters
+        ----------
+        
+        window_type : WindowType
+            specify the type of model fitting window
+
+        """
+        if type(window_type) != WindowType:
+            raise TypeError("Window type must be a WindowType enum value!")
+        
+        self.params_dict["window_type"] = window_type
+
+    def set_ellipse_scale(self, ellipse_scale):
+        """
+        Parameters
+        ----------
+        
+        ellipse_scale : double
+            specify scaling of elliptic footprints 
+
+        """
+        self.params_dict["ellipse_scale"] = ellipse_scale
 
 
 def print_model_fitting_info(group, show_params=False, prefix='', file=sys.stderr):