Repository: kevin218/eureka
Issue: New scatter_ppm parameter does not work with shared fits

As of right now the new scatter_ppm parameter in Stage 5 introduced in PR #131 does not work properly for shared fits. Only the scatter_ppm parameter gets used for all wavelengths while the scatter_ppm_# parameters are ignored.

Changed files:
eureka/S5_lightcurve_fitting/likelihood.py, eureka/S5_lightcurve_fitting/fitters.py, eureka/S5_lightcurve_fitting/plots_s5.py

diff --git a/eureka/S5_lightcurve_fitting/fitters.py b/eureka/S5_lightcurve_fitting/fitters.py
index 190f0e8c..3be4d397 100644
--- a/eureka/S5_lightcurve_fitting/fitters.py
+++ b/eureka/S5_lightcurve_fitting/fitters.py
@@ -48,6 +48,8 @@ def lsqfitter(lc, model, meta, log, calling_function='lsq', **kwargs):
         Also saving covariance matrix for later estimation of sampler step size.
     - January 7-22, 2022 Megan Mansfield
         Adding ability to do a single shared fit across all channels
+    - February 28-March 1, 2022 Caroline Piaulet
+        Adding scatter_ppm parameter
     """
     # Group the different variable types
     freenames, freepars, prior1, prior2, priortype, indep_vars = group_variables(model)
@@ -76,6 +78,12 @@ def lsqfitter(lc, model, meta, log, calling_function='lsq', **kwargs):
     best_model.components[0].update(fit_params, freenames)
 
     model.update(fit_params, freenames)
+    if "scatter_ppm" in freenames:
+        ind = [i for i in np.arange(len(freenames)) if freenames[i][0:11] == "scatter_ppm"]
+        lc.unc_fit = np.ones_like(lc.flux) * fit_params[ind[0]] * 1e-6        
+        if len(ind)>1:
+            for chan in np.arange(lc.flux.size//lc.time.size):
+                lc.unc_fit[chan*lc.time.size:(chan+1)*lc.time.size] = fit_params[ind[chan]] * 1e-6
     
     # Save the covariance matrix in case it's needed to estimate step size for a sampler
     model_lc = model.eval()
@@ -184,6 +192,10 @@ def emceefitter(lc, model, meta, log, **kwargs):
         Adding ability to do a single shared fit across all channels
     - February 23-25, 2022 Megan Mansfield
         Added log-uniform and Gaussian priors.
+    - February 28-March 1, 2022 Caroline Piaulet
+        Adding scatter_ppm parameter. Added statements to avoid some initial 
+        state issues.
+
     """
     if not hasattr(meta, 'lsq_first') or meta.lsq_first:
         # Only call lsq fitter first if asked or lsq_first option wasn't passed (allowing backwards compatibility)
@@ -228,6 +240,12 @@ def emceefitter(lc, model, meta, log, **kwargs):
         log.writelog('Warning: >=1 params hit the lower bound in the lsq fit. Setting to the middle of the interval.')
         freepars[ind_min] = pmid[ind_min]
     
+    ind_zero_step = np.where(step_size==0.)
+    if len(ind_zero_step[0]):
+        log.writelog('Warning: >=1 params would have a zero step. changing to 0.001 * prior range')
+        step_size[ind_zero_step] = 0.001*(pmax[ind_zero_step] - pmin[ind_zero_step])
+        
+    
     pos = np.array([freepars + np.array(step_size)*np.random.randn(ndim) for i in range(nwalkers)])
     uniformprior=np.where(priortype=='U')
     loguniformprior=np.where(priortype=='LU')
@@ -279,8 +297,11 @@ def emceefitter(lc, model, meta, log, **kwargs):
 
     model.update(fit_params, freenames)
     if "scatter_ppm" in freenames:
