Add sweep and coherent monitors of higher order

mbtrack2/tracking/__init__.py

@@ -31,3 +31,4 @@ from mbtrack2.tracking.wakepotential import (
     LongRangeResistiveWall,
     WakePotential,
 )
+from mbtrack2.tracking.sweep import Sweep
\ No newline at end of file

mbtrack2/tracking/monitors/monitors.py

@@ -1083,12 +1083,18 @@ class BunchSpectrumMonitor(Monitor):
         dict_buffer = {
             "incoherent": (3, self.n_fft // 2 + 1, buffer_size),
             "coherent": (3, self.n_fft // 2 + 1, buffer_size),
+            "coherent_q":(3, self.n_fft//2+1, buffer_size),
+            "coherent_s":(3, self.n_fft//2+1, buffer_size),
+            "coherent_o":(3, self.n_fft//2+1, buffer_size),
             "mean_incoherent": (3, buffer_size),
             "std_incoherent": (3, buffer_size)
         }
         dict_file = {
             "incoherent": (3, self.n_fft // 2 + 1, total_size),
             "coherent": (3, self.n_fft // 2 + 1, total_size),
+            "coherent_q":(3, self.n_fft//2+1, total_size),
+            "coherent_s":(3, self.n_fft//2+1, total_size),
+            "coherent_o":(3, self.n_fft//2+1, total_size),
             "mean_incoherent": (3, total_size),
             "std_incoherent": (3, total_size)
         }
@@ -1104,6 +1110,9 @@ class BunchSpectrumMonitor(Monitor):
         self.positions = np.zeros(
             (self.size_list, self.sample_size, self.save_every + 1))
         self.mean = np.zeros((self.size_list, self.save_every + 1))
+        self.std = np.zeros((self.size_list, self.save_every+1))
+        self.skew = np.zeros((self.size_list, self.save_every+1))
+        self.kurtosis = np.zeros((self.size_list, self.save_every+1))
 
         index = np.arange(0, int(mp_number))
         self.index_sample = sorted(random.sample(list(index),
@@ -1111,6 +1120,9 @@ class BunchSpectrumMonitor(Monitor):
 
         self.incoherent = np.zeros((3, self.n_fft // 2 + 1, self.buffer_size))
         self.coherent = np.zeros((3, self.n_fft // 2 + 1, self.buffer_size))
+        self.coherent_q = np.zeros((3, self.n_fft//2+1, self.buffer_size))
+        self.coherent_s = np.zeros((3, self.n_fft//2+1, self.buffer_size))
+        self.coherent_o = np.zeros((3, self.n_fft//2+1, self.buffer_size))
 
         self.file[self.group_name].create_dataset("freq",
                                                   data=self.frequency_samples)
@@ -1172,6 +1184,9 @@ class BunchSpectrumMonitor(Monitor):
                 self.positions[value, :, self.save_count] = np.nan
 
             self.mean[:, self.save_count] = bunch.mean[self.mean_index]
+            self.std[:, self.save_count] = bunch.std[self.mean_index]
+            self.skew[:, self.save_count] = bunch.skew[self.mean_index]
+            self.kurtosis[:, self.save_count] = bunch.kurtosis[self.mean_index]
 
             self.save_count += 1
 
@@ -1201,6 +1216,10 @@ class BunchSpectrumMonitor(Monitor):
             self.coherent[self.store_dict[key], :,
                           self.buffer_count] = self.get_coherent_spectrum(
                               self.mean[value])
+            self.coherent_q[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.std[value])
+            self.coherent_s[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.skew[value])
+            self.coherent_o[self.store_dict[key],:,self.buffer_count] = self.get_coherent_spectrum(self.kurtosis[value])
+
 
         self.buffer_count += 1
 
@@ -1236,6 +1255,15 @@ class BunchSpectrumMonitor(Monitor):
                                          self.buffer_size:(self.write_count +
                                                            1) *
                                          self.buffer_size] = self.coherent
+        self.file[self.group_name]["coherent_q"][:,:, 
+                self.write_count * self.buffer_size:(self.write_count+1) * 
+                self.buffer_size] = self.coherent_q
+        self.file[self.group_name]["coherent_s"][:,:, 
+                self.write_count * self.buffer_size:(self.write_count+1) * 
+                self.buffer_size] = self.coherent_s
+        self.file[self.group_name]["coherent_o"][:,:, 
+                self.write_count * self.buffer_size:(self.write_count+1) * 
+                self.buffer_size] = self.coherent_o
 
         self.file.flush()
         self.write_count += 1

mbtrack2/tracking/particles.py

@@ -9,7 +9,7 @@ import numpy as np
 import pandas as pd
 import seaborn as sns
 from scipy.constants import c, e, m_e, m_p
-
+from scipy.stats import moment
 
 class Particle:
     """
@@ -244,6 +244,24 @@ class Bunch:
         """
         std = [[self[name].std()] for name in self]
         return np.squeeze(np.array(std))
+    
+    @property
+    def skew(self):
+        """
+        Return the standard deviation of the position of alive 
+        particles for each coordinates.
+        """
+        skew = [[moment(self[name],3)] for name in self]
+        return np.squeeze(np.array(skew))
+    
+    @property
+    def kurtosis(self):
+        """
+        Return the standard deviation of the position of alive 
+        particles for each coordinates.
+        """
+        kurtosis = [[moment(self[name],4)] for name in self]
+        return np.squeeze(np.array(kurtosis))    
 
     @property
     def emit(self):

mbtrack2/tracking/sweep.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Mar 29 15:12:28 2024
+
+@author: gamelina
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+from mbtrack2.tracking.element import Element
+from scipy.signal import chirp
+
+class Sweep(Element):
+    def __init__(self, ring, f0, f1, t1, level):
+        self.ring = ring
+        self.t = np.arange(0, t1, ring.T0)
+        self.N = len(self.t)
+        self.count = 0
+        self.level = level
+        self.sweep = chirp(self.t, f0, t1, f1)
+
+    @Element.parallel
+    def track(self, bunch):
+        """
+        Tracking method for the element.
+        No bunch to bunch interaction, so written for Bunch objects and
+        @Element.parallel is used to handle Beam objects.
+        
+        Parameters
+        ----------
+        bunch : Bunch or Beam object
+        """
+        sweep_val = self.sweep[self.count]
+        bunch["delta"] += self.level / self.ring.E0 * sweep_val
+        self.count += 1
+        if self.count >= self.N:
+            self.count = 0
+
+    def plot(self):
+        fig, ax = plt.subplots()
+        ax.plot(self.t, self.sweep)
+        return fig
\ No newline at end of file