[Python] Why is reshape/ravel in NumPy faster (or about as fast) than explicit flat indexing...

I'm solving a 2D PDE using solve_ivp from SciPy and was trying to optimize my function by avoiding calls to .reshape() and .ravel() on every timestep. I rewrote my dudt function to operate purely in 1D using flat indexing and manually computed neighbor indices to apply the Laplacian stencil.

This leads to three questions:

1. 
   Is there a silly mistake I didn't catch and is my re-indexed version slow because of it.
   
   
2. 
   Is reshape() (and ravel()) really that fast in NumPy? Does it make a copy or just return a view?
   
   
3. 
   Why isn't the "fully vectorized flat indexing" version faster, even though it avoids reshaping altogether?
   

Any insights into how NumPy handles memory views, indexing, and performance in this context would be appreciated.



Here's my full code:

import os

import numpy as np
import pandas as pd
import scipy as sp
import matplotlib.pyplot as plt

# +++++++++++++++++++++++++++++

t_start = 0
t_end = 1
N_timesteps = 100
my_t = np.linspace(t_start, t_end, N_timesteps)

x_start = 0
y_start = 0
x_end = 2
y_end = 2
Nx_spatial = 80
Ny_spatial = 100
Ni = Nx_spatial
Nj = Ny_spatial
my_x = np.linspace(x_start, x_end, Nx_spatial)
my_y = np.linspace(x_start, x_end, Ny_spatial)
my_dx = my_x[1] - my_x[0]
my_dy = my_y[1] - my_y[0]
X, Y = np.meshgrid(my_x, my_y)

my_u0 = np.zeros_like(X)

def gaussian_2d(x, y, a, mx, my, sx, sy): return a * np.exp(-((x - mx)**2 / (2 * sx**2) + (y - my)**2 / (2 * sy**2)))
my_b = gaussian_2d(X, Y, a=100, mx=0.5, my=0.5, sx=0.1, sy=0.1) + chm.gaussian_2d(X, Y, a=-100, mx=1.5, my=1.0, sx=0.1, sy=0.1)

# my re-indexization matrices

K_center = np.arange(Ni*Nj).reshape(Nj, Ni)
K_right = np.roll(K_center, 1, axis = 1)
K_left = np.roll(K_center, -1, axis = 1)
K_up = np.roll(K_center, 1, axis = 0)
K_down = np.roll(K_center, -1, axis = 0)

center = K_center.ravel()
right = K_right.ravel()
left = K_left.ravel()
up = K_up.ravel()
down = K_down.ravel()

def dudt(t, u_flat, dx, dy, b):
u = u_flat.reshape(Ny_spatial, Nx_spatial)

# Neumann
u[ 0, :] = u[ 1, :]
u[-1, :] = u[-2, :]
u[:, 0] = u[:, 1]
u[:, -1] = u[:, -2]

d2u_dx2 = (1/dx**2) * (np.roll(u, 1, axis=0) - 2*u + np.roll(u, -1, axis=0))
d2u_dy2 = (1/dy**2) * (np.roll(u, 1, axis=1) - 2*u + np.roll(u, -1, axis=1))

# Hm... this re-indexed variant ain't actually faster.
# d2u_dy2 = (1/dy**2) * (u[up] - 2*u

+ u[down])
# d2u_dx2 = (1/dx**2) * (u

- 2*u

+ u

)
#
du_dt = (d2u_dx2 + d2u_dy2) + b

# Dirichlet
du_dt[:, 0] = 0
du_dt[:, -1] = 0
du_dt[ 0, :] = 0
du_dt[-1, :] = 0
#

return du_dt.ravel()
# return du_dt

sol = sp.integrate.solve_ivp(
fun = lambda t, u: dudt(t, u, dx=my_dx, dy=my_dy, b=my_b),
t_span = (t_start, t_end),
y0 = my_u0.ravel(),
method = "RK45",
t_eval = my_t,
)


print("ok I'm plotting now.")
from matplotlib.animation import FuncAnimation

if True:
fig, ax = plt.subplots()
im = ax.imshow(sol.y[:, 0].reshape(Ny_spatial, Nx_spatial),
extent=[x_start, x_end, y_start, y_end],
origin='lower',
cmap='viridis',
vmin=np.min(sol.y), vmax=np.max(sol.y)+0.1)
cbar = plt.colorbar(im, ax=ax)

# Create meshgrid for quiver (skip every 5 points)
x = np.linspace(x_start, x_end, Nx_spatial)
y = np.linspace(y_start, y_end, Ny_spatial)
X, Y = np.meshgrid(x, y)

skip = (slice(None, None, 5), slice(None, None, 5)) # every 5th point
quiv = ax.quiver(X[skip], Y[skip], np.zeros_like(X[skip]), np.zeros_like(Y[skip]), color='white', scale=100)

def update(frame):
u = sol.y[:, frame].reshape(Ny_spatial, Nx_spatial)
im.set_array(u)

# Compute gradient (dy, dx order!)
dy, dx = np.gradient(u, y, x)
quiv.set_UVC(-dx[skip], -dy[skip])

ax.set_title(f"t = {my_t[frame]:.2f}")
return [im, quiv]

ani = FuncAnimation(fig, update, frames=N_timesteps, interval=50, blit=False)
plt.show()

Continue reading...​