From de2c017f0a40e251cba94c98640eb730b4e31471 Mon Sep 17 00:00:00 2001
From: Yngve Levinsen <yngve.levinsen@ess.eu>
Date: Tue, 10 Nov 2020 19:45:58 +0100
Subject: [PATCH] added first iteration of ess.TraceWin.project
---
ess/TraceWin.py | 128 ++++++++++++++++++
.../tw_project_file_reverse_engineered.json | 97 +++++++++++++
setup.py | 1 +
3 files changed, 226 insertions(+)
create mode 100644 ess/data/tw_project_file_reverse_engineered.json
diff --git a/ess/TraceWin.py b/ess/TraceWin.py
index 445f812..eda0eaf 100644
--- a/ess/TraceWin.py
+++ b/ess/TraceWin.py
@@ -1492,3 +1492,131 @@ class field_map:
data = self.map.reshape(totmapshape)
for j in data:
fout.write("{}\n".format(j))
+
+
+class project:
+ def __init__(self, project_fname=None, settings_fname=None):
+ """
+ Read and modify project files
+
+ Example::
+ p = project('SPK.ini')
+ for diff in p.compare_to('MEBT.ini'):
+ print(diff)
+ p.set('main:beam1_energy', 89e6)
+ p.save()
+ """
+ import json
+ import pkg_resources
+
+ if settings_fname is None:
+ self._refdict = json.loads(pkg_resources.resource_string(__name__, "data/tw_project_file_reverse_engineered.json"))
+ else:
+ self._refdict = json.loads(open(self._settings_fname, "r").read())
+
+ self._settings_fname = settings_fname
+ self._project_fname = project_fname
+
+ self._dict = {}
+ if self._project_fname is not None:
+ self._read_settings()
+
+ def _read_settings(self):
+ import struct
+ import textwrap
+
+ with open(self._project_fname, "rb") as f:
+ hexlist = textwrap.wrap(f.read().hex(), 2)
+ for key in self._refdict:
+ o = self._refdict[key]
+ current = "".join(hexlist[o[0] : o[0] + o[2]])
+ if o[1] == "bool":
+ if current == o[3]:
+ self._dict[key] = False
+ elif current == o[4]:
+ self._dict[key] = True
+ else:
+ raise ValueError(f"Wrong setting found for {key}, found {current}, expected one of {o[3:]}")
+ elif o[1] in ["d", "f", "i"]:
+ self._dict[key] = struct.unpack(o[1], b"".fromhex(current))[0]
+
+ def print_settings(self, settings=None):
+ """
+ Print the settings given, or all by default
+
+ :param settings: List of the settings to print
+ """
+ if settings is None:
+ settings = self.keys()
+
+ for setting in settings:
+ print(setting, self._dict[setting])
+
+ def keys(self):
+ return self._dict.keys()
+
+ def get(self, parameter):
+ """
+ Get the setting of the parameter
+ """
+ return self._dict[parameter]
+
+ def set(self, parameter, value):
+ """
+ Set the new value for parameter
+ """
+ current = self.get(parameter)
+ if isinstance(current, bool):
+ if not isinstance(value, bool):
+ raise ValueError(f"{parameter} should be True or False")
+ elif isinstance(current, (float, int)):
+ if not isinstance(value, (float, int)):
+ raise ValueError(f"{parameter} should be a number")
+ self._dict[parameter] = value
+
+ def save(self, fname=None):
+ """
+ Save the project file
+
+ If fname not given, overwrite original file
+ """
+ import struct
+ from textwrap import wrap
+
+ if fname is None:
+ fname = self._project_fname
+
+ with open(self._project_fname, "rb") as f:
+ hexlist = wrap(f.read().hex(), 2)
+
+ for key in self._dict:
+ o = self._refdict[key]
+ v = self._dict[key]
+ if isinstance(v, bool):
+ if v:
+ v = [o[-1]]
+ else:
+ v = [o[-2]]
+ else:
+ v = wrap(struct.pack(o[1], v).hex(), 2)
+ for i in range(len(v)):
+ hexlist[o[0] + i] = v[i]
+
+ open(fname, "wb").write(bytes.fromhex("".join(hexlist)))
+
+ def compare_to(self, other):
+ """
+ Compare the settings of this file to a
+ different project file
+
+ :param other: project object, or file path to other project file
+ """
+ diffs = []
+
+ if isinstance(other, str):
+ other = project(other)
+
+ for key in self.keys():
+ if self.get(key) != other.get(key):
+ diffs.append([key, self.get(key), other.get(key)])
+ return diffs
diff --git a/ess/data/tw_project_file_reverse_engineered.json b/ess/data/tw_project_file_reverse_engineered.json
new file mode 100644
index 0000000..4e22f43
--- /dev/null
+++ b/ess/data/tw_project_file_reverse_engineered.json
@@ -0,0 +1,97 @@
+{
+ "partran": [
+ 109,
+ "bool",
+ 1,
+ "00",
+ "01"
+ ],
+ "toutatis": [
+ 135,
+ "bool",
+ 1,
+ "00",
+ "01"
+ ],
+ "match:input_matched_beam": [
+ 128,
+ "bool",
+ 1,
+ "00",
+ "01"
+ ],
+ "match:family_twiss": [
+ 113,
+ "bool",
+ 1,
+ "00",
+ "01"
+ ],
+ "match:diag": [
+ 118,
+ "bool",
+ 1,
+ "00",
+ "01"
+ ],
+ "match:limit_criterion": [
+ 12260,
+ "d",
+ 8
+ ],
+ "main:beam1_npart": [
+ 11984,
+ "i",
+ 4
+ ],
+ "main:beam2_npart": [
+ 11988,
+ "i",
+ 4
+ ],
+ "main:beam1_energy": [
+ 12100,
+ "d",
+ 8
+ ],
+ "main:beam1_freq": [
+ 12068,
+ "d",
+ 8
+ ],
+ "main:beam1_curr": [
+ 12084,
+ "d",
+ 8
+ ],
+ "main:beam1_dcycle": [
+ 12648,
+ "d",
+ 8
+ ],
+ "main:beam1_emit_xxp": [
+ 12116,
+ "d",
+ 8
+ ],
+ "main:beam1_emit_yyp": [
+ 12132,
+ "d",
+ 8
+ ],
+ "main:beam1_rms_energy_spread": [
+ 22752,
+ "d",
+ 8
+ ],
+ "main:beam1_twiss_alphaz": [
+ 12572,
+ "d",
+ 8
+ ],
+ "main:beam1_twiss_betaz": [
+ 12580,
+ "d",
+ 8
+ ]
+}
diff --git a/setup.py b/setup.py
index 1052449..b8dc58d 100644
--- a/setup.py
+++ b/setup.py
@@ -26,5 +26,6 @@ setup(
"Programming Language :: Python :: 3.8",
],
packages=["ess"],
+ package_data={"ess": ["data/*.json"]},
scripts=["scripts/twcli", "scripts/tracewin_errorstudy", "scripts/ibsimu2tw"],
)
--
GitLab