diff --git a/src/nqrduck_spectrometer/measurement.py b/src/nqrduck_spectrometer/measurement.py index a0d59dc..5923a6b 100644 --- a/src/nqrduck_spectrometer/measurement.py +++ b/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] \ No newline at end of file