From 419116aff956e3cbb007341fd09ff18b7897e2d3 Mon Sep 17 00:00:00 2001 From: jupfi Date: Thu, 23 May 2024 16:39:32 +0200 Subject: [PATCH] Refactoring and Linting. --- src/nqrduck_spectrometer/measurement.py | 229 +++++++++++------------- 1 file changed, 102 insertions(+), 127 deletions(-) diff --git a/src/nqrduck_spectrometer/measurement.py b/src/nqrduck_spectrometer/measurement.py index d473d28..39c32d3 100644 --- a/src/nqrduck_spectrometer/measurement.py +++ b/src/nqrduck_spectrometer/measurement.py @@ -1,9 +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 sympy.utilities.lambdify import lambdify from nqrduck.helpers.signalprocessing import SignalProcessing as sp from nqrduck.helpers.functions import Function @@ -30,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__( @@ -48,42 +47,37 @@ 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 - def add_fit(self, fit): + def add_fit(self, fit: "Fit") -> None: """Adds a fit to the measurement. Args: @@ -91,16 +85,15 @@ class Measurement: """ self.fits.append(fit) - def delete_fit(self, 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, name : str): + def edit_fit_name(self, fit: "Fit", name: str) -> None: """Edits the name of a fit. Args: @@ -110,37 +103,33 @@ class Measurement: logger.debug(f"Editing fit name to {name}.") fit.name = name - # Data saving and loading - - def to_json(self): - """Converts the measurement to a json-compatible format. + 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"]]) - obj = cls( + measurement = cls( json["name"], np.array(json["tdx"]), tdy, @@ -148,103 +137,98 @@ class Measurement: IF_frequency=json["IF_frequency"], ) - # Add fits - for fit in json["fits"]: - obj.add_fit(Fit.from_json(fit, obj)) + for fit_json in json["fits"]: + measurement.add_fit(Fit.from_json(fit_json, measurement)) - return obj + return measurement - # Measurement data + # 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): + def fits(self) -> list: """Fits of the measurement.""" return self._fits - + @fits.setter - def fits(self, value): + def fits(self, value: list) -> None: self._fits = value -class Fit(): + +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. """ + subclasses = [] - def __init_subclass__(cls, **kwargs): + 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: + 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): - """Fits the measurement data, sets the x and y data and sets the fit parameters and covariance. """ + 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 @@ -255,109 +239,100 @@ class Fit(): raise ValueError("Domain not recognized.") initial_guess = self.initial_guess() - parameters, covariance = curve_fit(self.fit_function, x, abs(y), p0=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, *parameters) - - self.parameters = parameters - self.covariance = covariance - - def get_fit_parameters_string(self): - """Get the fit parameters as a string. + self.y = self.fit_function(x, *self.parameters) - Returns: - str : The fit parameters as a string. - """ - return " ".join([f"{param:.2f}" for param in self.parameters]) - - def fit_function(self, x, *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. + *parameters: The fit parameters. Returns: - np.array : The y data. + np.array: The y data. """ raise NotImplementedError - - def initial_guess(self): + + def initial_guess(self) -> list: """Initial guess for the fit. Returns: - list : The initial guess. + list: The initial guess. """ raise NotImplementedError - - def to_json(self): - """Converts the fit to a json-compatible format. + def to_json(self) -> dict: + """Converts the fit to a JSON-compatible format. Returns: - dict : The fit in json-compatible format. + dict: The fit in JSON-compatible format. """ return { "name": self.name, "class": self.__class__.__name__, } - + @classmethod - def from_json(cls, data: dict, measurement : Measurement): - """Converts the json format to a fit. + 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. + data (dict): The fit in JSON-compatible format. + measurement (Measurement): The measurement. Returns: - Fit : The fit. + Fit: The fit. """ for subclass in cls.subclasses: - logger.debug(f"Keys data: {data.keys()}") if subclass.__name__ == data["class"]: - cls = subclass - break + return subclass(name=data["name"], measurement=measurement) + + raise ValueError(f"Subclass {data['class']} not found.") - return cls(measurement, data["name"]) - @property - def x(self): + def x(self) -> np.array: """The x data of the fit.""" return self._x - + @x.setter - def x(self, value): + def x(self, value: np.array) -> None: self._x = value @property - def y(self): + def y(self) -> np.array: """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, name = "T2*") -> None: - domain = "time" - measurement = measurement - super().__init__(name, domain, measurement) - - def fit(self): + @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() - # Create dict with fit parameters and covariance self.parameters = { "S0": self.parameters[0], "T2Star": self.parameters[1], - "covariance": self.covariance + "covariance": self.covariance, } - def fit_function (self, t, S0, T2Star): + 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): - return [1, 1] \ No newline at end of file + + def initial_guess(self) -> list: + """Initial guess for the T2* fit.""" + return [1, 1]