improved examples

examples/demo_comparison.py

@@ -0,0 +1,44 @@
+import numpy as np
+
+from sylt.tracking import Bunch, Ring, Tracker
+
+from properties import ring_properties, bunch_properties
+
+ring = Ring(**ring_properties)
+
+NUM_TURNS = 5_000       # number of tracked turns
+N = 200e10              # bunch intensity
+n = 30_000              # number of macro particles
+ng = n - 1_000          # number of ghost particles
+
+# Define Trackers
+trackers = {
+    'SPM': Tracker(Bunch(**bunch_properties, n=n), ring),
+    'SPM+SC': Tracker(Bunch(**bunch_properties, n=n, N=N), ring),
+    'SPM+SC+TM': Tracker(Bunch(**bunch_properties, n=n, N=N, eps=None), ring)
+}
+
+if __name__ == "__main__":
+
+    for key, tracker in trackers.items():
+
+        # Define "Ghost" Particles
+        tracker.UPDATE = False
+        tracker.bunch.w[:ng] = np.zeros(ng)
+        tracker.bunch.tau[:ng] = np.random.choice(
+            np.linspace(-1, 1, 300)*tracker.tau_hat, ng)
+
+        # Track Particles
+        data = np.empty((2, NUM_TURNS, n))
+        for turn in range(NUM_TURNS):
+            tracker.track()
+            data[:, turn] = tracker.bunch.tau, tracker.bunch.w
+
+        # Clean Distribution
+        data[:, :, tracker.clean(tracker.bunch.tau, tracker.bunch.w)] = np.nan
+
+        # Export Data
+        np.savez(f"./data/{key}", tau=data[0], w=data[1],
+                 N=tracker.bunch.N, eps=tracker.bunch.eps,
+                 **ring_properties, **bunch_properties, ng=ng,
+                 )

examples/demo_spm.py

@@ -0,0 +1,29 @@
+import matplotlib.pyplot as plt
+
+from sylt.tracking import Bunch, Ring, Tracker
+
+from properties import ring_properties, bunch_properties
+
+ring = Ring(**ring_properties)
+
+bunch = Bunch(
+    **bunch_properties,
+    n=50_000,   # number of macroparticles
+    # N=200e10,   # include intensity effects
+    eps=None,   # initialize transverse distribution
+)
+
+tracker = Tracker(bunch, ring)
+
+tracker.show('start')
+
+tracker.match(k=0.5)
+tracker.show('mismatched')
+
+for turn in range(1_000):
+    tracker.track()
+
+tracker.clean()
+tracker.show('stop')
+
+plt.savefig('./figs/demo.png')

examples/demo_tune.py → examples/demo_tune_blur.py

@@ -1,41 +1,30 @@
 import matplotlib.pyplot as plt
 import numpy as np
-from sylt.analysis import compute_synchrotron_frequency
 
 from sylt.tracking import Bunch, Ring, Tracker
-from sylt.plotting import plot_tune
 
-ring = Ring(
-    R=100,          # machine radius (m)
-    h=8,            # rf harmonic
-    Vg=80e3,        # gap voltage (V)
-    gamma_t=6.1,    # transition
+from sylt.analysis import compute_synchrotron_frequency
+from sylt.plotting import plot_tune
 
-    beta=17,
-    D=(2.66, 0),
-    b=(73e-3, 35e-3),
-)
+from properties import ring_properties, bunch_properties
 
-E = 2e9+938.272e6
-n = 500
-sig_w = 0e6
-sig_tau = 30e-9
+ring = Ring(**ring_properties)
 
-N = 400e10
-sig_eps = 1e-6
+N = 400e10      # bunch intensity
+n = 500         # number of macroparticles
 
 trackers = {
-    'SPM': Tracker(Bunch(E, sig_tau, sig_w, n), ring),
-    'SPM+SC': Tracker(Bunch(E, sig_tau, sig_w, n, N, sig_eps), ring),
-    'SPM+SC+TM': Tracker(Bunch(E, sig_tau, sig_w, n, N, sig_eps, eps=None), ring)
+    'SPM': Tracker(Bunch(**bunch_properties, n=n), ring),
+    'SPM+SC': Tracker(Bunch(**bunch_properties, n=n, N=N), ring),
+    'SPM+SC+TM': Tracker(Bunch(**bunch_properties, n=n, N=N, eps=None), ring)
 }
 
 DATA = {}
 for i, (key, tracker) in enumerate(trackers.items()):
     tracker.UPDATE = False
 
     tau = []
-    for turn in range(5_000):
+    for turn in range(500_000):
         tracker.track()
         tau.append(tracker.bunch.tau)
 
@@ -48,6 +37,9 @@
 
 plot_tune(DATA)
 
+for ext in ['png', 'pdf', 'svg']:
+    plt.savefig(f'./figs/tune.{ext}')
+
 plt.show()
 
 print('Finished !')

examples/plot_comparison.py

