diff --git a/ess/fieldmap.py b/ess/fieldmap.py
index 294c2646f3fa6997e4b2a53f0e21ab15b043147d..29ee0c5a7aa8be609e91d1cf93b2116b7d460c6b 100644
--- a/ess/fieldmap.py
+++ b/ess/fieldmap.py
@@ -303,58 +303,74 @@ def cut_fieldmap_length(
return Header, ShortenedMatrix
-def field_on_axis(Header, Field, fieldmap_dim: int):
+def field_on_axis(header, field, fieldmap_dim: int, x_on_axis: float = 0, y_on_axis: float = 0):
"""
parameters
----------
- Header: 1D numpy array of size 3, 5, 9 (for 1D, 2D and 3D fieldmaps)
- TraceWin fieldmap header info
- Field: 1D numpy array containing the nD field data
+ header: 1D numpy array of size 3, 5, 9 (for 1D, 2D and 3D field maps)
+ TraceWin field map header info
+ field: 1D numpy array containing the nD field data
fieldmap_dim: int
- dimesion of the fieldmap, 1, 2, 3 for 1D, 2D and 3D fieldmaps in TraceWin format
-
+ dimension of the field map, 1, 2, 3 for 1D, 2D and 3D field maps in TraceWin format
+ x_on_axis: float
+ Field is given for x=x_on_axis [default 0]
+ y_on_axis: float
+ Field is given for y=y_on_axis [default 0]
returns
-------
returns a 2D array containing the z coordinates and the field on axis in the z direction
- Take a 2D or 3D fieldmap as a 1D numpy array (as from TraceWin) in MV/m,
- dimension of the fieldmap,
+ Take a 2D or 3D field map as a 1D numpy array (as from TraceWin) in MV/m,
+ dimension of the field map,
returns the field on-axis in the z direction
example:
- field_on_axis_from_file('Data/FM/HB_W_coupler.edz', 3, 'b')
+ field_on_axis(header, field, 3)
"""
import numpy
if fieldmap_dim == 1:
- return Field
+ if x_on_axis or y_on_axis:
+ raise ValueError("Cannot define nonzero x/y for 1D field map")
+ return field
elif fieldmap_dim == 2:
- Nz, Nrx = Header[0], Header[2]
- Fieldmapmatrix = np.reshape(Field, (int(Nz + 1), int(Nrx + 1)))
+ Nz, Nrx = header[0], header[2]
+ field_map_matrix = np.reshape(field, (int(Nz + 1), int(Nrx + 1)))
midpoint = 0
elif fieldmap_dim == 3:
- Nz, Nrx, Ny = Header[0], Header[2], Header[5]
- Fieldmapmatrix = np.reshape(Field, (int(Nz + 1), int((Nrx + 1) * (Ny + 1))))
- midpoint = int(Nrx * Ny / 2)
-
- dz = Header[1] / Nz
+ Nz, Nrx, Ny = header[0], header[2], header[5]
+ field_map_matrix = np.reshape(field, (int(Nz + 1), int((Nrx + 1) * (Ny + 1))))
+ midpoint = int((Nrx + 1) * (Ny + 1) / 2)
+
+ dx_steps = 0
+ dy_steps = 0
+ if x_on_axis:
+ dx = (header[4] - header[3]) / (Nrx + 1)
+ dx_steps = int(x_on_axis / dx)
+
+ if y_on_axis:
+ dy = (header[7] - header[6]) / (Ny + 1)
+ dy_steps = int(y_on_axis / dy)
+ midpoint += int(dx_steps + (Nrx + 1) * dy_steps)
+
+ dz = header[1] / Nz
field_on_axis = numpy.zeros((2, int(Nz) + 1))
for i in range(int(Nz + 1)):
field_on_axis[0, i] = i * dz
- field_on_axis[1, :] = Fieldmapmatrix[:, midpoint]
+ field_on_axis[1, :] = field_map_matrix[:, midpoint]
return field_on_axis
-def field_on_axis_from_file(filepath, fieldmap_dim: int, file_type: str):
+def field_on_axis_from_file(file_path, fieldmap_dim: int, file_type: str):
"""
parameters
----------
- filepath: string
- name of the fieldmap file including its path and extension
+ file_path: string
+ name of the field map file including its path and extension
fieldmap_dim: int
- dimesion of the fieldmap, 1, 2, 3 for 1D, 2D and 3D fieldmaps in TraceWin format
+ dimension of the field map, 1, 2, 3 for 1D, 2D and 3D field maps in TraceWin format
file_type: string
- 'b' for binary and 'a' for ascii fieldmap files.
+ 'b' for binary and 'a' for ascii field map files.
returns
-------
@@ -367,11 +383,71 @@ def field_on_axis_from_file(filepath, fieldmap_dim: int, file_type: str):
field_on_axis_from_file('Data/FM/HB_W_coupler.edz', 3, 'b')
"""
- Header, Field = read_fieldmap(filepath, fieldmap_dim, file_type)
+ Header, Field = read_fieldmap(file_path, fieldmap_dim, file_type)
return field_on_axis(Header, Field, fieldmap_dim)
+def convert_3d_to_1d(path: str, map_name: str, map_type: str = "electric", file_type: str = "a", dist_from_axis: float = 0.0):
+ """
+ parameters
+ ---------
+ path: str
+ path to the field map
+ map_name: str
+ name of field map
+ map_type: str
+ electric [default] or magnetic
+ file_type: str
+ 'b' for binary and 'a' for ascii field map files.
+ dist_from_axis: float
+ In case of magnetic quadrupole field, the distance from axis where field is read (T/m). Unit m.
+ Ignored if map_type is electric.
+
+ returns
+ -------
+ Saves 1D field map to file. Currently overwriting automatically (renaming original) so careful.
+
+ Example:
+ fieldmap.convert_3d_to_1d('mebt', 'MEBT_Q', 'magnetic', 'a', 0.003)
+ """
+ if map_type == "electric":
+ dist_from_axis = 0
+ file_end = "edz"
+ file_end_out = "edz"
+ elif map_type == "magnetic":
+ if dist_from_axis:
+ file_end = "bsy"
+ else:
+ file_end = "bsz"
+ file_end_out = "bsz"
+ else:
+ raise ValueError(f"Wrong map_type: {map_type}")
+
+ file_path = ""
+ file_name = f"{map_name}.{file_end}"
+ for f in os.listdir(path):
+ if f == file_name:
+ file_path = os.path.join(path, f)
+ if not file_path:
+ raise ValueError(f"Could not find the map {map_name}")
+ header, field = read_fieldmap(file_path, 3, file_type)
+
+ if map_type == "electric":
+ dist_from_axis = 0
+
+ field_1d = field_on_axis(header, field, 3, x_on_axis=dist_from_axis)[1]
+ if dist_from_axis:
+ # Here we give T/m rather than T given for 3D field map
+ field_1d = field_1d / dist_from_axis
+
+ header_1d = np.array([header[0], header[1], header[-1]])
+ file_path_out = f"{file_path[:-4]}_1D.{file_end_out}"
+
+ save_to_ascii(header_1d, field_1d, file_path_out)
+ print(f"1D field map saved as {file_path_out}")
+
+
def save_to_ascii(Header, Field, output_filepath: str):
"""
Saves a TraceWin readable ASCII file
@@ -487,7 +563,7 @@ def batch_cut_fieldmap(
parameters
----------
- fieldmaps_path: string
+ fieldmaps_ path: string
path to the fieldmap folder (ASCII or binary) excluding the trailing / and filename
fieldmap_dim: int
dimension of the fieldmap, 1, 2 and 3 for 1D, 2D and 3D TraceWin formatted fieldmap
@@ -555,7 +631,7 @@ def create_tw_fieldmap(
---------
fieldmaps_path: str
path to the fieldmap
- tabular_file_with_path: file name
+ tabular_file_name: file name
tabular data with coordinates z, x and y or z and r or only z depending on dimension as first 1-3 cols.
n_col is the number of data columns except coordinate ones (1-6)
coordinates should be on a uniform grid on each axis (but different axes are independent)