Skip to content

Commit

Permalink
Updates after merge fixing and linting
Browse files Browse the repository at this point in the history
  • Loading branch information
@ColwynGulliford
ColwynGulliford committed Oct 18, 2024
1 parent cab6562 commit 36e9cdf
Show file tree
Hide file tree
Showing 3 changed files with 189 additions and 196 deletions.
20 changes: 10 additions & 10 deletions pmd_beamphysics/fields/fieldmesh.py
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@

import functools
import os
from copy import deepcopy
Expand All @@ -13,7 +12,10 @@
)
from pmd_beamphysics.fields.expansion import expand_fieldmesh_from_onaxis
from pmd_beamphysics.interfaces.ansys import read_ansys_ascii_3d_fields
from pmd_beamphysics.interfaces.cst import read_cst_ascii_3d_complex_fields, read_cst_ascii_3d_static_field
from pmd_beamphysics.interfaces.cst import (
read_cst_ascii_3d_complex_fields,
read_cst_ascii_3d_static_field,
)

from pmd_beamphysics.interfaces.astra import (
astra_1d_fieldmap_data,
Expand Down Expand Up @@ -532,8 +534,9 @@ def write_gpt(self, filePath, asci2gdf_bin=None, verbose=True):
Writes a GPT field file.
"""

return write_gpt_fieldmap(self, filePath, asci2gdf_bin=asci2gdf_bin, verbose=verbose)

return write_gpt_fieldmap(
self, filePath, asci2gdf_bin=asci2gdf_bin, verbose=verbose
)

@functools.wraps(write_impact_emfield_cartesian)
def write_impact_emfield_cartesian(self, filename):
Expand All @@ -550,7 +553,6 @@ def write_impact_emfield_cartesian(self, filename):

return write_impact_emfield_cartesian(self, filename)


# Superfish
@functools.wraps(write_superfish_t7)
def write_superfish(self, filePath, verbose=False):
Expand Down Expand Up @@ -600,17 +602,15 @@ def from_ansys_ascii_3d(cls, *, efile=None, hfile=None, frequency=None):

@classmethod
def from_cst_3d(cls, field_file1, field_file2=None, frequency=0):

if field_file2 is not None:
# field_file1 -> efile, field_file2 -> hfile
data = read_cst_ascii_3d_complex_fields(field_file1, field_file2, frequency=frequency, harmonic=1)
data = read_cst_ascii_3d_complex_fields(
field_file1, field_file2, frequency=frequency, harmonic=1
)
else:
data = read_cst_ascii_3d_static_field(field_file1)

return cls(data=data)




@classmethod
def from_astra_3d(cls, common_filename, frequency=0):
Expand Down
155 changes: 78 additions & 77 deletions pmd_beamphysics/interfaces/cst.py
View file
Original file line number Diff line number Diff line change
@@ -1,18 +1,17 @@

import numpy as np
import os

from scipy.constants import mu_0 as mu0

def get_scale(unit):

def get_scale(unit):
"""
Returns the scaling factor corresponding to the provided unit.
Parameters
----------
unit : str
The unit for which the scaling factor is requested.
The unit for which the scaling factor is requested.
Accepted values are:
- '[mm]' : millimeters
- '[V/m]' : volts per meter
Expand All @@ -21,26 +20,26 @@ def get_scale(unit):
Returns
-------
float
The scaling factor corresponding to the unit.
The scaling factor corresponding to the unit.
- 1e-3 for millimeters ('[mm]')
- 1 for volts per meter ('[V/m]')
- mu0 (vacuum permeability) for amperes per meter ('[A/m]')
"""

if(unit=='[mm]'):
if unit == "[mm]":
return 1e-3
elif(unit=='[V/m]'):
elif unit == "[V/m]":
return 1
elif(unit=='[A/m]'):
elif unit == "[A/m]":
return mu0

def get_vec(x):

def get_vec(x):
"""
Processes the input list or array `x` to compute key vector parameters.
The function computes the minimum and maximum values of `x`, the difference
between consecutive unique and sorted elements (`dx`), and the number of
The function computes the minimum and maximum values of `x`, the difference
between consecutive unique and sorted elements (`dx`), and the number of
unique elements in `x`.
Parameters
Expand All @@ -55,12 +54,12 @@ def get_vec(x):
xmax : float
The maximum value in `x`.
dx : float
The uniform difference between consecutive unique elements in `x`.
The uniform difference between consecutive unique elements in `x`.
Raises an assertion error if differences are not uniform.
nx : int
The number of unique elements in `x`.
"""

sx = set(x)
nx = len(sx)
xlist = np.array(sorted(list(sx)))
Expand All @@ -70,7 +69,6 @@ def get_vec(x):
return min(x), max(x), dx, nx


<<<<<<< HEAD
def read_cst_ascii_3d_field(filePath, n_header=2):
"""
Parses a 3D field file generated by CST, extracting field data and header information.
Expand All @@ -91,15 +89,15 @@ def read_cst_ascii_3d_field(filePath, n_header=2):
data : numpy.ndarray
A NumPy array containing the parsed 3D field data. The format of the data depends on
whether the field is static or time-varying. The array is ordered in Fortran ('F') format.
header_info : dict
A dictionary containing extracted information from the file header, such as units and
A dictionary containing extracted information from the file header, such as units and
field types (E or H field components).
Notes
-----
The header structure for static fields is expected to follow:
- x [units], y [units], z [units], Fx [units], Fy [units], Fz [units]
- Example:
```
Expand All @@ -120,60 +118,68 @@ def read_cst_ascii_3d_field(filePath, n_header=2):
The data in the file is assumed to be ordered in Fortran ('F') format, meaning column-major order.
"""
with open(filePath, 'r') as fid:

with open(filePath, "r") as fid:
header = fid.readline()

headers = header.split()

columns, units = headers[::2], headers[1::2]

#print(columns, units)
# print(columns, units)

field_columns = list(set([c[:2] for c in columns if c.startswith('E') or c.startswith('H')]))
field_columns = list(
set([c[:2] for c in columns if c.startswith("E") or c.startswith("H")])
)

if all([f.startswith('E') for f in field_columns]):
field_type = 'electric'
elif all([f.startswith('H') for f in field_columns]):
field_type = 'magnetic'
if all([f.startswith("E") for f in field_columns]):
field_type = "electric"
elif all([f.startswith("H") for f in field_columns]):
field_type = "magnetic"
else:
raise ValueError('Mixed CST mode not curretly supported.')
raise ValueError("Mixed CST mode not curretly supported.")

dat = np.loadtxt(filePath, skiprows=n_header)
X = dat[:,0]*get_scale(units[0])
Y = dat[:,1]*get_scale(units[1])
Z = dat[:,2]*get_scale(units[2])

X = dat[:, 0] * get_scale(units[0])
Y = dat[:, 1] * get_scale(units[1])
Z = dat[:, 2] * get_scale(units[2])

xmin, xmax, dx, nx = get_vec(X)
ymin, ymax, dy, ny = get_vec(Y)
zmin, zmax, dz, nz = get_vec(Z)

shape = (nx, ny, nz)

shape = (nx, ny, nz)

# Check if the field is complex:
if( len(columns)==9 ):

if len(columns) == 9:
# - sign to convert to exp(-i omega t)
Fx = (dat[:,3] - 1j*dat[:,4]).reshape(shape, order='F')*get_scale(units[3])
Fy = (dat[:,5] - 1j*dat[:,6]).reshape(shape, order='F')*get_scale(units[4])
Fz = (dat[:,7] - 1j*dat[:,8]).reshape(shape, order='F')*get_scale(units[5])

elif( len(columns)==6 ):

Fx = dat[:,3].reshape(shape, order='F')*get_scale(units[3])
Fy = dat[:,4].reshape(shape, order='F')*get_scale(units[4])
Fz = dat[:,5].reshape(shape, order='F')*get_scale(units[5])
Fx = (dat[:, 3] - 1j * dat[:, 4]).reshape(shape, order="F") * get_scale(
units[3]
)
Fy = (dat[:, 5] - 1j * dat[:, 6]).reshape(shape, order="F") * get_scale(
units[4]
)
Fz = (dat[:, 7] - 1j * dat[:, 8]).reshape(shape, order="F") * get_scale(
units[5]
)

elif len(columns) == 6:
Fx = dat[:, 3].reshape(shape, order="F") * get_scale(units[3])
Fy = dat[:, 4].reshape(shape, order="F") * get_scale(units[4])
Fz = dat[:, 5].reshape(shape, order="F") * get_scale(units[5])


attrs = {}
attrs['gridOriginOffset'] = (xmin, ymin, zmin)
attrs['gridSpacing'] = (dx, dy, dz)
attrs['gridSize'] = (nx, ny, nz)
attrs["gridOriginOffset"] = (xmin, ymin, zmin)
attrs["gridSpacing"] = (dx, dy, dz)
attrs["gridSize"] = (nx, ny, nz)

components = {
f"{field_type}Field/x": Fx,
f"{field_type}Field/y": Fy,
f"{field_type}Field/z": Fz,
}

components = {f'{field_type}Field/x':Fx, f'{field_type}Field/y':Fy, f'{field_type}Field/z':Fz}

return attrs, components


Expand All @@ -195,28 +201,26 @@ def read_cst_ascii_3d_static_field(ffile):
data : dict
FieldMesh ready data dict for the static field.
"""

attrs, components = read_cst_ascii_3d_field(ffile)

attrs['eleAnchorPt'] = 'center'
attrs['gridGeometry'] = 'rectangular'
attrs['axisLabels'] = ('x', 'y', 'z')
attrs['gridLowerBound'] = (0, 0, 0)
<<<<<<< HEAD
attrs['harmonic'] = 0
attrs['fundamentalFrequency'] = 0

attrs["eleAnchorPt"] = "center"
attrs["gridGeometry"] = "rectangular"
attrs["axisLabels"] = ("x", "y", "z")
attrs["gridLowerBound"] = (0, 0, 0)
attrs["harmonic"] = 0
attrs["fundamentalFrequency"] = 0

data = dict(attrs=attrs, components=components)

return data


def read_cst_ascii_3d_complex_fields(efile, hfile, frequency, harmonic=1):

"""
Parses a complete 3D field map from corresponding CST E and H field files for a complex electromagnetic mode.
This function reads and processes both the electric (E) and magnetic (H) field data from separate ASCII files
This function reads and processes both the electric (E) and magnetic (H) field data from separate ASCII files
generated by CST, and associates them with the specified frequency and harmonic of the electromagnetic mode.
Parameters
Expand All @@ -238,30 +242,27 @@ def read_cst_ascii_3d_complex_fields(efile, hfile, frequency, harmonic=1):
data : dict
FieldMesh style data dict containing the complex fields
"""


assert os.path.exists(efile), "Could not find electric field file"
assert os.path.exists(hfile), "Could not find magnetic field file"

e_attrs, e_components = read_cst_ascii_3d_field(efile)
b_attrs, b_components = read_cst_ascii_3d_field(hfile)

assert e_attrs['gridOriginOffset'] == b_attrs['gridOriginOffset']
assert e_attrs['gridSpacing'] == b_attrs['gridSpacing']
assert e_attrs['gridSize'] == b_attrs['gridSize']
assert e_attrs["gridOriginOffset"] == b_attrs["gridOriginOffset"]
assert e_attrs["gridSpacing"] == b_attrs["gridSpacing"]
assert e_attrs["gridSize"] == b_attrs["gridSize"]

components = {**e_components, **b_components}

attrs = e_attrs
attrs['eleAnchorPt'] = 'center'
attrs['gridGeometry'] = 'rectangular'
attrs['axisLabels'] = ('x', 'y', 'z')
attrs['gridLowerBound'] = (0, 0, 0)
attrs['harmonic'] = harmonic
attrs['fundamentalFrequency'] = frequency
attrs["eleAnchorPt"] = "center"
attrs["gridGeometry"] = "rectangular"
attrs["axisLabels"] = ("x", "y", "z")
attrs["gridLowerBound"] = (0, 0, 0)
attrs["harmonic"] = harmonic
attrs["fundamentalFrequency"] = frequency

data = dict(attrs=attrs, components=components)


return data

Loading

0 comments on commit 36e9cdf

Please sign in to comment.