pixelmap.h

//
//  pixelmap.h
//  GLAMER
//
//  Created by Robert Benton Metcalf on 26/04/24.
//
 
#ifndef pixelmap_h
#define pixelmap_h
 
#include <vector>
#include <tuple>
#include <type_traits>
 
#include "utilities_slsim.h"
#include "point.h"
#include "cpfits.h"
 
struct Grid;
struct GridMap;
class Source;
 

enum class PixelMapUnits {
  ndef     // not defined
  ,surfb   // ergs / s / cm**2
  ,count_per_sec
  ,ADU   // Analogue-to-Digital Units
  ,mass
  ,mass_density
};
 
std::string to_string(PixelMapUnits unit);

template <typename T=double>
class PixelMap
{
  
  static_assert( std::is_same<T, float>::value || std::is_same<T, double>::value
                , "PixeMap can only be instantated with a double or float template argument");
  
public:
  PixelMap(const PixelMap<T>& pmap, double res_ratio);
  PixelMap();
  PixelMap(const PixelMap<T>& other);
  PixelMap(PixelMap<T>&& other);
  PixelMap(const PixelMap<T>& pmap, const double* center, std::size_t Npixels);
  PixelMap(const PixelMap<T>& pmap,long nx,long ny, std::size_t Npixels);
  PixelMap(const double* center, std::size_t Npixels
          ,double resolution
          ,PixelMapUnits u = PixelMapUnits::ndef);
  PixelMap(const double* center, std::size_t Nx, std::size_t Ny
          ,double resolution
          ,PixelMapUnits u = PixelMapUnits::ndef);
  
  PixelMap(std::string fitsfilename   
          ,double my_res=-1         
          ,PixelMapUnits u = PixelMapUnits::ndef
          ,std::string extension = ""  
  );
  
  // for conversion from double to float
  template<typename OtherT>
  PixelMap(const PixelMap<OtherT>&);
  
  ~PixelMap(){
    map.resize(0);
  };
 
  
  PixelMap<T>& operator=(const PixelMap<T> &other);
  PixelMap<T>& operator=(PixelMap<T> &&other);
  
  void ChangeUnits(PixelMapUnits u){units=u;}
  PixelMapUnits getUnits() const {return units;}
  
  inline bool valid() const { return map.size(); };
  inline std::size_t size() const { return map.size(); };
  
  inline std::size_t getNx() const { return Nx; }
  inline std::size_t getNy() const { return Ny; }
  inline double getRangeX() const { return rangeX; }
  inline double getRangeY() const { return rangeY; }
  //inline double* getCenter() const{ return &center[0]; }
  void const getCenter(Point_2d &c) const{ c[0]=center[0]; c[1]=center[1];}
  Point_2d getCenter() const{
    Point_2d c(center[0],center[1]);
    return c;
  }
  inline double getResolution() const { return resolution; }
  
  // coordinates of lower left corner 
  inline Point_2d getLLBoundary() const{ return Point_2d(map_boundary_p1[0],map_boundary_p1[1]); }
  
  double getRA(){return RA;}
  double getDEC(){return DEC;}
  
  void setRAandDec(double RAin,double DECin){
    RA = RAin;
    DEC = DECin;
  }
  
  // Zero the whole map
  void Clean(){for(auto &a : map) a = 0;}
 
  void AddImages(ImageInfo *imageinfo,int Nimages,float rescale = 1.);
  void AddImages(std::vector<ImageInfo> &imageinfo,int Nimages,float rescale = 1.);
  void AddGridBrightness(Grid &grid);
  void AddGridMapBrightness(const GridMap &grid);
  void AddUniformImages(ImageInfo *imageinfo,int Nimages,T value);
  PosType AddSource(Source &source);
  PosType AddSource(Source &source,int oversample);
  void AddPointSource(const Point_2d &x,T flux);
  
  void copy_in(const PixelMap& pmap);
  
  void paste(const PixelMap& pmap);

  void paste(const PixelMap& pmap,long nx,long ny);
  
  PixelMap<T> convolve(const PixelMap<T> &kernel);
  PixelMap<T> convolve2(const PixelMap<T> &kernel);
 
  
  PixelMap<T> cutout(long xmin,long xmax,long ymin,long ymax);

  void duplicate(const PixelMap& pmap);

