Easily store einsum paths

[evanberkowitz]

For numeric reproducibility it would be good to store the einsum paths that different observables use. (It depends on the lattice size, nt, etc.)

Annoyingly the output of np.einsum_path can't just be dumped to h5.

import h5py as h5
import numpy as np

M = np.arange(1024).reshape(4,4,8,8)

path, summary = np.einsum_path('abde,acdf,bcef->',M,M,M,optimize="optimal")

# path = ['einsum_path', (0, 1), (0, 1)]
# print(summary) yields
#   Complete contraction:  abde,acdf,bcef->
#         Naive scaling:  6
#     Optimized scaling:  6
#      Naive FLOP count:  9.830e+04
#  Optimized FLOP count:  6.758e+04
#   Theoretical speedup:  1.455
#  Largest intermediate:  1.024e+03 elements
#--------------------------------------------------------------------------
#scaling                  current                                remaining
#--------------------------------------------------------------------------
#   6             acdf,abde->bcef                              bcef,bcef->
#   4                 bcef,bcef->                                       ->

# Another example
path, summary = np.einsum_path('abde,acdf,bcef,gghh->',M,M,M,M,optimize="optimal")

but because it's a mixed array path can't easily be written to an h5 path.

I tried pickling the whole shebang, but I couldn't figure out how to send the pickle.dump to an h5 field. There is an hdf5 pickle but I'd be afraid to use it and h5py at the same time.

How should this be accomplished? Sometimes many different contractions share an optimization (as in spin-spin correlators). Could be a string attribute (str(path) works).

[evanberkowitz]

Also found hdf5pickle but development stopped ~ 4 years ago, and may be dead.

[jl-wynen]

hickle is built on top of h5py, I would expect it to work with our other code.

If this is the only place where we would use it, we can avoid the dependency by doing it manually:

from io import BytesIO
import pickle

import h5py as h5
import numpy as np

M = np.arange(1024).reshape(4,4,8,8)

path, summary = np.einsum_path('abde,acdf,bcef->',M,M,M,optimize="optimal")

def saveEinsumPath(h5group, path):
    buffer = BytesIO()
    pickle.dump(path, buffer)
    h5group["einsum_path"] = np.void(buffer.getbuffer())

def loadEinsumPath(h5group):
    buffer = BytesIO(h5group["einsum_path"][()].tobytes())
    return pickle.load(buffer)

with h5.File("einsum_test.h5", "w") as f:
    saveEinsumPath(f, path)

with h5.File("einsum_test.h5", "r") as f:
    loaded = loadEinsumPath(f)

assert path == loaded

This is just something I quickly threw together from PyTorch related code and it still needs some safety checks...

[jl-wynen]

But do we actually need this? Optimizing the paths takes virtually no time. Or do you want to make sure we can reproduce the measurement with exactly the same numerical errors?

[evanberkowitz]

Yeah, the point is for exact reproducibility we need exactly the same einsum path. I'm not sure if the einsum optimization algorithm is deterministic, but it's a part of our algorithm that isn't really in our code.

(And if someone tried to reproduce using some other routine, they'd need to know the order.)