Filling an array with the same data at a regular interval without a for loop

Question

I have some code which generates the coordinates of a cylindrically-symmetric surface, with coordinates given as \$(r, \theta, \phi)\$. At the moment, I generate the coordinates of one \$\phi\$ slice, and store that in a \$2 \times N\$ array (for \$N\$ bins), and then in a for loop I copy this array for each value of \$\phi\$ from \$0\$ to \$2\pi\$:

import numpy as np

# this is the number of bins that my surface is chopped into
numbins = 50

# these are the values for r
r_vals = np.linspace(0.0001, 50, numbins, endpoint = True)

# these are the values for theta
theta_vals = np.linspace(0, np.pi / 2, numbins, endpoint = True)

# I combine the r and theta values into a 2xnumbins array for one "slice" of phi
phi_slice = np.vstack([r_vals,theta_vals]).T

# this is the array which will store all of the coordinates of the surface
surface = np.zeros((numbins**2,3))
lower_bound = 0
upper_bound = numbins

# this is the size of the bin for phi
dphi = (2 * np.pi) / numbins

# this is the for loop I'd like to eliminate.
# For every value of phi, it puts a copy of the phi_slice array into
# the surface array, so that the surface is cylindrical about phi.
for phi in np.linspace(0, (2 * np.pi) - dphi, numbins):
    surface[lower_bound:upper_bound, :2] = phi_slice
    surface[lower_bound:upper_bound,2] = phi
    lower_bound += numbins
    upper_bound += numbins

I'm calling this routine in a numerical integration of 1e6 or 1e7 steps, and while numbins is 50 in the example above, in practice it'll be in the thousands. This for loop is a choking point, and I'd really like to eliminate it to speed things up. Is there a good NumPythonic way to do the same thing as this for loop?

Answer 1

For starter, I would suggest making a function out of this pile of code so its easier to read an reason about:

def create_surface(bins_count=50, radius_start=0.0001, radius_end=50):
    radii = np.linspace(radius_start, radius_end, bins_count, endpoint=True)
    thetas = np.linspace(0, np.pi / 2, bins_count, endpoint=True)

    coordinates = np.vstack([radii, thetas]).T

    surface = np.zeros((bins_count**2, 3))
    lower_bound = 0
    upper_bound = numbins

    # this is the size of the bin for phi
    delta_phi = (2 * np.pi) / bins_count

    for phi in np.linspace(0, (2 * np.pi) - delta_phi, bins_count):
        surface[lower_bound:upper_bound, :2] = coordinates
        surface[lower_bound:upper_bound, 2] = phi
        lower_bound += bins_count
        upper_bound += bins_count
    return surface

Optionnally, throwing in a docstring here wouldn't hurt.

Now this function has two purposes: compute the coordinates and build a surface out of it, let's split things up.

def build_coordinates(bins_count=50, radius_start=0.0001, radius_end=50):
    radii = np.linspace(radius_start, radius_end, bins_count, endpoint=True)
    thetas = np.linspace(0, np.pi / 2, bins_count, endpoint=True)

    return np.vstack([radii, thetas]).T


def create_surface(coordinates):
    bins_count, two = coordinates.shape
    assert two == 2

    surface = np.zeros((bins_count**2, 3))
    lower_bound = 0
    upper_bound = numbins

    # this is the size of the bin for phi
    delta_phi = (2 * np.pi) / bins_count

    for phi in np.linspace(0, (2 * np.pi) - delta_phi, bins_count):
        surface[lower_bound:upper_bound, :2] = coordinates
        surface[lower_bound:upper_bound, 2] = phi
        lower_bound += bins_count
        upper_bound += bins_count
    return surface

Usage being phi_slice = build_coordinates(50); surface = create_surface(phi_slice). With the advantage of being able to build several surfaces out of the same coordinates.

Now let's start to look into removing that for loop. So basicaly, the surface you want to build should look like:

[[r0, Θ0, φ0],
 [r1, Θ1, φ0],
 ...
 [rn, Θn, φ0],
 [r0, Θ0, φ1],
 [r1, Θ1, φ1],
 ...
 [rn, Θn, φ1],
 ...
 [r0, Θ0, φn],
 [r1, Θ1, φn],
 ...
 [rn, Θn, φn]]

Where

[[r0, Θ0],
 [r1, Θ1],
 ...
 [rn, Θn]]

are the coordinates (your phi_slice) and

[φ0, ..., φn]

are np.linspace(0, (2 * np.pi) - dphi, numbins).

So better build surface by repeating phi_slice and the elevations enought time and then concatenating them together rather than copy/pasting them into a blank array.

This is rather simple since numpy provides np.repeat and np.concatenate. But both have their specifics:

np.repeat return a flat array when no axis is given and repeat each element the requested amount of time before starting to repeat the next one. This is exactly what is needed for the elevations:

>>> np.repeat(np.array([φ0, φ1, φ2]), 3)
array([φ0, φ0, φ0, φ1, φ1, φ1, φ2, φ2, φ2])

But not quite so for the coordinates:

>>> np.repeat(np.array([[r0, Θ0], [r1, Θ1], [r2, Θ2]]), 2, axis=0)
array([[r0, Θ0],
       [r0, Θ0],
       [r1, Θ1],
       [r1, Θ1],
       [r2, Θ2],
       [r2, Θ2]])

Instead, we would like to repeat the whole block of coordinates at once. So let's add a dimension to repeat along it and then reshape the result to remove the extra dimension:

>>> coordinates = np.array([[r0, Θ0], [r1, Θ1], [r2, Θ2]])
>>> np.repeat(coordinates[np.newaxis, :, :], 2, axis=0).reshape(6, 2)
array([[r0, Θ0],
       [r1, Θ1],
       [r2, Θ2],
       [r0, Θ0],
       [r1, Θ1],
       [r2, Θ2]])

Where 6 is the length of axis 0 times the amount of repetitions.

The last step being to combine both repeated arrays using np.concatenate. However concatenate requires that both array have the same amount of axis so we need to add one to the elevation vector:

>>> elevations = np.repeat(np.array([φ0, φ1, φ2]), 3)
>>> coordinates = np.array([[r0, Θ0], [r1, Θ1], [r2, Θ2]])
>>> repeated_coordinates = np.repeat(coordinates[np.newaxis, :, :], 3, axis=0).reshape(9, 2)
>>> np.concatenate((repeated_coordinates, elevations[:, np.newaxis]), axis=1)
array([[r0, Θ0, φ0],
       [r1, Θ1, φ0],
       [r2, Θ2, φ0],
       [r0, Θ0, φ1],
       [r1, Θ1, φ1],
       [r2, Θ2, φ1],
       [r0, Θ0, φ2],
       [r1, Θ1, φ2],
       [r2, Θ2, φ2]])

This yield:

def build_coordinates(bins_count=50, radius_start=0.0001, radius_end=50):
    radii = np.linspace(radius_start, radius_end, bins_count, endpoint=True)
    thetas = np.linspace(0, np.pi / 2, bins_count, endpoint=True)

    return np.vstack([radii, thetas]).T


def create_surface(coordinates):
    bins_count, two = coordinates.shape
    assert two == 2

    elevation = np.linspace(0, 2 * np.pi * (1 - 1/bins_count), bins_count)
    elevations = np.repeat(elevation, bins_count)
    coords = np.repeat(coordinates[np.newaxis, :, :], bins_count, axis=0).reshape(bins_count**2, 2)

    return np.concatenate((coords, elevations[:, np.newaxis]), axis=1)