-        ind = np.where(freenames == "scatter_ppm")
-        lc.unc_fit = medians[ind[0][0]]*1e-6
+        ind = [i for i in np.arange(len(freenames)) if freenames[i][0:11] == "scatter_ppm"]
+        lc.unc_fit = np.ones_like(lc.flux) * fit_params[ind[0]] * 1e-6        
+        if len(ind)>1:
+            for chan in np.arange(lc.flux.size//lc.time.size):
+                lc.unc_fit[chan*lc.time.size:(chan+1)*lc.time.size] = fit_params[ind[chan]] * 1e-6
 
     # Plot fit
     if meta.isplots_S5 >= 1:
@@ -293,7 +314,7 @@ def emceefitter(lc, model, meta, log, **kwargs):
     for freenames_i, fit_params_i in zip(freenames, fit_params):
         log.writelog('{0}: {1}'.format(freenames_i, fit_params_i))
     log.writelog('')
-
+    
     # Plot Allan plot
     if meta.isplots_S5 >= 3:
         plots.plot_rms(lc, model, meta, fitter='emcee')
@@ -340,6 +361,8 @@ def dynestyfitter(lc, model, meta, log, **kwargs):
         Adding ability to do a single shared fit across all channels
     - February 23-25, 2022 Megan Mansfield
         Added log-uniform and Gaussian priors.
+    - February 28-March 1, 2022 Caroline Piaulet
+        Adding scatter_ppm parameter. 
     """
     # Group the different variable types
     freenames, freepars, prior1, prior2, priortype, indep_vars = group_variables(model)
@@ -400,6 +423,14 @@ def dynestyfitter(lc, model, meta, log, **kwargs):
     best_model.components[0].update(fit_params, freenames)
 
     model.update(fit_params, freenames)
+    if "scatter_ppm" in freenames:
+        ind = [i for i in np.arange(len(freenames)) if freenames[i][0:11] == "scatter_ppm"]
+        lc.unc_fit = np.ones_like(lc.flux) * fit_params[ind[0]] * 1e-6        
+        if len(ind)>1:
+            for chan in np.arange(lc.flux.size//lc.time.size):
+                lc.unc_fit[chan*lc.time.size:(chan+1)*lc.time.size] = fit_params[ind[chan]] * 1e-6
+
+
     model_lc = model.eval()
     residuals = (lc.flux - model_lc) #/ lc.unc
 
@@ -457,6 +488,8 @@ def lmfitter(lc, model, meta, log, **kwargs):
 
     - December 29, 2021 Taylor Bell
         Updated documentation. Reduced repeated code.
+    - February 28-March 1, 2022 Caroline Piaulet
+        Adding scatter_ppm parameter. 
     """
         #TODO: Do something so that duplicate param names can all be handled (e.g. two Polynomail models with c0). Perhaps append something to the parameter name like c0_1 and c0_2?)
 
@@ -501,6 +534,13 @@ def lmfitter(lc, model, meta, log, **kwargs):
     # best_model.name = ', '.join(['{}:{}'.format(k, round(v[0], 2)) for k, v in params.dict.items()])
 
     model.update(fit_params, freenames)
+    if "scatter_ppm" in freenames:
+        ind = [i for i in np.arange(len(freenames)) if freenames[i][0:11] == "scatter_ppm"]
+        lc.unc_fit = np.ones_like(lc.flux) * fit_params[ind[0]] * 1e-6        
+        if len(ind)>1:
+            for chan in np.arange(lc.flux.size//lc.time.size):
+                lc.unc_fit[chan*lc.time.size:(chan+1)*lc.time.size] = fit_params[ind[chan]] * 1e-6
+
     # Plot fit
     if meta.isplots_S5 >= 1:
         plots.plot_fit(lc, model, meta, fitter='lmfitter')
diff --git a/eureka/S5_lightcurve_fitting/likelihood.py b/eureka/S5_lightcurve_fitting/likelihood.py
index 25c2544a..b8bc0564 100644
--- a/eureka/S5_lightcurve_fitting/likelihood.py
+++ b/eureka/S5_lightcurve_fitting/likelihood.py
@@ -36,8 +36,11 @@ def ln_like(theta, lc, model, freenames):
     model_lc = model.eval()
     residuals = (lc.flux - model_lc)
     if "scatter_ppm" in freenames:
-        ind = np.where(freenames=="scatter_ppm")
-        lc.unc_fit = theta[ind]*1e-6
+        ind = [i for i in np.arange(len(freenames)) if freenames[i][0:11] == "scatter_ppm"]
+        lc.unc_fit = np.ones_like(lc.flux) * theta[ind[0]] * 1e-6        
+        if len(ind)>1:
+            for chan in np.arange(lc.flux.size//lc.time.size):
+                lc.unc_fit[chan*lc.time.size:(chan+1)*lc.time.size] = theta[ind[chan]] * 1e-6
     else:
         lc.unc_fit = deepcopy(lc.unc)
     ln_like_val = (-0.5 * (np.sum((residuals / lc.unc_fit) ** 2+ np.log(2.0 * np.pi * (lc.unc_fit) ** 2))))
diff --git a/eureka/S5_lightcurve_fitting/plots_s5.py b/eureka/S5_lightcurve_fitting/plots_s5.py
index b67aadcb..b3ecd367 100644
--- a/eureka/S5_lightcurve_fitting/plots_s5.py
+++ b/eureka/S5_lightcurve_fitting/plots_s5.py
@@ -33,10 +33,13 @@ def plot_fit(lc, model, meta, fitter, isTitle=True):
         Moved plotting code to a separate function.
     - January 7-22, 2022 Megan Mansfield
         Adding ability to do a single shared fit across all channels
+    - February 28-March 1, 2022 Caroline Piaulet
+        Adding scatter_ppm parameter
     """
     if type(fitter)!=str:
         raise ValueError('Expected type str for fitter, instead received a {}'.format(type(fitter)))
 
+
     model_sys_full = model.syseval()
     model_phys_full, new_time = model.physeval(interp=meta.interp)
     model_lc = model.eval()
@@ -44,7 +47,7 @@ def plot_fit(lc, model, meta, fitter, isTitle=True):
     for i, channel in enumerate(lc.fitted_channels):
         flux = np.copy(lc.flux)
         if "unc_fit" in lc.__dict__.keys():
-            unc = deepcopy(lc.unc_fit)
+            unc = deepcopy(np.array(lc.unc_fit))
         else:
             unc = np.copy(lc.unc)
         model = np.copy(model_lc)
@@ -113,8 +116,6 @@ def plot_rms(lc, model, meta, fitter):
         Moved plotting code to a separate function.
     - January 7-22, 2022 Megan Mansfield
         Adding ability to do a single shared fit across all channels
-    - February 18, 2022 Caroline Piaulet
-        Using the fitted unc for plotted, if exists
     """
     if type(fitter)!=str:
         raise ValueError('Expected type str for fitter, instead received a {}'.format(type(fitter)))
@@ -124,14 +125,9 @@ def plot_rms(lc, model, meta, fitter):
 
     for channel in lc.fitted_channels:
         flux = np.copy(lc.flux)
-        if "unc_fit" in lc.__dict__.keys():
-            unc = np.copy(lc.unc_fit)
-        else:
-            unc = np.copy(lc.unc)
         model = np.copy(model_lc)
         if lc.share:
             flux = flux[channel*len(lc.time):(channel+1)*len(lc.time)]
-            unc = unc[channel*len(lc.time):(channel+1)*len(lc.time)]
             model = model[channel*len(lc.time):(channel+1)*len(lc.time)]
         
         residuals = flux - model
@@ -236,7 +232,7 @@ def plot_res_distr(lc, model, meta, fitter):
     for channel in lc.fitted_channels:
         flux = np.copy(lc.flux)
         if "unc_fit" in lc.__dict__.keys():
-            unc = np.copy(lc.unc_fit)
+            unc = np.copy(np.array(lc.unc_fit))
         else:
             unc = np.copy(lc.unc)
         model = np.copy(model_lc)