Refactoring and Linting.

This commit is contained in:
jupfi 2024-05-23 16:39:32 +02:00
parent f5b6f3a689
commit 419116aff9

View file

@ -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 logging
import numpy as np import numpy as np
from scipy.optimize import curve_fit from scipy.optimize import curve_fit
from sympy.utilities.lambdify import lambdify
from nqrduck.helpers.signalprocessing import SignalProcessing as sp from nqrduck.helpers.signalprocessing import SignalProcessing as sp
from nqrduck.helpers.functions import Function from nqrduck.helpers.functions import Function
@ -30,8 +29,8 @@ class Measurement:
target_frequency (float): Target frequency of the measurement. target_frequency (float): Target frequency of the measurement.
frequency_shift (float): Frequency shift of the measurement. frequency_shift (float): Frequency shift of the measurement.
IF_frequency (float): Intermediate frequency 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__( def __init__(
@ -48,42 +47,37 @@ class Measurement:
self.tdx = tdx self.tdx = tdx
self.tdy = tdy self.tdy = tdy
self.target_frequency = target_frequency 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.IF_frequency = IF_frequency
self.fdx, self.fdy = sp.fft(tdx, tdy, frequency_shift)
self.fits = [] self.fits = []
def apodization(self, function: Function): def apodization(self, function: Function) -> "Measurement":
"""Applies apodization to the measurement data. """Applies apodization to the measurement data.
Args: Args:
function (Function): Apodization function. 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 duration = (self.tdx[-1] - self.tdx[0]) * 1e-6
resolution = duration / len(self.tdx) resolution = duration / len(self.tdx)
logger.debug("Resolution: %s", resolution) logger.debug("Resolution: %s", resolution)
y_weight = function.get_pulse_amplitude(duration, resolution) y_weight = function.get_pulse_amplitude(duration, resolution)
tdy_apodized = self.tdy * y_weight
tdy_measurement = self.tdy * y_weight
apodized_measurement = Measurement( apodized_measurement = Measurement(
self.name, self.name,
self.tdx, self.tdx,
tdy_measurement, tdy_apodized,
target_frequency=self.target_frequency, target_frequency=self.target_frequency,
IF_frequency=self.IF_frequency, IF_frequency=self.IF_frequency,
) )
return apodized_measurement return apodized_measurement
def add_fit(self, fit): def add_fit(self, fit: "Fit") -> None:
"""Adds a fit to the measurement. """Adds a fit to the measurement.
Args: Args:
@ -91,16 +85,15 @@ class Measurement:
""" """
self.fits.append(fit) self.fits.append(fit)
def delete_fit(self, fit): def delete_fit(self, fit: "Fit") -> None:
"""Deletes a fit from the measurement. """Deletes a fit from the measurement.
Args: Args:
fit (Fit): The fit to delete. fit (Fit): The fit to delete.
""" """
self.fits.remove(fit) 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. """Edits the name of a fit.
Args: Args:
@ -110,37 +103,33 @@ class Measurement:
logger.debug(f"Editing fit name to {name}.") logger.debug(f"Editing fit name to {name}.")
fit.name = name fit.name = name
# Data saving and loading def to_json(self) -> dict:
"""Converts the measurement to a JSON-compatible format.
def to_json(self):
"""Converts the measurement to a json-compatible format.
Returns: Returns:
dict : The measurement in json-compatible format. dict: The measurement in JSON-compatible format.
""" """
return { return {
"name": self.name, "name": self.name,
"tdx": self.tdx.tolist(), "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, "target_frequency": self.target_frequency,
"IF_frequency": self.IF_frequency, "IF_frequency": self.IF_frequency,
"fits": [fit.to_json() for fit in self.fits], "fits": [fit.to_json() for fit in self.fits],
} }
@classmethod @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: Args:
json (dict) : The measurement in json-compatible format. json (dict): The measurement in JSON-compatible format.
Returns: Returns:
Measurement : The measurement. Measurement: The measurement.
""" """
tdy = np.array([complex(y[0], y[1]) for y in json["tdy"]]) tdy = np.array([complex(y[0], y[1]) for y in json["tdy"]])
obj = cls( measurement = cls(
json["name"], json["name"],
np.array(json["tdx"]), np.array(json["tdx"]),
tdy, tdy,
@ -148,103 +137,98 @@ class Measurement:
IF_frequency=json["IF_frequency"], IF_frequency=json["IF_frequency"],
) )
# Add fits for fit_json in json["fits"]:
for fit in json["fits"]: measurement.add_fit(Fit.from_json(fit_json, measurement))
obj.add_fit(Fit.from_json(fit, obj))
return obj return measurement
# Measurement data # Properties for encapsulation
@property @property
def name(self): def name(self) -> str:
"""Name of the measurement.""" """Name of the measurement."""
return self._name return self._name
@name.setter @name.setter
def name(self, value): def name(self, value: str) -> None:
self._name = value self._name = value
@property @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 return self._tdx
@tdx.setter @tdx.setter
def tdx(self, value): def tdx(self, value: np.array) -> None:
self._tdx = value self._tdx = value
@property @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 return self._tdy
@tdy.setter @tdy.setter
def tdy(self, value): def tdy(self, value: np.array) -> None:
self._tdy = value self._tdy = value
@property @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 return self._fdx
@fdx.setter @fdx.setter
def fdx(self, value): def fdx(self, value: np.array) -> None:
self._fdx = value self._fdx = value
@property @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 return self._fdy
@fdy.setter @fdy.setter
def fdy(self, value): def fdy(self, value: np.array) -> None:
self._fdy = value self._fdy = value
# Pulse parameters
@property @property
def target_frequency(self): def target_frequency(self) -> float:
"""Target frequency of the measurement.""" """Target frequency of the measurement."""
return self._target_frequency return self._target_frequency
@target_frequency.setter @target_frequency.setter
def target_frequency(self, value): def target_frequency(self, value: float) -> None:
self._target_frequency = value self._target_frequency = value
@property @property
def fits(self): def fits(self) -> list:
"""Fits of the measurement.""" """Fits of the measurement."""
return self._fits return self._fits
@fits.setter @fits.setter
def fits(self, value): def fits(self, value: list) -> None:
self._fits = value 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. """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. 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 = [] subclasses = []
def __init_subclass__(cls, **kwargs): def __init_subclass__(cls, **kwargs) -> None:
"""Adds the subclass to the list of subclasses.""" """Adds the subclass to the list of subclasses."""
super().__init_subclass__(**kwargs) super().__init_subclass__(**kwargs)
cls.subclasses.append(cls) 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.""" """Initializes the fit."""
self.name = name self.name = name
self.domain = domain self.domain = domain
self.measurement = measurement self.measurement = measurement
self.fit() self.fit()
def fit(self): def fit(self) -> None:
"""Fits the measurement data, sets the x and y data and sets the fit parameters and covariance. """ """Fits the measurement data and sets the fit parameters and covariance."""
if self.domain == "time": if self.domain == "time":
x = self.measurement.tdx x = self.measurement.tdx
y = self.measurement.tdy y = self.measurement.tdy
@ -255,48 +239,38 @@ class Fit():
raise ValueError("Domain not recognized.") raise ValueError("Domain not recognized.")
initial_guess = self.initial_guess() 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.x = x
self.y = self.fit_function(x, *parameters) self.y = self.fit_function(x, *self.parameters)
self.parameters = parameters def fit_function(self, x: np.array, *parameters) -> np.array:
self.covariance = 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. """The fit function.
Args: Args:
x (np.array): The x data. x (np.array): The x data.
*parameters : The fit parameters. *parameters: The fit parameters.
Returns: Returns:
np.array : The y data. np.array: The y data.
""" """
raise NotImplementedError raise NotImplementedError
def initial_guess(self): def initial_guess(self) -> list:
"""Initial guess for the fit. """Initial guess for the fit.
Returns: Returns:
list : The initial guess. list: The initial guess.
""" """
raise NotImplementedError raise NotImplementedError
def to_json(self) -> dict:
def to_json(self): """Converts the fit to a JSON-compatible format.
"""Converts the fit to a json-compatible format.
Returns: Returns:
dict : The fit in json-compatible format. dict: The fit in JSON-compatible format.
""" """
return { return {
"name": self.name, "name": self.name,
@ -304,60 +278,61 @@ class Fit():
} }
@classmethod @classmethod
def from_json(cls, data: dict, measurement : Measurement): def from_json(cls, data: dict, measurement: Measurement) -> "Fit":
"""Converts the json format to a fit. """Converts the JSON format to a fit.
Args: Args:
data (dict) : The fit in json-compatible format. data (dict): The fit in JSON-compatible format.
measurement (Measurement) : The measurement. measurement (Measurement): The measurement.
Returns: Returns:
Fit : The fit. Fit: The fit.
""" """
for subclass in cls.subclasses: for subclass in cls.subclasses:
logger.debug(f"Keys data: {data.keys()}")
if subclass.__name__ == data["class"]: if subclass.__name__ == data["class"]:
cls = subclass return subclass(name=data["name"], measurement=measurement)
break
return cls(measurement, data["name"]) raise ValueError(f"Subclass {data['class']} not found.")
@property @property
def x(self): def x(self) -> np.array:
"""The x data of the fit.""" """The x data of the fit."""
return self._x return self._x
@x.setter @x.setter
def x(self, value): def x(self, value: np.array) -> None:
self._x = value self._x = value
@property @property
def y(self): def y(self) -> np.array:
"""The y data of the fit.""" """The y data of the fit."""
return self._y return self._y
@y.setter @y.setter
def y(self, value): def y(self, value: np.array) -> None:
self._y = value self._y = value
class T2StarFit(Fit): class T2StarFit(Fit):
"""T2* fit for measurement data."""
def __init__(self, measurement: Measurement, name = "T2*") -> None: def __init__(self, measurement: Measurement, name: str = "T2*") -> None:
domain = "time" """Initializes the T2* fit."""
measurement = measurement super().__init__(name, "time", measurement)
super().__init__(name, domain, measurement)
def fit(self): def fit(self) -> None:
"""Fits the measurement data and sets the fit parameters and covariance."""
super().fit() super().fit()
# Create dict with fit parameters and covariance
self.parameters = { self.parameters = {
"S0": self.parameters[0], "S0": self.parameters[0],
"T2Star": self.parameters[1], "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) return S0 * np.exp(-t / T2Star)
def initial_guess(self): def initial_guess(self) -> list:
"""Initial guess for the T2* fit."""
return [1, 1] return [1, 1]