Source code for fbpic.openpmd_diag.boosted_field_diag

# Copyright 2016, FBPIC contributors
# Authors: Remi Lehe, Manuel Kirchen
# License: 3-Clause-BSD-LBNL
"""
This file defines the class BackTransformedFieldDiagnostic

Major features:
- The class reuses the existing methods of FieldDiagnostic
  as much as possible, through class inheritance
- The class implements memory buffering of the slices, so as
  not to write to disk at every timestep
- Parallel output is not implemented for the moment
"""
import os
import numpy as np
from scipy.constants import c
from .field_diag import FieldDiagnostic

# Check if CUDA is available, then import CUDA functions
from fbpic.utils.cuda import cuda_installed
if cuda_installed:
    import cupy
    from fbpic.utils.cuda import cuda, cuda_tpb_bpg_1d, compile_cupy

[docs] class BackTransformedFieldDiagnostic(FieldDiagnostic): """ Class that writes the fields *in the lab frame*, from a simulation in the boosted frame """ def __init__(self, zmin_lab, zmax_lab, v_lab, dt_snapshots_lab, Ntot_snapshots_lab, gamma_boost, period, fldobject, comm=None, fieldtypes=["E", "B"], write_dir=None, t_min_snapshots_lab=0., t_max_snapshots_lab=np.inf): """ Initialize diagnostics that retrieve the data in the lab frame, as a series of snapshot (one file per snapshot), within a virtual moving window defined by zmin_lab, zmax_lab, v_lab. The parameters defined below are specific to the back-transformed diagnostics. See the documentation of `FieldDiagnostic` for the other parameters. Parameters ---------- zmin_lab: float (in meters) The position of the left edge of the virtual moving window, *in the lab frame*, at t=0 zmax_lab: float (in meters) The position of the right edge of the virtual moving window, *in the lab frame*, at t=0 v_lab: float (in m.s^-1) Speed of the moving window *in the lab frame* dt_snapshots_lab: float (in seconds) Time interval *in the lab frame* between two successive snapshots Ntot_snapshots_lab: int Total number of snapshots that this diagnostic will produce period: int Number of iterations for which the data is accumulated in memory, before finally writing it to the disk. fieldtypes : a list of strings, optional The strings are either "rho", "E", "B" or "J" and indicate which field should be written. Default : only "E" and "B" are written. This is because the backward Lorentz transform is not as precise for "rho" and "J" as for "E" and "B" (because "rho" and "J" are staggered in time). The user can still output "rho" and "J" by changing `fieldtypes`, but has to be aware that there may errors in the backward transform. Moreover, writing rho/J slows down the simulation, as these fields are then brought from spectral to real space, at each iteration. t_min_snapshots_lab, t_max_snapshots_lab: floats (seconds) The time *in the lab frame* between which snapshots are created (`t_min_snapshots_lab` is inclusive, `t_max_snapshots_lab` is exclusive). """ # Do not leave write_dir as None, as this may conflict with # the default directory ('./diags') in which diagnostics in the # boosted frame are written. if write_dir is None: write_dir='lab_diags' # Initialize the normal attributes of a FieldDiagnostic FieldDiagnostic.__init__(self, period, fldobject, comm, fieldtypes, write_dir) # Register the boost quantities self.gamma_boost = gamma_boost self.inv_gamma_boost = 1./gamma_boost self.beta_boost = np.sqrt( 1. - self.inv_gamma_boost**2 ) self.inv_beta_boost = 1./self.beta_boost # Find the z resolution and size of the diagnostic *in the lab frame* # (Needed to initialize metadata in the openPMD file) dz_lab = c*self.fld.dt * self.inv_beta_boost*self.inv_gamma_boost Nz = int( (zmax_lab - zmin_lab)/dz_lab ) + 1 self.inv_dz_lab = 1./dz_lab # Get number of radial cells in the output # (if possible, remove damp cells) if self.comm is None: Nr = self.fld.interp[0].Nr else: Nr = self.comm.get_Nr(with_damp=False) # Create the list of LabSnapshot objects self.snapshots = [] for i in range( Ntot_snapshots_lab ): t_lab = i * dt_snapshots_lab if t_lab >= t_min_snapshots_lab and t_lab < t_max_snapshots_lab: snapshot = LabSnapshot( t_lab, zmin_lab + v_lab*t_lab, zmax_lab + v_lab*t_lab, self.write_dir, i, self.fld, Nr ) self.snapshots.append( snapshot ) # Initialize a corresponding empty file self.create_file_empty_meshes( snapshot.filename, i, snapshot.t_lab, Nr, Nz, snapshot.zmin_lab, dz_lab, self.fld.dt) # Create a slice handler, which will do all the extraction, Lorentz # transformation, etc for each slice to be registered in a LabSnapshot self.slice_handler = SliceHandler(self.gamma_boost, self.beta_boost, Nr) def write( self, iteration ): """ Redefines the method write of the parent class FieldDiagnostic """ # At each timestep, store a slices of the fields in memory buffers self.store_snapshot_slices( iteration ) # Every self.period, write the buffered slices to disk if iteration % self.period == 0: self.flush_to_disk() def store_snapshot_slices( self, iteration ): """ Store slices of the fields in the memory buffers of the corresponding lab snapshots Parameter --------- iteration : int The current iteration in the boosted frame simulation """ # If needed: Bring rho/J from spectral space (where they where # smoothed/corrected) to real space if "rho" in self.fieldtypes or "J" in self.fieldtypes: # Get 'rho_prev', since it correspond to rho at time n self.fld.spect2interp('rho_prev') self.fld.spect2interp('J') # Exchange rho and J if needed if (self.comm is not None) and (self.comm.size > 1): if not self.fld.exchanged_source['J']: self.comm.exchange_fields(self.fld.interp, 'J', 'add') if not self.fld.exchanged_source['rho_prev']: self.comm.exchange_fields(self.fld.interp, 'rho', 'add') # Find the limits of the local subdomain at this iteration if self.comm is None: zmin_boost = self.fld.interp[0].zmin zmax_boost = self.fld.interp[0].zmax else: # If a communicator is provided, remove guard and damp cells zmin_boost, zmax_boost = self.comm.get_zmin_zmax( local=True, with_damp=False, with_guard=False, rank=self.rank ) # Extract the current time in the boosted frame time = iteration * self.fld.dt # Loop through the labsnapshots for snapshot in self.snapshots: # Update the positions of the output slice of this snapshot # in the lab and boosted frame (current_z_lab and current_z_boost) snapshot.update_current_output_positions( time, self.inv_gamma_boost, self.inv_beta_boost ) # For this snapshot: # - check if the output position *in the boosted frame* # is in the current local domain # - check if the output position *in the lab frame* # is within the lab-frame boundaries of the current snapshot if ( (snapshot.current_z_boost > zmin_boost) and \ (snapshot.current_z_boost < zmax_boost) and \ (snapshot.current_z_lab > snapshot.zmin_lab) and \ (snapshot.current_z_lab < snapshot.zmax_lab) ): # In this case, extract the proper slice from the field array, # and store the results into snapshot.slice_array # (when running on the GPU, snapshot.slice_array # is a device array) self.slice_handler.extract_slice( self.fld, self.comm, snapshot.current_z_boost, zmin_boost, snapshot.slice_array ) # Register snapshot.slice_array in the list of buffers # (when running on the GPU, the slice to the CPU) snapshot.register_slice( self.inv_dz_lab ) def flush_to_disk( self ): """ Writes the buffered slices of fields to the disk Erase the buffered slices of the LabSnapshot objects """ # Loop through the labsnapshots and flush the data for snapshot in self.snapshots: # Compact the successive slices that have been buffered # over time into a single array, on each proc field_array, iz_min, iz_max = snapshot.compact_slices() # Perform the Lorentz transformation of the field values # *from the boosted frame to the lab frame*, on each proc if field_array is not None: self.slice_handler.transform_fields_to_lab_frame( field_array ) # Gather the slices on the first proc if self.comm is not None and self.comm.size > 1: global_field_array, global_iz_min, global_iz_max = \ self.gather_slices( field_array, iz_min, iz_max ) else: global_field_array = field_array global_iz_min = iz_min global_iz_max = iz_max # First proc writes the global array to disk (if it is not empty) if (self.rank==0) and (global_field_array is not None): # Write to disk self.write_slices( global_field_array, global_iz_min, global_iz_max, snapshot, self.slice_handler.field_to_index) # Erase the memory buffers snapshot.buffered_slices = [] snapshot.buffer_z_indices = [] def gather_slices( self, field_array, iz_min, iz_max ): """ Stitch together the field_array of the different processors Parameters: ----------- field_array: ndarray of reals, or None If the local proc has no slice data, this is None Otherwise, it is an array of shape (10, 2*Nm-1, Nr, nslice_local) iz_min, iz_max: ints or None If the local proc has no slice data, this is None Otherwise, it corresponds to the indices at which the data should written, in final dataset which is on disk Returns: -------- A tuple with: global_field_array: an array of shape (10, 2*Nm-1, Nr, nslice_global), or None if none of the procs had any data global_izmin, global_izmax: the indices at which the global_field_array should be written (or None) """ # Gather objects into lists (one element per proc) # Note: this is slow, as it uses the generic mpi4py routines gather. # (This is because for some proc field_array can be None.) mpi_comm = self.comm.mpi_comm field_array_list = mpi_comm.gather( field_array ) iz_min_list = mpi_comm.gather( iz_min ) iz_max_list = mpi_comm.gather( iz_max ) # First proc: merge the results if self.rank == 0: # Check whether any processor had some slices no_slices = True for f_array in field_array_list: if f_array is not None: no_slices = False n_modes = f_array.shape[1] # n_modes is 2*Nm - 1 Nr = f_array.shape[2] # If there are no slices, set global quantities to None if no_slices: global_field_array = None global_iz_min = None global_iz_max = None # If there are some slices, gather them else: # Find the global iz_min and global iz_max global_iz_min = min([n for n in iz_min_list if n is not None]) global_iz_max = max([n for n in iz_max_list if n is not None]) # Allocate a the global field array, with the proper size nslice = global_iz_max - global_iz_min data_shape = (10, n_modes, Nr, nslice) global_field_array = np.zeros( data_shape ) # Loop through all the processors # Fit the field arrays one by one into the global_field_array for i_proc in range(self.comm.size): # If this proc has no data, skip it if field_array_list[ i_proc ] is None: continue # Longitudinal indices within the array global_field_array s_min = iz_min_list[ i_proc ] - global_iz_min s_max = iz_max_list[ i_proc ] - global_iz_min # Copy the array to the proper position global_field_array[:,:,:, s_min:s_max] = \ field_array_list[i_proc] # The first proc returns the result return( global_field_array, global_iz_min, global_iz_max ) # Other processors return a dummy placeholder else: return( None, None, None ) def write_slices( self, field_array, iz_min, iz_max, snapshot, f2i ): """ For one given snapshot, write the slices of the different fields to an openPMD file Parameters ---------- field_array: array of reals Array of shape (10, 2*Nm-1, Nr, nslices) which contains field values iz_min, iz_max: integers The indices between which the slices will be written iz_min is inclusice and iz_max is exclusive snapshot: a LabSnaphot object f2i: dict Dictionary of correspondance between the field names and the integer index in the field_array """ # Open the file without parallel I/O in this implementation f = self.open_file( snapshot.filename ) field_path = "/data/%d/fields/" %snapshot.iteration field_grp = f[field_path] # Loop over the different quantities that should be written for fieldtype in self.fieldtypes: # Scalar field if fieldtype == "rho": data = field_array[ f2i[ "rho" ] ] self.write_field_slices( field_grp, data, "rho", iz_min, iz_max ) # Vector field elif fieldtype in ["E", "B", "J"]: for coord in self.coords: quantity = "%s%s" %(fieldtype, coord) path = "%s/%s" %(fieldtype, coord) data = field_array[ f2i[ quantity ] ] self.write_field_slices( field_grp, data, path, iz_min, iz_max ) # Close the file f.close() def write_field_slices( self, field_grp, data, path, iz_min, iz_max ): """ Writes the slices of a given field into the openPMD file """ dset = field_grp[ path ] dset[:, :, iz_min:iz_max ] = data
class LabSnapshot: """ Class that stores data relative to one given snapshot in the lab frame (i.e. one given *time* in the lab frame) """ def __init__(self, t_lab, zmin_lab, zmax_lab, write_dir, i, fld, Nr_output): """ Initialize a LabSnapshot Parameters ---------- t_lab: float (seconds) Time of this snapshot *in the lab frame* zmin_lab, zmax_lab: floats Longitudinal limits of this snapshot write_dir: string Absolute path to the directory where the data for this snapshot is to be written i: int Number of the file where this snapshot is to be written fld: a Fields object This is passed only in order to determine how to initialize the slice_array buffer (either on the CPU or GPU) Nr_output: int Number of cells in the r direction, in the final output (This typically excludes the radial damping cells) """ # Deduce the name of the filename where this snapshot writes self.filename = os.path.join( write_dir, 'hdf5/data%08d.h5' %i) self.iteration = i # Time and boundaries in the lab frame (constants quantities) self.zmin_lab = zmin_lab self.zmax_lab = zmax_lab self.t_lab = t_lab # Positions where the fields are to be registered # (Change at every iteration) self.current_z_lab = 0 self.current_z_boost = 0 # Buffered field slice and corresponding array index in z self.buffered_slices = [] self.buffer_z_indices = [] # Allocate a buffer for only one slice (avoids having to # reallocate arrays when running on the GPU) data_shape = (10, 2*fld.Nm-1, Nr_output) if fld.use_cuda is False: self.slice_array = np.empty( data_shape ) else: self.slice_array = cupy.empty( data_shape ) def update_current_output_positions( self, t_boost, inv_gamma, inv_beta ): """ Update the positions of output for this snapshot, so that if corresponds to the time t_boost in the boosted frame Parameters ---------- t_boost: float (seconds) Time of the current iteration, in the boosted frame inv_gamma, inv_beta: floats Inverse of the Lorentz factor of the boost, and inverse of the corresponding beta. """ t_lab = self.t_lab # This implements the Lorentz transformation formulas, # for a snapshot having a fixed t_lab self.current_z_boost = ( t_lab*inv_gamma - t_boost )*c*inv_beta self.current_z_lab = ( t_lab - t_boost*inv_gamma )*c*inv_beta def register_slice( self, inv_dz_lab ): """ Store the slice of fields represented by self.slice_array and also store the z index at which this slice should be written in the final lab frame array Parameters ---------- inv_dz_lab: float Inverse of the grid spacing in z, *in the lab frame* """ # Find the index of the slice in the lab frame if self.buffer_z_indices == []: # No previous index: caculate it from the absolute z_lab iz_lab = int( (self.current_z_lab - self.zmin_lab)*inv_dz_lab ) else: # By construction, this index shoud be the previous index - 1 # Handling integers avoids unstable roundoff errors, when # self.current_z_lab is very close to zmin_lab + iz_lab*dz_lab iz_lab = self.buffer_z_indices[-1] - 1 # Store the array and the index self.buffer_z_indices.append( iz_lab ) # Make a copy of the array if it is directly on the CPU if type(self.slice_array) is np.ndarray: self.buffered_slices.append( self.slice_array.copy() ) # or copy from the GPU else: self.buffered_slices.append( self.slice_array.get() ) def compact_slices(self): """ Compact the successive slices that have been buffered over time into a single array, and return the indices at which this array should be written. Returns ------- field_array: an array of reals of shape (10, 2*Nm-1, Nr_output, nslices) In the above nslices is the number of buffered slices iz_min, iz_max: integers The indices between which the slices should be written (iz_min is inclusive, iz_max is exclusive) Returns None if the slices are empty """ # Return None if the slices are empty if len(self.buffer_z_indices) == 0: return( None, None, None ) # Check that the indices of the slices are contiguous # (This should be a consequence of the transformation implemented # in update_current_output_positions, and of the calculation # of inv_dz_lab.) iz_old = self.buffer_z_indices[0] for iz in self.buffer_z_indices[1:]: if iz != iz_old - 1: raise UserWarning('In the boosted frame diagnostic, ' 'the buffered slices are not contiguous in z.\n' 'The boosted frame diagnostics may be inaccurate.') break iz_old = iz # Pack the different slices together # Reverse the order of the slices when stacking the array, # since the slices where registered for right to left field_array = np.stack( self.buffered_slices[::-1], axis=-1 ) # Get the first and last index in z # (Following Python conventions, iz_min is inclusive, # iz_max is exclusive) iz_min = self.buffer_z_indices[-1] iz_max = self.buffer_z_indices[0] + 1 return( field_array, iz_min, iz_max ) class SliceHandler: """ Class that extracts, Lorentz-transforms and writes slices of the fields """ def __init__( self, gamma_boost, beta_boost, Nr_output ): """ Initialize the SliceHandler object Parameters ---------- gamma_boost, beta_boost: floats The Lorentz factor of the boost and the corresponding beta Nr_output: int Number of cells in the r direction, in the final output (This typically excludes the radial damping cells) """ # Store the arguments self.gamma_boost = gamma_boost self.beta_boost = beta_boost self.Nr_output = Nr_output # Create a dictionary that contains the correspondance # between the field names and array index self.field_to_index = {'Er':0, 'Et':1, 'Ez':2, 'Br':3, 'Bt':4, 'Bz':5, 'Jr':6, 'Jt':7, 'Jz':8, 'rho':9} def extract_slice( self, fld, comm, z_boost, zmin_boost, slice_array ): """ Fills `slice_array` with a slice of the fields at z_boost (the fields returned are still in the boosted frame ; for performance, the Lorentz transform of the fields values is performed only when flushing to disk) Parameters ---------- fld: a Fields object The object from which to extract the fields comm: a BoundaryCommunicator object Contains information about the gard cells in particular z_boost: float (meters) Position of the slice in the boosted frame zmin_boost: float (meters) Position of the left end of physical part of the local subdomain (i.e. excludes guard cells) slice_array: either a numpy array or a cuda device array An array of reals that packs together the slices of the different fields (always on array on the CPU). The first index of this array corresponds to the field type (10 different field types), and the correspondance between the field type and integer index is given field_to_index The shape of this arrays is (10, 2*Nm-1, Nr_output) """ # Find the index of the slice in the boosted frame # and the corresponding interpolation shape factor dz = fld.interp[0].dz # Find the interpolation data in the z direction z_staggered_gridunits = ( z_boost - zmin_boost - 0.5*dz )/dz iz = int( z_staggered_gridunits ) Sz = iz + 1 - z_staggered_gridunits # Add the guard cells to the index iz if comm is not None: iz += comm.n_guard if comm.left_proc is None: iz += comm.nz_damp+comm.n_inject # Extract the slice directly on the CPU # Fill the pre-allocated CPU array slice_array if fld.use_cuda is False : # Extract a slice of the fields *in the boosted frame* # at z_boost, using interpolation, and store them in slice_array self.extract_slice_cpu( fld, iz, Sz, slice_array ) # Extract the slice on the GPU # Fill the pre-allocated GPU array slice_array else: # Prepare kernel call dim_grid_1d, dim_block_1d = cuda_tpb_bpg_1d( self.Nr_output ) # Extract the slices interp = fld.interp for m in range(fld.Nm): extract_slice_cuda[ dim_grid_1d, dim_block_1d ]( self.Nr_output, iz, Sz, slice_array, interp[m].Er, interp[m].Et, interp[m].Ez, interp[m].Br, interp[m].Bt, interp[m].Bz, interp[m].Jr, interp[m].Jt, interp[m].Jz, interp[m].rho, m) def extract_slice_cpu( self, fld, iz, Sz, slice_array ): """ Extract a slice of the fields at iz and iz+1, and interpolated between those two points using Sz and (1-Sz) Parameters ---------- fld: a Fields object iz: int Index at which to extract the fields Sz: float Interpolation shape factor used at iz slice_array: np.ndarray Array of shape (10, 2*Nm-1, Nr_output ) """ # Shortcut for the correspondance between field and integer index f2i = self.field_to_index # Loop through the fields, and extract the proper slice for each field for quantity in self.field_to_index.keys(): # Here typical values for `quantity` are e.g. 'Er', 'Bz', 'rho' # Interpolate the centered field in z slice_array[ f2i[quantity], :, : ] = Sz*self.get_dataset( fld, quantity, iz_slice=iz ) slice_array[ f2i[quantity], :, : ] += (1.-Sz) * self.get_dataset( fld, quantity, iz_slice=iz+1 ) def get_dataset( self, fld, quantity, iz_slice ): """ Extract a given quantity, at a given slice index, from the fields Parameters ---------- fld: a Fields object quantity: string Indicates the quantity to be extracted (e.g. 'Er', 'Bz', 'rho') iz_slice: int Indicates the position of the slice in the field array Returns ------- An array of reals, whose format is close to the final openPMD output. In particular, the array of fields is of shape ( 2*Nm-1, Nr) (real and imaginary part are separated for each mode) """ # Shortcut Nr = self.Nr_output # Allocate the array to be returned data_shape = (2*fld.Nm-1, Nr) output_array = np.empty( data_shape ) # Get the mode 0 : only the real part is non-zero output_array[0,:] = getattr(fld.interp[0], quantity)[iz_slice,:Nr].real # Get the higher modes # There is a factor 2 here so as to comply with the convention in # Lifschitz et al., which is also the convention adopted in FBPIC for m in range(1,fld.Nm): higher_mode_slice = 2*getattr(fld.interp[m], quantity)[iz_slice,:Nr] output_array[2*m-1, :] = higher_mode_slice.real output_array[2*m, :] = higher_mode_slice.imag return( output_array ) def transform_fields_to_lab_frame( self, fields ): """ Modifies the array `fields` in place, to transform the field values from the boosted frame to the lab frame. The transformation is a transformation with -beta_boost, thus the corresponding formulas are: - for the transverse part of E and B: $\vec{E}_{lab} = \gamma(\vec{E} - c\vec{\beta} \times\vec{B})$ $\vec{B}_{lab} = \gamma(\vec{B} + \vec{\beta}/c \times\vec{E})$ - for rho and Jz: $\rho_{lab} = \gamma(\rho + \beta J_{z}/c)$ $J_{z,lab} = \gamma(J_z + c\beta \rho)$ Parameter --------- fields: array of floats An array that packs together the slices of the different fields. The shape of this arrays is (10, 2*Nm-1, Nr_output, nslices) where nslices is the number of slices that have been buffered """ # Some shortcuts gamma = self.gamma_boost cbeta = c*self.beta_boost beta_c = self.beta_boost/c # Shortcut to give the correspondance between field name # (e.g. 'Ex', 'rho') and integer index in the array f2i = self.field_to_index # Lorentz transformations # For E and B # (NB: Ez and Bz are unchanged by the Lorentz transform) # Use temporary arrays when changing Er and Bt in place er_lab = gamma*( fields[f2i['Er']] + cbeta * fields[f2i['Bt']] ) bt_lab = gamma*( fields[f2i['Bt']] + beta_c * fields[f2i['Er']] ) fields[ f2i['Er'], ... ] = er_lab fields[ f2i['Bt'], ... ] = bt_lab # Use temporary arrays when changing Et and Br in place et_lab = gamma*( fields[f2i['Et']] - cbeta * fields[f2i['Br']] ) br_lab = gamma*( fields[f2i['Br']] - beta_c * fields[f2i['Et']] ) fields[ f2i['Et'], ... ] = et_lab fields[ f2i['Br'], ... ] = br_lab # For rho and J # (NB: the transverse components of J are unchanged) # Use temporary arrays when changing rho and Jz in place rho_lab = gamma*( fields[f2i['rho']] + beta_c * fields[f2i['Jz']] ) Jz_lab = gamma*( fields[f2i['Jz']] + cbeta * fields[f2i['rho']] ) fields[ f2i['rho'], ... ] = rho_lab fields[ f2i['Jz'], ... ] = Jz_lab if cuda_installed: @compile_cupy def extract_slice_cuda( Nr, iz, Sz, slice_arr, Er, Et, Ez, Br, Bt, Bz, Jr, Jt, Jz, rho, m ): """ Extract a slice of the fields at iz and iz+1, and interpolated between those two points using Sz and (1-Sz) Parameters ---------- Nr: int Number of cells transversally iz: int Index at which to extract the fields Sz: float Interpolation shape factor used at iz slice_arr: cupy.empty Array of floats of shape (10, 2*Nm-1, Nr) Er, Et, etc...: cupy.empty Array of complexs of shape (Nz, Nr), for the azimuthal mode m m: int Index of the azimuthal mode involved """ # One thread per radial position ir = cuda.grid(1) # Intermediate variables izp = iz+1 Szp = 1. - Sz if ir < Nr: # Interpolate the field in the longitudinal direction # and store it into pre-packed arrays # For the higher modes: # There is a factor 2 here so as to comply with the convention in # Lifschitz et al., which is also the convention of FBPIC # For performance, this is included in the shape factor. if m > 0: Sz = 2*Sz Szp = 2*Szp # Index at which the mode should be added # in the array `slice_arr` im = 2*m-1 else: im = 0 # Real part slice_arr[0,im,ir] = Sz*Er[iz,ir].real + Szp*Er[izp,ir].real slice_arr[1,im,ir] = Sz*Et[iz,ir].real + Szp*Et[izp,ir].real slice_arr[2,im,ir] = Sz*Ez[iz,ir].real + Szp*Ez[izp,ir].real slice_arr[3,im,ir] = Sz*Br[iz,ir].real + Szp*Br[izp,ir].real slice_arr[4,im,ir] = Sz*Bt[iz,ir].real + Szp*Bt[izp,ir].real slice_arr[5,im,ir] = Sz*Bz[iz,ir].real + Szp*Bz[izp,ir].real slice_arr[6,im,ir] = Sz*Jr[iz,ir].real + Szp*Jr[izp,ir].real slice_arr[7,im,ir] = Sz*Jt[iz,ir].real + Szp*Jt[izp,ir].real slice_arr[8,im,ir] = Sz*Jz[iz,ir].real + Szp*Jz[izp,ir].real slice_arr[9,im,ir] = Sz*rho[iz,ir].real + Szp*rho[izp,ir].real if m > 0: # Imaginary part slice_arr[0,im+1,ir] = Sz*Er[iz,ir].imag + Szp*Er[izp,ir].imag slice_arr[1,im+1,ir] = Sz*Et[iz,ir].imag + Szp*Et[izp,ir].imag slice_arr[2,im+1,ir] = Sz*Ez[iz,ir].imag + Szp*Ez[izp,ir].imag slice_arr[3,im+1,ir] = Sz*Br[iz,ir].imag + Szp*Br[izp,ir].imag slice_arr[4,im+1,ir] = Sz*Bt[iz,ir].imag + Szp*Bt[izp,ir].imag slice_arr[5,im+1,ir] = Sz*Bz[iz,ir].imag + Szp*Bz[izp,ir].imag slice_arr[6,im+1,ir] = Sz*Jr[iz,ir].imag + Szp*Jr[izp,ir].imag slice_arr[7,im+1,ir] = Sz*Jt[iz,ir].imag + Szp*Jt[izp,ir].imag slice_arr[8,im+1,ir] = Sz*Jz[iz,ir].imag + Szp*Jz[izp,ir].imag slice_arr[9,im+1,ir] = Sz*rho[iz,ir].imag + Szp*rho[izp,ir].imag # Alias, for backward compatibility BoostedFieldDiagnostic = BackTransformedFieldDiagnostic