add test example for damper 0

tests/test_damper0.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+from collections import deque
+
+import matplotlib.pyplot as plt
+import numpy as np
+from tqdm import tqdm
+
+# Create a CavityResonator object with chosen parameters.
+from mbtrack2 import (
+    Beam,
+    BeamMonitor,
+    BunchSpectrumMonitor,
+    ProportionalIntegralIQLoop,
+    ProportionalIntegralIQLoopMode0Damper,
+    ProportionalIntegralLoop,
+    plot_beamdata,
+)
+from mbtrack2.tracking import (
+    Bunch,
+    CavityResonator,
+    Electron,
+    LongitudinalMap,
+    RFCavity,
+    Synchrotron,
+    SynchrotronRadiation,
+    TransverseMap,
+)
+from mbtrack2.utilities import Optics
+
+np.random.seed(1)
+
+# Define ring parameters
+h = 20  # Harmonic number
+L = 100  # Circumference [m]
+E0 = 1.5e9  # Nominal energy [eV]
+particle = Electron()
+ac = 1e-3  # Momentum compaction factor
+U0 = 200e3  # Energy loss per turn [eV]
+tau = np.array([1e-3, 1e-3,
+                2e-3])  # Horizontal, vertical, longitudinal damping times [s]
+tune = np.array([12.2, 15.3])  # Horizontal and vertical tunes
+emit = np.array([10e-9,
+                 10e-12])  # Horizontal and vertical equilibrium emittance
+sigma_0 = 15e-12  # Natural bunch length [s]
+sigma_delta = 1e-3  # Energy spread
+chro = [2.0, 3.0]  # Chromaticities
+
+# Local optics at tracking point
+local_beta = np.array([3, 2])
+local_alpha = np.array([0, 0])
+local_dispersion = np.array([0, 0, 0, 0])
+optics = Optics(local_beta=local_beta,
+                local_alpha=local_alpha,
+                local_dispersion=local_dispersion)
+
+# Create Synchrotron object
+ring = Synchrotron(h=h,
+                   optics=optics,
+                   particle=particle,
+                   L=L,
+                   E0=E0,
+                   ac=ac,
+                   U0=U0,
+                   tau=tau,
+                   emit=emit,
+                   tune=tune,
+                   sigma_delta=sigma_delta,
+                   sigma_0=sigma_0,
+                   chro=chro)
+
+# Define tracking elements
+LongMap = LongitudinalMap(ring)
+RF = RFCavity(ring, m=1, Vc=1e6, theta=np.arccos(ring.U0 / 1e6))
+SR = SynchrotronRadiation(ring)
+TransMap = TransverseMap(ring)
+
+m = 1
+Rs = 1e6  # Shunt impedance [Ohm]
+Q = 10000
+QL = 5000
+detune = 100  # Detuning in Hz
+Ncav = 1
+Vc = 1e6  # Target cavity voltage [V]
+theta = np.arccos(ring.U0 / Vc)  # Total cavity phase in [rad].
+n_bin = 75
+
+cavity = CavityResonator(ring,
+                         m=m,
+                         Rs=Rs,
+                         Q=Q,
+                         QL=QL,
+                         detune=detune,
+                         Ncav=Ncav,
+                         Vc=Vc,
+                         theta=theta,
+                         n_bin=n_bin)
+
+PIMode0 = ProportionalIntegralIQLoopMode0Damper(ring,
+                                                cavity,
+                                                gain=[0.5, 1e4],
+                                                sample_num=5,
+                                                every=5,
+                                                delay=10,
+                                                IIR_cutoff=0,
+                                                FF=True,
+                                                enable_damper=True,
+                                                damper_gain=[0.5, 100])
+cavity.feedback.append(PIMode0)
+
+# Initialize a Beam using the filling_pattern method.
+# In mbtrack2, Beam is created with a filling pattern of length ring.h.
+# Here we create a filling pattern with True for the first 3 buckets and False for the rest.
+filling_pattern = np.zeros(ring.h, dtype=bool)
+filling_pattern[:] = True  # Only the first 3 buckets have bunches
+
+# Create a Beam object and initialize it with the filling pattern
+beam = Beam(ring)
+beam.init_beam(filling_pattern=filling_pattern,
+               current_per_bunch=1e-3,
+               mp_per_bunch=100)
+
+# Apply an initial longitudinal offset ("tau") to each non-empty bunch to generate a mode0 signal.
+# Here we use a different offset for each bunch.
+for i, bunch in enumerate(beam.not_empty):
+    bunch["delta"] += 1e-9 * (i+1)
+
+# filling cavity to accelerate simulation
+cavity.init_phasor(beam)
+print(cavity.cavity_voltage)
+
+n_turns = 8000  # turns
+cavity_phase_history = []  # record the phase of RF signal in cavity
+cavity_voltage_history = []  # record the voltage of RF signal in cavity
+beam_phase_history = []  # record mean beam phase（tau）
+beam_delta_history = []  # record mean beam  energy (delta)
+beam_phasor_amp_history = []  # record magnitude of beam phasor
+beam_phasor_phase_history = []  # record phase of beam phasor
+
+for turn in range(n_turns):
+    LongMap.track(beam)
+    TransMap.track(beam)
+    SR.track(beam)
+
+    if turn == 100:
+        print("setpoint voltage", cavity.Vc, "phase:", cavity.theta)
+    # before cavity.track(beam)，save beam to cavity for feedback
+    cavity._last_beam = beam
+    cavity.track(beam)
+
+    # Record cavity parameters for analysis
+    cavity_phase_history.append(cavity.theta)
+    cavity_voltage_history.append(cavity.cavity_voltage)
+    '''
+    # Record cavity beam phasor amplitude and phase
+    beam_phasor_amp_history.append(np.abs(cavity.cavity_phasor_record))
+    beam_phasor_phase_history.append(np.angle(cavity.cavity_phasor_record))
+    '''
+    # Compute average beam phase (tau) and average delta over non-empty bunches
+    not_empty = list(beam.not_empty)
+    if len(not_empty) > 0:
+        avg_beam_tau = np.mean([b.mean[4] for b in not_empty])
+        avg_beam_delta = np.mean([b["delta"] for b in not_empty])
+    else:
+        avg_beam_tau = 0
+        avg_beam_delta = 0
+    beam_phase_history.append(avg_beam_tau)
+    beam_delta_history.append(avg_beam_delta)
+
+# cavity magnitude and phase
+turns_array = np.arange(n_turns)
+
+plt.figure()
+plt.plot(turns_array, cavity_phase_history, label="Cavity Phase (rad)")
+plt.xlabel("Turn")
+plt.ylabel("Cavity Phase (rad)")
+plt.title("Cavity Phase Evolution")
+plt.legend()
+plt.show()
+
+plt.figure()
+plt.plot(turns_array, cavity_voltage_history, label="Cavity Voltage (V)")
+plt.xlabel("Turn")
+plt.ylabel("Cavity Voltage (V)")
+plt.title("Cavity Voltage Evolution")
+plt.legend()
+plt.show()
+'''
+# %% [code]
+# magnitude and phase of beam phasor
+plt.figure()
+plt.plot(turns_array, beam_phasor_amp_history, label="Beam Phasor Amplitude (V)")
+plt.xlabel("Turn")
+plt.ylabel("Amplitude (V)")
+plt.title("Beam Phasor Amplitude Evolution")
+plt.legend()
+plt.show()
+
+plt.figure()
+plt.plot(turns_array, beam_phasor_phase_history, label="Beam Phasor Phase (rad)")
+plt.xlabel("Turn")
+plt.ylabel("Phase (rad)")
+plt.title("Beam Phasor Phase Evolution")
+plt.legend()
+plt.show()
+'''
+# %% [code]
+# mean tau and delta of Beam
+plt.figure()
+plt.plot(turns_array, beam_phase_history, label="Average Beam Phase (tau, s)")
+plt.xlabel("Turn")
+plt.ylabel("Average Beam Phase (s)")
+plt.title("Beam Phase Evolution")
+plt.legend()
+plt.show()
+
+plt.figure()
+plt.plot(turns_array, beam_delta_history, label="Average Beam Delta")
+plt.xlabel("Turn")
+plt.ylabel("Average Beam Delta")
+plt.title("Beam Delta Evolution")
+plt.legend()
+plt.show()