Merge pull request #102 from jrmaddison/jrmaddison/MPI

Use mpi4py.MPI instead of dolfin.MPI

Author: bearecinos

aux/Uobs_from_momsolve.py

@@ -15,9 +15,10 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with tlm_adjoint.  If not, see <https://www.gnu.org/licenses/>.
 
-from fenics_ice.backend import Function, HDF5File, MPI, Mesh, Point, \
+from fenics_ice.backend import Function, HDF5File, Mesh, Point, \
     RectangleMesh, VectorFunctionSpace, project
 
+import mpi4py.MPI as MPI  # noqa: N817
 import numpy as np
 from pathlib import Path
 import h5py
@@ -51,7 +52,7 @@ def main(dd, infile, outfile, noise_sdev, L, seed=0, ls=None):
     assert Path(outfile).suffix == ".h5"
     assert noise_sdev > 0.0
 
-    infile = HDF5File(MPI.comm_world, str(Path(dd)/infile), 'r')
+    infile = HDF5File(MPI.COMM_WORLD, str(Path(dd)/infile), 'r')
 
     # Get mesh from file
     mesh = Mesh()

aux/gen_rect_mesh.py

@@ -22,7 +22,7 @@
 
 from fenics_ice.backend import File, Point, RectangleMesh
 
-from mpi4py import MPI
+import mpi4py.MPI as MPI  # noqa: N817
 import argparse
 
 def gen_rect_mesh(nx, ny, xmin, xmax, ymin, ymax, outfile, direction='right'):

aux/plotting/plot_dq_ts.py

@@ -15,8 +15,9 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with tlm_adjoint.  If not, see <https://www.gnu.org/licenses/>.
 
-from fenics_ice.backend import Function, FunctionSpace, HDF5File, Mesh, MPI
+from fenics_ice.backend import Function, FunctionSpace, HDF5File, Mesh
 
+import mpi4py.MPI as MPI  # noqa: N817
 import sys
 import pickle
 import numpy as np
@@ -93,7 +94,7 @@
     t    = mesh.cells()
 
     hdffile = str(next(outdir.glob("*dQ_ts.h5")))
-    hdf5data = HDF5File(MPI.comm_world, hdffile, 'r')
+    hdf5data = HDF5File(MPI.COMM_WORLD, hdffile, 'r')
     hdf5data.read(dQ, f'dQ/vector_{n_sens}')
 
 

aux/plotting/plot_leading_eigenfuncs.py

@@ -15,8 +15,9 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with tlm_adjoint.  If not, see <https://www.gnu.org/licenses/>.
 
-from fenics_ice.backend import Function, FunctionSpace, HDF5File, MPI, Mesh
+from fenics_ice.backend import Function, FunctionSpace, HDF5File, Mesh
 
+import mpi4py.MPI as MPI  # noqa: N817
 import sys
 import pickle
 import numpy as np
@@ -91,7 +92,7 @@
     for j in range(4):
         k = j + e_offset
         vr_file = str(next((run_dir / 'output').glob("*vr.h5")))
-        hdf5data = HDF5File(MPI.comm_world, vr_file, 'r')
+        hdf5data = HDF5File(MPI.COMM_WORLD, vr_file, 'r')
         hdf5data.read(eigenfunc, f'v/vector_{k}')
 
         sind = j+1+i*4

fenics_ice/inout.py

@@ -19,10 +19,11 @@
 Module to handle model input & output
 """
 
-from .backend import File, Function, HDF5File, MPI, XDMFFile, \
+from .backend import File, Function, HDF5File, XDMFFile, \
     configure_checkpointing
 from tlm_adjoint.fenics.backend import backend_Function
 
+import mpi4py.MPI as MPI  # noqa: N817
 import sys
 import time
 import csv
@@ -48,7 +49,7 @@ class Writer(ABC):
     unnamed_re = re.compile("f_[0-9]+")
     suffix = None       # the subclass specific file extension
 
-    def __init__(self, fpath, comm=MPI.comm_world):
+    def __init__(self, fpath, comm=MPI.COMM_WORLD):
         assert comm is not None, "Need an MPI communicator"
         self._fpath = Path(fpath)
         self.comm = comm
@@ -245,7 +246,7 @@ def write_dqval(dQ_ts, cntrl_names, params):
     outdir_f = Path(outdir)/phase_name/phase_suffix
     # TODO add this file to diags once Dan makes his pull request
     vtkfile = File(str((diagdir_f/h5_filename).with_suffix(".pvd")))
-    hdf5out = HDF5File(MPI.comm_world, str(outdir_f/h5_filename), 'w')
+    hdf5out = HDF5File(MPI.COMM_WORLD, str(outdir_f/h5_filename), 'w')
     n = 0.0
 
     # Loop dQ sample times ('num_sens')
@@ -304,15 +305,15 @@ def write_variable(var, params, name=None, outdir=None, phase_name='', phase_suf
         xml_fname = str(outfname.with_suffix(".xml"))
         File(xml_fname) << outvar
     if 'h5' in output_var_format:
-        hdf5out = HDF5File(MPI.comm_world, str(outfname.with_suffix(".h5")), 'w')
+        hdf5out = HDF5File(MPI.COMM_WORLD, str(outfname.with_suffix(".h5")), 'w')
         hdf5out.write(outvar, name)
         hdf5out.close()
     if 'all' in output_var_format:
         vtk_fname = str(outfname.with_suffix(".pvd"))
         xml_fname = str(outfname.with_suffix(".xml"))
         File(vtk_fname) << outvar
         File(xml_fname) << outvar
-        hdf5out = HDF5File(MPI.comm_world, str(outfname.with_suffix(".h5")), 'w')
+        hdf5out = HDF5File(MPI.COMM_WORLD, str(outfname.with_suffix(".h5")), 'w')
         hdf5out.write(outvar, name)
         hdf5out.close()
 

fenics_ice/mesh.py

@@ -19,11 +19,12 @@
 Module to handle all things mesh to avoid code repetition in run scripts.
 """
 
