Suggestion: analytic_grid initialization speedup using NumPy vectorization
- 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