488 lines
15 KiB
Python
488 lines
15 KiB
Python
import logging
|
|
import numpy as np
|
|
import sympy
|
|
from decimal import Decimal
|
|
from pathlib import Path
|
|
from PyQt6.QtGui import QPixmap
|
|
from nqrduck.contrib.mplwidget import MplWidget
|
|
from nqrduck.helpers.signalprocessing import SignalProcessing as sp
|
|
from nqrduck.assets.icons import PulseParamters
|
|
from .base_spectrometer_model import BaseSpectrometerModel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Function:
|
|
name: str
|
|
parameters: list
|
|
expression: str | sympy.Expr
|
|
resolution: Decimal
|
|
start_x: float
|
|
end_x: float
|
|
|
|
def __init__(self, expr) -> None:
|
|
self.parameters = []
|
|
self.expr = expr
|
|
self.resolution = Decimal(1 / 30.72e6)
|
|
self.start_x = -1
|
|
self.end_x = 1
|
|
|
|
def get_time_points(self, pulse_length: Decimal) -> np.ndarray:
|
|
"""Returns the time domain points for the function with the given pulse length.
|
|
|
|
Args:
|
|
pulse_length (Decimal): The pulse length in seconds.
|
|
|
|
Returns:
|
|
np.ndarray: The time domain points.
|
|
"""
|
|
# Get the time domain points
|
|
n = int(pulse_length / self.resolution)
|
|
t = np.linspace(0, float(pulse_length), n)
|
|
return t
|
|
|
|
def evaluate(self, pulse_length: Decimal, resolution : Decimal = None) -> np.ndarray:
|
|
"""Evaluates the function for the given pulse length.
|
|
|
|
Args:
|
|
pulse_length (Decimal): The pulse length in seconds.
|
|
resolution (Decimal, optional): The resolution of the function in seconds. Defaults to None.
|
|
|
|
Returns:
|
|
np.ndarray: The evaluated function.
|
|
"""
|
|
if resolution is None:
|
|
resolution = self.resolution
|
|
n = int(pulse_length / resolution)
|
|
t = np.linspace(self.start_x, self.end_x, n)
|
|
x = sympy.symbols("x")
|
|
|
|
found_variables = dict()
|
|
# Create a dictionary of the parameters and their values
|
|
for parameter in self.parameters:
|
|
found_variables[parameter.symbol] = parameter.value
|
|
|
|
final_expr = self.expr.subs(found_variables)
|
|
# If the expression is a number (does not depend on x), return an array of that number
|
|
if final_expr.is_number:
|
|
return np.full(t.shape, float(final_expr))
|
|
|
|
f = sympy.lambdify([x], final_expr, "numpy")
|
|
|
|
return f(t)
|
|
|
|
def get_tdx(self, pulse_length: float) -> np.ndarray:
|
|
"""Returns the time domain points and the evaluated function for the given pulse length.
|
|
|
|
Args:
|
|
pulse_length (float): The pulse length in seconds.
|
|
|
|
Returns:
|
|
np.ndarray: The time domain points.
|
|
"""
|
|
n = int(pulse_length / self.resolution)
|
|
t = np.linspace(self.start_x, self.end_x, n)
|
|
return t
|
|
|
|
def frequency_domain_plot(self, pulse_length: Decimal) -> MplWidget:
|
|
"""Plots the frequency domain of the function for the given pulse length.
|
|
|
|
Args:
|
|
pulse_length (Decimal): The pulse length in seconds.
|
|
|
|
Returns:
|
|
MplWidget: The matplotlib widget containing the plot.
|
|
"""
|
|
mpl_widget = MplWidget()
|
|
td = self.get_time_points(pulse_length)
|
|
yd = self.evaluate(pulse_length)
|
|
xdf, ydf = sp.fft(td, yd)
|
|
mpl_widget.canvas.ax.plot(xdf, abs(ydf))
|
|
mpl_widget.canvas.ax.set_xlabel("Frequency in Hz")
|
|
mpl_widget.canvas.ax.set_ylabel("Magnitude")
|
|
mpl_widget.canvas.ax.grid(True)
|
|
return mpl_widget
|
|
|
|
def time_domain_plot(self, pulse_length: Decimal) -> MplWidget:
|
|
"""Plots the time domain of the function for the given pulse length.
|
|
|
|
Args:
|
|
pulse_length (Decimal): The pulse length in seconds.
|
|
|
|
Returns:
|
|
MplWidget: The matplotlib widget containing the plot.
|
|
"""
|
|
mpl_widget = MplWidget()
|
|
td = self.get_time_points(pulse_length)
|
|
mpl_widget.canvas.ax.plot(td, abs(self.evaluate(pulse_length)))
|
|
mpl_widget.canvas.ax.set_xlabel("Time in s")
|
|
mpl_widget.canvas.ax.set_ylabel("Magnitude")
|
|
mpl_widget.canvas.ax.grid(True)
|
|
return mpl_widget
|
|
|
|
def get_pulse_amplitude(self, pulse_length: Decimal, resolution : Decimal = None) -> np.array:
|
|
"""Returns the pulse amplitude in the time domain.
|
|
|
|
Args:
|
|
pulse_length (Decimal): The pulse length in seconds.
|
|
resolution (Decimal, optional): The resolution of the function in seconds. Defaults to None.
|
|
|
|
Returns:
|
|
np.array: The pulse amplitude.
|
|
"""
|
|
return self.evaluate(pulse_length, resolution=resolution)
|
|
|
|
def add_parameter(self, parameter: "Function.Parameter") -> None:
|
|
"""Adds a parameter to the function.
|
|
|
|
Args:
|
|
parameter (Function.Parameter): The parameter to add."""
|
|
self.parameters.append(parameter)
|
|
|
|
def to_json(self) -> dict:
|
|
"""Returns a json representation of the function.
|
|
|
|
Returns:
|
|
dict: The json representation of the function.
|
|
"""
|
|
return {
|
|
"name": self.name,
|
|
"parameters": [parameter.to_json() for parameter in self.parameters],
|
|
"expression": str(self.expr),
|
|
"resolution": self.resolution,
|
|
"start_x": self.start_x,
|
|
"end_x": self.end_x,
|
|
}
|
|
|
|
@classmethod
|
|
def from_json(cls, data: dict) -> "Function":
|
|
"""Creates a function from a json representation.
|
|
|
|
Args:
|
|
data (dict): The json representation of the function.
|
|
|
|
Returns:
|
|
Function: The function.
|
|
"""
|
|
for subclass in cls.__subclasses__():
|
|
if subclass.name == data["name"]:
|
|
cls = subclass
|
|
break
|
|
|
|
obj = cls()
|
|
obj.expr = data["expression"]
|
|
obj.name = data["name"]
|
|
obj.resolution = data["resolution"]
|
|
obj.start_x = data["start_x"]
|
|
obj.end_x = data["end_x"]
|
|
|
|
obj.parameters = []
|
|
for parameter in data["parameters"]:
|
|
obj.add_parameter(Function.Parameter.from_json(parameter))
|
|
|
|
return obj
|
|
|
|
@property
|
|
def expr(self):
|
|
"""The sympy expression of the function."""
|
|
return self._expr
|
|
|
|
@expr.setter
|
|
def expr(self, expr):
|
|
if isinstance(expr, str):
|
|
try:
|
|
self._expr = sympy.sympify(expr)
|
|
except:
|
|
logger.error("Could not convert %s to a sympy expression", expr)
|
|
raise SyntaxError("Could not convert %s to a sympy expression" % expr)
|
|
elif isinstance(expr, sympy.Expr):
|
|
self._expr = expr
|
|
|
|
@property
|
|
def resolution(self):
|
|
"""The resolution of the function in seconds."""
|
|
return self._resolution
|
|
|
|
@resolution.setter
|
|
def resolution(self, resolution):
|
|
try:
|
|
self._resolution = Decimal(resolution)
|
|
except:
|
|
logger.error("Could not convert %s to a decimal", resolution)
|
|
raise SyntaxError("Could not convert %s to a decimal" % resolution)
|
|
|
|
@property
|
|
def start_x(self):
|
|
"""The x value where the evalution of the function starts."""
|
|
return self._start_x
|
|
|
|
@start_x.setter
|
|
def start_x(self, start_x):
|
|
try:
|
|
self._start_x = float(start_x)
|
|
except:
|
|
logger.error("Could not convert %s to a float", start_x)
|
|
raise SyntaxError("Could not convert %s to a float" % start_x)
|
|
|
|
@property
|
|
def end_x(self):
|
|
"""The x value where the evalution of the function ends."""
|
|
return self._end_x
|
|
|
|
@end_x.setter
|
|
def end_x(self, end_x):
|
|
try:
|
|
self._end_x = float(end_x)
|
|
except:
|
|
logger.error("Could not convert %s to a float", end_x)
|
|
raise SyntaxError("Could not convert %s to a float" % end_x)
|
|
|
|
def get_pixmap(self):
|
|
"""This is the default pixmap for every function. If one wants to have a custom pixmap, this method has to be overwritten.
|
|
|
|
Returns:
|
|
QPixmap -- The default pixmap for every function"""
|
|
pixmap = PulseParamters.TXCustom()
|
|
return pixmap
|
|
|
|
class Parameter:
|
|
def __init__(self, name: str, symbol: str, value: float) -> None:
|
|
self.name = name
|
|
self.symbol = symbol
|
|
self.value = value
|
|
self.default = value
|
|
|
|
def set_value(self, value: float) -> None:
|
|
"""Sets the value of the parameter.
|
|
|
|
Args:
|
|
value (float): The new value of the parameter.
|
|
"""
|
|
self.value = value
|
|
logger.debug("Parameter %s set to %s", self.name, self.value)
|
|
|
|
def to_json(self) -> dict:
|
|
"""Returns a json representation of the parameter.
|
|
|
|
Returns:
|
|
dict: The json representation of the parameter.
|
|
"""
|
|
return {
|
|
"name": self.name,
|
|
"symbol": self.symbol,
|
|
"value": self.value,
|
|
"default": self.default,
|
|
}
|
|
|
|
@classmethod
|
|
def from_json(cls, data):
|
|
"""Creates a parameter from a json representation.
|
|
|
|
Args:
|
|
data (dict): The json representation of the parameter.
|
|
|
|
Returns:
|
|
Function.Parameter: The parameter.
|
|
"""
|
|
obj = cls(data["name"], data["symbol"], data["value"])
|
|
obj.default = data["default"]
|
|
return obj
|
|
|
|
|
|
class RectFunction(Function):
|
|
"""The rectangular function."""
|
|
|
|
name = "Rectangular"
|
|
|
|
def __init__(self) -> None:
|
|
expr = sympy.sympify("1")
|
|
super().__init__(expr)
|
|
|
|
def get_pixmap(self):
|
|
pixmap = PulseParamters.TXRect()
|
|
return pixmap
|
|
|
|
|
|
class SincFunction(Function):
|
|
"""The sinc function."""
|
|
|
|
name = "Sinc"
|
|
|
|
def __init__(self) -> None:
|
|
expr = sympy.sympify("sin(x * l)/ (x * l)")
|
|
super().__init__(expr)
|
|
self.add_parameter(Function.Parameter("Scale Factor", "l", 2))
|
|
self.start_x = -np.pi
|
|
self.end_x = np.pi
|
|
|
|
def get_pixmap(self):
|
|
pixmap = PulseParamters.TXSinc()
|
|
return pixmap
|
|
|
|
|
|
class GaussianFunction(Function):
|
|
"""The Gaussian function."""
|
|
|
|
name = "Gaussian"
|
|
|
|
def __init__(self) -> None:
|
|
expr = sympy.sympify("exp(-0.5 * ((x - mu) / sigma)**2)")
|
|
super().__init__(expr)
|
|
self.add_parameter(Function.Parameter("Mean", "mu", 0))
|
|
self.add_parameter(Function.Parameter("Standard Deviation", "sigma", 1))
|
|
self.start_x = -np.pi
|
|
self.end_x = np.pi
|
|
|
|
def get_pixmap(self):
|
|
pixmap = PulseParamters.TXGauss()
|
|
return pixmap
|
|
|
|
|
|
# class TriangleFunction(Function):
|
|
# def __init__(self) -> None:
|
|
# expr = sympy.sympify("triang(x)")
|
|
# super().__init__(lambda x: triang(x))
|
|
|
|
|
|
class CustomFunction(Function):
|
|
"""A custom function."""
|
|
|
|
name = "Custom"
|
|
|
|
def __init__(self) -> None:
|
|
expr = sympy.sympify(" 2 * x**2 + 3 * x + 1")
|
|
super().__init__(expr)
|
|
|
|
|
|
class Option:
|
|
"""Defines options for the pulse parameters which can then be set accordingly."""
|
|
|
|
def __init__(self, name: str, value) -> None:
|
|
self.name = name
|
|
self.value = value
|
|
|
|
def set_value(self):
|
|
raise NotImplementedError
|
|
|
|
def to_json(self):
|
|
return {"name": self.name, "value": self.value, "type": self.TYPE}
|
|
|
|
@classmethod
|
|
def from_json(cls, data) -> "Option":
|
|
for subclass in cls.__subclasses__():
|
|
if subclass.TYPE == data["type"]:
|
|
cls = subclass
|
|
break
|
|
|
|
# Check if from_json is implemented for the subclass
|
|
if cls.from_json.__func__ == Option.from_json.__func__:
|
|
obj = cls(data["name"], data["value"])
|
|
else:
|
|
obj = cls.from_json(data)
|
|
|
|
return obj
|
|
|
|
|
|
class BooleanOption(Option):
|
|
"""Defines a boolean option for a pulse parameter option."""
|
|
|
|
TYPE = "Boolean"
|
|
|
|
def set_value(self, value):
|
|
self.value = value
|
|
|
|
|
|
class NumericOption(Option):
|
|
"""Defines a numeric option for a pulse parameter option."""
|
|
|
|
TYPE = "Numeric"
|
|
|
|
def set_value(self, value):
|
|
self.value = float(value)
|
|
|
|
|
|
class FunctionOption(Option):
|
|
"""Defines a selection option for a pulse parameter option.
|
|
It takes different function objects."""
|
|
|
|
TYPE = "Function"
|
|
|
|
def __init__(self, name, functions) -> None:
|
|
super().__init__(name, functions[0])
|
|
self.functions = functions
|
|
|
|
def set_value(self, value):
|
|
self.value = value
|
|
|
|
def get_function_by_name(self, name):
|
|
for function in self.functions:
|
|
if function.name == name:
|
|
return function
|
|
raise ValueError("Function with name %s not found" % name)
|
|
|
|
def to_json(self):
|
|
return {"name": self.name, "value": self.value.to_json(), "type": self.TYPE}
|
|
|
|
@classmethod
|
|
def from_json(cls, data):
|
|
functions = [function() for function in Function.__subclasses__()]
|
|
obj = cls(data["name"], functions)
|
|
obj.value = Function.from_json(data["value"])
|
|
return obj
|
|
|
|
def get_pixmap(self):
|
|
return self.value.get_pixmap()
|
|
|
|
|
|
class TXPulse(BaseSpectrometerModel.PulseParameter):
|
|
RELATIVE_AMPLITUDE = "Relative TX Amplitude"
|
|
TX_PHASE = "TX Phase"
|
|
TX_PULSE_SHAPE = "TX Pulse Shape"
|
|
|
|
def __init__(self, name) -> None:
|
|
super().__init__(name)
|
|
self.add_option(NumericOption(self.RELATIVE_AMPLITUDE, 0))
|
|
self.add_option(NumericOption(self.TX_PHASE, 0))
|
|
self.add_option(
|
|
FunctionOption(
|
|
self.TX_PULSE_SHAPE,
|
|
[RectFunction(), SincFunction(), GaussianFunction(), CustomFunction()],
|
|
),
|
|
)
|
|
|
|
def get_pixmap(self):
|
|
if self.get_option_by_name(self.RELATIVE_AMPLITUDE).value > 0:
|
|
return self.get_option_by_name(self.TX_PULSE_SHAPE).get_pixmap()
|
|
else:
|
|
pixmap = PulseParamters.TXOff()
|
|
return pixmap
|
|
|
|
|
|
class RXReadout(BaseSpectrometerModel.PulseParameter):
|
|
RX = "RX"
|
|
|
|
def __init__(self, name) -> None:
|
|
super().__init__(name)
|
|
self.add_option(BooleanOption(self.RX, False))
|
|
|
|
def get_pixmap(self):
|
|
if self.get_option_by_name(self.RX).value == False:
|
|
pixmap = PulseParamters.RXOff()
|
|
else:
|
|
pixmap = PulseParamters.RXOn()
|
|
return pixmap
|
|
|
|
|
|
class Gate(BaseSpectrometerModel.PulseParameter):
|
|
GATE_STATE = "Gate State"
|
|
|
|
def __init__(self, name) -> None:
|
|
super().__init__(name)
|
|
self.add_option(BooleanOption(self.GATE_STATE, False))
|
|
|
|
def get_pixmap(self):
|
|
if self.get_option_by_name(self.GATE_STATE).value == False:
|
|
pixmap = PulseParamters.GateOff()
|
|
else:
|
|
pixmap = PulseParamters.GateOn()
|
|
return pixmap
|