Merge branch 'feature-fitting' into development

src/nqrduck_spectrometer/measurement.py

@@ -1,7 +1,8 @@
-"""Class for handling measurement data."""
+"""This module defines the measurement data structure and the fit class for measurement data."""
 
 import logging
 import numpy as np
+from scipy.optimize import curve_fit
 from nqrduck.helpers.signalprocessing import SignalProcessing as sp
 from nqrduck.helpers.functions import Function
 
@@ -28,8 +29,8 @@ class Measurement:
         target_frequency (float): Target frequency of the measurement.
         frequency_shift (float): Frequency shift of the measurement.
         IF_frequency (float): Intermediate frequency of the measurement.
-        xf (np.array): Frequency axis for the x axis of the measurement data.
+        fdx (np.array): Frequency axis for the x axis of the measurement data.
-        yf (np.array): Frequency axis for the y axis of the measurement data.
+        fdy (np.array): Frequency axis for the y axis of the measurement data.
     """
 
     def __init__(
@@ -46,69 +47,89 @@ class Measurement:
         self.tdx = tdx
         self.tdy = tdy
         self.target_frequency = target_frequency
-        self.fdx, self.fdy = sp.fft(tdx, tdy, frequency_shift)
+        self.frequency_shift = frequency_shift
         self.IF_frequency = IF_frequency
+        self.fdx, self.fdy = sp.fft(tdx, tdy, frequency_shift)
+        self.fits = []
 
-    def apodization(self, function: Function):
+    def apodization(self, function: Function) -> "Measurement":
         """Applies apodization to the measurement data.
 
         Args:
             function (Function): Apodization function.
 
-        returns:
+        Returns:
-            Measurement : The apodized measurement.
+            Measurement: The apodized measurement.
         """
-        # Get the y data weights from the function
         duration = (self.tdx[-1] - self.tdx[0]) * 1e-6
-
         resolution = duration / len(self.tdx)
-
         logger.debug("Resolution: %s", resolution)
 
         y_weight = function.get_pulse_amplitude(duration, resolution)
-
+        tdy_apodized = self.tdy * y_weight
-        tdy_measurement = self.tdy * y_weight
 
         apodized_measurement = Measurement(
             self.name,
             self.tdx,
-            tdy_measurement,
+            tdy_apodized,
             target_frequency=self.target_frequency,
             IF_frequency=self.IF_frequency,
         )
-
         return apodized_measurement
 
-    # Data saving and loading
+    def add_fit(self, fit: "Fit") -> None:
+        """Adds a fit to the measurement.
 
-    def to_json(self):
+        Args:
-        """Converts the measurement to a json-compatible format.
+            fit (Fit): The fit to add.
+        """
+        self.fits.append(fit)
+
+    def delete_fit(self, fit: "Fit") -> None:
+        """Deletes a fit from the measurement.
+
+        Args:
+            fit (Fit): The fit to delete.
+        """
+        self.fits.remove(fit)
+
+    def edit_fit_name(self, fit: "Fit", name: str) -> None:
+        """Edits the name of a fit.
+
+        Args:
+            fit (Fit): The fit to edit.
+            name (str): The new name.
+        """
+        logger.debug(f"Editing fit name to {name}.")
+        fit.name = name
+
+    def to_json(self) -> dict:
+        """Converts the measurement to a JSON-compatible format.
 
         Returns:
-            dict : The measurement in json-compatible format.
+            dict: The measurement in JSON-compatible format.
         """
         return {
             "name": self.name,
             "tdx": self.tdx.tolist(),
-            "tdy": [
+            "tdy": [[x.real, x.imag] for x in self.tdy],
-                [x.real, x.imag] for x in self.tdy
-            ],  # Convert complex numbers to list
             "target_frequency": self.target_frequency,
             "IF_frequency": self.IF_frequency,
+            "fits": [fit.to_json() for fit in self.fits],
         }
 
     @classmethod
-    def from_json(cls, json: dict):
+    def from_json(cls, json: dict) -> "Measurement":
-        """Converts the json format to a measurement.
+        """Converts the JSON format to a measurement.
 
         Args:
-            json (dict) : The measurement in json-compatible format.
+            json (dict): The measurement in JSON-compatible format.
 
         Returns:
