FIX: adding fg_scale and using it when not using normalization

include/chi2.cuh

@@ -23,7 +23,7 @@ class Chi2 : public Fi {
 
  private:
   VirtualImageProcessor* ip;
-  float fg_scale;
+  float fg_scale = 1.0;
   int imageToAdd;
   float* result_dchi2;
 };

include/classes/fi.cuh

@@ -37,7 +37,7 @@ class Fi {
   std::string setName(std::string name) { this->name = name; };
 
   float get_fivalue() { return this->fi_value; };
-  bool get_normalize() { return this->normalize; };
+  bool getNormalize() { return this->normalize; };
   float getPenalizationFactor() { return this->penalization_factor; };
   void set_fivalue(float fi) { this->fi_value = fi; };
   void setPenalizationFactor(float p) { this->penalization_factor = p; };

include/classes/synthesizer.cuh

@@ -85,7 +85,7 @@ class Synthesizer {
   Error* error = NULL;
   int griddingThreads = 0;
   bool gridding = false;
-  float fg_scale = 0.0;
+  float fg_scale = 1.0;
   float vis_noise = 0.0;
   void (*Order)(Optimizer* o, Image* I) = NULL;
   int imagesChanged = 0;

include/classes/virtualimageprocessor.cuh

@@ -15,9 +15,10 @@ class VirtualImageProcessor {
                           float xobs,
                           float yobs,
                           float freq,
-                          int primary_beam) = 0;
+                          int primary_beam,
+                          float fg_scale) = 0;
   virtual void calculateInu(cufftComplex* image, float* I, float freq) = 0;
-  virtual void chainRule(float* I, float freq) = 0;
+  virtual void chainRule(float* I, float freq, float fg_scale) = 0;
   virtual void configure(int i) = 0;
   virtual CKernel* getCKernel() { return this->ckernel; };
   virtual void setCKernel(CKernel* ckernel) { this->ckernel = ckernel; };

include/functions.cuh

@@ -145,14 +145,18 @@ __host__ T deviceReduce(T* in, long N, int input_threads);
 __host__ float deviceMaxReduce(float* in, long N, int input_threads);
 __host__ float deviceMinReduce(float* in, long N, int input_threads);
 __host__ float simulate(float* I, VirtualImageProcessor* ip);
-__host__ float chi2(float* I, VirtualImageProcessor* ip, bool normalize);
+__host__ float chi2(float* I,
+                    VirtualImageProcessor* ip,
+                    bool normalize,
+                    float fg_scale);
 __host__ void linkRestartDGi(float* dgi);
 __host__ void linkAddToDPhi(float* dphi, float* dgi, int index);
 __host__ void dchi2(float* I,
                     float* dxi2,
                     float* result_dchi2,
                     VirtualImageProcessor* ip,
-                    bool normalize);
+                    bool normalize,
+                    float fg_scale);
 __host__ void defaultNewP(float* p, float* xi, float xmin, int image);
 __host__ void particularNewP(float* p, float* xi, float xmin, int image);
 __host__ void defaultEvaluateXt(float* xt,
@@ -172,10 +176,11 @@ __host__ void linkApplyBeam2I(cufftComplex* image,
                               float xobs,
                               float yobs,
                               float freq,
-                              int primary_beam);
+                              int primary_beam,
+                              float fg_scale);
 __host__ void linkClipWNoise2I(float* I);
 __host__ void linkCalculateInu2I(cufftComplex* image, float* I, float freq);
-__host__ void linkChain2I(float* chain, float freq, float* I);
+__host__ void linkChain2I(float* chain, float freq, float* I, float fg_scale);
 __host__ void normalizeImage(float* image, float normalization_factor);
 __host__ void linkClip(float* I);
 __host__ void DEntropy(float* I,

include/imageProcessor.cuh

@@ -15,9 +15,10 @@ class ImageProcessor : public VirtualImageProcessor {
                   float xobs,
                   float yobs,
                   float freq,
-                  int primary_beam);
+                  int primary_beam,
+                  float fg_scale);
   void calculateInu(cufftComplex* image, float* I, float freq);
-  void chainRule(float* I, float freq);
+  void chainRule(float* I, float freq, float fg_scale);
   void configure(int i);
 };
 

src/chi2.cu

