提升np.fromfunction性能及解决高维度数组内存报错的技术问询

Hey Juan, let's tackle your two problems head-on—both speed and memory issues are common when dealing with massive arrays like this, and there are practical fixes for both.

1. Boosting Execution Speed #

Your original np.fromfunction approach is slow because it relies on a Python lambda being called element-wise for every single entry in the array. For large dimensions, this Python-level loop becomes a massive bottleneck. Instead, use NumPy's vectorized operations (implemented in optimized C) to avoid per-element Python overhead.

Replace fromfunction with Vectorized Broadcasting #

Here's a much faster alternative for your uint8-based bit-extraction task:

import numpy as np

dim = 32
# Generate all i values as a column vector
i = np.arange(2**dim, dtype=np.uint32).reshape(-1, 1)
# Generate all j values as a row vector
j = np.arange(dim, dtype=np.uint32).reshape(1, -1)
# Vectorized shift-and-mask operation (broadcasted automatically)
y_shift = ((i >> j) & 1).astype(np.uint8)

Why this works better:

• NumPy handles the shifting and masking operations in bulk, avoiding Python loop overhead.
• Broadcasting aligns the column vector i and row vector j to compute the entire array in one go, rather than calling a lambda millions/billions of times.

2. Avoiding "zsh: killed python3" (Memory Overload) #

When dim=32, your target array has 2^32 = 4,294,967,296 rows and 32 columns. Even with uint8 dtype, that's 128 GB of memory—way more than most machines can handle. The "killed" error happens because your OS terminates the Python process when it runs out of available RAM.

Fix 1: Use Memory-Mapped Files (np.memmap) #

Store the array on disk instead of keeping it entirely in memory. This lets you work with the array as if it's in RAM, but data is loaded/saved in chunks:

import numpy as np

dim = 32
shape = (2**dim, dim)
dtype = np.uint8

# Create a memory-mapped file (stored on disk)
y_shift = np.memmap('high_dim_array.dat', dtype=dtype, mode='w+', shape=shape)

# Process in chunks to avoid memory overload
block_size = 2**20  # Adjust based on your available RAM (e.g., 1M rows per block)
for start in range(0, shape[0], block_size):
    end = min(start + block_size, shape[0])
    # Generate a chunk of i values
    i_block = np.arange(start, end, dtype=np.uint32).reshape(-1, 1)
    j = np.arange(dim, dtype=np.uint32).reshape(1, -1)
    # Write the chunk to the memmap file
    y_shift[start:end] = ((i_block >> j) & 1).astype(dtype)

# Later, to read the array back:
# y_shift = np.memmap('high_dim_array.dat', dtype=dtype, mode='r', shape=shape)

Fix 2: Reevaluate Whether You Need the Full Array #

Ask yourself: Do you really need to store every entry explicitly? For many use cases (e.g., machine learning, bitwise computations), you can compute values on-the-fly instead of pregenerating the entire array. For example, if you're iterating over rows, compute each row's bits when you need it, rather than storing all rows upfront.

Fix 3: Use More Compact Storage (If Applicable) #

If your dim is a multiple of 8, you can pack 8 bits into a single uint8 value to reduce memory usage by 8x. For dim=32, this would shrink the array to (2^32, 4) (16 GB total). You'd need to handle bit extraction manually when accessing values, but it's a tradeoff for reduced storage:

# Example for dim=32 (pack 8 bits per byte)
i = np.arange(2**32, dtype=np.uint32).reshape(-1, 1)
# Pack 8 bits into each uint8 column
y_compact = ((i >> np.arange(0, 32, 8).reshape(1, -1)) & 0xFF).astype(np.uint8)

内容的提问来源于stack exchange，提问作者Juan