llnl / llnl/prorad

Suggestion: analytic_grid initialization speedup using NumPy vectorization

Open
#3 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
9
Forks
6
PR merge metrics
No merged PRs in 30d

Description

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:

```python
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:
```python
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):
```python
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
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in prorad.py at the existing analytic grid initialization and compare its inputs and Grid construction with the proposed analytic_grid2 section. Review the vectorized fieldsvec example and determine how the new format should coexist with fields. Done means equivalent grid output with the vectorized initialization and a demonstrated speed improvement.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.