"""
v0.1:
Now that TraceWin can also convert binary fieldmaps to ASCII the main usage is reduced to trimming fieldmaps
-- M. Eshraqi 2018 July 09
v0.2:
Added function to read ASCII fieldmaps and added descriptions
Added function save_to_ascii
Added function batch_convert_to_Ascci
-- M. Eshraqi 2018 July 23
v0.3:
minor bug fixed in 2D fields
added function unify_grid_size
v0.4:
added function create_tw_fieldmap
added function uniform_mesh
fixed a bug in function save_to_ascii which was writing Nx, Ny and Nz as floats
-- M. Eshraqi 2018 Aug 09
---------------------------------------------------------------------------------------------------------
From TraceWin manual:
---------------------------------------------------------------------------------------------------------
The field map file syntax is the following in the BINARY format:
- Dimension 1 ::
nz (integer 4 bytes) zmax (double 8 bytes)
Norm (double 8 bytes)
# Nz Zmax
# Norm
for k=0 to nz
Fz(k.zmax/nz) (float 4 bytes)
- Dimension 2 ::
nz (integer 4 bytes) zmax (double 8 bytes)
nx (integer 4 bytes) xmax (double 8 bytes)
Norm (double 8 bytes)
# Nz Zmax
# Nx Xmax
# Norm
for k=0 to nz
for i=0 to nr
Fz(k*zmax/nz, i*rmax/nr) (float 4 bytes)
- Dimension 3 ::
(Be careful about the dimension order)
nz (integer 4 bytes) zmax (double 8 bytes)
nx (integer 4 bytes) xmin (double 8 bytes) xmax (double 8 bytes)
ny (integer 4 bytes) ymin (double 8 bytes) ymax (double 8 bytes)
Norm (double 8 bytes)
Be careful about the dimension order
# Nz Zmax
# Nx Xmin Xmax
# Ny Ymin Ymax
# Norm
for k=0 to nz
for j=0 to ny
for i=0 to nx
Fz(k*zmax/nz, ymin+j*(ymax-ymin)/ny, xmin+i*(xmax-xmin)/nx) (float 4 bytes)
Warning: The lattice has to be regular.
The normalization factor is equal to ke/Norm or kb/Norm.
Fz are in MV/m for electric field or in T for magnetic field. The dimensions are in meter.
---------------------------------------------------------------------------------------------------------
"""
__all__ = [
"read_fieldmap",
"cut_fieldmap_length",
"field_on_axis",
"field_on_axis_from_file",
"convert_3d_to_1d",
"save_to_ascii",
"batch_convert_to_ascii",
"unify_grid_size",
"batch_cut_fieldmap",
"uniform_mesh",
"create_tw_fieldmap",
]
# class FieldMapy:
# '''
# class contains functions needed for manipulating fieldmaps for TraceWin.
# '''
__version__ = "0.4"
import os
import numpy as np
def _read_binary_fieldmap(filepath: str, fieldmap_dim: int):
"""
parameters
----------
filepath: string
path to the binary fieldmap file including the path, file name and extension
fieldmap_dim: int
dimension of the fieldmap, 1, 2 and 3 for 1D, 2D and 3D TraceWin formatted fieldmap
returns
-------
a list containing Header and Field
Header: 1D 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
"""
field_File = open(filepath, "r")
if fieldmap_dim == 3:
Header_Type = np.dtype(
[
("Nz", "int32"),
("Zmax", "float64"),
("Nx", "int32"),
("Xmin", "float64"),
("Xmax", "float64"),
("Ny", "int32"),
("Ymin", "float64"),
("Ymax", "float64"),
("Norm", "float64"),
]
)
THeader = np.fromfile(field_File, dtype=Header_Type, count=1)
Nz = THeader["Nz"][0]
Zmax = THeader["Zmax"][0]
Nx = THeader["Nx"][0]
Xmin = THeader["Xmin"][0]
Xmax = THeader["Xmax"][0]
Ny = THeader["Ny"][0]
Ymin = THeader["Ymin"][0]
Ymax = THeader["Ymax"][0]
Norm = THeader["Norm"][0]
Header = np.zeros(len(THeader))
Header = np.array([Nz, Zmax, Nx, Xmin, Xmax, Ny, Ymin, Ymax, Norm])
elif fieldmap_dim == 2:
Header_Type = np.dtype([("Nz", "int32"), ("Zmax", "float64"), ("Nx", "int32"), ("Xmax", "float64"), ("Norm", "float64")])
THeader = np.fromfile(field_File, dtype=Header_Type, count=1)
Nz = THeader["Nz"][0]
Zmax = THeader["Zmax"][0]
Nx = THeader["Nx"][0]
Xmax = THeader["Xmax"][0]
Norm = THeader["Norm"][0]
Header = np.zeros(len(THeader))
Header = np.array([Nz, Zmax, Nx, Xmax, Norm])
elif fieldmap_dim == 1:
Header_Type = np.dtype([("Nz", "int32"), ("Zmax", "float64"), ("Norm", "float32")])
Header = np.fromfile(field_File, dtype=Header_Type, count=1)
Nz = THeader["Nz"][0]
Zmax = THeader["Zmax"][0]
Norm = THeader["Norm"][0]
Header = np.zeros(len(THeader))
Field_Type = np.dtype("float32")
Nf = int(Nz + 1)
if fieldmap_dim == 2:
Nf *= int(Nx + 1)
elif fieldmap_dim == 3:
Nf *= int(Ny + 1) * int(Nx + 1)
Field = np.fromfile(field_File, dtype=Field_Type, count=Nf)
return Header, Field
def _read_ascii_fieldmap(filepath, interpolate=None):
"""
parameters
----------
filepath: string
path to the ASCII fieldmap file including the path, file name and extension
returns
-------
a list containing Header and Field
Header: 1D 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
"""
from ess import TraceWin
fmap = TraceWin.field_map(filepath)
if interpolate:
fmap.interpolate(interpolate)
Fieldmap = np.array(fmap.get_flat_fieldmap())
Header = np.array(fmap.header)
return Header, Fieldmap
[docs]def read_fieldmap(filepath: str, fieldmap_dim: int, file_type: str):
"""
parameters
----------
filepath: string
path to the binary fieldmap file including the path, file name and extension
fieldmap_dim: int
dimension of the fieldmap, 1, 2 and 3 for 1D, 2D and 3D TraceWin formatted fieldmap
file_type: str
'b' for binary and 'a' for ascii fieldmap files.
returns
-------
a list containing Header and Field
Header: 1D 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
"""
if file_type == "b":
header, field = _read_binary_fieldmap(filepath, fieldmap_dim)
elif file_type == "a":
header, field = _read_ascii_fieldmap(filepath)
else:
raise ValueError("file_type should be a for ascii or b for binary")
return header, field
[docs]def cut_fieldmap_length(
filepath: str,
fieldmap_dim: int,
entrance_rows_cut: int,
exit_rows_cut: int,
file_type: str,
):
"""
parameters
----------
filepath: string
path to the fieldmap file (ASCII or binary) including the path, file name and extension
fieldmap_dim: int
dimension of the fieldmap, 1, 2 and 3 for 1D, 2D and 3D TraceWin formatted fieldmap
entrance_rows_cut: int
number of data rows to be removed from array from the entrance side of the field
exit_rows_cut: int
number of data rows to be removed from array from the exit side of the field
file_type: str
'b' for binary and 'a' for ascii fieldmap files.
returns
-------
a list containing Header and Field
Header: 1D 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 with only the remaining rows
raises
------
prints a message if the file_type is not specified to be a (ascii) or b (binary)
todo
-------
cutting the field in other directions
"""
Header, Field = read_fieldmap(filepath, fieldmap_dim, file_type)
lowx = np.arange(0, int(entrance_rows_cut)) # list of rows to be removed from the entrance side
highx = np.arange(int(Header[0] + 1 - exit_rows_cut), int(Header[0] + 1)) # list of rows to be removed from the exit side
newNz = (Header[0] + 1) - entrance_rows_cut - exit_rows_cut # new number of field data points after removal of rows
stepZ = Header[1] / Header[0]
# depending the dimension of the field the fieldmap is reshaped for line removal
if fieldmap_dim == 1:
newlength = newNz
elif fieldmap_dim == 2:
Fieldmapmatrix = np.reshape(Field, (int(Header[0] + 1), int(Header[2] + 1)))
newlength = newNz * (Header[2] + 1)
elif fieldmap_dim == 3:
Fieldmapmatrix = np.reshape(Field, (int(Header[0] + 1), int((Header[2] + 1) * (Header[5] + 1))))
newlength = newNz * (Header[2] + 1) * (Header[5] + 1)
if entrance_rows_cut > 0 and exit_rows_cut > 0:
x = np.delete(Fieldmapmatrix, (lowx, highx), axis=0)
elif entrance_rows_cut == 0:
x = np.delete(Fieldmapmatrix, (highx), axis=0)
elif exit_rows_cut == 0:
x = np.delete(Fieldmapmatrix, (lowx), axis=0)
# reshaping the array back to a 1D array
ShortenedMatrix = np.reshape(x, int(newlength))
# creating new header data (if other directions are added to the cut fuction the header for those have
# to be rewritten too)
Header[0] = newNz - 1
Header[1] = stepZ * (newNz - 1)
return Header, ShortenedMatrix
[docs]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 field maps)
TraceWin field map header info
field: 1D numpy array containing the nD field data
fieldmap_dim: int
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 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(header, field, 3)
"""
import numpy
if fieldmap_dim == 1:
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]
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]
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, :] = field_map_matrix[:, midpoint]
return field_on_axis
[docs]def field_on_axis_from_file(file_path, fieldmap_dim: int, file_type: str):
"""
parameters
----------
file_path: string
name of the field map file including its path and extension
fieldmap_dim: int
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 field map files.
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,
returns the field on-axis in the z direction
example:
field_on_axis_from_file('Data/FM/HB_W_coupler.edz', 3, 'b')
"""
Header, Field = read_fieldmap(file_path, fieldmap_dim, file_type)
return field_on_axis(Header, Field, fieldmap_dim)
[docs]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 - field_on_axis(header, field, 3)[1]) / 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}")
[docs]def save_to_ascii(Header, Field, output_filepath: str):
"""
Saves a TraceWin readable ASCII file
Parameters
----------
Header: array like
1D array of size (3, 5 or 9) for 1D, 2D or 3D fieldmap header info, TraceWin format
Field: array like
1D array of size Nz fieldmap containing field data
output_filepath: string
path to the output file
"""
# The dimension is defined from the size of the Header:
if len(Header) == 3:
fieldmap_dim = 1
elif len(Header) == 5:
fieldmap_dim = 2
elif len(Header) == 9:
fieldmap_dim = 3
else:
raise ValueError("Wrong length of header, cannot write {}".format(output_filepath))
with open(output_filepath, "w") as out_file:
# stepZ = Header[1] / Header[0] # stepZ = Zmax / Nz
out_file.write("{0[0]:.0f} {0[1]}\n".format(Header))
if fieldmap_dim == 1:
out_file.write("{}\n".format(Header[2]))
elif fieldmap_dim == 2:
out_file.write("{0[2]:.0f} {0[3]}\n".format(Header))
out_file.write("{}\n".format(Header[4]))
elif fieldmap_dim == 3:
out_file.write("{0[2]:.0f} {0[3]} {0[4]}\n".format(Header))
out_file.write("{0[5]:.0f} {0[6]} {0[7]}\n".format(Header))
out_file.write("{}\n".format(Header[8]))
for fval in Field:
out_file.write("{:g}\n".format(fval))
[docs]def batch_convert_to_ascii(binary_field_path: str, fieldmap_dim: int):
"""
parameters
----------
binary_field_path: string
path to the binary files to be converted to ASCII format
fieldmap_dim: int
dimension of the fieldmap to be converted
"""
onlyfiles = [f for f in os.listdir(binary_field_path) if os.path.isfile(os.path.join(binary_field_path, f))]
for f in onlyfiles:
if f[0] != ".":
Header, Field = _read_binary_fieldmap(os.path.join(binary_field_path, f), fieldmap_dim)
save_to_ascii(Header, Field, os.path.join(binary_field_path, "ASCII", f))
return
[docs]def unify_grid_size(old_x: list, old_y: list, new_x: list, kind="linear", save_to_file=None):
"""
Takes a mesh as x values and the corresponding values of a function.
creates a spline interpolation for those data sets.
Using a new_x it returns the correcposnding new_y.
Could be used to convert an irregular mesh to a regular mesh or
to increase the number of mesh points
parameters
----------
old_x: list,
old_y: list,
new_x: list
kind: str
An integer specifying the order of the spline interpolator to use. Default is ‘linear’.
- ‘linear’, ‘nearest’, ‘zero’, ‘slinear’, ‘quadratic’, ‘cubic’: refer to a spline
interpolation of zeroth, first, second or third order
- ‘previous’ or ‘next’: return the previous or next value of the point
save_to_file: str
If given, file name where this should be saved.
Returns
-------
new_y: list
"""
from scipy import interpolate
f = interpolate.interp1d(old_x, old_y, kind)
new_y = f(new_x)
if save_to_file is not None:
field = np.zeros((2, len(new_x)))
field[0, :] = new_x
field[1, :] = new_y
np.savetxt(save_to_file, np.transpose(field), delimiter="\t")
return new_y
[docs]def batch_cut_fieldmap(
fieldmaps_path: str,
fieldmap_dim: int,
entrance_rows_cut: int,
exit_rows_cut: int,
file_type: str,
):
"""
Saves a TraceWin readable ASCII file
parameters
----------
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
entrance_rows_cut: int
number of data rows to be removed from array from the entrance side of the field
exit_rows_cut: int
number of data rows to be removed from array from the exit side of the field
file_type: str
'b' for binary and 'a' for ascii fieldmap files.
"""
onlyfiles = [f for f in os.listdir(fieldmaps_path) if os.path.isfile(os.path.join(fieldmaps_path, f))]
out_path = "ASCII"
full_out_path = os.path.join(fieldmaps_path, out_path)
if not os.path.exists(full_out_path):
os.makedirs(full_out_path)
for f in onlyfiles:
if f[0] != ".":
filepath = os.path.join(fieldmaps_path, f)
if entrance_rows_cut + exit_rows_cut > 0:
Header, Field = cut_fieldmap_length(filepath, fieldmap_dim, entrance_rows_cut, exit_rows_cut, file_type)
save_to_ascii(Header, Field, os.path.join(fieldmaps_path, "ASCII", f))
[docs]def create_tw_fieldmap(
fieldmaps_path: str,
tabular_file_name: str,
fieldmap_dim: int = 1,
n_col: int = 1,
skiprows: int = 0,
comments: str = "#",
extensions=["F1", "F2", "F3", "F4", "F5", "F6"],
xyz=[1, 2, 3],
delimiter=None,
):
"""
parameters
---------
fieldmaps_path: str
path to the fieldmap
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)
fieldmap_dim: int
dimension of the fieldmap data
n_col: int
number of data columns in fieldmap
skiprows: int
number of header lines to skip
comments: str
Comment lines start with this character
extensions: list
list of column names to be used for naming the files, default is [F1, F2, F3, F4, F5, F6]
xyz: list
order of columns, xyz = [col_number(x), col), col_number(y), col_number(z)]
e.g.if column 1 is x, col 2 is y and col 3 is z, xyz list = [1,2,3] (default)
delimiter: string
Delimiter character that separates columns. Default is whitespace
"""
if len(xyz) < fieldmap_dim:
print("xyz array is shorter than the number of dimensions")
return
tabular_file_with_path = os.path.join(fieldmaps_path, tabular_file_name)
field_array = np.loadtxt(tabular_file_with_path, delimiter=delimiter, skiprows=skiprows, comments=comments)
items_per_row = np.array([[2, 1, 0, 0], [2, 2, 1, 0], [2, 3, 3, 1]]) # for 1D, 2D and 3D fields in an ASCII fieldmap in TraceWin format
l_header = np.sum(items_per_row[fieldmap_dim - 1, :], axis=0) # getting the length of the header from the items
Header = np.zeros(l_header)
# sort first on x coordinate, then y and then z for TraceWin fieldmap
for i in range(0, fieldmap_dim):
coor = field_array[:, xyz[i] - 1]
chk_mesh, d_mesh, n_mesh = uniform_mesh(coor)
if chk_mesh:
if i == fieldmap_dim - 1:
Header[0] = n_mesh
Header[1] = max(coor) - min(coor)
# \n
if fieldmap_dim == 1:
Header[2] = 1
# \n
elif fieldmap_dim == 2:
if i == fieldmap_dim - 2:
Header[2] = n_mesh
Header[3] = max(coor) - min(coor)
# \n
Header[4] = 1
# \n
elif fieldmap_dim == 3:
if i == fieldmap_dim - 2:
Header[2] = n_mesh
Header[3] = min(coor)
Header[4] = max(coor)
# \n
elif i == fieldmap_dim - 3:
Header[5] = n_mesh
Header[6] = min(coor)
Header[7] = max(coor)
# \n
Header[8] = 1
# \n
else: # fieldmap_dim > 3 or fieldmap_dim < 1:
print("Invalid dimension, should be 0, 1 or 2")
return
ordr = np.argsort(coor, kind="mergesort")
field_array = field_array[ordr, :]
filename_wo_ext = os.path.splitext(tabular_file_name)[0]
for i in range(fieldmap_dim, fieldmap_dim + n_col):
save_to_ascii(
Header,
field_array[:, i],
os.path.join(fieldmaps_path, filename_wo_ext + "." + extensions[i - fieldmap_dim]),
)