-from .backend import FunctionSpace, MPI, Mesh, MeshFunction, \
-    MeshValueCollection, VectorFunctionSpace, XDMFFile, parameters
+from .backend import FunctionSpace, Mesh, MeshFunction, MeshValueCollection, \
+    VectorFunctionSpace, XDMFFile, parameters
 
 from . import model
 
+import mpi4py.MPI as MPI  # noqa: N817
 import os
 import numpy as np
 from pathlib import Path
@@ -50,7 +51,7 @@ def get_mesh(params):
 
     elif filetype == '.xdmf':
         mesh_in = Mesh()
-        mesh_xdmf = XDMFFile(MPI.comm_world, str(meshfile))
+        mesh_xdmf = XDMFFile(MPI.COMM_WORLD, str(meshfile))
         mesh_xdmf.read(mesh_in)
 
     else:
@@ -68,10 +69,10 @@ def get_mesh_length(mesh):
 
     comm = mesh.mpi_comm()
 
-    xmin = MPI.min(comm, xmin)
-    xmax = MPI.max(comm, xmax)
-    ymin = MPI.min(comm, ymin)
-    ymax = MPI.max(comm, ymax)
+    xmin = comm.allreduce(xmin, op=MPI.MIN)
+    xmax = comm.allreduce(xmax, op=MPI.MAX)
+    ymin = comm.allreduce(ymin, op=MPI.MIN)
+    ymax = comm.allreduce(ymax, op=MPI.MAX)
 
     L1 = xmax - xmin
     L2 = ymax - ymin
@@ -128,7 +129,7 @@ def get_ff_from_file(params, model, fill_val=0):
 
     # Read the MeshValueCollection (sparse)
     ff_mvc = MeshValueCollection("size_t", model.mesh, dim=dim-1)
-    ff_xdmf = XDMFFile(MPI.comm_world, str(ff_file))
+    ff_xdmf = XDMFFile(MPI.COMM_WORLD, str(ff_file))
     ff_xdmf.read(ff_mvc)
 
     # Create FacetFunction filled w/ default

fenics_ice/model.py

