Suggestion: analytic_grid initialization speedup using NumPy vectorization

[sfeister]

The section of prorad.py where it initializes the grid by looping over values in serial or even parallel is a choke point in speed:

            print("'params.grid_nthreads' not specified. Initializing grid in serial.")
            for i in field_xrange_idx:
                for j in field_yrange_idx:
                    for k in field_zrange_idx:
                        x = X_coords[i]
                        y = Y_coords[j]
                        z = Z_coords[k]
                        gridvals[i,j,k,:] = params.fields((x,y,z))   

There is a faster way to initialize the analytic grid, taking advantage of NumPy's vectorization capabilities. However, it requires more familiarity with the NumPy syntax by the user who writes the "fields()" function. I've written a local branch which leverages a new, vectorized "fields()" function in my input deck and then implements a new format in prorad.py called "analytic_grid2"; I can push it into this repo as an example if there is interest. I make a 300x300x300 analytic grid in a few seconds using the vectorized technique, but it a much smaller grid takes many times longer to generate with the original method.

Suggested addition to prorad.py:

    if fformat == "analytic_grid2":
        # Load user-defined analytic fields, Scott Feister's version using vectorized grid
        
        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

        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)
        
        X, Y, Z = np.meshgrid(X_coords, Y_coords, Z_coords)

        print("Initializing grid in vectorized serial.")
        gridvals = params.fieldsvec((X, Y, Z), NUM_FIELDS)
        
        gridspacings = (lx/ngridx, ly/ngridy, lz/ngridz)
 
        grid = Grid(gridvals, gridspacings, (xoffset,yoffset,zoffset), (lx,ly,lz), cyl_coords=cyl_coords)

Corresponding example of definition of fields in input deck (a radial blob of E & B):

rad = 0.01 # Radius of blob

def fields(coord): # REQUIRED FOR ANALYTIC_GRID
    x,y,z = coord
    r = (x**2+y**2+z**2)
 
    if r < rad: # Create strongest field in the center
        Evec = (1 - r / rad) * np.array([1e4, 1e4, 1e4])
        Bvec = (1 - r / rad) * np.array([-4e4, -4e4, -4e4])
    else:
        Evec = np.array([0, 0, 0])
        Bvec = np.array([0, 0, 0])

    return (Evec[0], Evec[1], Evec[2], Bvec[0], Bvec[1], Bvec[2], 0.0,0.0,0.0)

def fieldsvec(coord, NUM_FIELDS): # REQUIRED FOR ANALYTIC_GRID2
    X, Y, Z = coord
    R = X**2 + Y**2 + Z**2
   
    Evec = np.zeros((X.shape[0],X.shape[1],X.shape[2],3))
    Bvec = np.zeros((X.shape[0],X.shape[1],X.shape[2],3))
    gridvals = np.zeros((X.shape[0],X.shape[1],X.shape[2],NUM_FIELDS))
    
    ct = R < rad
    Evec[ct,:] = np.expand_dims((1 - R[ct] / rad), axis=-1) * np.array([1e4, 1e4, 1e4])
    Bvec[ct,:] = np.expand_dims((1 - R[ct] / rad), axis=-1) * np.array([-4e4, -4e4, -4e4])
    
    gridvals[:,:,:,0:3] = Evec
    gridvals[:,:,:,3:6] = Bvec
    gridvals[:,:,:,6:9] = 0.0
    return gridvals