<jupyter_start><jupyter_text>Target:- get something of shape [Nt, Na, 3]- compute autocorrelation for each combination of (I, a), (J, b)- return as [Nt, Na, 3, Na, 3]<jupyter_code># arb. shaped array:
Nt = 4
Na = 2

array = np.ones([Nt, Na, 3])
Nt, *shape = array.shape
Nt, shape
# create mock data
# fill some data
for (t, I, a) in np.ndindex((Nt, *shape)):
    array[t, I, a] = t + a + I

array
# reshape to (X, Nt)
data = array.reshape(Nt, -1).swapaxes(0, 1)

data.shape, data
# array returns row for row
for x in data:
    print(x)

# product returns all combinations of rows
for X in product(data[:3], data[:3]):
    print(X)
# apply correlate to each row
def correlate(X, **kwargs):
    Nt = len(X[0])
    return sl.correlate(X[0], X[1], **kwargs)[Nt - 1 :]

# correlate with `fft=True`
def c(X):
    return correlate(X, method='fft')
# reference:
ref_result = np.zeros([Nt, *shape, *shape])
for (I, a) in np.ndindex(*shape):
    for (J, b) in np.ndindex(*shape):
        X = (array[:, I, a], array[:, J, b])
        ref_result[:, I, a, J, b] = c(X)

ref_result.shape
%%time
result1 = np.zeros((data.shape[0]**2, data.shape[1]))

for ii, X in enumerate(product(data, data)):
    result1[ii] = c(X)
%%time
# parallelize
with multiprocessing.Pool() as p:
    result2 = np.array(p.map(c, product(data, data)))
    
# to see a speedup, one needs larger data
np.allclose(result1, result2)
# shape back
result = np.moveaxis(result2.reshape((*shape, *shape, Nt)), -1, 0)
result.shape
np.allclose(result, ref_result)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17551813471502592, 'fraction_numerical': 0.01683937823834197, 'mean_word_length': 3.9838709677419355, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '524411', 'license': '0BSD', 'n_tokens_mistral': 630, 'n_tokens_neox': 611, 'n_words': 188, 'path': 'numerics/numpy_multiprocess.ipynb', 'repo_name': 'flokno/python_recipes', 'size': '1544'}