@@ -13,6 +13,7 @@ extern int nPenalizators;
 Chi2::Chi2() {
   this->ip = new ImageProcessor();
   this->name = "Chi2";
+  this->normalize = false;
 };
 
 void Chi2::configure(int penalizatorIndex,
@@ -42,13 +43,13 @@ void Chi2::configure(int penalizatorIndex,
 
 float Chi2::calcFi(float* p) {
   float result = 0.0f;
-  this->set_fivalue(chi2(p, ip, this->normalize));
+  this->set_fivalue(chi2(p, ip, this->normalize, this->fg_scale));
   result = (penalization_factor) * (this->get_fivalue());
   return result;
 };
 
 void Chi2::calcGi(float* p, float* xi) {
-  dchi2(p, xi, result_dchi2, ip, this->normalize);
+  dchi2(p, xi, result_dchi2, ip, this->normalize, this->fg_scale);
 };
 
 void Chi2::restartDGi() {

src/functions.cu

@@ -2303,6 +2303,7 @@ __global__ void apply_beam2I(float antenna_diameter,
                              long N,
                              float xobs,
                              float yobs,
+                             float fg_scale,
                              float freq,
                              double DELTAX,
                              double DELTAY,
@@ -2313,7 +2314,8 @@ __global__ void apply_beam2I(float antenna_diameter,
   float atten = attenuation(antenna_diameter, pb_factor, pb_cutoff, freq, xobs,
                             yobs, DELTAX, DELTAY, primary_beam);
 
-  image[N * i + j] = make_cuFloatComplex(image[N * i + j].x * atten, 0.0f);
+  image[N * i + j] =
+      make_cuFloatComplex(image[N * i + j].x * atten * fg_scale, 0.0f);
 }
 
 __global__ void apply_beam2I(float antenna_diameter,
@@ -3647,6 +3649,7 @@ __global__ void DChi2(float* noise,
                       float* w,
                       long N,
                       long numVisibilities,
+                      float fg_scale,
                       float noise_cut,
                       float ref_xobs,
                       float ref_yobs,
@@ -3691,7 +3694,7 @@ __global__ void DChi2(float* noise,
       dchi2 += w[v] * ((Vr[v].x * cosk) + (Vr[v].y * sink));
     }
 
-    dchi2 *= atten;
+    dchi2 *= fg_scale * atten;
     if (normalize)
       dchi2 /= numVisibilities;
 
@@ -3707,6 +3710,7 @@ __global__ void DChi2(float* noise,
                       float* w,
                       long N,
                       long numVisibilities,
+                      float fg_scale,
                       float noise_cut,
                       float ref_xobs,
                       float ref_yobs,
@@ -3750,7 +3754,7 @@ __global__ void DChi2(float* noise,
       dchi2 += w[v] * ((Vr[v].x * cosk) + (Vr[v].y * sink));
     }
 
-    dchi2 *= atten * gcf_i;
+    dchi2 *= fg_scale * atten * gcf_i;
 
     if (normalize)
       dchi2 /= numVisibilities;
@@ -3843,6 +3847,7 @@ __global__ void DChi2_total_alpha(float* noise,
                                   float nu,
                                   float nu_0,
                                   float noise_cut,
+                                  float fg_scale,
                                   float threshold,
                                   long N,
                                   long M) {
@@ -3856,7 +3861,7 @@ __global__ void DChi2_total_alpha(float* noise,
   alpha = I[N * M + N * i + j];
 
   dI_nu_0 = powf(nudiv, alpha);
-  dalpha = I_nu_0 * dI_nu_0 * logf(nudiv);
+  dalpha = I_nu_0 * dI_nu_0 * fg_scale * logf(nudiv);
 
   if (noise[N * i + j] < noise_cut) {
     if (I_nu_0 > threshold) {
@@ -3899,6 +3904,7 @@ __global__ void chainRule2I(float* chain,
                             float nu,
                             float nu_0,
                             float noise_cut,
+                            float fg_scale,
                             long N,
                             long M) {
   const int j = threadIdx.x + blockDim.x * blockIdx.x;
@@ -3911,7 +3917,7 @@ __global__ void chainRule2I(float* chain,
   alpha = I[N * M + N * i + j];
 
   dI_nu_0 = powf(nudiv, alpha);
-  dalpha = I_nu_0 * dI_nu_0 * logf(nudiv);
+  dalpha = I_nu_0 * dI_nu_0 * fg_scale * logf(nudiv);
 
   chain[N * i + j] = dI_nu_0;
   chain[N * M + N * i + j] = dalpha;
@@ -4060,7 +4066,7 @@ __global__ void noise_reduction(float* noise_I, long N, long M) {
     noise_I[N * M + N * i + j] = 0.0f;
 }
 
-__host__ float simulate(float* I, VirtualImageProcessor* ip) {
+__host__ float simulate(float* I, VirtualImageProcessor* ip, float fg_scale) {
   cudaSetDevice(firstgpu);
 
   float resultchi2 = 0.0f;
@@ -4090,7 +4096,7 @@ __host__ float simulate(float* I, VirtualImageProcessor* ip) {
                        datasets[d].fields[f].ref_xobs_pix,
                        datasets[d].fields[f].ref_yobs_pix,
                        datasets[d].fields[f].nu[i],
-                       datasets[d].antennas[0].primary_beam);
+                       datasets[d].antennas[0].primary_beam, fg_scale);
 
         if (NULL != ip->getCKernel()) {
           apply_GCF<<<numBlocksNN, threadsPerBlockNN>>>(
@@ -4186,7 +4192,10 @@ __host__ float simulate(float* I, VirtualImageProcessor* ip) {
   return 0.5f * resultchi2;
 };
 
-__host__ float chi2(float* I, VirtualImageProcessor* ip, bool normalize) {
+__host__ float chi2(float* I,
+                    VirtualImageProcessor* ip,
+                    bool normalize,
+                    float fg_scale) {
   cudaSetDevice(firstgpu);
 
   float reduced_chi2 = 0.0f;
@@ -4216,7 +4225,7 @@ __host__ float chi2(float* I, VirtualImageProcessor* ip, bool normalize) {
                        datasets[d].fields[f].ref_xobs_pix,
                        datasets[d].fields[f].ref_yobs_pix,
                        datasets[d].fields[f].nu[i],
-                       datasets[d].antennas[0].primary_beam);
+                       datasets[d].antennas[0].primary_beam, fg_scale);
 
         if (NULL != ip->getCKernel()) {
           apply_GCF<<<numBlocksNN, threadsPerBlockNN>>>(
@@ -4317,7 +4326,8 @@ __host__ void dchi2(float* I,
                     float* dxi2,
                     float* result_dchi2,
                     VirtualImageProcessor* ip,
-                    bool normalize) {
+                    bool normalize,
+                    float fg_scale) {
   cudaSetDevice(firstgpu);
 
   for (int d = 0; d < nMeasurementSets; d++) {
@@ -4352,7 +4362,7 @@ __host__ void dchi2(float* I,
                     datasets[d].fields[f].device_visibilities[i][s].uvw,
                     datasets[d].fields[f].device_visibilities[i][s].weight, N,
                     datasets[d].fields[f].numVisibilitiesPerFreqPerStoke[i][s],
-                    noise_cut, datasets[d].fields[f].ref_xobs_pix,
+                    fg_scale, noise_cut, datasets[d].fields[f].ref_xobs_pix,
                     datasets[d].fields[f].ref_yobs_pix,
                     datasets[d].fields[f].phs_xobs_pix,
                     datasets[d].fields[f].phs_yobs_pix, DELTAX, DELTAY,
@@ -4368,7 +4378,7 @@ __host__ void dchi2(float* I,
                     datasets[d].fields[f].device_visibilities[i][s].uvw,
                     datasets[d].fields[f].device_visibilities[i][s].weight, N,
                     datasets[d].fields[f].numVisibilitiesPerFreqPerStoke[i][s],
-                    noise_cut, datasets[d].fields[f].ref_xobs_pix,
+                    fg_scale, noise_cut, datasets[d].fields[f].ref_xobs_pix,
                     datasets[d].fields[f].ref_yobs_pix,
                     datasets[d].fields[f].phs_xobs_pix,
                     datasets[d].fields[f].phs_yobs_pix, DELTAX, DELTAY,
@@ -4402,8 +4412,8 @@ __host__ void dchi2(float* I,
                   DChi2_total_alpha<<<numBlocksNN, threadsPerBlockNN>>>(
                       device_noise_image, result_dchi2,
                       vars_gpu[gpu_idx].device_dchi2, I,
-                      datasets[d].fields[f].nu[i], nu_0, noise_cut, threshold,
-                      N, M);
+                      datasets[d].fields[f].nu[i], nu_0, noise_cut, fg_scale,
+                      threshold, N, M);
                 checkCudaErrors(cudaDeviceSynchronize());
               }
             }
@@ -4464,10 +4474,11 @@ __host__ void linkApplyBeam2I(cufftComplex* image,
                               float xobs,
                               float yobs,
                               float freq,
-                              int primary_beam) {
+                              int primary_beam,
+                              float fg_scale) {
   apply_beam2I<<<numBlocksNN, threadsPerBlockNN>>>(
-      antenna_diameter, pb_factor, pb_cutoff, image, N, xobs, yobs, freq,
-      DELTAX, DELTAY, primary_beam);
+      antenna_diameter, pb_factor, pb_cutoff, image, N, xobs, yobs, fg_scale,
+      freq, DELTAX, DELTAY, primary_beam);
   checkCudaErrors(cudaDeviceSynchronize());
 };
 
@@ -4477,9 +4488,9 @@ __host__ void linkCalculateInu2I(cufftComplex* image, float* I, float freq) {
   checkCudaErrors(cudaDeviceSynchronize());
 };
 
-__host__ void linkChain2I(float* chain, float freq, float* I) {
-  chainRule2I<<<numBlocksNN, threadsPerBlockNN>>>(chain, device_noise_image, I,
-                                                  freq, nu_0, noise_cut, N, M);
+__host__ void linkChain2I(float* chain, float freq, float* I, float fg_scale) {
+  chainRule2I<<<numBlocksNN, threadsPerBlockNN>>>(
+      chain, device_noise_image, I, freq, nu_0, noise_cut, fg_scale, N, M);
   checkCudaErrors(cudaDeviceSynchronize());
 };
 

src/imageProcessor.cu

@@ -17,15 +17,16 @@ void ImageProcessor::apply_beam(cufftComplex* image,
                                 float xobs,
                                 float yobs,
                                 float freq,
-                                int primary_beam) {
+                                int primary_beam,
+                                float fg_scale) {
   if (image_count == 2)
     linkApplyBeam2I(image, antenna_diameter, pb_factor, pb_cutoff, xobs, yobs,
-                    freq, primary_beam);
+                    freq, primary_beam, fg_scale);
 };
 
-void ImageProcessor::chainRule(float* I, float freq) {
+void ImageProcessor::chainRule(float* I, float freq, float fg_scale) {
   if (image_count == 2)
-    linkChain2I(chain, freq, I);
+    linkChain2I(chain, freq, I, fg_scale);
 };
 
 void ImageProcessor::clipWNoise(float* I) {

src/mfs.cu

@@ -895,10 +895,6 @@ void MFS::setDevice() {
     printf("noise (Jy/pix) = %e\n", noise_jypix);
   }
 
-  Fi* chi2 = optimizer->getObjectiveFunction()->getFiByName("Chi2");
-  if (NULL != chi2)
-    chi2->setFgScale(this->fg_scale);
-
   std::vector<float> u_mask;
   if (variables.user_mask != "NULL") {
     u_mask = ioImageHandler->read_data_float_FITS(variables.user_mask);
@@ -958,8 +954,15 @@ void MFS::clearRun() {
 void MFS::run() {
   optimizer->getObjectiveFunction()->setIo(ioImageHandler);
 
+  Fi* chi2 = optimizer->getObjectiveFunction()->getFiByName("Chi2");
+
+  if (NULL != chi2)
+    chi2->setFgScale(this->fg_scale);
+
+  if (NULL != chi2 && chi2->getNormalize())
+    chi2->setFgScale(1.0f);
+
   if (this->gridding) {
-    Fi* chi2 = optimizer->getObjectiveFunction()->getFiByName("Chi2");
     if (NULL != chi2)
       chi2->setCKernel(this->ckernel);
   }
@@ -987,7 +990,7 @@ void MFS::run() {
   float chi2_final = 0.0f;
   float final_S = 0.0f;
   float lambda_S = 0.0f;
-  Fi* chi2 = optimizer->getObjectiveFunction()->getFiByName("Chi2");
+
   if (NULL != chi2) {
     chi2_final = chi2->get_fivalue();
   }
@@ -1050,8 +1053,8 @@ void MFS::writeImages() {
   printf("Saving final image to disk\n");
   if (IoOrderEnd == NULL) {
     ioImageHandler->printNotPathImage(image->getImage(), "JY/PIXEL",
-                                      optimizer->getCurrentIteration(), 0, 1.0,
-                                      true);
+                                      optimizer->getCurrentIteration(), 0,
+                                      this->fg_scale, true);
     if (print_images)
       ioImageHandler->printNotNormalizedImage(
           image->getImage(), "alpha.fits", "", optimizer->getCurrentIteration(),