  template <typename S>
  PosType AddSource_parallel(S &source,int oversample){
    Point_2d s_center;
    source.getTheta(s_center);
    
    if( (s_center[0] + source.getRadius() ) < map_boundary_p1[0] ) return 0.0;
    if( (s_center[0] - source.getRadius() ) > map_boundary_p2[0] ) return 0.0;
    if( (s_center[1] + source.getRadius() ) < map_boundary_p1[1] ) return 0.0;
    if( (s_center[1] - source.getRadius() ) > map_boundary_p2[1] ) return 0.0;
    
    int nthreads = Utilities::GetNThreads();
    PosType totals[nthreads];
    std::vector<std::thread> thr;
    
    size_t block = map.size()/nthreads;
    for(int i = 0; i < nthreads ;++i){
      thr.push_back(std::thread(&PixelMap<T>::addsource_<S>,this
                      ,i*block,std::min((i+1)*block-1,map.size()-1)
                      ,oversample,std::ref(source)
                      ,std::ref(totals[i])));
    }
 
    for(int ii=0;ii < nthreads;++ii) thr[ii].join();
    
    PosType total =0;
    for(int ii=0;ii < nthreads;++ii) total += totals[ii];
    
    return total;
  }
 
  void AddCurve(ImageInfo *curve,T value);
  void AddCurve(Kist<Point> *imagekist,T value);
  void AddCurve(std::vector<Point_2d> &curve,T value);
  void AddCurve(std::vector<RAY> &curve,T value);
  
  void drawline(double x1[],double x2[],T value,bool add);
  void DrawLine(long x0,long x1,long y0,long y1,T value,bool add);
  void DrawLineGS(long x0,long x1,long y0,long y1,T value,bool add);
  void drawcircle(PosType r_center[],PosType radius,PosType value);
  void drawdisk(PosType r_center[],PosType radius,PosType value,int Nstrip);
  void AddGrid(const Grid &grid,T value = 1.0);
  void AddGrid(const Grid &grid,LensingVariable val);
  
  void Renormalize(T factor);
  void AddValue(std::size_t i, T value);
  void AssignValue(std::size_t i, T value);
  void printASCII() const;
  void printASCIItoFile(std::string filename) const;
  void printFITS(std::string filename,bool Xflip = false, bool ctype = false, bool verbose = false);
  void printFITS(std::string filename,std::vector< std::tuple<std::string,double,std::string> > &extra_header_info, bool verbose);
  
  void printFITS(std::string filename  
                ,std::vector<std::string> &headercards 
                 );
 
  void smooth(double sigma);
 
  inline T getValue(std::size_t i) const { return map[i]; }
  inline T & operator[](std::size_t i) { return map[i]; };
  inline T operator[](std::size_t i) const { return map[i]; };
  inline T & operator()(std::size_t i) { return map[i]; };
  inline T operator()(std::size_t i) const { return map[i]; };
  inline T operator()(std::size_t i,std::size_t j) const { return map[i + Nx*j]; };
  inline T & operator()(std::size_t i,std::size_t j) { return map[i + Nx*j]; };
  
  PixelMap& operator+=(const PixelMap& rhs);
  void operator+=(float f){map +=f;};
  void operator+=(double f){map +=f;};
  //friend PixelMap operator+(const PixelMap&, const PixelMap&);
  PixelMap operator+(const PixelMap&) const;
 
  PixelMap& operator-=(const PixelMap& rhs);
  //friend PixelMap operator-(const PixelMap&, const PixelMap&);
  PixelMap operator-(const PixelMap&) const;
  void operator-=(float f){map -=f;};
  void operator-=(double f){map -=f;};
 
  PixelMap& operator*=(const PixelMap& rhs);
  //friend PixelMap operator*(const PixelMap&, const PixelMap&);
  PixelMap operator*(const PixelMap& a) const;
  PixelMap operator/(const PixelMap& a) const;
 
  PixelMap& operator*=(PosType b);
  PixelMap operator*(PosType b) const;
 
  std::valarray<T>& data() { return map; }
  
  bool agrees(const PixelMap& other) const;
  
  //friend void swap(PixelMap&, PixelMap&);
  static void swap(PixelMap<T>&, PixelMap<T>&);
  
  PosType ave() const {return map.sum()/map.size();}
  PosType sum() const {return map.sum();};
  size_t size(){return map.size();}
  T max() const{ return map.max(); }
  T min() const{ return map.min(); }
 
  void FindArc(PosType &radius,PosType *xc,PosType *arc_center,PosType &arclength,PosType &width
                         ,PosType threshold);
  
  long find_index(PosType const x[],long &ix,long &iy) const;
  long find_index(PosType const x[]) const;
  
