From f2943135f11d6f6904df231f503c9260b597d056 Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Thu, 18 Aug 2022 15:59:47 +0200
Subject: [PATCH] Add new methods to the WakeField class
WakeField.drop allows to delete a component or a list of component from the WakeField.
WakeField.save and WakeField.load add a shortcut to pickle WakeField objects.
---
mbtrack2/impedance/impedance_model.py | 6 ++
mbtrack2/impedance/wakefield.py | 80 +++++++++++++++++++++++++++
2 files changed, 86 insertions(+)
diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index f4a557e..2985cb7 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -531,6 +531,9 @@ class ImpedanceModel(Element):
def save(self, file):
"""
Save impedance model to file.
+
+ The same pandas version is needed on both saving and loading computer
+ for the pickle to work.
Parameters
----------
@@ -550,6 +553,9 @@ class ImpedanceModel(Element):
def load(self, file):
"""
Load impedance model from file.
+
+ The same pandas version is needed on both saving and loading computer
+ for the pickle to work.
Parameters
----------
diff --git a/mbtrack2/impedance/wakefield.py b/mbtrack2/impedance/wakefield.py
index f21b317..5b6f92a 100644
--- a/mbtrack2/impedance/wakefield.py
+++ b/mbtrack2/impedance/wakefield.py
@@ -5,6 +5,7 @@ wake functions.
"""
import warnings
+import pickle
import pandas as pd
import numpy as np
import scipy as sc
@@ -667,6 +668,12 @@ class WakeField:
Add two WakeFields taking into account beta functions.
add_several_wakefields(wakefields, beta)
Add a list of WakeFields taking into account beta functions.
+ drop(component)
+ Delete a component or a list of component from the WakeField.
+ save(file)
+ Save WakeField element to file.
+ load(file)
+ Load WakeField element from file.
"""
def __init__(self, structure_list=None, name=None):
@@ -736,6 +743,79 @@ class WakeField:
"Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad"]
return np.array([comp for comp in dir(self) if comp in valid])
+ def drop(self, component):
+ """
+ Delete a component or a list of component from the WakeField.
+
+ Parameters
+ ----------
+ component : str or list of str
+ Component or list of components to drop.
+ If "Z", drop all impedance components.
+ If "W"", drop all wake function components.
+
+ Returns
+ -------
+ None.
+
+ """
+ if component == "Z":
+ component = self.impedance_components
+ elif component == "W":
+ component = self.wake_components
+
+ if isinstance(component, str):
+ delattr(self, component)
+ elif isinstance(component, list) or isinstance(component, np.ndarray):
+ for comp in component:
+ delattr(self, comp)
+ else:
+ raise TypeError("component should be a str or list of str.")
+
+ def save(self, file):
+ """
+ Save WakeField element to file.
+
+ The same pandas version is needed on both saving and loading computer
+ for the pickle to work.
+
+ Parameters
+ ----------
+ file : str
+ File where the WakeField element is saved.
+
+ Returns
+ -------
+ None.
+
+ """
+ with open(file,"wb") as f:
+ pickle.dump(self, f)
+
+ @staticmethod
+ def load(file):
+ """
+ Load WakeField element from file.
+
+ The same pandas version is needed on both saving and loading computer
+ for the pickle to work.
+
+ Parameters
+ ----------
+ file : str
+ File where the WakeField element is saved.
+
+ Returns
+ -------
+ wakefield : WakeField
+ Loaded wakefield.
+
+ """
+ with open(file, 'rb') as f:
+ wakefield = pickle.load(f)
+
+ return wakefield
+
@staticmethod
def add_wakefields(wake1, beta1, wake2, beta2):
"""
--
GitLab