Merge pull request #69 from CHIMEFRB/chimefrb-updates

Chimefrb updates

fitburst/analysis/fitter.py

@@ -10,6 +10,7 @@
 from scipy.optimize import least_squares
 import fitburst.routines.derivative as deriv
 import numpy as np
+import sys
 
 class LSFitter:
     """
@@ -60,14 +61,97 @@ def __init__(self, data: float, model: object, good_freq: bool, weighted_fit: bo
 
         # set parameters for fitter configuration.
         self.good_freq = good_freq
-        self.global_parameters = ["dm", "scattering_timescale"]
+        self.global_parameters = ["dm", "dm_index", "scattering_timescale", "scattering_index"]
         self.success = None
         self.weights = None
         self.weighted_fit = weighted_fit
 
         # before running fit, determine per-channel weights.
         self._set_weights()
 
+    def compute_hessian(self, data: float, parameter_list: list) -> float:
+        """
+        Computes the Jacobian matrix for the scipy.optimize.least_squares solver.
+
+        Parameters
+        ----------
+        parameter_list : list
+            A list of names for fit parameters.
+        
+        Returns
+        -------
+        jacobian : float
+            The Jacobian matrix for the residuals vector supplied to least_squares()
+
+        Notes
+        -----
+        The parameter_list argument is not actually used in this method, but is 
+        supplied in order to conform to the definition of the callable needed by 
+        least_squares() for exact calculation of the Jacobian in terms of derivatives.
+        """
+
+        # load all parameter values into a dictionary.
+        parameter_dict = self.model.get_parameters_dict()
+
+        # before calculating, compute residual.
+        residual = data - self.model.compute_model(data = data)
+
+        # define the scale of the Hessian matrix and its labels.
+        par_labels_output = []
+        par_labels = []
+        num_par = 0
+        print("here are the fit parameters:", self.fit_parameters)
+
+        for current_par in self.fit_parameters:
+            if np.any([current_par == x for x in self.global_parameters]):
+                par_labels_output += [current_par]
+                par_labels += [current_par]
+                num_par += 1
+
+            else:
+                par_labels_output += [f"{current_par}{idx+1}" for idx in range(self.model.num_components)]
+                par_labels += ([current_par] * self.model.num_components)
+                num_par += (self.model.num_components)
+
+        hessian = np.zeros((num_par, num_par), dtype=float)
+
+        # now loop over all fit parameters and compute derivatives.
+        for current_par_idx_1, current_par_1 in zip(range(num_par), par_labels):
+            current_par_deriv_1 = getattr(deriv, f"deriv_model_wrt_{current_par_1}")
+            current_deriv_1 = current_par_deriv_1(
+                parameter_dict, self.model, component=(current_par_idx_1 % self.model.num_components)
+            )
+
+            # for efficient calculation, only compute one half of the matrix and fill the other half appropriately.
+            for current_par_idx_2, current_par_2 in zip(range(current_par_idx_1, num_par), par_labels[current_par_idx_1:]):
+
+                # also compute a given derivative for all burst components.
+                current_par_deriv_2 = getattr(deriv, f"deriv_model_wrt_{current_par_2}")
+                current_deriv_2 = current_par_deriv_2(
+                    parameter_dict, self.model, component=(current_par_idx_2 % self.model.num_components)
+                )
+
+                # correct name ordering of mixed partial derivative, if necessary.
+                try:
+                    current_mixed_deriv = getattr(deriv, f"deriv2_model_wrt_{current_par_1}_{current_par_2}")
+
+                except AttributeError:
+                    current_mixed_deriv = getattr(deriv, f"deriv2_model_wrt_{current_par_2}_{current_par_1}")
+
+                # only computed mixed derivative for parameters that describe the same component.
+                current_deriv_mixed = 0
+
+                if (current_par_idx_1 % self.model.num_components) == (current_par_idx_2 % self.model.num_components):
+                    current_deriv_mixed = current_mixed_deriv(
+                        parameter_dict, self.model, component=(current_par_idx_2 % self.model.num_components)
+                    )
+
+                # finally, compute the hessian here.
+                current_hes = 2 * current_deriv_1 * current_deriv_2 - residual * current_deriv_mixed
+                hessian[current_par_idx_1, current_par_idx_2] = np.sum(current_hes * self.weights[:, None])
+                hessian[current_par_idx_2, current_par_idx_1] = hessian[current_par_idx_1, current_par_idx_2]
+
+        return hessian, par_labels_output
 
     def compute_jacobian(self, parameter_list: list) -> float:
         """
@@ -102,8 +186,14 @@ def compute_jacobian(self, parameter_list: list) -> float:
         for current_parameter in self.fit_parameters:
             current_deriv = getattr(deriv, f"deriv_model_wrt_{current_parameter}")
 