-            Measurement : The measurement.
+            Measurement: The measurement.
         """
         tdy = np.array([complex(y[0], y[1]) for y in json["tdy"]])
-        return cls(
+        measurement = cls(
             json["name"],
             np.array(json["tdx"]),
             tdy,
@@ -116,58 +137,227 @@ class Measurement:
             IF_frequency=json["IF_frequency"],
         )
 
-    # Measurement data
+        for fit_json in json["fits"]:
+            measurement.add_fit(Fit.from_json(fit_json, measurement))
+
+        return measurement
+
+    # Properties for encapsulation
     @property
-    def name(self):
+    def name(self) -> str:
         """Name of the measurement."""
         return self._name
 
     @name.setter
-    def name(self, value):
+    def name(self, value: str) -> None:
         self._name = value
 
     @property
-    def tdx(self):
+    def tdx(self) -> np.array:
-        """Time axis for the x axis of the measurement data."""
+        """Time domain data for the measurement (x)."""
         return self._tdx
 
     @tdx.setter
-    def tdx(self, value):
+    def tdx(self, value: np.array) -> None:
         self._tdx = value
 
     @property
-    def tdy(self):
+    def tdy(self) -> np.array:
-        """Time axis for the y axis of the measurement data."""
+        """Time domain data for the measurement (y)."""
         return self._tdy
 
     @tdy.setter
-    def tdy(self, value):
+    def tdy(self, value: np.array) -> None:
         self._tdy = value
 
     @property
-    def fdx(self):
+    def fdx(self) -> np.array:
-        """Frequency axis for the x axis of the measurement data."""
+        """Frequency domain data for the measurement (x)."""
         return self._fdx
 
     @fdx.setter
-    def fdx(self, value):
+    def fdx(self, value: np.array) -> None:
         self._fdx = value
 
     @property
-    def fdy(self):
+    def fdy(self) -> np.array:
-        """Frequency axis for the y axis of the measurement data."""
+        """Frequency domain data for the measurement (y)."""
         return self._fdy
 
     @fdy.setter
-    def fdy(self, value):
+    def fdy(self, value: np.array) -> None:
         self._fdy = value
 
-    # Pulse parameters
     @property
-    def target_frequency(self):
+    def target_frequency(self) -> float:
         """Target frequency of the measurement."""
         return self._target_frequency
 
     @target_frequency.setter
-    def target_frequency(self, value):
+    def target_frequency(self, value: float) -> None:
         self._target_frequency = value
+
+    @property
+    def fits(self) -> list:
+        """Fits of the measurement."""
+        return self._fits
+
+    @fits.setter
+    def fits(self, value: list) -> None:
+        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.
+    """
+
+    subclasses = []
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        """Adds the subclass to the list of subclasses."""
+        super().__init_subclass__(**kwargs)
+        cls.subclasses.append(cls)
+
+    def __init__(self, name: str, domain: str, measurement: Measurement) -> None:
+        """Initializes the fit."""
+        self.name = name
+        self.domain = domain
+        self.measurement = measurement
+        self.fit()
+
+    def fit(self) -> None:
+        """Fits the measurement data and sets 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()
+        self.parameters, self.covariance = curve_fit(
+            self.fit_function, x, abs(y), p0=initial_guess
+        )
+
+        self.x = x
+        self.y = self.fit_function(x, *self.parameters)
+
+    def fit_function(self, x: np.array, *parameters) -> np.array:
+        """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) -> list:
+        """Initial guess for the fit.
+
+        Returns:
+            list: The initial guess.
+        """
+        raise NotImplementedError
+
+    def to_json(self) -> dict:
+        """Converts the fit to a JSON-compatible format.
+
+        Returns:
+            dict: The fit in JSON-compatible format.
+        """
+        return {
+            "name": self.name,
+            "class": self.__class__.__name__,
+        }
+
+    @classmethod
+    def from_json(cls, data: dict, measurement: Measurement) -> "Fit":
+        """Converts the JSON format to a fit.
+
+        Args:
+            data (dict): The fit in JSON-compatible format.
+            measurement (Measurement): The measurement.
+
+        Returns:
+            Fit: The fit.
+        """
+        for subclass in cls.subclasses:
+            if subclass.__name__ == data["class"]:
+                return subclass(name=data["name"], measurement=measurement)
+
+        raise ValueError(f"Subclass {data['class']} not found.")
+
+    @property
+    def x(self) -> np.array:
+        """The x data of the fit."""
+        return self._x
+
+    @x.setter
+    def x(self, value: np.array) -> None:
+        self._x = value
+
+    @property
+    def y(self) -> np.array:
+        """The y data of the fit."""
+        return self._y
+
+    @y.setter
+    def y(self, value: np.array) -> None:
+        self._y = value
+
+
+class T2StarFit(Fit):
+    """T2* fit for measurement data."""
+
+    def __init__(self, measurement: Measurement, name: str = "T2*") -> None:
+        """Initializes the T2* fit."""
+        super().__init__(name, "time", measurement)
+
+    def fit(self) -> None:
+        """Fits the measurement data and sets the fit parameters and covariance."""
+        super().fit()
+        self.parameters = {
+            "S0": self.parameters[0],
+            "T2Star": self.parameters[1],
+            "covariance": self.covariance,
+        }
+
+    def fit_function(self, t: np.array, S0: float, T2Star: float) -> np.array:
+        """The T2* fit function used for curve fitting."""
+        return S0 * np.exp(-t / T2Star)
+
+    def initial_guess(self) -> list:
+        """Initial guess for the T2* fit."""
+        return [1, 1]
+    
+class LorentzianFit(Fit):
+    """Lorentzian fit for measurement data."""
+
+    def __init__(self, measurement: Measurement, name: str = "Lorentzian") -> None:
+        """Initializes the Lorentzian fit."""
+        super().__init__(name, "frequency", measurement)
+
+    def fit(self) -> None:
+        """Fits the measurement data and sets the fit parameters and covariance."""
+        super().fit()
+        self.parameters = {
+            "S0": self.parameters[0],
+            "T2Star": self.parameters[1],
+            "covariance": self.covariance,
+        }
+        logger.debug("Lorentzian fit parameters: %s", self.parameters)
+
+    def fit_function(self, f: np.array, S0: float, T2Star: float) -> np.array:
+        """The Lorentzian fit function used for curve fitting."""
+        return S0 / (1 + (2 * np.pi * f * T2Star) ** 2)
+
+    def initial_guess(self) -> list:
+        """Initial guess for the Lorentzian fit."""
+        return [1, 1]