Resonance sweep.py

"""
Automatic resonance sweep for RTB2004 + Voltcraft FG-1302.

Python commands the generator frequency (confirmed command: SOUR1:FREQ),
and reads the response of both scope channels at each step. Because the
frequency is COMMANDED (not guessed from the signal), it is known exactly,
so there is no FFT-resolution problem.

Setup:
  - Generator USB -> PC ; Scope USB -> PC (both confirmed).
  - Generator Out1 -> amplifier -> actuator piezo.
  - CH1 = piezo sensor ; CH2 = accelerometer (via 480E09).
  - SET THE DRIVE AMPLITUDE BY HAND on the generator BEFORE running, modest
    level (resonances will amplify it). Make sure nothing clips at resonance.

Scope settings:
  - FFT OFF. Vertical scales set so the LARGEST resonance response does not clip.
  - Timebase: a few cycles of the LOWEST sweep frequency should fit, but since
    we measure amplitude at a known frequency, exact timebase is not critical;
    ~10 ms/div is a reasonable general choice for 20-1500 Hz. Keep it fixed.

It sweeps START_HZ -> STOP_HZ in STEP_HZ steps, measuring both channels,
then saves Excel + plot (auto-numbered) and prints the peak (resonance) of each.
"""

import pyvisa
import numpy as np
import os, glob, re, time
import matplotlib
# Live plotting needs an interactive backend; only force Agg when not live.
import matplotlib.pyplot as plt

# ---------------- SETTINGS YOU MAY CHANGE ----------------
SCOPE_ADDR = "TCPIP::169.254.153.1::INSTR"
GEN_ADDR   = "USB0::0x5345::0x1235::25131357::INSTR"

START_HZ   = 200.0
STOP_HZ    = 400.0
STEP_HZ    = 10
SETTLE_S   = 0.4      # wait after setting each frequency, for plate to settle
N_AVG      = 2        # captures averaged per channel at each frequency
DRIVE_VPP  = 1.0   # generator amplitude in Vpp (BEFORE x30 amplifier)
CYCLES_ON_SCREEN = 6  # how many cycles to capture per frequency (sets timebase)
LIVE_PLOT  = True     # show a live-updating plot while sweeping (set False to disable)

CH1_LABEL  = "piezo_sensor"
CH2_LABEL  = "accelerometer"
# EXPECTED_F removed - no reference line drawn (testing a blade, not the plate)

SAVE_DIR   = r"C:\Users\thot_ni\Documents\Post processing\Excel data"
PLOT_DIR   = r"C:\Users\thot_ni\Documents\Post processing\plots"
BASE_NAME  = "auto sweep"
# ---------------------------------------------------------


def set_timebase_for(scope, f):
    """Set time/div so the screen (10 divisions) shows ~CYCLES_ON_SCREEN cycles
    of frequency f. This keeps the capture well-resolved at every frequency."""
    if f <= 0:
        return
    total_time = CYCLES_ON_SCREEN / f      # seconds to show that many cycles
    tdiv = total_time / 10.0               # 10 divisions across the screen
    # clamp to the scope's sensible range
    tdiv = max(1e-7, min(tdiv, 5.0))
    scope.write(f"TIM:SCAL {tdiv:.6e}")

def acquire_fresh(scope):
    """Force ONE fresh acquisition and wait for it to finish, so the data we
    then read is current (not a stale buffer)."""
    scope.write("SING")          # arm a single acquisition
    try:
        scope.query("*OPC?")     # wait until the acquisition is complete
    except Exception:
        time.sleep(0.3)          # fallback wait if *OPC? not supported

def read_channel(scope, ch):
    scope.write(f"CHAN{ch}:STAT ON")
    head = scope.query(f"CHAN{ch}:DATA:HEAD?").strip()
    x0, x1, n, _ = [float(v) for v in head.split(",")]
    n = int(n)
    dt = (x1 - x0) / (n - 1)
    scope.write("FORM ASC")
    raw = scope.query(f"CHAN{ch}:DATA?")
    v = np.array([float(x) for x in raw.strip().split(",")])
    return dt, v


def amp_at(volts, dt, target_f):
    """Response amplitude (zero-to-peak, volts).
    The plate is driven at one known frequency, so the captured AC waveform IS
    that frequency. We use a robust peak-to-peak of the AC signal and halve it.
    Using percentiles (99th - 1st) instead of absolute max/min rejects the odd
    noise spike, so it stays clean. This avoids any FFT-resolution issue at low
    frequencies."""
    v = volts - np.mean(volts)
    hi = np.percentile(v, 99.0)
    lo = np.percentile(v, 1.0)
    pp = hi - lo
    return pp / 2.0      # zero-to-peak amplitude