  long find_index(PosType x,PosType y,long &ix,long &iy) const;
  long find_index(PosType x,PosType y) const;
  
  void find_position(PosType x[],std::size_t const index) const;
  void find_position(PosType x[],std::size_t const ix,std::size_t const iy) const;
  
  PixelMap<T> rotate(
                  PosType theta  
                  ,T scale=1  
                  );
  
  T linear_interpolate(PosType x[]);
  
  void drawgrid(int N,PosType value);
  void drawPoints(std::vector<Point *> points,PosType size,PosType value);
  
  void drawPoints(std::vector<Point> points,PosType size,PosType value);
  
  void drawCurve(std::vector<Point *> points,PosType value){
    for(int i=0;i<points.size()-1;++i) drawline(points[i]->x,points[i+1]->x,value,false);
  }
  void drawCurve(std::vector<Point> points,PosType value){
    for(int i=0;i<points.size()-1;++i) drawline(points[i].x,points[i+1].x,value,false);
  }
  void drawPoints(std::vector<Point_2d> points,PosType size,PosType value);
  void drawCurve(std::vector<Point_2d> points,PosType value){
    for(int i=0;i<points.size()-1;++i) drawline(points[i].x,points[i+1].x,value,false);
  }
  void drawSquare(PosType p1[],PosType p2[],PosType value);
  void drawBox(PosType p1[],PosType p2[],PosType value,int Nstrip);

  void PowerSpectrum(std::vector<PosType> &power_spectrum   
                     ,std::vector<PosType> &lvec            
                     ,bool overwrite = true                 
                     );

  void PowerSpectrum(std::vector<PosType> &power_spectrum   
                     ,const std::vector<PosType> &lbins            
                     ,std::vector<PosType> &lave            
                     ,bool overwrite = true                 
                     );
 
  void AdaptiveSmooth(PosType value);
  
  void find_contour(T level
                    ,std::vector<std::vector<Point_2d> > &points
                    ,std::vector<bool> &hits_edge
                    ) const;
  
  void find_islands_holes(T level,
                    std::vector<std::vector<size_t> > &points
                    ) const;
  
  void lens_definition(T min_sn_per_image                 
                       ,T pixel_threshold                 
                       ,int &Nimages                           
                       ,T &total_sig_noise_source         
                       ,std::vector<size_t> &maxima_indexes    
                       ,std::vector<std::vector<size_t> > &image_points
                       ,bool &lens_TF                          
                       ,T &level                          
                       ,size_t &n_pix_in_source                 
                       ,bool verbose = false
                       );
 
  
  std::vector<size_t> maxima(T minlevel) const;
  
  //int count_islands(std::list<size_t> &pixel_index,std::vector<std::list<size_t>::iterator> &heads) const;
  int count_islands(std::vector<size_t> &pixel_index) const;

  size_t threshold(std::list<size_t> &pixel_index,PosType value){
    for(size_t i=0;i<map.size();++i) if(value < map[i]) pixel_index.push_back(i);
    return pixel_index.size();
  }
  
  void flipY(){
    for(size_t i=0;i<Nx;++i){
      for(size_t j=0;j<(Ny-1)/2;++j){
        std::swap( map[i + Nx*j],map[i + Nx*(Ny-j-1)] );
      }
    }
  }

  void flipX(){
    for(size_t i=0;i<(Nx-1)/2;++i){
      for(size_t j=0;j<Ny;++j){
        std::swap( map[i + Nx*j],map[Nx-i-1 + Nx*j] );
      }
    }
  }

  void doubleFlip(){
    flipX();
    flipY();
  }
  
  
  void recenter(PosType newcenter[2] 
                 );
  void recenter(Point_2d newcenter 
                 );
  
  /* \brief convolve the image with a kernel.
   
   It is assumed that the size of the kernel is much smaller than the image and
   that the kernal has the same pixel size as the image.
   **/
  //void convolve(PixelMap &kernel,long center_x = 0,long center_y = 0);
 
  PixelMap<T> downsize(int n 
                    );

  PixelMap interpolate(int n);

