numpy.load access performance

Note to self: When loading datasets stored as (compressed) .npz files and iterating through them across some data dimension, it is much faster to first load the full array into memory rather than accessing via index repeatedly. I come across this loading an atomistic dataset (sn2), where this change lead to an almost 5000x speedup!

In this synthetic example, on my MacBook Pro, it's still about 80x faster:

import numpy as np
from time import monotonic

n = 10000
data = np.random.random((n, 6, 3))

np.savez_compressed("data.npz", data=data)

data = np.load("data.npz")

start = monotonic()

stuff = []
for i in range(n):
    stuff.append(data["data"][i])

print(f"{(monotonic()-start)*1e6/n} µs/sample")
# -> 2897.6689249975607 µs/sample

data = np.array(data["data"])

start = monotonic()

stuff = []
for i in range(n):
    stuff.append(data[i])

print(f"{(monotonic()-start)*1e6/n} µs/sample")
# -> 34.760837513022125 µs/sample