@@ -46,7 +46,7 @@ def __init__(self, mesh_in, input_data, param_in, init_fields=True,
         self.params = param_in
         self.input_data = input_data
         self.solvers = []
-        self.parallel = MPI.size(mesh_in.mpi_comm()) > 1
+        self.parallel = mesh_in.mpi_comm().size > 1
 
         # Generate Domain and Function Spaces
         self.mesh = mesh_in

fenics_ice/solver.py

@@ -23,7 +23,7 @@
     line_search_rank0_scipy_scalar_search_wolfe1 as line_search_rank0
 
 import logging
-import mpi4py.MPI as MPI
+import mpi4py.MPI as MPI  # noqa: N817
 import numpy as np
 from pathlib import Path
 import time

runs/run_errorprop.py

@@ -15,12 +15,13 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with tlm_adjoint.  If not, see <https://www.gnu.org/licenses/>.
 
-from fenics_ice.backend import Function, HDF5File, MPI
+from fenics_ice.backend import Function, HDF5File
 
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["OPENBLAS_NUM_THREADS"] = "1"
 
+import mpi4py.MPI as MPI  # noqa: N817
 from pathlib import Path
 import pickle
 import numpy as np
@@ -102,7 +103,7 @@ def run_errorprop(config_file):
     # and eigenvectors from .h5 file
     eps = params.constants.float_eps
     W = []
-    with HDF5File(MPI.comm_world, str(outdir_e/vecfile), 'r') as hdf5data:
+    with HDF5File(MPI.COMM_WORLD, str(outdir_e/vecfile), 'r') as hdf5data:
         for i in range(nlam):
             w = Function(space)
             hdf5data.read(w, f'v/vector_{i}')
@@ -120,7 +121,7 @@ def run_errorprop(config_file):
 
     # File containing dQoi_dCntrl (i.e. Jacobian of parameter to observable (Qoi))
     outdir_qoi = Path(outdir)/phase_time/phase_suffix_qoi
-    hdf5data = HDF5File(MPI.comm_world, str(outdir_qoi/dqoi_h5file), 'r')
+    hdf5data = HDF5File(MPI.COMM_WORLD, str(outdir_qoi/dqoi_h5file), 'r')
 
     dQ_cntrl = Function(space, space_type="conjugate_dual")
 
@@ -181,7 +182,7 @@ def run_errorprop(config_file):
     diag_dir = Path(params.io.diagnostics_dir)/phase_err/phase_suffix_err
     outdir_err = Path(params.io.output_dir)/phase_err/phase_suffix_err
 
-    # if(MPI.comm_world.rank == 0):
+    # if(MPI.COMM_WORLD.rank == 0):
     plt.semilogy(sigma_steps, sigma_conv)
     plt.title("Convergence of sigmaQoI")
     plt.ylabel("sigma QoI")

runs/run_invsigma.py

@@ -16,12 +16,13 @@
 # along with tlm_adjoint.  If not, see <https://www.gnu.org/licenses/>.
 
 from fenics_ice.backend import FiniteElement, Function, FunctionSpace, \
-    HDF5File, MPI, TestFunction, assemble, assign, inner, dx
+    HDF5File, TestFunction, assemble, assign, inner, dx
 
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["OPENBLAS_NUM_THREADS"] = "1"
 
+import mpi4py.MPI as MPI  # noqa: N817
 from pathlib import Path
 import pickle
 import numpy as np
@@ -43,8 +44,8 @@ def patch_fun(mesh_in, params):
     import random
     from scipy.spatial import KDTree
 
-    comm = MPI.comm_world
-    rank = MPI.rank(comm)
+    comm = MPI.COMM_WORLD
+    rank = comm.rank
     root = rank == 0
 
     # Test DG function
@@ -114,7 +115,7 @@ def patch_fun(mesh_in, params):
 def run_invsigma(config_file):
     """Compute control sigma values from eigendecomposition"""
 
-    comm = MPI.comm_world
+    comm = MPI.COMM_WORLD
 
     # Read run config file
     params = ConfigParser(config_file)

runs/run_sample.py

@@ -15,12 +15,13 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with tlm_adjoint.  If not, see <https://www.gnu.org/licenses/>.
 
-from fenics_ice.backend import Function, HDF5File, MPI, project
+from fenics_ice.backend import Function, HDF5File, project
 
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["OPENBLAS_NUM_THREADS"] = "1"
 
+import mpi4py.MPI as MPI  # noqa: N817
 import sys
 import numpy as np
 import pickle
@@ -120,7 +121,7 @@ def run_sample(config_file):
             lam = lam[:max_lam] 
 
         y = Function(space)
-        with HDF5File(MPI.comm_world, str(outdir_e/vecfile), 'r') as hdf5data:
+        with HDF5File(MPI.COMM_WORLD, str(outdir_e/vecfile), 'r') as hdf5data:
             for i in range(len(lam)):
                 w = Function(space)
                 hdf5data.read(w, f'v/vector_{i}')

tests/conftest.py

@@ -129,7 +129,7 @@ def pytest_configure(config):
     pytest.parallel = MPI.COMM_WORLD.size > 1
 
     comm = MPI.COMM_WORLD
-    rank = comm.Get_rank()
+    rank = comm.rank
 
     # Clone the input data once even when test are run in parallel
     if rank == 0:

tests/test_config.py

@@ -17,7 +17,6 @@
 
 import fenics_ice.backend as fe
 
-from mpi4py import MPI
 import pytest
 import numpy as np
 import fenics_ice as fice

tests/test_runs.py

@@ -24,7 +24,7 @@
 from runs import run_inv, run_forward, run_eigendec, run_errorprop, run_invsigma
 from fenics_ice import config
 import shutil
-from mpi4py import MPI
+
 
 def EQReset():
     """Take care of tlm_adjoint EquationManager"""
@@ -396,4 +396,4 @@ def test_run_smith_error_prop(temp_model, monkeypatch):
 
     pytest.check_float_result(Q_sigma,
                               Q_sigma_expected_at_50_eival,
-                              work_dir, 'expected_Q_sigma')
+                              work_dir, 'expected_Q_sigma')