From 92be75554a1b511060dc06c59b63d2bb5208f057 Mon Sep 17 00:00:00 2001
From: Alexis Gamelin <alexis.gamelin@synchrotron-soleil.fr>
Date: Fri, 29 Mar 2024 15:25:19 +0100
Subject: [PATCH] Add sweep and coherent monitors of higher order
---
mbtrack2/tracking/__init__.py | 1 +
mbtrack2/tracking/monitors/monitors.py | 28 +++++++++++++++++
mbtrack2/tracking/particles.py | 20 +++++++++++-
mbtrack2/tracking/sweep.py | 42 ++++++++++++++++++++++++++
4 files changed, 90 insertions(+), 1 deletion(-)
create mode 100644 mbtrack2/tracking/sweep.py
diff --git a/mbtrack2/tracking/__init__.py b/mbtrack2/tracking/__init__.py
index 02843e2..a4ea3ff 100644
--- a/mbtrack2/tracking/__init__.py
+++ b/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
diff --git a/mbtrack2/tracking/monitors/monitors.py b/mbtrack2/tracking/monitors/monitors.py
index 151fe2a..f5abcbb 100644
--- a/mbtrack2/tracking/monitors/monitors.py
+++ b/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
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index 7ba97ff..bff2263 100644
--- a/mbtrack2/tracking/particles.py
+++ b/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):
diff --git a/mbtrack2/tracking/sweep.py b/mbtrack2/tracking/sweep.py
new file mode 100644
index 0000000..61d0c3d
--- /dev/null
+++ b/mbtrack2/tracking/sweep.py
@@ -0,0 +1,42 @@
+#!/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
--
GitLab