diff --git a/prorad.py b/prorad.py index 77f12c8..b961605 100755 --- a/prorad.py +++ b/prorad.py @@ -42,6 +42,7 @@ from scipy import interpolate import matplotlib as mpl import os + if not os.environ.has_key('DISPLAY'): mpl.use('Agg') import matplotlib.pyplot as plt @@ -50,35 +51,39 @@ import importlib from multiprocessing import Pool from timeit import default_timer as timer + sys.dont_write_bytecode = True # Physical constants in cgs -proton_charge = 4.8032e-10 # statcoul -charge_SI = 1.6022e-19 # coulombs -proton_mass = 1836.2*9.1094e-28 # g -c = 2.9979e10 # cm/sec -c_SI = c*1e-2 -E_SItoCGS = 1e6/c +proton_charge = 4.8032e-10 # statcoul +charge_SI = 1.6022e-19 # coulombs +proton_mass = 1836.2 * 9.1094e-28 # g +c = 2.9979e10 # cm/sec +c_SI = c * 1e-2 +E_SItoCGS = 1e6 / c MeV_to_ergs = 1.602e-6 MeV_to_Joules = 1.602e-13 # Number of field values stored at each grid point (Ex,Ey,Ez,Bx,By,Bz,nu,Z,N) NUM_FIELDS = 9 + def main(): # Load input parameters - try: import_params(sys.argv[1]) + try: + import_params(sys.argv[1]) except (ImportError, IndexError): - print("ERROR: You must supply as a command line argument the name of a valid input parameter file contained in 'inputs/'") + print( + "ERROR: You must supply as a command line argument the name of a valid input parameter file contained in 'inputs/'") print(" (omit the path and '.py' extension)") return # Load in the grid - grid = load_grid(fformat=params.fformat,fname=params.fname) + grid = load_grid(fformat=params.fformat, fname=params.fname) if grid is None: quit() - + # Push the protons through the grid to obtain final locations on film plane - film_x,film_y,Ep,traces = push_protons(grid) + film_x, film_y, Ep, traces = push_protons(grid) # Default plotting options plot_fluence = True @@ -87,22 +92,33 @@ def main(): save_images = False # Overwrite plotting options with any that are specified in the input file - try: plot_fluence = params.plot_fluence - except AttributeError: pass - try: plot_traces = params.plot_traces - except AttributeError: pass - try: plot_quiver = params.plot_quiver - except AttributeError: pass - try: save_images = params.save_images - except AttributeError: pass - - cbar, fluence_fig, traces_fig, quiver_fig = plot_results(grid, film_x, film_y, Ep, traces, plot_fluence=plot_fluence, plot_traces=plot_traces, plot_quiver=plot_quiver, save_images=save_images) + try: + plot_fluence = params.plot_fluence + except AttributeError: + pass + try: + plot_traces = params.plot_traces + except AttributeError: + pass + try: + plot_quiver = params.plot_quiver + except AttributeError: + pass + try: + save_images = params.save_images + except AttributeError: + pass + + cbar, fluence_fig, traces_fig, quiver_fig = plot_results(grid, film_x, film_y, Ep, traces, + plot_fluence=plot_fluence, plot_traces=plot_traces, + plot_quiver=plot_quiver, save_images=save_images) plt.show() # If the user has defined 'fileout', save final x and y proton positions to a file with that name in 'outputs' directory. try: - save_results(film_x,film_y,params.fileout) - except AttributeError: pass + save_results(film_x, film_y, params.fileout) + except AttributeError: + pass def import_params(inputfile): @@ -111,55 +127,48 @@ def import_params(inputfile): Normally, this will be called by main and passed whatever input name is supplied by the user on the command line. If prorad is imported in an interactive python shell, the user must call this functon to - import the desired input parameter file. - + import the desired input parameter file. Parameters _________ inputfile : string - Name of an input file contained in the 'inputs' directory, + Name of an input file contained in the 'inputs' directory, omitting the '.py' or '.pyc' suffix. - """ - globals()['params'] = importlib.import_module('.'+inputfile, 'inputs') + globals()['params'] = importlib.import_module('.' + inputfile, 'inputs') class Grid(object): """ - - Stores field values on a grid in 3 dimensions. + Stores field values on a grid in 3 dimensions. Parameters ---------- gridvals : 4D ndarray with shape (nx,ny,nz,NUM_FIELDS) All 9 field values (Ex,Ey,Ez,Bx,By,Bz,nu,Z,N) defined on a 3D grid - gridspacings : tuple of floats (dx,dy,dz) or (dR,dTheta,dz) Spacing between grid points along each dimension - gridcorner : tuple of floats (xoffset,yoffset,zoffset) Spatial coordinates of where the corner of the grid (gridvals[0,0,0,:]) is located (typically (-lx/2,-ly/2,0.0) for cartesian grid, (0.0,0.0,0.0) for cylindrical) - gridextent : tuple of floats (lx, ly, lz) Length, width and depth of grid in Cartesian coordinates. For cylindrical grid, lx and ly should be at least the diameter. - """ + def __init__(self, gridvals, gridspacings, gridcorner, gridextent, cyl_coords=False): self.vals = gridvals self.xoffset, self.yoffset, self.zoffset = gridcorner self.dx, self.dy, self.dz = gridspacings - self.nx = len(gridvals[:,0,0,0]) - self.ny = len(gridvals[0,:,0,0]) - self.nz = len(gridvals[0,0,:,0]) + self.nx = len(gridvals[:, 0, 0, 0]) + self.ny = len(gridvals[0, :, 0, 0]) + self.nz = len(gridvals[0, 0, :, 0]) self.lx, self.ly, self.lz = gridextent self.cyl_coords = cyl_coords -def load_grid(fformat=None,fname=None,ngridx=None,ngridy=None,ngridz=None): +def load_grid(fformat=None, fname=None, ngridx=None, ngridy=None, ngridz=None): """ Produce a Grid object using either analytic definitions or the output file of an MHD simulation. - Parameters ---------- fformat : string @@ -176,85 +185,101 @@ def load_grid(fformat=None,fname=None,ngridx=None,ngridy=None,ngridz=None): ngridz : int, optional Desired number of grid points along z dimension. Must be either passed as an argument OR defined in the input file. Can be left undefined for FLASH grids. - Returns ------- grid : Grid object populated with field values, grid spacings, grid offset, and grid dimensions. - """ start_time = timer() - + # If grid dimensions not passed as arguments, try to get them from input file if ngridx is None: - try: ngridx = params.ngridx - except AttributeError: pass + try: + ngridx = params.ngridx + except AttributeError: + pass if ngridy is None: - try: ngridy = params.ngridy - except AttributeError: pass + try: + ngridy = params.ngridy + except AttributeError: + pass if ngridz is None: - try: ngridz = params.ngridz - except AttributeError: pass + try: + ngridz = params.ngridz + except AttributeError: + pass grid = None if fformat is None: - try: fformat = params.fformat + try: + fformat = params.fformat except AttributeError: print("ERROR: No file format supplied. Aborting.") return None if fname is None and (fformat != 'analytic_grid' and fformat != 'LSP'): - try: fname = params.fname + try: + fname = params.fname except AttributeError: print("ERROR: File name required for fformat='" + fformat + "'. Aborting.") - if fformat == "analytic_grid": # Load user-defined analytic fields - - print("Generating "+str(ngridx)+"x"+str(ngridy)+"x"+str(ngridz)+" grid from analytic fields...") - + + print("Generating " + str(ngridx) + "x" + str(ngridy) + "x" + str(ngridz) + " grid from analytic fields...") + cyl_coords = False - try: cyl_coords = params.cyl_coords - except AttributeError: pass + try: + cyl_coords = params.cyl_coords + except AttributeError: + pass lx, ly, lz = params.lx, params.ly, params.lz xoffset, yoffset, zoffset = params.gridcorner - X_coords = np.linspace(xoffset, xoffset+lx, ngridx) - Y_coords = np.linspace(yoffset, yoffset+ly, ngridy) - Z_coords = np.linspace(zoffset, zoffset+lz, ngridz) - - # Which grid indices to populate. Omitted indices will be left as zeros. + X_coords = np.linspace(xoffset, xoffset + lx, ngridx) + Y_coords = np.linspace(yoffset, yoffset + ly, ngridy) + Z_coords = np.linspace(zoffset, zoffset + lz, ngridz) + + # Which grid indices to populate. Omitted indices will be left as zeros. # All indices are populated unless specified otherwise in the input file. # Leaving out indices that the user knows will have zero fields can speed things up. field_xrange_idx = range(ngridx) field_yrange_idx = range(ngridy) field_zrange_idx = range(ngridz) - - try: field_xrange_idx = params.field_xrange_idx - except AttributeError: pass - try: field_yrange_idx = params.field_yrange_idx - except AttributeError: pass - try: field_zrange_idx = params.field_zrange_idx - except AttributeError: pass - - gridvals = np.zeros((ngridx,ngridy,ngridz,NUM_FIELDS)) - - try: + + try: + field_xrange_idx = params.field_xrange_idx + except AttributeError: + pass + try: + field_yrange_idx = params.field_yrange_idx + except AttributeError: + pass + try: + field_zrange_idx = params.field_zrange_idx + except AttributeError: + pass + + gridvals = np.zeros((ngridx, ngridy, ngridz, NUM_FIELDS)) + + try: # If grid_nthreads is defined in input file, initialize the specified grid elements in parallel. - # Using Python's multiprocessing library rather than mpi, so actual parallelism is currently limited + # Using Python's multiprocessing library rather than mpi, so actual parallelism is currently limited # to the number of processors on a single node. p = Pool(params.grid_nthreads) print("Using " + str(params.grid_nthreads) + " threads to initialize grid") - coords = np.meshgrid(X_coords[field_xrange_idx],Y_coords[field_yrange_idx],Z_coords[field_zrange_idx],indexing='ij') - coords = itertools.product(X_coords[field_xrange_idx],Y_coords[field_yrange_idx],Z_coords[field_zrange_idx]) - initialized_vals = np.array(p.map(params.fields,coords)) - initialized_vals = initialized_vals.reshape(len(field_xrange_idx),len(field_yrange_idx),len(field_zrange_idx),NUM_FIELDS) - idx1,idx2,idx3 = np.meshgrid(field_xrange_idx,field_yrange_idx,field_zrange_idx,indexing='ij') - gridvals[idx1,idx2,idx3] = initialized_vals - except AttributeError: + coords = np.meshgrid(X_coords[field_xrange_idx], Y_coords[field_yrange_idx], Z_coords[field_zrange_idx], + indexing='ij') + coords = itertools.product(X_coords[field_xrange_idx], Y_coords[field_yrange_idx], + Z_coords[field_zrange_idx]) + initialized_vals = np.array(p.map(params.fields, coords)) + initialized_vals = initialized_vals.reshape(len(field_xrange_idx), len(field_yrange_idx), + len(field_zrange_idx), NUM_FIELDS) + idx1, idx2, idx3 = np.meshgrid(field_xrange_idx, field_yrange_idx, field_zrange_idx, indexing='ij') + gridvals[idx1, idx2, idx3] = initialized_vals + except: print("'params.grid_nthreads' not specified. Initializing grid in serial.") for i in field_xrange_idx: for j in field_yrange_idx: @@ -262,11 +287,11 @@ def load_grid(fformat=None,fname=None,ngridx=None,ngridy=None,ngridz=None): x = X_coords[i] y = Y_coords[j] z = Z_coords[k] - gridvals[i,j,k,:] = params.fields((x,y,z)) + gridvals[i, j, k, :] = params.fields((x, y, z)) - gridspacings = (lx/ngridx, ly/ngridy, lz/ngridz) - - grid = Grid(gridvals, gridspacings, (xoffset,yoffset,zoffset), (lx,ly,lz), cyl_coords=cyl_coords) + gridspacings = (lx / ngridx, ly / ngridy, lz / ngridz) + + grid = Grid(gridvals, gridspacings, (xoffset, yoffset, zoffset), (lx, ly, lz), cyl_coords=cyl_coords) elif fformat == "FLASH": print("Loading FLASH grid...") @@ -284,22 +309,22 @@ def load_grid(fformat=None,fname=None,ngridx=None,ngridy=None,ngridz=None): magx = uniform_data[u'magx'].in_cgs().to_ndarray() magy = uniform_data[u'magy'].in_cgs().to_ndarray() magz = uniform_data[u'magz'].in_cgs().to_ndarray() - + right_edge = ds.domain_right_edge.in_cgs().to_ndarray() left_edge = ds.domain_left_edge.in_cgs().to_ndarray() - gridspacings = (right_edge-left_edge)/magx.shape + gridspacings = (right_edge - left_edge) / magx.shape ngridx, ngridy, ngridz = magx.shape - - gridvals = np.zeros((ngridx,ngridy,ngridz,NUM_FIELDS)) + + gridvals = np.zeros((ngridx, ngridy, ngridz, NUM_FIELDS)) for i in range(ngridx): for j in range(ngridy): for k in range(ngridz): # TODO: Calculate electric fields too - gridvals[i,j,k,3:6] = [magx[i,j,k], magy[i,j,k], magz[i,j,k]] - - grid = Grid(gridvals, gridspacings, left_edge, right_edge-left_edge) + gridvals[i, j, k, 3:6] = [magx[i, j, k], magy[i, j, k], magz[i, j, k]] + + grid = Grid(gridvals, gridspacings, left_edge, right_edge - left_edge) elif fformat == "HYDRA": print("Loading HYDRA grid...") @@ -311,64 +336,65 @@ def load_grid(fformat=None,fname=None,ngridx=None,ngridy=None,ngridz=None): print("ERROR: You need to install and configure Yorick and the hyddjs tools to load HYDRA grids.") print("Contact Dave Strozzi (strozzi2@llnl.gov) for info.") return None - + H = stdY.h2open(fname) - varL = ['x', 'y', 'z', 'Bx', 'By', 'Bz', 'p', 'eden', 'zb','ireg','tmat'] - stg = DJH.h2stgchk(H,varL) + varL = ['x', 'y', 'z', 'Bx', 'By', 'Bz', 'p', 'eden', 'zb', 'ireg', 'tmat'] + stg = DJH.h2stgchk(H, varL) x = stg['x'][0] y = stg['y'][0] z = stg['z'][0] - #zb = stg['zb'][0] + # zb = stg['zb'][0] xmin, xmax = 0.0, np.amax(x) zmin, zmax = np.amin(z), np.amax(z) try: - xmin,xmax = params.hyd_xrange - except AttributeError: pass + xmin, xmax = params.hyd_xrange + except AttributeError: + pass try: - zmin,zmax = params.hyd_zrange - except AttributeError: pass + zmin, zmax = params.hyd_zrange + except AttributeError: + pass - R = np.linspace(0, xmax, ngridx) - Theta = np.linspace(0, 2*np.pi-2*np.pi/ngridy, ngridy) + Theta = np.linspace(0, 2 * np.pi - 2 * np.pi / ngridy, ngridy) Z = np.linspace(zmin, zmax, ngridz) dz = Z[1] - Z[0] dR = R[1] - R[0] dTheta = Theta[1] - Theta[0] - xz_vals = np.zeros((ngridx,ngridz,NUM_FIELDS)) - X_2D, Z_2D = np.meshgrid(R,Z,indexing='ij') - + xz_vals = np.zeros((ngridx, ngridz, NUM_FIELDS)) + X_2D, Z_2D = np.meshgrid(R, Z, indexing='ij') + # Calculate electric field from electron temperature and density gradient - edens = DJH.h2interp(H,stg,'eden',X_2D,Z_2D)*1e6 # m^-3 - epres = DJH.h2interp(H,stg,'p',X_2D,Z_2D)*1e11 # Pa - #etemp = DJH.h2interp(H,stg,'tmat',X_2D,Z_2D)*1e-3 # eV - np.seterr(divide='ignore',invalid='ignore') # Temporarily ignore errors because we could have divide by zeros - #Exz = (etemp/proton_charge)*np.nan_to_num(np.divide(np.gradient(epres,dR,dz),epres)) - Exz = -np.nan_to_num(np.divide(np.gradient(epres,dR/100.0,dz/100.0),charge_SI*edens))*E_SItoCGS - np.seterr(divide='warn',invalid='warn') # Restore to normal error behavior - - xz_vals[:,:,0], xz_vals[:,:,2] = Exz - xz_vals[:,:,3] = DJH.h2interp(H,stg,'Bx',X_2D,Z_2D) - xz_vals[:,:,4] = DJH.h2interp(H,stg,'By',X_2D,Z_2D) - xz_vals[:,:,5] = DJH.h2interp(H,stg,'Bz',X_2D,Z_2D) - #xz_vals[:,:,7] = DJH.h2interp(H,stg,'zb',X_2D,Z_2D) - - gridvals = np.zeros((ngridx,ngridy,ngridz,NUM_FIELDS)) + edens = DJH.h2interp(H, stg, 'eden', X_2D, Z_2D) * 1e6 # m^-3 + epres = DJH.h2interp(H, stg, 'p', X_2D, Z_2D) * 1e11 # Pa + # etemp = DJH.h2interp(H,stg,'tmat',X_2D,Z_2D)*1e-3 # eV + np.seterr(divide='ignore', invalid='ignore') # Temporarily ignore errors because we could have divide by zeros + # Exz = (etemp/proton_charge)*np.nan_to_num(np.divide(np.gradient(epres,dR,dz),epres)) + Exz = -np.nan_to_num(np.divide(np.gradient(epres, dR / 100.0, dz / 100.0), charge_SI * edens)) * E_SItoCGS + np.seterr(divide='warn', invalid='warn') # Restore to normal error behavior + + xz_vals[:, :, 0], xz_vals[:, :, 2] = Exz + xz_vals[:, :, 3] = DJH.h2interp(H, stg, 'Bx', X_2D, Z_2D) + xz_vals[:, :, 4] = DJH.h2interp(H, stg, 'By', X_2D, Z_2D) + xz_vals[:, :, 5] = DJH.h2interp(H, stg, 'Bz', X_2D, Z_2D) + # xz_vals[:,:,7] = DJH.h2interp(H,stg,'zb',X_2D,Z_2D) + + gridvals = np.zeros((ngridx, ngridy, ngridz, NUM_FIELDS)) for j in range(ngridy): theta = Theta[j] - gridvals[:,j,:,0] = xz_vals[:,:,0]*np.cos(theta) - gridvals[:,j,:,1] = xz_vals[:,:,0]*np.sin(theta) - gridvals[:,j,:,2] = xz_vals[:,:,2] - gridvals[:,j,:,3] = xz_vals[:,:,3]*np.cos(theta) - xz_vals[:,:,4]*np.sin(theta) - gridvals[:,j,:,4] = xz_vals[:,:,3]*np.sin(theta) + xz_vals[:,:,4]*np.cos(theta) - gridvals[:,j,:,5] = xz_vals[:,:,5] - - grid = Grid(gridvals, (dR,dTheta,dz), (0.0,0.0,0.0), (2*xmax,2*xmax,zmax-zmin), cyl_coords=True) + gridvals[:, j, :, 0] = xz_vals[:, :, 0] * np.cos(theta) + gridvals[:, j, :, 1] = xz_vals[:, :, 0] * np.sin(theta) + gridvals[:, j, :, 2] = xz_vals[:, :, 2] + gridvals[:, j, :, 3] = xz_vals[:, :, 3] * np.cos(theta) - xz_vals[:, :, 4] * np.sin(theta) + gridvals[:, j, :, 4] = xz_vals[:, :, 3] * np.sin(theta) + xz_vals[:, :, 4] * np.cos(theta) + gridvals[:, j, :, 5] = xz_vals[:, :, 5] + + grid = Grid(gridvals, (dR, dTheta, dz), (0.0, 0.0, 0.0), (2 * xmax, 2 * xmax, zmax - zmin), cyl_coords=True) elif fformat == "LSP": print("Loading LSP grid...") @@ -377,81 +403,82 @@ def load_grid(fformat=None,fname=None,ngridx=None,ngridy=None,ngridz=None): except ImportError: print("ERROR: You need the xdrlib tool to read LSP grids. Contact Drew Higginson for access.") return None - - FILE = lsp.flds(params.lsp_dirname,step=int(params.lsp_step)) - (X,Y,Z,t) = FILE.get_XYZt() - - (Ex,Name,Unit) = FILE.get_VarNameUnit(name='E',xyz='x') - (Ey,Name,Unit) = FILE.get_VarNameUnit(name='E',xyz='y') - (Ez,Name,Unit) = FILE.get_VarNameUnit(name='E',xyz='z') - (Bx,Name,Unit) = FILE.get_VarNameUnit(name='B',xyz='x') - (By,Name,Unit) = FILE.get_VarNameUnit(name='B',xyz='y') - (Bz,Name,Unit) = FILE.get_VarNameUnit(name='B',xyz='z') + + FILE = lsp.flds(params.lsp_dirname, step=int(params.lsp_step)) + (X, Y, Z, t) = FILE.get_XYZt() + + (Ex, Name, Unit) = FILE.get_VarNameUnit(name='E', xyz='x') + (Ey, Name, Unit) = FILE.get_VarNameUnit(name='E', xyz='y') + (Ez, Name, Unit) = FILE.get_VarNameUnit(name='E', xyz='z') + (Bx, Name, Unit) = FILE.get_VarNameUnit(name='B', xyz='x') + (By, Name, Unit) = FILE.get_VarNameUnit(name='B', xyz='y') + (Bz, Name, Unit) = FILE.get_VarNameUnit(name='B', xyz='z') ngridx = len(X) ngridy = len(Y) ngridz = len(Z) - gridvals = np.zeros((ngridx,ngridy,ngridz,NUM_FIELDS)) - - dx = X[1]-X[0] - dz = Z[1]-Z[0] - + gridvals = np.zeros((ngridx, ngridy, ngridz, NUM_FIELDS)) + + dx = X[1] - X[0] + dz = Z[1] - Z[0] + if ngridy == 1: # 2D grid - + # Extrude distance dy based on the wavelength of periodic features in x # (probably will not want to do this in general) - By_fft = np.fft.fft(By[:,0]) - x_periods = 1.0/np.fft.fftfreq(len(By_fft),dx) + By_fft = np.fft.fft(By[:, 0]) + x_periods = 1.0 / np.fft.fftfreq(len(By_fft), dx) dy = abs(x_periods[np.argmax(By_fft)]) - #gridvals[:,0,:,0] = Ex*1e5*E_SItoCGS - #gridvals[:,0,:,1] = Ey*1e5*E_SItoCGS - #gridvals[:,0,:,2] = Ez*1e5*E_SItoCGS - gridvals[:,0,:,3] = Bx - gridvals[:,0,:,4] = By - gridvals[:,0,:,5] = Bz - + # gridvals[:,0,:,0] = Ex*1e5*E_SItoCGS + # gridvals[:,0,:,1] = Ey*1e5*E_SItoCGS + # gridvals[:,0,:,2] = Ez*1e5*E_SItoCGS + gridvals[:, 0, :, 3] = Bx + gridvals[:, 0, :, 4] = By + gridvals[:, 0, :, 5] = Bz + # If params.lsp_ntile is defined, tile the grid that number of times in x and z try: - gridvals = np.tile(gridvals,(params.lsp_ntile,1,params.lsp_ntile,1)) - except AttributeError: pass + gridvals = np.tile(gridvals, (params.lsp_ntile, 1, params.lsp_ntile, 1)) + except AttributeError: + pass else: # 3D grid # TODO: Test 3D LSP grid - dy = Y[1]-Y[0] - - #gridvals[:,:,:,0] = Ex*1e5*E_SItoCGS - #gridvals[:,:,:,1] = Ey*1e5*E_SItoCGS - #gridvals[:,:,:,2] = Ez*1e5*E_SItoCGS - gridvals[:,:,:,3] = Bx - gridvals[:,:,:,4] = By - gridvals[:,:,:,5] = Bz - - lx = dx*len(gridvals[:,0,0,0]) - ly = dy*len(gridvals[0,:,0,0]) - lz = dz*len(gridvals[0,0,:,0]) - - grid = Grid(gridvals, (dx,dy,dz), (-lx/2.0,-ly/2.0,0.0), (lx,ly,lz)) - + dy = Y[1] - Y[0] + + # gridvals[:,:,:,0] = Ex*1e5*E_SItoCGS + # gridvals[:,:,:,1] = Ey*1e5*E_SItoCGS + # gridvals[:,:,:,2] = Ez*1e5*E_SItoCGS + gridvals[:, :, :, 3] = Bx + gridvals[:, :, :, 4] = By + gridvals[:, :, :, 5] = Bz + + lx = dx * len(gridvals[:, 0, 0, 0]) + ly = dy * len(gridvals[0, :, 0, 0]) + lz = dz * len(gridvals[0, 0, :, 0]) + + grid = Grid(gridvals, (dx, dy, dz), (-lx / 2.0, -ly / 2.0, 0.0), (lx, ly, lz)) + else: - print('"'+fformat+'"'+'is not a recognized file format. Aborting.') + print('"' + fformat + '"' + 'is not a recognized file format. Aborting.') return None end_time = timer() - print("Time elapsed during grid generation: " + str(end_time-start_time) + " s") + print("Time elapsed during grid generation: " + str(end_time - start_time) + " s") if grid.cyl_coords: print("Grid dR, dTheta, dz: " + str(grid.dx) + " cm, " + str(grid.dy) + " rad, " + str(grid.dz) + " cm") print("Grid nR, nTheta, nz: " + str(grid.nx) + ", " + str(grid.ny) + ", " + str(grid.nz)) - print("Grid lR, lTheta, lz: " + str(grid.lx)+" cm, " + str(grid.ly) + " rad, " + str(grid.lz) + " cm") + print("Grid lR, lTheta, lz: " + str(grid.lx) + " cm, " + str(grid.ly) + " rad, " + str(grid.lz) + " cm") else: print("Grid dx, dy, dz: " + str(grid.dx) + " cm, " + str(grid.dy) + " cm, " + str(grid.dz) + " cm") print("Grid nx, ny, nz: " + str(grid.nx) + ", " + str(grid.ny) + ", " + str(grid.nz)) - print("Grid lx, ly, lz: " + str(grid.lx)+" cm, " + str(grid.ly) + " cm, " + str(grid.lz) + " cm") + print("Grid lx, ly, lz: " + str(grid.lx) + " cm, " + str(grid.ly) + " cm, " + str(grid.lz) + " cm") return grid @@ -459,12 +486,10 @@ def load_grid(fformat=None,fname=None,ngridx=None,ngridy=None,ngridz=None): def push_protons(grid): """ Initialize arrays of proton positions, velocities, and energies, then simulate proton propogation through the grid. - Parameters ---------- grid : Grid Grid object produced by load_grid() - Returns ------- film_x : array of size NP @@ -472,32 +497,31 @@ def push_protons(grid): film_y : array of size NP 'y' positions of protons on film plane relative to film axes (in cm) Ep : array of size NP - Energies of protons at image plane (in MeV). + Energies of protons at image plane (in MeV). traces : ndarray of shape (ntraces, nz_prot+2, 3) Record of particle traces for random subset of protons. First dimension is the proton being tracked, second dimension is the index of the step, third dimension is the coordinates (x,y,z) at each location - """ NP = params.NP ntraces = params.ntraces nsteps = params.nsteps - + # Array to store trajectories of a subset of the protons. - # Three extra indices for when the protons are at the source, at the beginning of the grid, + # Three extra indices for when the protons are at the source, at the beginning of the grid, # and at the film. The 3 is for x,y,z coords of each proton. - traces = np.zeros((ntraces, nsteps+3, 3)) + traces = np.zeros((ntraces, nsteps + 3, 3)) mass = proton_mass charge = proton_charge - try: + try: mass = params.particle_mass print('Using user-defined particle mass.') except AttributeError: pass - try: + try: charge = params.particle_charge print('Using user-defined particle charge.') except AttributeError: @@ -507,14 +531,14 @@ def push_protons(grid): # The user can choose to supply their own initial x,y,z positions and velocities. # (can supply either all of these or none of them, but not just some) # If these are provided, don't need to define params.source_loc, params.prop_dir, - # params.l_s2start, params.E0, params.spread_angle, or params.r_source. Note that - # protons should all be moving in a roughly similar direction, otherwise - # large errors will be introduced because of how the equations of motion are + # params.l_s2start, params.E0, params.spread_angle, or params.r_source. Note that + # protons should all be moving in a roughly similar direction, otherwise + # large errors will be introduced because of how the equations of motion are # implemented in the pusher. - x,y,z = params.x, params.y, params.z - vx,vy,vz = params.vx, params.vy, params.vz + x, y, z = params.x, params.y, params.z + vx, vy, vz = params.vx, params.vy, params.vz - prop_dir = np.array([np.average(vx),np.average(vy),np.average(vz)]) + prop_dir = np.array([np.average(vx), np.average(vy), np.average(vz)]) prop_dir /= linalg.norm(prop_dir) print('Using user-defined initial positions and velocities.') @@ -524,140 +548,145 @@ def push_protons(grid): start_time = timer() # Protons "start" at source (meaning velocities are set as if they did) - traces[:,0,0:3] = params.source_loc + traces[:, 0, 0:3] = params.source_loc - v0 = np.sqrt(2.0*params.E0*MeV_to_ergs/mass) + v0 = np.sqrt(2.0 * params.E0 * MeV_to_ergs / mass) # 3D position of proton source source_loc = np.array(params.source_loc) - + # unit vector defining propogation direction - prop_dir = np.array(params.prop_dir)/linalg.norm(params.prop_dir) + prop_dir = np.array(params.prop_dir) / linalg.norm(params.prop_dir) # Distance from source to where protons actually start their propogation l_s2start = params.l_s2start - + # Angle between prop_dir and the outer edge of the cone defining the initial proton spread - phi_max = params.spread_angle*np.pi/180.0 + phi_max = params.spread_angle * np.pi / 180.0 - # Together, theta and phi give an isotropic distribution of protons on the surface of a - # unit sphere sector of angle phi_max. phis are polar angles from the distribution + # Together, theta and phi give an isotropic distribution of protons on the surface of a + # unit sphere sector of angle phi_max. phis are polar angles from the distribution # center, and thetas are azimuthal angles. - phis = np.arccos((1-np.cos(phi_max))*np.random.random_sample(NP) + np.cos(phi_max)) - thetas = 2*np.pi*np.random.random_sample(NP) + phis = np.arccos((1 - np.cos(phi_max)) * np.random.random_sample(NP) + np.cos(phi_max)) + thetas = 2 * np.pi * np.random.random_sample(NP) # Initialize at a starting plane as if source is at [0,0,0] pointing along z axis - positions = l_s2start*(np.array([0,0,1])[:,np.newaxis] + np.tan(phis)*np.array([np.cos(thetas),np.sin(thetas),np.zeros(NP)])) + positions = l_s2start * (np.array([0, 0, 1])[:, np.newaxis] + np.tan(phis) * np.array( + [np.cos(thetas), np.sin(thetas), np.zeros(NP)])) # Then rotate to correct direction and add in actual source position - rot_axis = np.cross([0,0,1],prop_dir) + rot_axis = np.cross([0, 0, 1], prop_dir) if linalg.norm(rot_axis) == 0: - rot_axis = [0,1,0] + rot_axis = [0, 1, 0] rot_axis /= linalg.norm(rot_axis) - rot_angle = np.arccos(np.dot([0,0,1],prop_dir)) + rot_angle = np.arccos(np.dot([0, 0, 1], prop_dir)) # New favorite way to define a rotation matrix: - rot_matrix = linalg.expm(np.cross(np.eye(3), rot_axis*rot_angle)) - positions = np.dot(rot_matrix,positions) - positions += source_loc[:,np.newaxis] - + rot_matrix = linalg.expm(np.cross(np.eye(3), rot_axis * rot_angle)) + positions = np.dot(rot_matrix, positions) + positions += source_loc[:, np.newaxis] + # Calculate velocity vector based on starting position - s2pos = positions-source_loc[:,np.newaxis] - velocities = v0*s2pos/np.linalg.norm(s2pos,axis=0) + s2pos = positions - source_loc[:, np.newaxis] + velocities = v0 * s2pos / np.linalg.norm(s2pos, axis=0) + + x, y, z = positions + vx, vy, vz = velocities - x,y,z = positions - vx,vy,vz = velocities - try: - # Add random noise to account for finite source. + # Add random noise to account for finite source. # NOTE: this means that protons might not necessarily align with # the starting end ending planes on either side of the grid, so # be careful if you have, for example, a very thin sheet of metal # at one end of your grid, that your l_prop will still push them # all through it. r_source = params.r_source - x += (2*np.random.random_sample(NP)-1)*r_source - y += (2*np.random.random_sample(NP)-1)*r_source - z += (2*np.random.random_sample(NP)-1)*r_source - except AttributeError: pass + x += (2 * np.random.random_sample(NP) - 1) * r_source + y += (2 * np.random.random_sample(NP) - 1) * r_source + z += (2 * np.random.random_sample(NP) - 1) * r_source + except AttributeError: + pass try: - # Add random noise to account for finite source. + # Add random noise to account for finite source. # NOTE: this means that protons might not necessarily align with # the starting end ending planes on either side of the grid, so # be careful if you have, for example, a very thin sheet of metal # at one end of your grid, that your l_prop will still push them # all through it. source_fwhm = params.source_fwhm - source_variance = (source_fwhm/2.355)**2 - gaussian_noise = np.random.multivariate_normal([0,0,0], np.eye(3)*source_variance, NP) - x += gaussian_noise[:,0] - y += gaussian_noise[:,1] - z += gaussian_noise[:,2] - except AttributeError: pass + source_variance = (source_fwhm / 2.355) ** 2 + gaussian_noise = np.random.multivariate_normal([0, 0, 0], np.eye(3) * source_variance, NP) + x += gaussian_noise[:, 0] + y += gaussian_noise[:, 1] + z += gaussian_noise[:, 2] + except AttributeError: + pass end_time = timer() - print("Time elapsed during proton initialization: " + str(end_time-start_time) + " s") - + print("Time elapsed during proton initialization: " + str(end_time - start_time) + " s") + # Proton step size - ds_prot = params.l_prop/nsteps + ds_prot = params.l_prop / nsteps print("Pushing " + str(NP) + " protons...") # Push protons using compiled module written in C start_time = timer() - _pmover.pmover(mass, charge, ds_prot, grid.dx, grid.dy, grid.dz, grid.xoffset, grid.yoffset, grid.zoffset, NP, nsteps,\ - grid.nx, grid.ny, grid.nz, ntraces, grid.cyl_coords, x, y, z, vx, vy, vz, prop_dir, grid.vals, traces) - - traces[:,nsteps+1,0] = x[0:ntraces] - traces[:,nsteps+1,1] = y[0:ntraces] - traces[:,nsteps+1,2] = z[0:ntraces] + _pmover.pmover(mass, charge, ds_prot, grid.dx, grid.dy, grid.dz, grid.xoffset, grid.yoffset, grid.zoffset, NP, + nsteps, \ + grid.nx, grid.ny, grid.nz, ntraces, grid.cyl_coords, x, y, z, vx, vy, vz, prop_dir, grid.vals, + traces) + + traces[:, nsteps + 1, 0] = x[0:ntraces] + traces[:, nsteps + 1, 1] = y[0:ntraces] + traces[:, nsteps + 1, 2] = z[0:ntraces] end_time = timer() - print("Time elapsed during proton push through grid: " + str(end_time-start_time) + " s") + print("Time elapsed during proton push through grid: " + str(end_time - start_time) + " s") - # final propagation to film plane start_time = timer() - - film_axis1 = np.array(params.film_axis1)/linalg.norm(params.film_axis1) - film_axis2 = np.array(params.film_axis2)/linalg.norm(params.film_axis2) - film_perp = np.cross(film_axis1,film_axis2) # Film normal vector + film_axis1 = np.array(params.film_axis1) / linalg.norm(params.film_axis1) + film_axis2 = np.array(params.film_axis2) / linalg.norm(params.film_axis2) + + film_perp = np.cross(film_axis1, film_axis2) # Film normal vector film_loc = np.array(params.film_loc) - positions = np.array([x,y,z]) - velocities = np.array([vx,vy,vz]) - positions += velocities*(np.dot((film_loc[:,np.newaxis]-positions).T,film_perp)/np.dot(velocities.T,film_perp)) - + positions = np.array([x, y, z]) + velocities = np.array([vx, vy, vz]) + positions += velocities * ( + np.dot((film_loc[:, np.newaxis] - positions).T, film_perp) / np.dot(velocities.T, film_perp)) + # Project final positions onto basis vectors that span the film - film_x = np.dot((positions-film_loc[:,np.newaxis]).T,film_axis1) - film_y = np.dot((positions-film_loc[:,np.newaxis]).T,film_axis2) + film_x = np.dot((positions - film_loc[:, np.newaxis]).T, film_axis1) + film_y = np.dot((positions - film_loc[:, np.newaxis]).T, film_axis2) + + x, y, z = positions - x,y,z = positions - end_time = timer() - print("Time elapsed during final propogation and projection onto film plane: " + str(end_time-start_time) + " s") - + print("Time elapsed during final propogation and projection onto film plane: " + str(end_time - start_time) + " s") + # Report final energy in MeV - Ep = 0.5*mass*(vx**2+vy**2+vz**2)/MeV_to_ergs - - traces[:,nsteps+2,0] = x[0:ntraces] - traces[:,nsteps+2,1] = y[0:ntraces] - traces[:,nsteps+2,2] = z[0:ntraces] - + Ep = 0.5 * mass * (vx ** 2 + vy ** 2 + vz ** 2) / MeV_to_ergs + + traces[:, nsteps + 2, 0] = x[0:ntraces] + traces[:, nsteps + 2, 1] = y[0:ntraces] + traces[:, nsteps + 2, 2] = z[0:ntraces] + film_x = np.nan_to_num(film_x) film_y = np.nan_to_num(film_y) Ep = np.nan_to_num(Ep) traces = np.nan_to_num(traces) - - return film_x,film_y,Ep,traces + + return film_x, film_y, Ep, traces -def plot_results(grid,film_x,film_y,Ep,traces,plot_fluence=True,plot_quiver=False,plot_traces=True,save_images=False): +def plot_results(grid, film_x, film_y, Ep, traces, plot_fluence=True, plot_quiver=False, plot_traces=True, + save_images=False): """ Plot proton fluence, vector field slice, and/or 3D particle traces - Parameters ---------- grid : Grid @@ -675,13 +704,12 @@ def plot_results(grid,film_x,film_y,Ep,traces,plot_fluence=True,plot_quiver=Fals Whether to plot proton fluence as a 2D histogram plot_quiver : boolean, optional Whether to plot a 2D x-y slice of the magnetic or electric fields using a quiver plot. - Also plots material Z using imshow. + Also plots material Z using imshow. plot_traces : boolean, optional Whether to plot particle traces in 3D using mpl_toolkits.mplot3d (if installed) save_images : boolean, optional Whether to save fluence and/or quiver plots in the 'outputs' directory. Will be saved as png images, and named according to the name of the input parameter file. - """ quiver_fig = None @@ -689,34 +717,34 @@ def plot_results(grid,film_x,film_y,Ep,traces,plot_fluence=True,plot_quiver=Fals fluence_fig = None cbar = None - if plot_quiver: + if plot_quiver: # TODO: Clean up this section - + # Plot material Z and field cross section quiver_fig = plt.figure(1) plt.axes().set_aspect('equal') - + # Material Z # if grid.cyl_coords is False: # plt.imshow(grid.vals[:,:,grid.nz/2,7].T, origin='lower', cmap='Reds', interpolation='nearest') # plot material Z - + # x-y electric fields # quiverU = grid.vals[:,:,grid.nz/2,0] # quiverV = grid.vals[:,:,grid.nz/2,1] # x-y magnetic fields - quiverU = grid.vals[:,:,grid.nz/3,3] - quiverV = grid.vals[:,:,grid.nz/3,4] - quiverX = range(len(quiverU[:,0])) - quiverY = range(len(quiverU[0,:])) - quiverX,quiverY = np.meshgrid(quiverX, quiverY, indexing='ij') + quiverU = grid.vals[:, :, grid.nz / 3, 3] + quiverV = grid.vals[:, :, grid.nz / 3, 4] + quiverX = range(len(quiverU[:, 0])) + quiverY = range(len(quiverU[0, :])) + quiverX, quiverY = np.meshgrid(quiverX, quiverY, indexing='ij') if grid.cyl_coords: # X-Y view - quiverR = np.linspace(grid.dx/2,grid.dx*grid.nx+grid.dx/2,grid.nx) - quiverTheta = np.linspace(0.0,2*np.pi-grid.dy,grid.ny) - quiverR,quiverTheta = np.meshgrid(quiverR,quiverTheta, indexing='ij') - quiverX = quiverR*np.cos(quiverTheta) - quiverY = quiverR*np.sin(quiverTheta) + quiverR = np.linspace(grid.dx / 2, grid.dx * grid.nx + grid.dx / 2, grid.nx) + quiverTheta = np.linspace(0.0, 2 * np.pi - grid.dy, grid.ny) + quiverR, quiverTheta = np.meshgrid(quiverR, quiverTheta, indexing='ij') + quiverX = quiverR * np.cos(quiverTheta) + quiverY = quiverR * np.sin(quiverTheta) # Z-X view (uncomment to activate) # quiverX = np.linspace(grid.zoffset, grid.zoffset+grid.lz, grid.nz) @@ -730,21 +758,21 @@ def plot_results(grid,film_x,film_y,Ep,traces,plot_fluence=True,plot_quiver=Fals plt.quiver(quiverX[quiverMask], quiverY[quiverMask], quiverU[quiverMask], quiverV[quiverMask]) plt.xlabel("x") plt.ylabel("y") - plt.xlim([min(quiverX.flatten()),max(quiverX.flatten())]) - plt.ylim([min(quiverY.flatten()),max(quiverY.flatten())]) + plt.xlim([min(quiverX.flatten()), max(quiverX.flatten())]) + plt.ylim([min(quiverY.flatten()), max(quiverY.flatten())]) if save_images: - plt.savefig('outputs/'+sys.argv[1]+'_quiver.png', bbox_inches='tight', transparent=True, dpi=200) - + plt.savefig('outputs/' + sys.argv[1] + '_quiver.png', bbox_inches='tight', transparent=True, dpi=200) + if plot_traces: # Plot particle traces in 3D try: from mpl_toolkits.mplot3d import Axes3D - traces_fig = plt.figure(2,figsize=(7,7)) + traces_fig = plt.figure(2, figsize=(7, 7)) ax = traces_fig.gca(projection='3d') for i in range(params.ntraces): - #ax.plot(traces[i,0:params.nsteps+2,0],traces[i,0:params.nsteps+2,1],traces[i,0:params.nsteps+2,2]) - ax.plot(traces[i,:,0],traces[i,:,1],traces[i,:,2]) + # ax.plot(traces[i,0:params.nsteps+2,0],traces[i,0:params.nsteps+2,1],traces[i,0:params.nsteps+2,2]) + ax.plot(traces[i, :, 0], traces[i, :, 1], traces[i, :, 2]) # Plot grid bounds. # NOTE: DO NOT MISTAKE GRID BOUNDS FOR WHERE YOUR HOHLRAUM IS, ESPECIALLY IN CYLINDRICAL CASE @@ -763,25 +791,28 @@ def plot_results(grid,film_x,film_y,Ep,traces,plot_fluence=True,plot_quiver=Fals # art3d.pathpatch_2d_to_3d(cyl_bottom, z=grid.zoffset, zdir="z") # Cylinder sides - X, Z = np.meshgrid(np.linspace(-radius, radius, 20), np.linspace(grid.zoffset, grid.zoffset+grid.lz, 20)) - Y = np.sqrt(radius**2 - X**2) # Pythagorean theorem + X, Z = np.meshgrid(np.linspace(-radius, radius, 20), + np.linspace(grid.zoffset, grid.zoffset + grid.lz, 20)) + Y = np.sqrt(radius ** 2 - X ** 2) # Pythagorean theorem ax.plot_surface(X, Y, Z, linewidth=1, color='k', alpha=0.2) ax.plot_surface(X, -Y, Z, linewidth=1, color='k', alpha=0.2) else: # Plot edges of 3d box as dotted black lines - for s, e in itertools.combinations(np.array(list(itertools.product([params.gridcorner[0],params.gridcorner[0]+grid.lx],[params.gridcorner[1],params.gridcorner[1]+grid.ly], [params.gridcorner[2],params.gridcorner[2]+grid.lz]))), 2): - if np.sum(np.abs(s-e)) in (grid.lx,grid.ly,grid.lz): + for s, e in itertools.combinations(np.array(list( + itertools.product([params.gridcorner[0], params.gridcorner[0] + grid.lx], + [params.gridcorner[1], params.gridcorner[1] + grid.ly], + [params.gridcorner[2], params.gridcorner[2] + grid.lz]))), 2): + if np.sum(np.abs(s - e)) in (grid.lx, grid.ly, grid.lz): ax.plot3D(*zip(s, e), color='k', linestyle='--') - + if grid.cyl_coords: - ax.set_xlim([grid.xoffset-grid.lx,grid.xoffset+grid.lx]) - ax.set_ylim([grid.xoffset-grid.lx,grid.xoffset+grid.lx]) - ax.set_zlim([grid.zoffset,grid.zoffset+grid.lz]) + ax.set_xlim([grid.xoffset - grid.lx, grid.xoffset + grid.lx]) + ax.set_ylim([grid.xoffset - grid.lx, grid.xoffset + grid.lx]) + ax.set_zlim([grid.zoffset, grid.zoffset + grid.lz]) else: - ax.set_xlim([grid.xoffset,grid.xoffset+grid.lx]) - ax.set_ylim([grid.yoffset,grid.yoffset+grid.ly]) - ax.set_zlim([grid.zoffset,grid.zoffset+grid.lz]) - + ax.set_xlim([grid.xoffset, grid.xoffset + grid.lx]) + ax.set_ylim([grid.yoffset, grid.yoffset + grid.ly]) + ax.set_zlim([grid.zoffset, grid.zoffset + grid.lz]) ax.set_xlabel('x') ax.set_ylabel('y') @@ -799,22 +830,25 @@ def plot_results(grid,film_x,film_y,Ep,traces,plot_fluence=True,plot_quiver=Fals maxfluence = 150.0 try: maxfluence = params.hist_maxfluence - except AttributeError: pass - myhist = plt.hist2d(film_x*10,film_y*10, bins=300, cmap='gray_r', range=[[-xmax*10,xmax*10],[-ymax*10,ymax*10]], vmin=0.0, vmax=maxfluence) + except AttributeError: + pass + myhist = plt.hist2d(film_x * 10, film_y * 10, bins=300, cmap='gray_r', + range=[[-xmax * 10, xmax * 10], [-ymax * 10, ymax * 10]], vmin=0.0, vmax=maxfluence) plt.xlabel('mm') plt.ylabel('mm') - #plt.title('9.5 MeV deuteron') + # plt.title('9.5 MeV deuteron') cbar = plt.colorbar(format='%05.2f') try: # Include interactive draggable colorbar, if available from libraries import draggable_cbar - cbar = draggable_cbar.DraggableColorbar(cbar,myhist[3]) + cbar = draggable_cbar.DraggableColorbar(cbar, myhist[3]) cbar.connect() - except ImportError: pass + except ImportError: + pass if save_images: - plt.savefig('outputs/'+sys.argv[1]+'_fluence.png', bbox_inches='tight') + plt.savefig('outputs/' + sys.argv[1] + '_fluence.png', bbox_inches='tight') return cbar, fluence_fig, traces_fig, quiver_fig @@ -822,7 +856,6 @@ def plot_results(grid,film_x,film_y,Ep,traces,plot_fluence=True,plot_quiver=Fals def save_results(x_arr, y_arr, fileout): """ Save final proton positions to a csv output file - Parameters ---------- x_arr : array of size NP @@ -834,11 +867,12 @@ def save_results(x_arr, y_arr, fileout): """ if fileout == '': return - f = open('outputs/' + fileout,'w') + f = open('outputs/' + fileout, 'w') for i in range(params.NP): - x,y = x_arr[i], y_arr[i] + x, y = x_arr[i], y_arr[i] f.write(str(x) + ', ' + str(y) + '\n') f.close() + if __name__ == '__main__': - main() + main() \ No newline at end of file