From d6f87adeaadb6bbfcf90b741c184b245aa27486d Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Fri, 12 Aug 2022 16:18:27 +0200
Subject: [PATCH] Add save and load function to ImpedanceModel
---
mbtrack2/impedance/impedance_model.py | 47 +++++++++++++++++++++++++++
1 file changed, 47 insertions(+)
diff --git a/mbtrack2/impedance/impedance_model.py b/mbtrack2/impedance/impedance_model.py
index 2bf0f9d..f4a557e 100644
--- a/mbtrack2/impedance/impedance_model.py
+++ b/mbtrack2/impedance/impedance_model.py
@@ -5,6 +5,7 @@ Module where the ImpedanceModel class is defined.
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
+import pickle
from scipy.integrate import trapz
from scipy.interpolate import interp1d
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
@@ -61,6 +62,10 @@ class ImpedanceModel(Element):
Sum all the elements with the same name in the model into sum_name.
plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None)
Plot the contributions of different kind of WakeFields.
+ save(file)
+ Save impedance model to file.
+ load(file)
+ Load impedance model from file.
"""
def __init__(self, ring, wakefield_list=None, wakefiled_positions=None):
@@ -522,3 +527,45 @@ class ImpedanceModel(Element):
ax.set_ylabel("Power loss [W]")
return pf0, power_loss
+
+ def save(self, file):
+ """
+ Save impedance model to file.
+
+ Parameters
+ ----------
+ file : str
+ File where the impedance model is saved.
+
+ Returns
+ -------
+ None.
+
+ """
+ to_save = {"wakefields":self.wakefields,
+ "positions":self.positions}
+ with open(file,"wb") as f:
+ pickle.dump(to_save, f)
+
+ def load(self, file):
+ """
+ Load impedance model from file.
+
+ Parameters
+ ----------
+ file : str
+ File where the impedance model is saved.
+
+ Returns
+ -------
+ None.
+
+ """
+
+ with open(file, 'rb') as f:
+ to_load = pickle.load(f)
+
+ self.wakefields = to_load["wakefields"]
+ self.positions = to_load["positions"]
+ self.sum_elements()
+ self.sum_by_name_all()
--
GitLab