   void addheader(std::string label,long value,std::string comment){
     bool found = false;
     for(auto  &a : headers_long){
       if(std::get<0>(a) == label){
         std::get<1>(a)=value;
         std::get<2>(a)=comment;
         found = true;
         break;
       }
     }
     if(!found) headers_long.push_back(std::make_tuple(label,value,comment));
   }
   void addheader(std::string label,size_t value,std::string comment){
     bool found = false;
     for(auto  &a : headers_long){
       if(std::get<0>(a) == label){
         std::get<1>(a)=value;
         std::get<2>(a)=comment;
         found = true;
         break;
       }
     }
     if(!found) headers_long.push_back(std::make_tuple(label,value,comment));
   }
   void addheader(std::string label,float value,std::string comment){
     bool found = false;
     for(auto  &a : headers_float){
       if(std::get<0>(a) == label){
         std::get<1>(a)=value;
         std::get<2>(a)=comment;
         found=true;
         break;
       }
     }
     if(!found) headers_float.push_back(std::make_tuple(label,value,comment));
   }
  void addheader(std::string label,double value,std::string comment){
    bool found = false;
    for(auto  &a : headers_float){
      if(std::get<0>(a) == label){
        std::get<1>(a)=value;
        std::get<2>(a)=comment;
        found=true;
        break;
      }
    }
    if(!found) headers_float.push_back(std::make_tuple(label,value,comment));
  }
  void addheader(std::string label,std::string value,std::string comment){
    bool found = false;
    for(auto  &a : headers_string){
      if(std::get<0>(a) == label){
        std::get<1>(a)=value;
        std::get<2>(a)=comment;
        found=true;
        break;
      }
    }
    if(!found) headers_string.push_back(std::make_tuple(label,value,comment));
  }
 
private:
  std::vector<std::tuple<std::string,float,std::string> > headers_float;
  std::vector<std::tuple<std::string,long,std::string> > headers_long;
  std::vector<std::tuple<std::string,std::string,std::string> > headers_string;
 
  std::valarray<T> map;
  long Nx;
  long Ny;
  double resolution,rangeX,rangeY,center[2];
  double RA=0,DEC=0; // optional coordinates of center
  double map_boundary_p1[2],map_boundary_p2[2];
  PixelMapUnits units= PixelMapUnits::ndef;
  
  void AddGrid_(const PointList &list,LensingVariable val);
 
  double LeafPixelArea(IndexType i,Branch * branch1);
  void PointsWithinLeaf(Branch * branch1, std::list <unsigned long> &neighborlist);
  bool inMapBox(Branch * branch1) const;
  bool inMapBox(double * branch1) const;
  
  bool pixels_are_neighbors(size_t i,size_t j) const;
  void _count_islands_(size_t current,std::list<size_t> &reservoir
                       ,std::list<size_t>::iterator &group) const;
  
  //void addsource_(size_t i1,size_t i2,int oversample,Source source,PosType &total);
  //void addsource_(size_t i1,size_t i2,int oversample,Po,
  //                   PosType &total);
 
  template <typename S>
    void addsource_(size_t i1,size_t i2,int oversample,
                    S &source,
                    PosType &total){
    PosType tmp_res = resolution*1.0/oversample;
    PosType tmp = tmp_res*tmp_res;
    PosType bl = resolution /2 - 0.5*tmp_res;
    PosType y[2],x[2];
    
    total = 0;
    
    for(size_t index = i1 ;index <= i2; ++index){
      find_position(y,index);
      y[0] -= bl;
      y[1] -= bl;
      for(int i = 0 ; i < oversample ; ++i){
        x[0] = y[0] + i*tmp_res;
        for(int j=0; j < oversample;++j){
          x[1] = y[1] + j*tmp_res;
          map[index] += source.SurfaceBrightness(x)*tmp;
          total += source.SurfaceBrightness(x)*tmp;
        }
      }
    }
  }
  
  // find if pixel k is in the curve which must be in pixel units, meant to be used in conjunction with Utilities::find_boundaries<>
  bool incurve(long k,std::vector<Point_2d> &curve) const;
};
 
template<typename T>
void swap(PixelMap<T>& x, PixelMap<T>& y)
{
  using std::swap;
  
  swap(x.map,y.map);
  
  swap(x.Nx, y.Nx);
  swap(x.Ny, y.Ny);
  swap(x.resolution, y.resolution);
  swap(x.rangeX, y.rangeX);
  swap(x.rangeY, y.rangeY);
  
  swap(x.center[0], y.center[0]);
  swap(x.center[1], y.center[1]);
  
  swap(x.map_boundary_p1[0], y.map_boundary_p1[0]);
  swap(x.map_boundary_p1[1], y.map_boundary_p1[1]);
  swap(x.map_boundary_p2[0], y.map_boundary_p2[0]);
  swap(x.map_boundary_p2[1], y.map_boundary_p2[1]);
}
 
#endif /* pixelmap_h */