From 01cd4c32dc1fec3d94eeac32c6bcc834dda54eef Mon Sep 17 00:00:00 2001
From: LemurPwned <jakub.mojsiejuk@gmail.com>
Date: Sun, 22 Dec 2024 21:18:57 +0100
Subject: [PATCH] adding a back scan

---
 view/helpers.py        | 21 +++++++++++++++++----
 view/simulation_fns.py |  9 +++++++--
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/view/helpers.py b/view/helpers.py
index 4f7d135..47027a2 100644
--- a/view/helpers.py
+++ b/view/helpers.py
@@ -51,7 +51,20 @@ def plot_optim(hvals, target, simulation, title="Optimisation"):
         st.pyplot(fig)
 
 
-def plot_data(Hscan, freqs, spec, mag=None, title="Resonance spectrum"):
+def plot_data(
+    Hscan, freqs, spec, mag=None, title: str = "Resonance spectrum", mode: str = "none"
+):
+    Hscan_spec = Hscan
+    if st.session_state["Hreturn"] and mode.lower() == "pimm":
+        print("Adding spectrum")
+        Hlen = len(Hscan) // 2
+        Hscan_spec = Hscan[:Hlen]
+        spec_left_right = spec[:Hlen, :]
+        spec_right_left = spec[Hlen:, :]
+        # invert the right side
+        spec_right_left = spec_right_left[::-1, :]
+        spec = spec_left_right + spec_right_left
+
     with plt.style.context(["dark_background"]):
         if mag is not None:
             fig = plt.figure(dpi=300)
@@ -77,7 +90,7 @@ def plot_data(Hscan, freqs, spec, mag=None, title="Resonance spectrum"):
         else:
             fig, ax1 = plt.subplots(dpi=300)
         ax1.pcolormesh(
-            Hscan / 1e3,
+            Hscan_spec / 1e3,
             freqs / 1e9,
             10 * np.log10(np.abs(spec.T)),
             shading="auto",
@@ -123,7 +136,7 @@ def simulate_vsd():
             Hoex_mag=st.session_state.Hoex_mag,
         )
         spec = Filters.detrend_axis(spec, axis=0)
-    plot_data(Hscan, freqs, spec, title="VSD spectrum")
+    plot_data(Hscan, freqs, spec, title="VSD spectrum", mode='vsd')
 
 
 def simulate_pimm():
@@ -137,4 +150,4 @@ def simulate_pimm():
             sim_time=st.session_state.sim_time * 1e-9,
         )
     mag = np.asarray(output["m_avg"])
-    plot_data(Hscan, freqs, spec, mag=mag, title="PIMM spectrum")
+    plot_data(Hscan, freqs, spec, mag=mag, title="PIMM spectrum", mode='pimm')
diff --git a/view/simulation_fns.py b/view/simulation_fns.py
index ff162da..3f931f5 100644
--- a/view/simulation_fns.py
+++ b/view/simulation_fns.py
@@ -143,9 +143,14 @@ def get_pimm_data(
     htheta, hphi = get_axis_angles(H_axis)
     hmin, hmax = min([Hmin, Hmax]), max([Hmin, Hmax])  # fix user input
     Hscan, Hvecs = FieldScan.amplitude_scan(
-        hmin, hmax, Hsteps, htheta, hphi, back=st.session_state["Hreturn"]
+        hmin, hmax, Hsteps, htheta, hphi
     )
+    if st.session_state["Hreturn"]:
+        Hscan = np.concatenate((Hscan, Hscan[::-1]))
+        Hvecs = np.concatenate((Hvecs, Hvecs[::-1]))
     j, rparams = prepare_simulation()
+    # avoid wait if the user sim's too short
+    wtime = 4e-9 if sim_time >= 6e-9 else 0.0
     spec, freqs, out = PIMM_procedure(
         j,
         Hvecs=Hvecs,
@@ -156,7 +161,7 @@ def get_pimm_data(
         max_frequency=st.session_state["max_freq"] * 1e9,
         simulation_duration=sim_time,
         disturbance=1e-6,
-        wait_time=4e-9,
+        wait_time=wtime,
     )
     return spec, freqs, out, Hscan