Filling an array with the same data at a regular interval without a for loop - Code Review Stack Exchange most recent 30 from codereview.stackexchange.com 2019-09-19T03:18:34Z https://codereview.stackexchange.com/feeds/question/156780 https://creativecommons.org/licenses/by-sa/4.0/rdf https://codereview.stackexchange.com/q/156780 7 Filling an array with the same data at a regular interval without a for loop Alex https://codereview.stackexchange.com/users/132509 2017-03-03T05:44:11Z 2017-03-03T17:03:34Z <p>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\$:</p> <pre><code>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 </code></pre> <p>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?</p> https://codereview.stackexchange.com/questions/156780/-/156805#156805 7 Answer by Mathias Ettinger for Filling an array with the same data at a regular interval without a for loop Mathias Ettinger https://codereview.stackexchange.com/users/84718 2017-03-03T13:16:08Z 2017-03-03T17:03:34Z <p>For starter, I would suggest making a function out of this pile of code so its easier to read an reason about:</p> <pre><code>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 </code></pre> <p>Optionnally, throwing in a docstring here wouldn't hurt.</p> <p>Now this function has two purposes: compute the coordinates and build a surface out of it, let's split things up.</p> <pre><code>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 </code></pre> <p>Usage being <code>phi_slice = build_coordinates(50); surface = create_surface(phi_slice)</code>. With the advantage of being able to build several surfaces out of the same coordinates.</p> <p>Now let's start to look into removing that <code>for</code> loop. So basicaly, the <code>surface</code> you want to build should look like:</p> <pre><code>[[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]] </code></pre> <p>Where</p> <pre><code>[[r0, Θ0], [r1, Θ1], ... [rn, Θn]] </code></pre> <p>are the coordinates (your <code>phi_slice</code>) and</p> <pre><code>[φ0, ..., φn] </code></pre> <p>are <code>np.linspace(0, (2 * np.pi) - dphi, numbins)</code>.</p> <p>So better build <code>surface</code> by repeating <code>phi_slice</code> and the elevations enought time and then concatenating them together rather than copy/pasting them into a blank array.</p> <p>This is rather simple since <code>numpy</code> provides <code>np.repeat</code> and <code>np.concatenate</code>. But both have their specifics:</p> <p><code>np.repeat</code> return a flat array when no <code>axis</code> 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:</p> <pre><code>&gt;&gt;&gt; np.repeat(np.array([φ0, φ1, φ2]), 3) array([φ0, φ0, φ0, φ1, φ1, φ1, φ2, φ2, φ2]) </code></pre> <p>But not quite so for the coordinates:</p> <pre><code>&gt;&gt;&gt; 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]]) </code></pre> <p>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:</p> <pre><code>&gt;&gt;&gt; coordinates = np.array([[r0, Θ0], [r1, Θ1], [r2, Θ2]]) &gt;&gt;&gt; np.repeat(coordinates[np.newaxis, :, :], 2, axis=0).reshape(6, 2) array([[r0, Θ0], [r1, Θ1], [r2, Θ2], [r0, Θ0], [r1, Θ1], [r2, Θ2]]) </code></pre> <p>Where 6 is the length of axis 0 times the amount of repetitions.</p> <p>The last step being to combine both repeated arrays using <code>np.concatenate</code>. However <code>concatenate</code> requires that both array have the same amount of axis so we need to add one to the elevation vector:</p> <pre><code>&gt;&gt;&gt; elevations = np.repeat(np.array([φ0, φ1, φ2]), 3) &gt;&gt;&gt; coordinates = np.array([[r0, Θ0], [r1, Θ1], [r2, Θ2]]) &gt;&gt;&gt; repeated_coordinates = np.repeat(coordinates[np.newaxis, :, :], 3, axis=0).reshape(9, 2) &gt;&gt;&gt; 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]]) </code></pre> <p>This yield:</p> <pre><code>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) </code></pre>