I just posted an example for that in a different discussion: MSEVB using Python API
What the -partition flag does in LAMMPS is that it takes the global communicator (MPI_COMM_WORLD) and splits it into 4. Then each of the MPI ranks creates a LAMMPS instance with the subcommunicator passed as argument (by default LAMMPS uses the global one).
For demonstration purposes, here is an example that runs two LAMMPS instances on two processors each, i.e. the equivalent of -partition 2x2, if you do mpirun -np 4:
import sys
from mpi4py import MPI
from lammps import lammps
npartition = 2
me_global = MPI.COMM_WORLD.Get_rank()
np_global = MPI.COMM_WORLD.Get_size()
np_per_comm = np_global / npartition
if np_global != npartition*np_per_comm:
sys.exit("inconsistent number of partitions")
if me_global == 0:
print("Running on %d partitions of %d processors" % (npartition, np_per_comm))
color = me_global // np_per_comm
comm = MPI.COMM_WORLD.Split(color=color)
lmp = lammps(cmdargs=['-nocite', '-log', 'log.lammps-'+str(me_global), '-screen', 'none'], comm=comm)
lmp.file('in.lammps')
lmp.close()
Of course, this is only a small part of what would need to be done. You also need to create communicators to exchange information between the different replica. This is where it gets complicated and beyond the scope of what can be explained here.