temp

Author: majosm

grudge/discretization.py

@@ -74,7 +74,6 @@
 
 
 PartID = Tuple[VolumeTag, int]
-ConvertibleToPartID = Union[Tuple[()], VolumeTag, int, PartID]
 
 
 # {{{ part ID normalization
@@ -190,7 +189,7 @@ def __init__(self, array_context: ArrayContext,
                 Mapping[DiscretizationTag, ElementGroupFactory]] = None,
             mpi_communicator: Optional["mpi4py.MPI.Intracomm"] = None,
             inter_part_connections: Optional[
-                Mapping[Tuple[ConvertibleToPartID, ConvertibleToPartID],
+                Mapping[Tuple[PartID, PartID],
                     DiscretizationConnection]] = None,
             ) -> None:
         """
@@ -287,32 +286,7 @@ def promote_to_part_id(key):
                         "DiscretizationCollection constructor is called in "
                         "'modern' mode")
 
-            if VTAG_ALL not in volume_discrs.keys():
-                # Multi-volume
-                if mpi_communicator is not None:
-                    def promote_to_part_id(key):
-                        return key
-                else:
-                    def promote_to_part_id(key):
-                        return (key, None)
-            else:
-                # Single-volume
-                if mpi_communicator is not None:
-                    def promote_to_part_id(key):
-                        return (VTAG_ALL, key)
-                else:
-                    def promote_to_part_id(key):
-                        return (VTAG_ALL, None)
-
-            volume_discrs = {
-                vtag: Discretization(
-                    array_context, _normalize_mesh_part_ids(volume_discr.mesh),
-                    self.group_factory_for_discretization_tag(DISCR_TAG_BASE))
-                for vtag, volume_discr in volume_discrs.items()}
-
-            self._inter_part_connections = {
-                promote_to_part_id(key): value
-                for key, value in inter_part_connections.items()}
+            self._inter_part_connections = inter_part_connections
 
         self._volume_discrs = volume_discrs
 
@@ -979,16 +953,41 @@ def make_discretization_collection(
 
     del order
 
-    volume_discrs = {
-            vtag: (
-                Discretization(
-                    array_context, mesh_or_discr,
-                    discr_tag_to_group_factory[DISCR_TAG_BASE])
-                if isinstance(mesh_or_discr, Mesh) else mesh_or_discr)
-            for vtag, mesh_or_discr in volumes.items()}
-
     mpi_communicator = getattr(array_context, "mpi_communicator", None)
 
+    if mpi_communicator is not None:
+        rank = mpi_communicator.Get_rank()
+    else:
+        rank = None
+
+    if VTAG_ALL not in volumes.keys():
+        # Multi-volume
+        if mpi_communicator is not None:
+            def promote_to_part_id(key):
+                return key
+        else:
+            def promote_to_part_id(key):
+                return (key, None)
+    else:
+        # Single-volume
+        if mpi_communicator is not None:
+            def promote_to_part_id(key):
+                return (VTAG_ALL, key)
+        else:
+            def promote_to_part_id(key):
+                return (VTAG_ALL, None)
+
+    if any(
+            isinstance(mesh_or_discr, Discretization)
+            for mesh_or_discr in volumes.values()):
+        raise NotImplementedError("Doesn't work at the moment")
+
+    volume_discrs = {
+        vtag: Discretization(
+            array_context, _normalize_mesh_part_ids(mesh, promote_to_part_id),
+            discr_tag_to_group_factory[DISCR_TAG_BASE])
+        for vtag, mesh in volumes.items()}
+
     return DiscretizationCollection(
             array_context=array_context,
             volume_discrs=volume_discrs,