def next_paths():
    os.makedirs(SAVE_DIR, exist_ok=True)
    os.makedirs(PLOT_DIR, exist_ok=True)
    ex = glob.glob(os.path.join(SAVE_DIR, f"{BASE_NAME} *.xlsx"))
    nums = [int(m.group(1)) for p in ex
            if (m := re.search(re.escape(BASE_NAME) + r" (\d+)\.xlsx$", os.path.basename(p)))]
    run = (max(nums) + 1) if nums else 1
    return run, os.path.join(SAVE_DIR, f"{BASE_NAME} {run}.xlsx"), \
                os.path.join(PLOT_DIR, f"{BASE_NAME} {run}.png")


# ---- connect both instruments ----
rm = pyvisa.ResourceManager()
scope = rm.open_resource(SCOPE_ADDR); scope.timeout = 20000
gen   = rm.open_resource(GEN_ADDR);   gen.timeout = 5000
print("Scope:    ", scope.query("*IDN?").strip())
print("Generator:", gen.query("*IDN?").strip())

# --- prepare the generator: sine, amplitude, OUTPUT ON ---
gen.write("SOUR1:FUNC SIN")
gen.write(f"SOUR1:VOLT {DRIVE_VPP}")
gen.write("OUTP1 ON")
time.sleep(0.5)
try:
    out_state = gen.query("OUTP1?").strip()
    print(f"Generator output state: {out_state}  (1 = ON)")
    if out_state.strip() not in ("1", "ON"):
        print("WARNING: generator output is not ON - the plate will not be driven!")
except Exception:
    pass

run_number, xlsx_path, png_path = next_paths()
print(f"This sweep -> {BASE_NAME} {run_number}.xlsx")

freqs = np.arange(START_HZ, STOP_HZ + STEP_HZ, STEP_HZ)
print(f"Sweeping {START_HZ:.0f} -> {STOP_HZ:.0f} Hz in {STEP_HZ:.0f} Hz steps "
      f"({len(freqs)} points). Ctrl+C to abort.")
print("-" * 64)
print(f"{'set Hz':>8} {'got Hz':>9} {CH1_LABEL[:9]:>10} {CH2_LABEL[:9]:>11}")
print("-" * 64)

f_list, a1_list, a2_list = [], [], []

# --- set up live plot window if enabled ---
if LIVE_PLOT:
    plt.ion()                                  # interactive mode on
    live_fig, live_ax = plt.subplots(figsize=(9, 5))
    (line1,) = live_ax.plot([], [], "-", color="C0", lw=1.3, label=CH1_LABEL)
    (line2,) = live_ax.plot([], [], "-", color="C1", lw=1.3, label=CH2_LABEL)
    live_ax.set_xlabel("Drive frequency (Hz)")
    live_ax.set_ylabel("Response amplitude (mV)")
    live_ax.set_title("Resonance sweep (live)")
    live_ax.grid(True, alpha=0.3)
    live_ax.legend(loc="upper right")
    live_ax.set_xlim(START_HZ, STOP_HZ)
    plt.show(block=False)

try:
    for f_set in freqs:
        gen.write(f"SOUR1:FREQ {f_set:.3f}")
        time.sleep(SETTLE_S)
        # read back what the generator actually set
        try:
            f_got = float(gen.query("SOUR1:FREQ?").strip())
        except Exception:
            f_got = f_set

        set_timebase_for(scope, f_got)    # NEW: timebase suited to this frequency
        c1s, c2s = [], []
        for _ in range(N_AVG):
            acquire_fresh(scope)          # force a fresh capture each time
            dt1, v1 = read_channel(scope, 1)
            dt2, v2 = read_channel(scope, 2)
            c1s.append(amp_at(v1, dt1, f_got))
            c2s.append(amp_at(v2, dt2, f_got))
        a1 = float(np.mean(c1s)); a2 = float(np.mean(c2s))

        # debug: on the first 3 steps, show what we actually captured
        if len(f_list) < 3:
            print(f"   [debug] CH1 samples={len(v1)} pp={v1.max()-v1.min():.4f}V | "
                  f"CH2 samples={len(v2)} pp={v2.max()-v2.min():.4f}V")

        f_list.append(f_got); a1_list.append(a1); a2_list.append(a2)
        print(f"{f_set:8.1f} {f_got:9.2f} {a1*1000:10.3f} {a2*1000:11.3f}")

        # update the live plot as data comes in
        if LIVE_PLOT:
            line1.set_data(f_list, [x*1000 for x in a1_list])
            line2.set_data(f_list, [x*1000 for x in a2_list])
            live_ax.relim(); live_ax.autoscale_view(scaley=True)
            live_fig.canvas.draw_idle()
            live_fig.canvas.flush_events()

except KeyboardInterrupt:
    print("\n(aborted by user - saving what was collected so far)")
finally:
    try:
        gen.write("OUTP1 OFF")   # stop driving the plate when done
    except Exception:
        pass
    try:
        scope.write("RUN")       # put scope back to continuous running
    except Exception:
        pass
    gen.close(); scope.close()

if LIVE_PLOT:
    plt.ioff()
    try:
        plt.close(live_fig)
    except Exception:
        pass

if len(f_list) < 2:
    print("Not enough data collected."); raise SystemExit

