From 36e9cdf38475188129c0b8b1362fee61966670da Mon Sep 17 00:00:00 2001 From: Colwyn Gulliford <36416205+ColwynGulliford@users.noreply.github.com> Date: Fri, 18 Oct 2024 14:00:39 -0400 Subject: [PATCH] Updates after merge fixing and linting --- pmd_beamphysics/fields/fieldmesh.py | 20 +-- pmd_beamphysics/interfaces/cst.py | 155 ++++++++++---------- pmd_beamphysics/interfaces/gpt.py | 210 +++++++++++++--------------- 3 files changed, 189 insertions(+), 196 deletions(-) diff --git a/pmd_beamphysics/fields/fieldmesh.py b/pmd_beamphysics/fields/fieldmesh.py index 9761874..29c98b2 100644 --- a/pmd_beamphysics/fields/fieldmesh.py +++ b/pmd_beamphysics/fields/fieldmesh.py @@ -1,4 +1,3 @@ - import functools import os from copy import deepcopy @@ -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, @@ -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): @@ -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): @@ -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): diff --git a/pmd_beamphysics/interfaces/cst.py b/pmd_beamphysics/interfaces/cst.py index a81813d..0168b4d 100644 --- a/pmd_beamphysics/interfaces/cst.py +++ b/pmd_beamphysics/interfaces/cst.py @@ -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 @@ -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 @@ -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))) @@ -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. @@ -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: ``` @@ -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 @@ -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 @@ -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 - diff --git a/pmd_beamphysics/interfaces/gpt.py b/pmd_beamphysics/interfaces/gpt.py index c14d81b..61d6d57 100644 --- a/pmd_beamphysics/interfaces/gpt.py +++ b/pmd_beamphysics/interfaces/gpt.py @@ -86,100 +86,93 @@ def run_asci2gdf(outfile, asci2gdf_bin, verbose=False): print("Written GDF file:", outfile) return outfile - - - -def write_gpt_fieldmap(fm, - outfile, - asci2gdf_bin=None, - verbose=False): + + +def write_gpt_fieldmap(fm, outfile, asci2gdf_bin=None, verbose=False): """ Writes a GPT fieldmap file from a FieldMesh object. - + Requires cylindrical geometry for now. """ - - if fm.geometry == 'cylindrical' and fm.coord_vec('r')[0] == fm.coord_vec('r')[-1]: - return write_gpt_1d_fieldmap(fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose) - elif fm.geometry == 'cylindrical': - return write_gpt_2d_fieldmap(fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose) + if fm.geometry == "cylindrical" and fm.coord_vec("r")[0] == fm.coord_vec("r")[-1]: + return write_gpt_1d_fieldmap( + fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose + ) - elif fm.geometry == 'rectangular': - return write_gpt_3d_fieldmap(fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose) - + elif fm.geometry == "cylindrical": + return write_gpt_2d_fieldmap( + fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose + ) + + elif fm.geometry == "rectangular": + return write_gpt_3d_fieldmap( + fm, outfile, asci2gdf_bin=asci2gdf_bin, verbose=verbose + ) else: - raise ValueError(f'Unknown geometry {fm.geometry}') - + raise ValueError(f"Unknown geometry {fm.geometry}") -def write_gpt_1d_fieldmap(fm, - outfile, - asci2gdf_bin=None, - verbose=False): +def write_gpt_1d_fieldmap(fm, outfile, asci2gdf_bin=None, verbose=False): """ Writes a GPT fieldmap file from a FieldMesh object. - + Requires cylindrical geometry for now. """ - - assert fm.geometry == 'cylindrical', f'Geometry: {fm.geometry} not implemented' - - assert fm.shape[1] == 1, 'Cylindrical symmetry required' - assert fm.coord_vec('r')[0]==0, 'r[0] must equal 0' + assert fm.geometry == "cylindrical", f"Geometry: {fm.geometry} not implemented" + + assert fm.shape[1] == 1, "Cylindrical symmetry required" + + assert fm.coord_vec("r")[0] == 0, "r[0] must equal 0" dat = {} - dat['Z'] = fm.coord_vec('z') + dat["Z"] = fm.coord_vec("z") - keys = ['Z'] + keys = ["Z"] if fm.is_static: if fm.is_pure_magnetic: - keys = ['Z', 'Bz'] - dat['Bz'] = np.real(fm['Bz'][0,0,:]) + keys = ["Z", "Bz"] + dat["Bz"] = np.real(fm["Bz"][0, 0, :]) elif fm.is_pure_electric: - keys = ['Z', 'Ez'] - dat['Er'] = np.real(fm['Er'][0,0,:]) - dat['Ez'] = np.real(fm['Ez'][0,0,:]) + keys = ["Z", "Ez"] + dat["Er"] = np.real(fm["Er"][0, 0, :]) + dat["Ez"] = np.real(fm["Ez"][0, 0, :]) else: - raise ValueError('Mixed static field TODO') - + raise ValueError("Mixed static field TODO") + else: - # Use internal Superfish routine - keys = ['Z', 'Ez'] - _, dat['Ez'], _, _ = fish_complex_to_real_fields(fm, verbose=verbose) - dat['Ez'] = dat['Ez'][0,0,:] - - - # Flatten dat - gptdata = np.array([dat[k].flatten() for k in keys]).T - - # Write file. + # Use internal Superfish routine + keys = ["Z", "Ez"] + _, dat["Ez"], _, _ = fish_complex_to_real_fields(fm, verbose=verbose) + dat["Ez"] = dat["Ez"][0, 0, :] + + # Flatten dat + gptdata = np.array([dat[k].flatten() for k in keys]).T + + # Write file. # Hack to delete final newline # https://stackoverflow.com/questions/28492954/numpy-savetxt-stop-newline-on-final-line - with open(outfile, 'w') as fout: - NEWLINE_SIZE_IN_BYTES = 1 # 2 on Windows? - np.savetxt(fout, gptdata, header=' '.join(keys), comments='') - fout.seek(0, os.SEEK_END) # Go to the end of the file. + with open(outfile, "w") as fout: + NEWLINE_SIZE_IN_BYTES = 1 # 2 on Windows? + np.savetxt(fout, gptdata, header=" ".join(keys), comments="") + fout.seek(0, os.SEEK_END) # Go to the end of the file. # Go backwards one byte from the end of the file. fout.seek(fout.tell() - NEWLINE_SIZE_IN_BYTES, os.SEEK_SET) - fout.truncate() # Truncate the file to this point. - - + fout.truncate() # Truncate the file to this point. + if asci2gdf_bin: run_asci2gdf(outfile, asci2gdf_bin, verbose=verbose) - elif verbose: - print(f'ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}') - - return outfile + elif verbose: + print( + f"ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}" + ) + return outfile -def write_gpt_2d_fieldmap(fm, - outfile, - asci2gdf_bin=None, - verbose=False): +def write_gpt_2d_fieldmap(fm, outfile, asci2gdf_bin=None, verbose=False): """ Writes a GPT fieldmap file from a FieldMesh object. @@ -232,68 +225,67 @@ def write_gpt_2d_fieldmap(fm, if asci2gdf_bin: run_asci2gdf(outfile, asci2gdf_bin, verbose=verbose) - elif verbose: - print(f'ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}') - - return outfile + elif verbose: + print( + f"ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}" + ) + return outfile -def write_gpt_3d_fieldmap(fm, - outfile, - asci2gdf_bin=None, - verbose=False): + +def write_gpt_3d_fieldmap(fm, outfile, asci2gdf_bin=None, verbose=False): """ Writes a 3D GPT fieldmap file from a FieldMesh object. """ - - assert fm.geometry == 'rectangular', f'Geometry: {fm.geometry} not implemented' - + + assert fm.geometry == "rectangular", f"Geometry: {fm.geometry} not implemented" + dat = {} - dat['X'], dat['Y'], dat['Z'] = np.meshgrid(fm.coord_vec('x'), fm.coord_vec('y'), fm.coord_vec('z'), indexing='ij') - - keys = ['X', 'Y', 'Z'] + dat["X"], dat["Y"], dat["Z"] = np.meshgrid( + fm.coord_vec("x"), fm.coord_vec("y"), fm.coord_vec("z"), indexing="ij" + ) + + keys = ["X", "Y", "Z"] if fm.is_static: if fm.is_pure_magnetic: - keys = keys + ['Bx', 'By', 'Bz'] - dat['Bx'] = np.real(fm['Bx']) - dat['By'] = np.real(fm['By']) - dat['Bz'] = np.real(fm['Bz']) - + keys = keys + ["Bx", "By", "Bz"] + dat["Bx"] = np.real(fm["Bx"]) + dat["By"] = np.real(fm["By"]) + dat["Bz"] = np.real(fm["Bz"]) + elif fm.is_pure_electric: - keys = keys + ['Ex', 'Ey', 'Ez'] - dat['Ex'] = np.real(fm['Ex']) - dat['Ey'] = np.real(fm['Ey']) - dat['Ez'] = np.real(fm['Ez']) + keys = keys + ["Ex", "Ey", "Ez"] + dat["Ex"] = np.real(fm["Ex"]) + dat["Ey"] = np.real(fm["Ey"]) + dat["Ez"] = np.real(fm["Ez"]) else: - raise ValueError('Mixed static field TODO') - + raise ValueError("Mixed static field TODO") + else: - # Use internal Superfish routine - raise ValueError('Complex 3D Fields not implement yet!') - #keys = ['X', 'Y', 'Z', + ['Ex', 'Ey', 'Ez']'Bx', 'By', 'Bz'] - #dat['Er'], dat['Ez'], dat['Bphi'], _ = fish_complex_to_real_fields(fm, verbose=verbose) - - - # Flatten dat - gptdata = np.array([dat[k].flatten() for k in keys]).T - - # Write file. + # Use internal Superfish routine + raise ValueError("Complex 3D Fields not implement yet!") + # keys = ['X', 'Y', 'Z', + ['Ex', 'Ey', 'Ez']'Bx', 'By', 'Bz'] + # dat['Er'], dat['Ez'], dat['Bphi'], _ = fish_complex_to_real_fields(fm, verbose=verbose) + + # Flatten dat + gptdata = np.array([dat[k].flatten() for k in keys]).T + + # Write file. # Hack to delete final newline # https://stackoverflow.com/questions/28492954/numpy-savetxt-stop-newline-on-final-line - with open(outfile, 'w') as fout: - NEWLINE_SIZE_IN_BYTES = 1 # 2 on Windows? - np.savetxt(fout, gptdata, header=' '.join(keys), comments='') - fout.seek(0, os.SEEK_END) # Go to the end of the file. + with open(outfile, "w") as fout: + NEWLINE_SIZE_IN_BYTES = 1 # 2 on Windows? + np.savetxt(fout, gptdata, header=" ".join(keys), comments="") + fout.seek(0, os.SEEK_END) # Go to the end of the file. # Go backwards one byte from the end of the file. fout.seek(fout.tell() - NEWLINE_SIZE_IN_BYTES, os.SEEK_SET) - fout.truncate() # Truncate the file to this point. - - + fout.truncate() # Truncate the file to this point. + if asci2gdf_bin: run_asci2gdf(outfile, asci2gdf_bin, verbose=verbose) - elif verbose: - print(f'ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}') - - return outfile + elif verbose: + print( + f"ASCII field data written. Convert to GDF using: asci2df -o field.gdf {outfile}" + ) - + return outfile