Implemented basic fitting.

src/nqrduck_spectrometer/measurement.py

@@ -2,6 +2,8 @@
 
 import logging
 import numpy as np
+from scipy.optimize import curve_fit
+from sympy.utilities.lambdify import lambdify
 from nqrduck.helpers.signalprocessing import SignalProcessing as sp
 from nqrduck.helpers.functions import Function
 
@@ -49,6 +51,8 @@ class Measurement:
         self.fdx, self.fdy = sp.fft(tdx, tdy, frequency_shift)
         self.IF_frequency = IF_frequency
 
+        self.fits = []
+
     def apodization(self, function: Function):
         """Applies apodization to the measurement data.
 
@@ -79,6 +83,14 @@ class Measurement:
 
         return apodized_measurement
 
+    def add_fit(self, fit):
+        """Adds a fit to the measurement.
+
+        Args:
+            fit (Fit): The fit to add.
+        """
+        self.fits.append(fit)
+
     # Data saving and loading
 
     def to_json(self):
@@ -171,3 +183,142 @@ class Measurement:
     @target_frequency.setter
     def target_frequency(self, value):
         self._target_frequency = value
+
+    @property
+    def fits(self):
+        """Fits of the measurement."""
+        return self._fits
+    
+    @fits.setter
+    def fits(self, value):
+        self._fits = value
+
+class Fit():
+    """The fit class for measurement data. A fit can be performed on either the frequency or time domain data.
+    
+    A measurement can have multiple fits.
+
+    Examples for fits in time domain would be the T2* relaxation time, while in frequency domain it could be the line width.
+
+    A fit has a name, a nqrduck function and a strategy for the algorithm to use.
+    """
+
+    def __init__(self, name: str, domain: str, measurement : Measurement) -> None:
+        """Initializes the fit."""
+        self.name = name
+        self.domain = domain
+        self.measurement = measurement
+
+    def fit(self):
+        """Fits the measurement data, sets the x and y data and returns the fit parameters and covariance. """
+        if self.domain == "time":
+            x = self.measurement.tdx
+            y = self.measurement.tdy
+        elif self.domain == "frequency":
+            x = self.measurement.fdx
+            y = self.measurement.fdy
+        else:
+            raise ValueError("Domain not recognized.")
+
+        initial_guess = self.initial_guess()
+        parameters, covariance = curve_fit(self.fit_function, x, abs(y), p0=initial_guess)
+
+        self.x = x
+        self.y = self.fit_function(x, *parameters)
+
+        
+        return parameters, covariance
+    
+    def get_fit_parameters_string(self):
+        """Get the fit parameters as a string.
+
+        Returns:
+            str : The fit parameters as a string.
+        """
+        return " ".join([f"{param:.2f}" for param in self.parameters])
+    
+    def fit_function(self, x, *parameters):
+        """The fit function.
+
+        Args:
+            x (np.array): The x data.
+            *parameters : The fit parameters.
+
+        Returns:
+            np.array : The y data.
+        """
+        raise NotImplementedError
+    
+    def initial_guess(self):
+        """Initial guess for the fit.
+
+        Returns:
+            list : The initial guess.
+        """
+        raise NotImplementedError
+        
+
+    def to_json(self):
+        """Converts the fit to a json-compatible format.
+
+        Returns:
+            dict : The fit in json-compatible format.
+        """
+        return {
+            "name": self.name,
+            "function": self.function.to_json(),
+            "strategy": self.strategy,
+        }
+    
+    @classmethod
+    def from_json(cls, json: dict):
+        """Converts the json format to a fit.
+
+        Args:
+            json (dict) : The fit in json-compatible format.
+
+        Returns:
+            Fit : The fit.
+        """
+        function = Function.from_json(json["function"])
+        return cls(json["name"], function, json["strategy"])
+    
+    @property
+    def x(self):
+        """The x data of the fit."""
+        return self._x
+    
+    @x.setter
+    def x(self, value):
+        self._x = value
+
+    @property
+    def y(self):
+        """The y data of the fit."""
+        return self._y
+    
+    @y.setter
+    def y(self, value):
+        self._y = value
+    
+class T2StarFit(Fit):
+
+    def __init__(self, measurement: Measurement) -> None:
+        self.name = "T2*"
+        self.domain = "time"
+        self.measurement = measurement
+    
+    def fit(self):
+        parameters, covariance = super().fit()
+        # Create dict with fit parameters and covariance
+        self.parameters = {
+            "S0": parameters[0],
+            "T2Star": parameters[1],
+            "covariance": covariance
+        }
+
+    def fit_function (self, t, S0, T2Star):
+        return S0 * np.exp(-t / T2Star)
+    
+    def initial_guess(self):
+        return [1, 1]