f  = np.array(f_list)
a1 = np.array(a1_list) * 1000.0  # mV
a2 = np.array(a2_list) * 1000.0
f1_peak = f[np.argmax(a1)]
f2_peak = f[np.argmax(a2)]

def find_modes(freq, amp, label):
    """Find ALL resonance peaks in a response curve, not just the tallest.
    Returns a list of (frequency, amplitude) for each detected mode.
    A peak must rise clearly above its surroundings (prominence) and be a local
    maximum, so noise wiggles and clipped plateaus are not counted as modes."""
    try:
        from scipy.signal import find_peaks
        # prominence = how far a peak stands out above the local baseline.
        # 0.15 of the data range is strict enough to skip small shoulder bumps
        # on the sides of big peaks, while still catching genuine smaller modes.
        prom = 0.08 * (np.max(amp) - np.min(amp))
        idx, props = find_peaks(amp, prominence=prom, distance=3)
        cand = [(float(freq[i]), float(amp[i])) for i in idx]
        # drop "shoulder" peaks: a small peak sitting very close (<40 Hz) to a
        # much taller peak (less than 1/3 its height) is almost certainly a bump
        # on that bigger peak's slope, not a separate mode.
        peaks = []
        for fp, ap in cand:
            is_shoulder = False
            for fp2, ap2 in cand:
                if fp2 != fp and abs(fp - fp2) < 60 and ap < 0.5 * ap2:
                    is_shoulder = True
                    break
            if not is_shoulder:
                peaks.append((fp, ap))
    except Exception:
        # fallback if scipy missing: simple local-maximum check
        peaks = []
        for i in range(1, len(amp) - 1):
            if amp[i] > amp[i-1] and amp[i] > amp[i+1]:
                if amp[i] > np.min(amp) + 0.10 * (np.max(amp) - np.min(amp)):
                    peaks.append((float(freq[i]), float(amp[i])))
    # sort by frequency
    peaks.sort(key=lambda p: p[0])
    print(f"\nModes detected on {label}:")
    if not peaks:
        print("  (none clearly detected)")
    for fp, ap in peaks:
        print(f"  {fp:8.1f} Hz   ({ap:.2f} mV)")
    return peaks

modes_ch1 = find_modes(f, a1, CH1_LABEL)
modes_ch2 = find_modes(f, a2, CH2_LABEL)

# --- save Excel ---
from openpyxl import Workbook
wb = Workbook(); ws = wb.active; ws.title = "auto_sweep"
ws.append(["frequency_Hz", f"{CH1_LABEL}_mV", f"{CH2_LABEL}_mV"])
for fi, b1, b2 in zip(f, a1, a2):
    ws.append([round(float(fi),3), round(float(b1),5), round(float(b2),5)])
ws.append([]); ws.append(["peak_"+CH1_LABEL+"_Hz", round(float(f1_peak),3)])
ws.append(["peak_"+CH2_LABEL+"_Hz", round(float(f2_peak),3)])
ws.append([]); ws.append([f"--- detected modes ({CH1_LABEL}) ---"])
for fp, ap in modes_ch1:
    ws.append([round(fp,2), round(ap,4)])
ws.append([]); ws.append([f"--- detected modes ({CH2_LABEL}) ---"])
for fp, ap in modes_ch2:
    ws.append([round(fp,2), round(ap,4)])
for i in range(1,4): ws.column_dimensions[chr(64+i)].width = 22
wb.save(xlsx_path)

# --- plot ---
fig, ax = plt.subplots(figsize=(9,5))
ax.plot(f, a1, "-", color="C0", lw=1.3, label=f"{CH1_LABEL} (peak {f1_peak:.1f} Hz)")
ax.plot(f, a2, "-", color="C1", lw=1.3, label=f"{CH2_LABEL} (peak {f2_peak:.1f} Hz)")
# Clean annotation: a thin dashed vertical line at each detected mode, and a
# single tidy text box listing all the mode frequencies in a corner. This is
# far more readable than scattering labels on top of the peaks.
for fp, ap in modes_ch1:
    ax.axvline(fp, color="C0", ls=":", lw=1.0, alpha=0.6)

if modes_ch1:
    mode_text = "Detected modes (Hz):\n" + "\n".join(f"  {fp:.0f}" for fp, _ in modes_ch1)
    ax.text(0.985, 0.97, mode_text, transform=ax.transAxes,
            ha="right", va="top", fontsize=9, color="C0",
            bbox=dict(boxstyle="round", facecolor="white", edgecolor="C0", alpha=0.9))

ax.set_xlabel("Drive frequency (Hz)")
ax.set_ylabel("Response amplitude (mV)")
ax.set_title("Automatic resonance sweep")
ax.grid(True, alpha=0.3); ax.legend()
fig.tight_layout(); fig.savefig(png_path, dpi=200)

print("-" * 64)
print(f"Data: {xlsx_path}")
print(f"Plot: {png_path}")
print(f"Peak {CH1_LABEL}: {f1_peak:.2f} Hz | Peak {CH2_LABEL}: {f2_peak:.2f} Hz")