Problems using MPI capabilities

[RomanArmenta]

Hello everyone. I have been working with Netket for some time now, but I have never used its MPI capabilities because I'm unfamiliar with it. If I run the following code in the terminal, I get an error.

# J1-J2 model
# Version with complex Hamiltonian
#
################################################################################

from dis import show_code
import os
import sys

# detect MPI rank
from mpi4py import MPI
rank = MPI.COMM_WORLD.Get_rank()

import jax
import netket as nk
import numpy as np
import json
import mpi4jax
if (rank==0):
    print("NetKet version: {}".format(nk.__version__))
    print("Jax devices: {}".format(jax.devices()))
    print("Jax version: {}".format(jax.__version__))
    print("MPI utils available: {}".format(nk.utils.mpi.available))

# Parameters
L  = 12   # length
J1 = 1.0  # nearest-neighbours exchange interaction
J2 = .50   # next-nearest-neighbours exchange interaction

MSR = 1  # Marshall sign rule (+1/-1)

# ## Hamiltonian

# Hilbert space
g = nk.graph.Chain(L, pbc=True)
hilbert = nk.hilbert.Spin(s=0.5, total_sz=0.0, N=g.n_nodes)
if (rank==0):
    print("Number of graph nodes: {:d}".format(g.n_nodes))
    print("Hilbert size: {:d}".format(hilbert.size))

# Pauli matrices
def sigx(i):
    return nk.operator.spin.sigmax(hilbert, i, dtype=np.complex128)

def sigy(i):
    return nk.operator.spin.sigmay(hilbert, i, dtype=np.complex128)

def sigz(i):
    return nk.operator.spin.sigmaz(hilbert, i, dtype=np.complex128)

def heisenberg(i, j, sgn=1):
    """Heisenberg two spin interaction including Marshall sign rule."""
    return sgn * (sigx(i)*sigx(j) + sigy(i)*sigy(j)) + sigz(i)*sigz(j)

# setup local Hamiltonian
Ha = nk.operator.LocalOperator(hilbert, dtype=np.complex128)  # Hamiltonian

# nearest neighbours
for i in range(L - 1):
    Ha += J1 * heisenberg(i, i+1, MSR)

Ha += J1 * heisenberg(L-1, 0, MSR)

# next nearest neighbours
for i in range(L - 2):
    Ha += J2 * heisenberg(i, i+2)

Ha += J2 * ( heisenberg(L-1, 1) + heisenberg(L-2, 0) )
if (rank==0):
    # check Hamiltonian
    print("Hamiltonian is hemitian: {}".format(Ha.is_hermitian))
    print("Number of local operators: {:d}".format(len(Ha.operators)))
    print("Hamiltonian size: {}".format(Ha.to_dense().shape))

# ## Exact diagonalization

ED = nk.exact.lanczos_ed(Ha, compute_eigenvectors=False)
E0 = ED[0]
if (rank==0):
    print("Exact ground state energy: {:.5f}".format(E0))

# ## Restricted Boltzmann Machine

# setup model
model = nk.models.RBM(alpha=1,dtype=np.complex128)

sa = nk.sampler.MetropolisExchange(hilbert, graph=g, d_max=2, n_chains_per_rank=1)
vs = nk.vqs.MCState(sa, model, n_samples=3000)
opt = nk.optimizer.Sgd(learning_rate=0.01)
sr = nk.optimizer.SR(diag_shift=0.01)

gs = nk.VMC(hamiltonian=Ha, optimizer=opt, variational_state=vs, preconditioner=sr)

# run simulations
output_files = "J1J2cplx_L{:d}_J{:.2f}_rbm".format(L, J2)
gs.run(out=output_files, n_iter=300,show_progress=True)

# print energy
Data = json.load(open("{:s}.log".format(output_files)))
E0rbm = np.mean(Data["Energy"]["Mean"]["real"][-500:-1])
print("RBM ground state energy: {:.5f}".format(E0rbm))
print("Hamiltonian is hemitian: {}".format(Ha.is_hermitian))
print("Number of local operators: {:d}".format(len(Ha.operators)))
print("Hamiltonian size: {}".format(Ha.to_dense().shape))

# ## Exact diagonalization

ED = nk.exact.lanczos_ed(Ha, compute_eigenvectors=False)
E0 = ED[0]
if (rank==0):
    print("Exact ground state energy: {:.5f}".format(E0))

# ## Restricted Boltzmann Machine

# setup model
model = nk.models.RBM(alpha=1,dtype=np.complex128)

sa = nk.sampler.MetropolisExchange(hilbert, graph=g, d_max=2, n_chains_per_rank=1)
vs = nk.vqs.MCState(sa, model, n_samples=3000)
opt = nk.optimizer.Sgd(learning_rate=0.01)
sr = nk.optimizer.SR(diag_shift=0.01)

gs = nk.VMC(hamiltonian=Ha, optimizer=opt, variational_state=vs, preconditioner=sr)

# run simulations
output_files = "J1J2cplx_L{:d}_J{:.2f}_rbm".format(L, J2)
gs.run(out=output_files, n_iter=300,show_progress=True)

# print energy
Data = json.load(open("{:s}.log".format(output_files)))
E0rbm = np.mean(Data["Energy"]["Mean"]["real"][-500:-1])
print("RBM ground state energy: {:.5f}".format(E0rbm))

If I run this code with mpirun -np 2 python3 my_code.py, I get the following error:

AttributeError: 'Config' object has no attribute 'omnistaging_enabled'

So this arises some questions:

1- Am I doing the coding right or it is an installation problem?
2- Can I take advantage of MPI capabilities to run two - or more - optimization processes? I mean, if I want to get the ground state of two different set of parameters of the J1-J2 model, can I do it at the same time with MPI?

[PhilipVinc]

I believe you have a too new version of jax with an older version of netket or mpi4jax.
If you update the relevant packages to the last version it will stop throwing this error.

pip install -U netket mpi4jax flax chex optax

should fix it

[PhilipVinc]

2- Can I take advantage of MPI capabilities to run two - or more - optimization processes? I mean, if I want to get the ground state of two different set of parameters of the J1-J2 model, can I do it at the same time with MPI?

No. What MPI does is that it splits your samples among the MPI processes. So if you have lots of samples (And only in that case) you will see a ~linear speedup in the number of nodes/ranks. If you want to optimise different sets of parameters you should just run the optimisation twice, regardless of MPI.