diff --git a/src/nqrduck_spectrometer/measurement.py b/src/nqrduck_spectrometer/measurement.py index a0d59dc..61d04d1 100644 --- a/src/nqrduck_spectrometer/measurement.py +++ b/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. - yf (np.array): Frequency axis for the y axis of the measurement data. + fdx (np.array): Frequency axis for the x 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: - Measurement : The apodized measurement. + Returns: + 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_measurement = self.tdy * y_weight + tdy_apodized = 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): - """Converts the measurement to a json-compatible format. + Args: + 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": [ - [x.real, x.imag] for x in self.tdy - ], # Convert complex numbers to list + "tdy": [[x.real, x.imag] for x in self.tdy], "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): - """Converts the json format to a measurement. + def from_json(cls, json: dict) -> "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): - """Time axis for the x axis of the measurement data.""" + def tdx(self) -> np.array: + """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): - """Time axis for the y axis of the measurement data.""" + def tdy(self) -> np.array: + """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): - """Frequency axis for the x axis of the measurement data.""" + def fdx(self) -> np.array: + """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): - """Frequency axis for the y axis of the measurement data.""" + def fdy(self) -> np.array: + """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]