From dca1c6816f0697ca3c6827fd07a0236a3189b922 Mon Sep 17 00:00:00 2001
From: jupfi <git@jupfi.me>
Date: Thu, 23 May 2024 17:04:58 +0200
Subject: [PATCH] Added fitting of frequency domain data.

---
 src/nqrduck_measurement/model.py              |  2 +-
 .../signalprocessing_options.py               | 19 ++-----------------
 src/nqrduck_measurement/view.py               |  8 +++++++-
 3 files changed, 10 insertions(+), 19 deletions(-)

diff --git a/src/nqrduck_measurement/model.py b/src/nqrduck_measurement/model.py
index 871a99b..9ba0e56 100644
--- a/src/nqrduck_measurement/model.py
+++ b/src/nqrduck_measurement/model.py
@@ -45,7 +45,7 @@ class MeasurementModel(ModuleModel):
 
     FILE_EXTENSION = "meas"
     # This constants are used to determine which view is currently displayed.
-    FFT_VIEW = "fft"
+    FFT_VIEW = "frequency"
     TIME_VIEW = "time"
 
     displayed_measurement_changed = pyqtSignal(Measurement)
diff --git a/src/nqrduck_measurement/signalprocessing_options.py b/src/nqrduck_measurement/signalprocessing_options.py
index 8f7dad4..4768e5c 100644
--- a/src/nqrduck_measurement/signalprocessing_options.py
+++ b/src/nqrduck_measurement/signalprocessing_options.py
@@ -2,7 +2,7 @@
 
 import logging
 import sympy
-from nqrduck_spectrometer.measurement import Measurement, Fit, T2StarFit
+from nqrduck_spectrometer.measurement import Measurement, Fit, T2StarFit, LorentzianFit
 from nqrduck.helpers.functions import Function, GaussianFunction, CustomFunction
 from nqrduck.helpers.formbuilder import (
     DuckFormBuilder,
@@ -27,22 +27,6 @@ class FIDFunction(Function):
 
         self.add_parameter(Function.Parameter("T2star (microseconds)", "T2star", 10))
 
-
-class LorentzianFunction(Function):
-    """The Lorentzian function."""
-
-    name = "Lorentzian"
-
-    def __init__(self) -> None:
-        """Lorentzian function."""
-        expr = sympy.sympify("1 / (1 + (x / T2star)^2)")
-        super().__init__(expr)
-        self.start_x = 0
-        self.end_x = 30
-
-        self.add_parameter(Function.Parameter("T2star (microseconds)", "T2star", 10))
-
-
 class Apodization(DuckFormBuilder):
     """Apodization parameter.
 
@@ -100,6 +84,7 @@ class Fitting(DuckFormBuilder):
 
         fits = {}
         fits["T2*"] = T2StarFit(self.measurement)
+        fits["Lorentzian"] = LorentzianFit(self.measurement)
 
         selection_field = DuckFormDropdownField(
             text=None,
diff --git a/src/nqrduck_measurement/view.py b/src/nqrduck_measurement/view.py
index 237b3ab..4702019 100644
--- a/src/nqrduck_measurement/view.py
+++ b/src/nqrduck_measurement/view.py
@@ -251,10 +251,16 @@ class MeasurementView(ModuleView):
             return
 
         for fit in measurement.fits:
-            logger.debug(f"Plotting fit {fit.name}.")
             if fit.domain == self.module.model.view_mode:
+                logger.debug(f"Plotting {fit.name} fit in domain {fit.domain}.")
                 x = fit.x
                 y = fit.y
+                # Shift the x values if the view mode is FFT
+                if fit.domain == self.module.model.FFT_VIEW:
+                    x = x + float(
+                        measurement.target_frequency - measurement.IF_frequency
+                    ) * 1e-6
+
                 self._ui_form.plotter.canvas.ax.plot(
                     x, y, label=f"{fit.name} Fit", linestyle="--"
                 )