Search code examples
pythonnumpyvectorpolar-coordinatescartesian-coordinates

Simulate speakers around a sphere using superposition - speed improvments needed


Note: Drastic speed improvements since posting, see edits at bottom.

I have some working code by it over utilizes loops and I'm pretty sure there should be a faster way of doing it. The size of the output array ends up being pretty large and so when I try to make other arrays the same size of the output, I run out of memory rather quickly.

I am simulating many speakers placed around a sphere all pointing toward the center. I have a simulation of a single speaker and I would like to leverage this single simulation by using the principle of superposition. Basically I want to sum up rotated copies of the single transducer simulation to get my final result.

I have an axisymmetric simulation of acoustic pressure data in cylindrical coordinates ("polar_coord_r", "polar_coord_z"). The pressure field from the simulation is unique at each R and Z value and completely described by a 2D array ("P_real_RZ").

I want to sum together multiple, rotated copies of the this pressure field onto a 3D Cartesian output array. Each copy is rotated to a different location on the sphere. Currently, I am specifying the rotation with an x,y,z point because it allows me to do vector math (spherical coordinates wouldn't allow me to do this as elegantly). The output array is rather large (770 × 770 × 804).

I have some working code to get the output from a single copy of the speaker ("transducer"). It takes about 12 seconds for each slice so it would take over two hours to add each new speaker!! I want to have a dozen or so copies of the speaker so this will take way to long.

The code takes a slice with constant X and computes the R and Z positions at each location in the that slice. "r_distance" is a 2D array containing the radial distance from a line passing between the origin and a point ("point"). Similarlity, "z_distance" is a 2D array containing the distance along that same line.

To get the pressure for the slice, I find the indices of the closest matching "polar_coord_r" and "polar_coord_z" to the computed R distances and Z distances. I use these indices to find what value of pressure (from P_real_RZ) to place at each value in the output.

Some definitions:

  • xx, yy, and zz are 1D arrays of describing the slices through the output volume
  • XXX, YYY, and ZZZ are 3D arrays produced by numpy.meshgrid
  • point is a point which defines the direction that the speaker is rotated. Basically it's just a position vector of the speakers center.
  • P_real_RZ is a 2D array which specifies the real pressure at each unique R and Z value.
  • polar_coord_r and polar_coord_z are 1D arrays which define the unique values of R and Z on which P_real_RZ is defined.
  • current_transducer (only one so far represented in this code) is the pressure values computer for the current point.
  • output is the result from summing many speakers/transducers together.

Any suggestions to speed up this code is greatly appreciated.

Working loop:

for i, x in enumerate(xx):
    # Creates a unit vector from origin to a point
    vector = normalize(point)

    # Create a slice of the cartesian space with constant X
    xyz_slice = np.array([x*np.ones_like(XXX[i,:,:]), YYY[i,:,:], ZZZ[i,:,:]])

    # Projects the position vector of each point of the slice onto the unit vector.
    projection = np.array(list(map(np.dot, xyz_slice, vector )))

    # Normalizes the projection which results in the Z distance along the line passing through the point
    #z_distance = np.apply_along_axis(np.linalg.norm, 0, projection) # this is the slow bit
    z_distance = np.linalg.norm(projection, axis=0) # I'm an idiot

    # Uses vector math to determine the distance from the line
    # Each point in the XYZ slice is the sum of vector along the line and the vector away from the line (radial vector).
    # By extension the position of the xyz point minus the projection of the point against the unit vector, results in the radial vector
    # Norm the radial vector to get the R value for everywhere in the slice
    #r_distance = np.apply_along_axis(np.linalg.norm, 0, xyz_slice - projection) # this is the slow bit
    r_distance = np.linalg.norm(xyz_slice - projection, axis=0) # I'm an idiot

    # Map the pressure data to each point in the slice using the R and Z distance with the RZ pressure slice.
    # look for a more efficient way to do this perhaps. currently takes about 12 seconds per slice
    r_indices = r_map_v(r_distance) # 1.3 seconds by itself
    z_indices = z_map_v(z_distance)
    r_indices[r_indices>384] = 384 # find and remove indicies above the max for r_distance
    z_indices[r_indices>803] = 803
    current_transducer[i,:,:] = P_real_RZ[z_indices, r_indices]

    # Sum the mapped pressure data into the output.
    output += current_transducer

I have also tried to work with the simulation data in the form of a 3D Cartesian array. That is the pressure data from the simulation for all x, y, and z values the same size as the output.I can rotate this 3D array in one direction (not two rotations needed for speakers arranged on a sphere). This takes up waaaay too much memory and is still painfully slow. I end up getting memory errors with this approach.

Edit: I found a slightly simpler way to do it but it is still slow. I've updated the code above so that there are no longer nested loops.

I ran a line profiler and the slowest lines by far were the two containing np.apply_along_axis(). I'm afraid I might have to rethink how I do this completely.

Final Edit: I initially had a nested loop which I assumed to be the issue. I don't know what made me think I needed to use apply_along_axis with linalg.norm. In any case that was the issue.


Solution

  • I haven't looked for all the ways that you could optimize this code, but this issue jumped out at me: "I ran a line profiler and the slowest lines by far were the two containing np.apply_along_axis()." np.linalg.norm accepts an axis argument. You can replace the line

    z_distance = np.apply_along_axis(np.linalg.norm, 0, projection)
    

    with

    z_distance = np.linalg.norm(projection, axis=0)
    

    (and likewise for the other use of np.apply_along_axis and np.linalg.norm). That should improve the performance a bit.