mpi: Add utility to get number of ranks on a single node

devito/mpi/distributed.py

@@ -256,6 +256,16 @@ def nprocs(self):
         else:
             return 1
 
+    @property
+    def nprocs_local(self):
+        if self.comm is not MPI.COMM_NULL:
+            local_comm = MPI.Comm.Split_type(self.comm, MPI.COMM_TYPE_SHARED)
+            node_size = local_comm.Get_size()
+            local_comm.Free()
+            return node_size
+        else:
+            return 1
+
     @property
     def topology(self):
         return self._topology

examples/seismic/preset_models.py

@@ -7,6 +7,16 @@
 __all__ = ['demo_model']
 
 
+def Gardners(vp, normalize=True):
+    """
+    Gardner's relation for vp in km/s
+    """
+    b = 1 / (0.31 * (1e3*vp)**0.25)
+    if normalize:
+        b[vp < 1.51] = 1.0
+    return b
+
+
 def demo_model(preset, **kwargs):
     """
     Utility function to create preset `Model` objects for
@@ -49,6 +59,7 @@ def demo_model(preset, **kwargs):
     vp = kwargs.pop('vp', 1.5)
     nlayers = kwargs.pop('nlayers', 3)
     fs = kwargs.pop('fs', False)
+    density = kwargs.pop('density', False)
 
     if preset.lower() in ['constant-elastic']:
         # A constant single-layer model in a 2D or 3D domain
@@ -76,7 +87,8 @@ def demo_model(preset, **kwargs):
     if preset.lower() in ['constant-isotropic']:
         # A constant single-layer model in a 2D or 3D domain
         # with velocity 1.5 km/s.
-
+        if density:
+            kwargs['b'] = 1
         return SeismicModel(space_order=space_order, vp=vp, origin=origin, shape=shape,
                             dtype=dtype, spacing=spacing, nbl=nbl, fs=fs, **kwargs)
 
@@ -99,7 +111,8 @@ def demo_model(preset, **kwargs):
         phi = None
         if len(shape) > 2 and preset.lower() not in ['constant-tti-noazimuth']:
             phi = .35
-
+        if density:
+            kwargs['b'] = 1
         return SeismicModel(space_order=space_order, vp=vp, origin=origin, shape=shape,
                             dtype=dtype, spacing=spacing, nbl=nbl, epsilon=epsilon,
                             delta=delta, theta=theta, phi=phi, bcs="damp", **kwargs)
@@ -119,6 +132,9 @@ def demo_model(preset, **kwargs):
         for i in range(1, nlayers):
             v[..., i*int(shape[-1] / nlayers):] = vp_i[i]  # Bottom velocity
 
+        if density:
+            kwargs['b'] = Gardners(v)
+
         return SeismicModel(space_order=space_order, vp=v, origin=origin, shape=shape,
                             dtype=dtype, spacing=spacing, nbl=nbl, bcs="damp",
                             fs=fs, **kwargs)
@@ -139,8 +155,7 @@ def demo_model(preset, **kwargs):
             v[..., i*int(shape[-1] / nlayers):] = vp_i[i]  # Bottom velocity
 
         vs = 0.5 * v[:]
-        b = 1 / (0.31 * (1e3*v)**0.25)
-        b[v < 1.51] = 1.0
+        b = Gardners(v)
         vs[v < 1.51] = 0.0
 
         return SeismicModel(space_order=space_order, vp=v, vs=vs, b=b,
@@ -214,6 +229,9 @@ def demo_model(preset, **kwargs):
         if len(shape) > 2 and preset.lower() not in ['layers-tti-noazimuth']:
             phi = .25*(v - 1.5)
 
+        if density:
+            kwargs['b'] = Gardners(v)
+
         model = SeismicModel(space_order=space_order, vp=v, origin=origin, shape=shape,
                              dtype=dtype, spacing=spacing, nbl=nbl, epsilon=epsilon,
                              delta=delta, theta=theta, phi=phi, bcs="damp",
@@ -245,6 +263,9 @@ def demo_model(preset, **kwargs):
         y, x = np.ogrid[-a:shape[0]-a, -b:shape[1]-b]
         v[x*x + y*y <= r*r] = vp
 
+        if density:
+            kwargs['b'] = Gardners(v)
+
         return SeismicModel(space_order=space_order, vp=v, origin=origin, shape=shape,
                             dtype=dtype, spacing=spacing, nbl=nbl, bcs="damp",
                             fs=fs, **kwargs)
@@ -347,7 +368,7 @@ def demo_model(preset, **kwargs):
 
         qp[:] = 3.516*((vp[:]*1000.)**2.2)*10**(-6)  # Li's empirical formula
 
-        b = 1 / (0.31*(vp[:]*1000.)**0.25)  # Gardner's relation
+        b = Gardners(vp, normalize=False)
 
         return SeismicModel(space_order=space_order, vp=vp, qp=qp, b=b, nbl=nbl,
                             dtype=dtype, origin=origin, shape=shape,

tests/test_mpi.py

@@ -35,6 +35,7 @@ def test_partitioning(self):
         }
         assert f.shape == expected[distributor.nprocs][distributor.myrank]
         assert f.size_global == 225
+        assert distributor.nprocs_local == distributor.nprocs
 
     @pytest.mark.parallel(mode=[2, 4])
     def test_partitioning_fewer_dims(self):