+            # before computing derivatives, account for global parameters.
+            num_derivs = self.model.num_components
+
+            if current_parameter in self.global_parameters:
+                num_derivs = 1
+
             # also compute a given derivative for all burst components.
-            for current_component in range(self.model.num_components):
+            for current_component in range(num_derivs):
                 current_jac = -current_deriv(
                     parameter_dict, self.model, component=current_component
                 )
@@ -205,7 +295,7 @@ def fit(self, exact_jacobian: bool = True) -> None:
 
         except Exception as exc:
             print("ERROR: solver encountered a failure! Debug!")
-            print(exc)
+            print(sys.exc_info()[2])
 
     def fix_parameter(self, parameter_list: list) -> None:
         """
@@ -368,6 +458,7 @@ def _compute_fit_statistics(self, spectrum_observed: float, fit_result: object)
             covariance = np.linalg.inv(hessian) * chisq_final_reduced
             uncertainties = [float(x) for x in np.sqrt(np.diag(covariance)).tolist()]
 
+            self.covariance = covariance
             self.fit_statistics["bestfit_uncertainties"] = self.load_fit_parameters_list(
                 uncertainties)
             self.fit_statistics["bestfit_covariance"] = None # return the full matrix at some point?

fitburst/analysis/model.py

@@ -351,7 +351,7 @@ def get_parameters_dict(self) -> dict:
 
         return parameter_dict
 
-    def update_parameters(self, model_parameters: dict) -> None:
+    def update_parameters(self, model_parameters: dict, global_parameters: list = ["dm", "scattering_timescale"]) -> None:
         """
         Overloads parameter values stored in object with those supplied by the user.
 
@@ -373,4 +373,8 @@ def update_parameters(self, model_parameters: dict) -> None:
 
             # if number of components is 2 or greater, update num_components attribute.
             if len(model_parameters[current_parameter]) > 1:
-                setattr(self, "num_components", len(model_parameters[current_parameter]))
+                num_components = len(model_parameters[current_parameter])
+                setattr(self, "num_components", num_components)
+
+                if current_parameter in global_parameters:
+                    setattr(self, current_parameter, [model_parameters[current_parameter][0]] * num_components)

fitburst/analysis/peak_finder.py

@@ -44,7 +44,7 @@ class FindPeak:
     #This is for creating the name of output files same as the input files.
     #name_of_file=str(input_file).replace('.npz','')
 
-    def __init__(self,data,time,freq,rms=None):    
+    def __init__(self, data: float, time: float, freq : float, rms: float = None):
         """
         Initializes FindPeak class with data(data is input file provided),time,freq as a parameters.
 
@@ -64,12 +64,12 @@ def __init__(self,data,time,freq,rms=None):
             this is the percentage by which we want to shift(increase) the original rms value 
 
         """     
-        self.data=data
-        self.freq=freq
-        self.time=time
-        self.rms=rms
+        self.data = data
+        self.freq = freq
+        self.time = time
+        self.rms = rms
 
-    def find_peak(self):
+    def find_peak(self, distance: int = 5):
         """
         Finds peak positions in a spectrum and the approximate temporal width of each peak. Prints out
         the peak positions and corresponding temporal width in a data frame.
@@ -79,7 +79,7 @@ def find_peak(self):
         """
         self.mean_intensity_freq=self.data.mean(axis=0)
         plt.plot(self.time*1000,self.mean_intensity_freq)
-        peaks_location=find_peaks(self.mean_intensity_freq, distance=5)
+        peaks_location=find_peaks(self.mean_intensity_freq, distance=distance)
 
         self.peak_times=self.time[peaks_location[0]]
         self.peak_mean_intensities=self.mean_intensity_freq[peaks_location[0]]
@@ -91,7 +91,8 @@ def find_peak(self):
         #Shift RMS line by given percentage if asked in the input.
         if self.rms is not None:
             #Increase the value of rms intensity by given rms shift percent.
-            shifted_rms_intensity=self.rms_intensity+(self.rms/100)*self.rms_intensity
+            #shifted_rms_intensity=self.rms_intensity+(self.rms/100)*self.rms_intensity
+            shifted_rms_intensity = self.rms * np.max(self.mean_intensity_freq)
 
             #Find the peaks greater than shifted_rms values.
             index_peaks_greater_rms=np.where(self.peak_mean_intensities>shifted_rms_intensity)
@@ -169,7 +170,8 @@ def get_parameters_dict(self, original_dict: dict, update_width=False):
         Returns: Dictionary with values in List
         --------
         """
-        self.find_peak()
+
+        #self.find_peak()
         mul_factor=len(self.time_of_arrivals)
         burst_parameters = {}