@@ -0,0 +1,68 @@
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+
+import numpy as np
+from glob import glob
+
+import os
+
+from demo_comparison import ng, trackers
+
+files = sorted(glob('./data/*.npz'))[::-1]
+
+# Load Data into Memory
+data = {}
+for file in files:
+    path, ext = os.path.splitext(file)
+    key = os.path.basename(path)
+    tmp = np.load(file)
+    data[key] = {'tau': tmp['tau']*1e9, 'w': tmp['w']*1e-6}
+
+NUM_TURNS, n = data['SPM']['tau'].shape
+
+fig, ax = plt.subplots(num='comparison', constrained_layout=True)
+
+artists = []
+
+
+def init(turn=0):
+
+    for i, (key, tracker) in enumerate(trackers.items()):
+        ax.plot([], [], f"C{i}", label=key)
+
+        tau = data[key]['tau'][turn]
+        w = data[key]['w'][turn]
+
+        artists.append([
+            ax.plot(tau[:ng], w[:ng], f'C{i}.', zorder=5-i)[0],
+            ax.plot(tau[ng:], w[ng:], f'C{i},', zorder=2-i)[0],
+        ])
+
+        ts, ws = tracker.separatrix()
+        ax.plot(ts*1e9, +ws*1e-6, 'm-')
+        ax.plot(ts*1e9, -ws*1e-6, 'm-')
+
+    ax.set_xlabel(r"$\tau$ (ns)")
+    ax.set_ylabel(r"$w$ (MeV)")
+    ax.legend()
+
+
+init(2_500)
+
+fig.savefig('./figs/comparison.png')
+
+
+def update(turn):
+    for key, (p1, p2) in zip(trackers.keys(), artists):
+        tau = data[key]['tau'][turn]
+        w = data[key]['w'][turn]
+        p1.set_data(tau[0:ng], w[0:ng])
+        p2.set_data(tau[ng:], w[ng:])
+    return sum(artists, [])
+
+
+ani = FuncAnimation(fig, update, frames=range(
+    0, NUM_TURNS, 20), blit=True, interval=0)
+
+plt.show()
+ani.save('./figs/comparison.gif', fps=60)

examples/properties.py

@@ -0,0 +1,17 @@
+ring_properties = {
+    'R': 100,              # machine radius (m)
+    'h': 8,                # rf harmonic
+    'Vg': 80e3,            # gap voltage (V)
+    'gamma_t': 6.1,        # transition
+
+    'beta': 17,            # nominal beta function
+    'D': (2.66, 0),        # nominal dispersion function
+    'b': (73e-3, 35e-3),   # aperture radius (m)
+}
+
+bunch_properties = {
+    'E': 2.938272e9,        # beam energy (eV)
+    'sig_w': 5e6,           # energy spread (eV)
+    'sig_tau': 30e-9,       # bunch length (s)
+    'sig_eps': 1e-6,        # transverse emittance (m)
+}

sylt/tracking.py

@@ -11,18 +11,15 @@
 Z_0 = value('characteristic impedance of vacuum')
 
 
-
-
-
 @dataclass
 class Bunch:
-    E: float    # particle energy
-    sig_tau: float  # width relative time
+    E: float            # particle energy
+    sig_tau: float      # width relative time
     sig_w: float = 0    # width relative energy
-    n: int = 10_000      # macroparticle number
+    n: int = 10_000     # macroparticle number
 
-    N: float = 0    # bunch intensity
-    sig_eps: float = None  # width emittance
+    N: float = 0        # bunch intensity
+    sig_eps: float = np.nan  # width emittance
 
     eps: float = 0       # assume particle emittance is zero
 
@@ -161,6 +158,9 @@ def __post_init__(self):
         self.kappa = self.eta / (self.bunch.beta**2*self.bunch.E)
         self.nu = self.Omega/self.omega
 
+        self.tau_hat = self.T/self.ring.h/2
+        # self.tau_hat = (np.pi-self.ring.vphi_s)/(self.ring.h*self.omega)
+
     def H(self, tau, w):
         """return particle hamiltonion"""
         phi = self.ring.h*self.omega*tau
@@ -206,11 +206,11 @@ def track(self):
 
     def clean(self, tau=None, w=None):
         """assign NaN to particle's who's longitudinal position is external to the separatrix"""
-        tau = self.tau if tau is None else tau
-        w = self.w if w is None else w
-        
-        tau_hat = (np.pi-self.ring.vphi_s)/(self.ring.h*self.omega)
-        lost = (self.H(tau_hat, 0) - self.H(tau, w)) > 0
+        tau = self.bunch.tau if tau is None else tau
+        w = self.bunch.w if w is None else w
+
+        lost = ((self.H(self.tau_hat, 0) - self.H(tau, w)) > 0) + \
+            (np.abs(tau) > self.tau_hat)
 
         self.bunch.tau[lost] = np.nan
         self.bunch.w[lost] = np.nan
@@ -290,13 +290,13 @@ def match(self, sig_w=None, sig_tau=None, k=0):
 
         if sig_tau is not None:
             sig_phi = ring.h*self.omega*sig_tau
-            sig_w = np.sqrt(bunch.q*ring.Vg/(np.abs(self.kappa)*np.pi*ring.h)
-                            * (1-np.cos(sig_phi)))
+            bunch.sig_w = np.sqrt(bunch.q*ring.Vg/(np.abs(self.kappa)*np.pi*ring.h)
+                                  * (1-np.cos(sig_phi)))
         if sig_w is not None:
             pass
 
-        bunch.sig_tau = sig_tau*(1+k)
-        bunch.sig_w = sig_w*(1-k)
+        bunch.sig_tau = bunch.sig_tau*(1+k)
+        bunch.sig_w = bunch.sig_w*(1-k)
         bunch.__post_init__()
 
         text = [