From 06e8469b2536a702935930155c11770ea69da7ab Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Wed, 1 Jun 2022 02:29:45 +0100 Subject: [PATCH 1/3] io: enable saving/loading distribution and permutation --- firedrake/checkpointing.py | 398 ++++++++++++------ firedrake/mesh.py | 77 +++- ...test_io_freeze_distribution_permutation.py | 81 ++++ 3 files changed, 414 insertions(+), 142 deletions(-) create mode 100644 tests/output/test_io_freeze_distribution_permutation.py diff --git a/firedrake/checkpointing.py b/firedrake/checkpointing.py index d029f1dc02..968adec713 100644 --- a/firedrake/checkpointing.py +++ b/firedrake/checkpointing.py @@ -8,7 +8,7 @@ from firedrake.cython import hdf5interface as h5i from firedrake.cython import dmcommon from firedrake.petsc import PETSc, OptionsManager -from firedrake.mesh import MeshTopology, ExtrudedMeshTopology, DEFAULT_MESH_NAME, make_mesh_from_coordinates +from firedrake.mesh import MeshTopology, ExtrudedMeshTopology, DEFAULT_MESH_NAME, make_mesh_from_coordinates, DistributedMeshOverlapType from firedrake.functionspace import FunctionSpace from firedrake import functionspaceimpl as impl from firedrake.functionspacedata import get_global_numbering, create_element @@ -46,6 +46,13 @@ r"""The prefix attached to the DG function resulting from projecting the original function to the embedding DG space.""" +# This is the distribution_parameters and reorder that one must use when +# distribution and permutation are loaded. +distribution_parameters_noop = {"partition": False, + "overlap_type": (DistributedMeshOverlapType.NONE, 0)} +reorder_noop = None + + class DumbCheckpoint(object): r"""A very dumb checkpoint object. @@ -537,9 +544,10 @@ def save_mesh(self, mesh): :arg mesh: the mesh to save. """ + mesh.init() # Handle extruded mesh tmesh = mesh.topology - if isinstance(tmesh, ExtrudedMeshTopology): + if mesh.extruded: # -- Save mesh topology -- base_tmesh = mesh._base_mesh.topology self._save_mesh_topology(base_tmesh) @@ -604,7 +612,6 @@ def save_mesh(self, mesh): path = self._path_to_mesh(tmesh.name, mesh.name) self.require_group(path) # Save Firedrake coodinates. - mesh.init() self.set_attr(path, PREFIX + "_coordinate_element", self._pickle(mesh._coordinates.function_space().ufl_element())) self.set_attr(path, PREFIX + "_coordinates", mesh._coordinates.name()) self._save_function_topology(mesh._coordinates) @@ -617,8 +624,12 @@ def save_mesh(self, mesh): @PETSc.Log.EventDecorator("SaveMeshTopology") def _save_mesh_topology(self, tmesh): # -- Save DMPlex -- + tmesh.init() topology_dm = tmesh.topology_dm tmesh_name = topology_dm.getName() + distribution_name = tmesh._distribution_name + perm_is = tmesh._plex_renumbering + permutation_name = tmesh._permutation_name if tmesh_name in self.require_group(self._path_to_topologies()): # Check if the global number of DMPlex points and # the global sum of DMPlex cone sizes are consistent. @@ -630,12 +641,50 @@ def _save_mesh_topology(self, tmesh): raise ValueError(f"Mesh ({tmesh_name}) already exists in {self.filename}, but the global number of DMPlex points is inconsistent: {order_array_size1} ({self.filename}) != {order_array_size} ({tmesh_name})") if ornt_array_size1 != ornt_array_size: raise ValueError(f"Mesh ({tmesh_name}) already exists in {self.filename}, but the global sum of all DMPlex cone sizes is inconsistent: {ornt_array_size1} ({self.filename}) != {ornt_array_size} ({tmesh_name})") + # We assume that each (conceptually the same) mesh topology (plex) + # is uniquely named (this is users' responsibility). + # With the current setup, "distributions" folder will always contain + # one and only one distribution. This can be amended once we make + # the global point numbers independent of the distributions. + path = self._path_to_distributions(tmesh_name) + if path not in self.h5pyfile: + # backward campat.: old files do not have "distributions" folder. + # Just save mesh + dist + perm. + do_save = True + else: + # There should be one and only one distribution. + _distribution_name, = self.h5pyfile[path].keys() + if _distribution_name != distribution_name: + raise RuntimeError(f"Can only save one distribution: found {_distribution_name} in file, but trying to save {distribution_name}") + else: + path = self._path_to_distribution(tmesh_name, distribution_name) + _comm_size = self.h5pyfile[path]["chart_sizes"].size + if _comm_size != topology_dm.comm.size: + raise RuntimeError(f"Trying to save distribution of the same name ({distribution_name}), but it was previously saved with {_comm_size} processes, and now you are using {topology_dm.comm.size} processes") + path = self._path_to_permutations(tmesh_name, distribution_name) + # There should be one and only one permutation. + _permutation_name, = self.h5pyfile[path].keys() + if _permutation_name != permutation_name: + raise RuntimeError(f"Can only save one permutation: found {_permutation_name} in file, but trying to save {permutation_name}") + # Regard tmesh + distribution + permutation has already been saved. + do_save = False else: + do_save = True + if do_save: self.viewer.pushFormat(format=ViewerHDF5.Format.HDF5_PETSC) with self.opts.inserted_options(): + topology_dm.distributionSetName(distribution_name) topology_dm.topologyView(viewer=self.viewer) + topology_dm.distributionSetName(None) topology_dm.labelsView(viewer=self.viewer) self.viewer.popFormat() + path = self._path_to_permutation(tmesh_name, distribution_name, permutation_name) + self.require_group(path) + self.viewer.pushGroup(path) + perm_is.setName("permutation") + perm_is.view(self.viewer) + perm_is.setName(None) + self.viewer.popGroup() @PETSc.Log.EventDecorator("SaveFunctionSpace") def _save_function_space(self, V): @@ -694,32 +743,31 @@ def _save_function_space_topology(self, tV): # create a symbolic link as /topologies/tmesh.name/topology <- /topologies/base_tmesh.name/topology # to have a full structure under tmesh.name, but at least for now we don't need to. base_tmesh_name = topology_dm.getName() - topology_dm.setName(tmesh.name) with self.opts.inserted_options(): + topology_dm.setName(tmesh.name) topology_dm.sectionView(self.viewer, dm) - topology_dm.setName(base_tmesh_name) + topology_dm.setName(base_tmesh_name) @PETSc.Log.EventDecorator("SaveFunction") def save_function(self, f, idx=None, name=None): r"""Save a :class:`~.Function`. :arg f: the :class:`~.Function` to save. - :arg idx: optional timestepping index. A function can + :kwarg idx: optional timestepping index. A function can either be saved in timestepping mode or in normal mode (non-timestepping); for each function of interest, this method must always be called with the idx parameter set or never be called with the idx parameter set. - :arg name: optional alternative name to save the function under + :kwarg name: optional alternative name to save the function under. """ + V = f.function_space() + mesh = V.mesh() if name: - g = Function(f.function_space(), val=f.dat, name=name) + g = Function(V, val=f.dat, name=name) return self.save_function(g, idx=idx) - # -- Save function space -- - V = f.function_space() self._save_function_space(V) # -- Save function -- - mesh = V.mesh() V_name = self._generate_function_space_name(V) if isinstance(V.topological, impl.MixedFunctionSpace): base_path = self._path_to_mixed_function(mesh.name, V_name, f.name()) @@ -788,10 +836,10 @@ def _save_function_topology(self, tf, idx=None): with tf.dat.vec_ro as vec: vec.setName(tf.name()) base_tmesh_name = topology_dm.getName() - topology_dm.setName(tmesh.name) with self.opts.inserted_options(): + topology_dm.setName(tmesh.name) topology_dm.globalVectorView(self.viewer, dm, vec) - topology_dm.setName(base_tmesh_name) + topology_dm.setName(base_tmesh_name) if idx is not None: self.viewer.popTimestepping() @@ -805,78 +853,94 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter distributing the mesh; see :func:`~.Mesh`. :returns: the loaded mesh. """ - if reorder is None: - reorder = parameters["reorder_meshes"] - if distribution_parameters is None: - distribution_parameters = {} - mesh_key = self._generate_mesh_key(name, reorder, distribution_parameters) - if mesh_key in self._mesh_cache: - return self._mesh_cache[mesh_key] tmesh_name = self._get_mesh_name_topology_name_map()[name] path = self._path_to_topology_extruded(tmesh_name) if path in self.h5pyfile: # -- Load mesh topology -- base_tmesh_name = self.get_attr(path, PREFIX_EXTRUDED + "_base_mesh") base_tmesh = self._load_mesh_topology(base_tmesh_name, reorder, distribution_parameters) - variable_layers = self.get_attr(path, PREFIX_EXTRUDED + "_variable_layers") - if variable_layers: - cell = base_tmesh.ufl_cell() - element = ufl.VectorElement("DP" if cell.is_simplex() else "DQ", cell, 0, dim=2) - _ = self._load_function_space_topology(base_tmesh, element) - base_tmesh_key = self._generate_mesh_key(base_tmesh.name, base_tmesh._did_reordering, base_tmesh._distribution_parameters) - sd_key = self._get_shared_data_key_for_checkpointing(base_tmesh, element) - _, _, lsf = self._function_load_utils[base_tmesh_key + sd_key] - nroots, _, _ = lsf.getGraph() - layers_a = np.empty(nroots, dtype=utils.IntType) - layers_a_iset = PETSc.IS().createGeneral(layers_a, comm=self.viewer.comm) - layers_a_iset.setName("_".join([PREFIX_EXTRUDED, "layers_iset"])) - self.viewer.pushGroup(path) - layers_a_iset.load(self.viewer) - self.viewer.popGroup() - layers_a = layers_a_iset.getIndices() - layers = np.empty((base_tmesh.cell_set.total_size, 2), dtype=utils.IntType) - unit = MPI._typedict[np.dtype(utils.IntType).char] - lsf.bcastBegin(unit, layers_a, layers, MPI.REPLACE) - lsf.bcastEnd(unit, layers_a, layers, MPI.REPLACE) + base_tmesh.init() + tmesh_key = self._generate_mesh_key_from_names(tmesh_name, + base_tmesh._distribution_name, + base_tmesh._permutation_name) + if tmesh_key in self._tmesh_cache: + tmesh = self._tmesh_cache[tmesh_key] else: - layers = self.get_attr(path, PREFIX_EXTRUDED + "_layers") - tmesh = ExtrudedMeshTopology(base_tmesh, layers, name=tmesh_name) + variable_layers = self.get_attr(path, PREFIX_EXTRUDED + "_variable_layers") + if variable_layers: + cell = base_tmesh.ufl_cell() + element = ufl.VectorElement("DP" if cell.is_simplex() else "DQ", cell, 0, dim=2) + _ = self._load_function_space_topology(base_tmesh, element) + base_tmesh_key = self._generate_mesh_key_from_names(base_tmesh.name, + base_tmesh._distribution_name, + base_tmesh._permutation_name) + sd_key = self._get_shared_data_key_for_checkpointing(base_tmesh, element) + _, _, lsf = self._function_load_utils[base_tmesh_key + sd_key] + nroots, _, _ = lsf.getGraph() + layers_a = np.empty(nroots, dtype=utils.IntType) + layers_a_iset = PETSc.IS().createGeneral(layers_a, comm=self.viewer.comm) + layers_a_iset.setName("_".join([PREFIX_EXTRUDED, "layers_iset"])) + self.viewer.pushGroup(path) + layers_a_iset.load(self.viewer) + self.viewer.popGroup() + layers_a = layers_a_iset.getIndices() + layers = np.empty((base_tmesh.cell_set.total_size, 2), dtype=utils.IntType) + unit = MPI._typedict[np.dtype(utils.IntType).char] + lsf.bcastBegin(unit, layers_a, layers, MPI.REPLACE) + lsf.bcastEnd(unit, layers_a, layers, MPI.REPLACE) + else: + layers = self.get_attr(path, PREFIX_EXTRUDED + "_layers") + tmesh = ExtrudedMeshTopology(base_tmesh, layers, name=tmesh_name) + self._tmesh_cache[tmesh_key] = tmesh # -- Load mesh -- - path = self._path_to_mesh(tmesh_name, name) - coord_element = self._unpickle(self.get_attr(path, PREFIX + "_coordinate_element")) - coord_name = self.get_attr(path, PREFIX + "_coordinates") - coordinates = self._load_function_topology(tmesh, coord_element, coord_name) - mesh = make_mesh_from_coordinates(coordinates, name) - if self.has_attr(path, PREFIX + "_radial_coordinates"): - radial_coord_element = self._unpickle(self.get_attr(path, PREFIX + "_radial_coordinate_element")) - radial_coord_name = self.get_attr(path, PREFIX + "_radial_coordinates") - radial_coordinates = self._load_function_topology(tmesh, radial_coord_element, radial_coord_name) - tV_radial_coord = impl.FunctionSpace(tmesh, radial_coord_element) - V_radial_coord = impl.WithGeometry.create(tV_radial_coord, mesh) - radial_coord_function_name = self.get_attr(path, PREFIX + "_radial_coordinate_function") - mesh.radial_coordinates = Function(V_radial_coord, val=radial_coordinates, name=radial_coord_function_name) - # The followings are conceptually redundant, but needed. - path = os.path.join(self._path_to_mesh(tmesh_name, name), PREFIX_EXTRUDED) - base_mesh_name = self.get_attr(path, PREFIX_EXTRUDED + "_base_mesh") - mesh._base_mesh = self.load_mesh(base_mesh_name) + mesh_key = self._generate_mesh_key_from_names(name, + base_tmesh._distribution_name, + base_tmesh._permutation_name) + if mesh_key in self._mesh_cache: + mesh = self._mesh_cache[mesh_key] + else: + path = self._path_to_mesh(tmesh_name, name) + coord_element = self._unpickle(self.get_attr(path, PREFIX + "_coordinate_element")) + coord_name = self.get_attr(path, PREFIX + "_coordinates") + coordinates = self._load_function_topology(tmesh, coord_element, coord_name) + mesh = make_mesh_from_coordinates(coordinates, name) + if self.has_attr(path, PREFIX + "_radial_coordinates"): + radial_coord_element = self._unpickle(self.get_attr(path, PREFIX + "_radial_coordinate_element")) + radial_coord_name = self.get_attr(path, PREFIX + "_radial_coordinates") + radial_coordinates = self._load_function_topology(tmesh, radial_coord_element, radial_coord_name) + tV_radial_coord = impl.FunctionSpace(tmesh, radial_coord_element) + V_radial_coord = impl.WithGeometry.create(tV_radial_coord, mesh) + radial_coord_function_name = self.get_attr(path, PREFIX + "_radial_coordinate_function") + mesh.radial_coordinates = Function(V_radial_coord, val=radial_coordinates, name=radial_coord_function_name) + # The followings are conceptually redundant, but needed. + path = os.path.join(self._path_to_mesh(tmesh_name, name), PREFIX_EXTRUDED) + base_mesh_name = self.get_attr(path, PREFIX_EXTRUDED + "_base_mesh") + mesh._base_mesh = self.load_mesh(base_mesh_name) + self._mesh_cache[mesh_key] = mesh else: utils._init() # -- Load mesh topology -- tmesh = self._load_mesh_topology(tmesh_name, reorder, distribution_parameters) - # -- Load coordinates -- - # tmesh.topology_dm has already been redistributed. - path = self._path_to_mesh(tmesh_name, name) - # Load firedrake coordinates directly. - # When implementing checkpointing for MeshHierarchy in the future, - # we will need to postpone calling tmesh.init(). - tmesh.init() - coord_element = self._unpickle(self.get_attr(path, PREFIX + "_coordinate_element")) - coord_name = self.get_attr(path, PREFIX + "_coordinates") - coordinates = self._load_function_topology(tmesh, coord_element, coord_name) - mesh = make_mesh_from_coordinates(coordinates, name) - # Load plex coordinates for a complete representation of plex. - tmesh.topology_dm.coordinatesLoad(self.viewer, tmesh.sfXC) - self._mesh_cache[mesh_key] = mesh + mesh_key = self._generate_mesh_key_from_names(name, + tmesh._distribution_name, + tmesh._permutation_name) + if mesh_key in self._mesh_cache: + mesh = self._mesh_cache[mesh_key] + else: + # -- Load coordinates -- + # tmesh.topology_dm has already been redistributed. + path = self._path_to_mesh(tmesh_name, name) + # Load firedrake coordinates directly. + # When implementing checkpointing for MeshHierarchy in the future, + # we will need to postpone calling tmesh.init(). + tmesh.init() + coord_element = self._unpickle(self.get_attr(path, PREFIX + "_coordinate_element")) + coord_name = self.get_attr(path, PREFIX + "_coordinates") + coordinates = self._load_function_topology(tmesh, coord_element, coord_name) + mesh = make_mesh_from_coordinates(coordinates, name) + # Load plex coordinates for a complete representation of plex. + tmesh.topology_dm.coordinatesLoad(self.viewer, tmesh.sfXC) + self._mesh_cache[mesh_key] = mesh return mesh @PETSc.Log.EventDecorator("LoadMeshTopology") @@ -890,44 +954,96 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters): :returns: The loaded :class:`~.MeshTopology`. """ # -- Load DMPlex -- - tmesh_key = self._generate_mesh_key(tmesh_name, reorder, distribution_parameters) + path = self._path_to_distributions(tmesh_name) + load_distribution_permutation = False + if path in self.h5pyfile: + _distribution_name, = self.h5pyfile[path].keys() + path = self._path_to_distribution(tmesh_name, _distribution_name) + _comm_size = self.get_attr(path, "comm_size") + if _comm_size == self.viewer.comm.size and \ + distribution_parameters is None and reorder is None: + load_distribution_permutation = True + if load_distribution_permutation: + path = self._path_to_distributions(tmesh_name) + distribution_name, = self.h5pyfile[path].keys() + path = self._path_to_permutations(tmesh_name, distribution_name) + permutation_name, = self.h5pyfile[path].keys() + distribution_parameters = distribution_parameters_noop + reorder = reorder_noop + else: + # Here mimic what Mesh function does. + if reorder is None: + reorder = parameters["reorder_meshes"] + if distribution_parameters is None: + distribution_parameters = {} + distribution_name = None + permutation_name = None + perm_is = None + # This is only to return the same tmesh object if the same set of arguments are given. + # Multiple tmesh_key might end up having the same value, but it is hard to process + # all distribution and reorder options at this stage (many things happen in MeshTopology constructor). + tmesh_key = self._generate_mesh_key(tmesh_name, distribution_name, permutation_name, reorder, distribution_parameters) if tmesh_key in self._tmesh_cache: - return self._tmesh_cache[tmesh_key] - plex = PETSc.DMPlex() - plex.create(comm=self.viewer.comm) - plex.setName(tmesh_name) - # Check format - path = os.path.join(self._path_to_topology(tmesh_name), "topology") - if any(d not in self.h5pyfile for d in [os.path.join(path, "cells"), - os.path.join(path, "cones"), - os.path.join(path, "order"), - os.path.join(path, "orientation")]): - raise RuntimeError(f"Unsupported PETSc ViewerHDF5 format used in {self.filename}") - format = ViewerHDF5.Format.HDF5_PETSC - self.viewer.pushFormat(format=format) - sfXB = plex.topologyLoad(self.viewer) - self.viewer.popFormat() - # -- Construct Mesh (Topology) -- - tmesh = MeshTopology(plex, name=plex.getName(), reorder=reorder, - distribution_parameters=distribution_parameters, sfXB=sfXB) - self.viewer.pushFormat(format=format) - # tmesh.topology_dm has already been redistributed. - sfXCtemp = tmesh.sfXB.compose(tmesh.sfBC) if tmesh.sfBC is not None else tmesh.sfXB - plex.labelsLoad(self.viewer, sfXCtemp) - self.viewer.popFormat() - # These labels are distribution dependent. - # We should be able to save/load labels selectively. - plex.removeLabel("pyop2_core") - plex.removeLabel("pyop2_owned") - plex.removeLabel("pyop2_ghost") - self._tmesh_cache[tmesh_key] = tmesh + tmesh = self._tmesh_cache[tmesh_key] + else: + plex = PETSc.DMPlex() + plex.create(comm=self.viewer.comm) + plex.setName(tmesh_name) + # Check format + path = os.path.join(self._path_to_topology(tmesh_name), "topology") + if any(d not in self.h5pyfile for d in [os.path.join(path, "cells"), + os.path.join(path, "cones"), + os.path.join(path, "order"), + os.path.join(path, "orientation")]): + raise RuntimeError(f"Unsupported PETSc ViewerHDF5 format used in {self.filename}") + format = ViewerHDF5.Format.HDF5_PETSC + self.viewer.pushFormat(format=format) + plex.distributionSetName(distribution_name) + sfXB = plex.topologyLoad(self.viewer) + plex.distributionSetName(None) + self.viewer.popFormat() + if load_distribution_permutation: + chart_size = np.empty(1, dtype=utils.IntType) + chart_sizes_iset = PETSc.IS().createGeneral(chart_size, comm=self.viewer.comm) + chart_sizes_iset.setName("chart_sizes") + path = self._path_to_distribution(tmesh_name, distribution_name) + self.viewer.pushGroup(path) + chart_sizes_iset.load(self.viewer) + self.viewer.popGroup() + chart_size = chart_sizes_iset.getIndices().item() + perm = np.empty(chart_size, dtype=utils.IntType) + perm_is = PETSc.IS().createGeneral(perm, comm=self.viewer.comm) + path = self._path_to_permutation(tmesh_name, distribution_name, permutation_name) + self.viewer.pushGroup(path) + perm_is.setName("permutation") + perm_is.load(self.viewer) + perm_is.setName(None) + self.viewer.popGroup() + else: + perm_is = None + # -- Construct Mesh (Topology) -- + tmesh = MeshTopology(plex, name=plex.getName(), reorder=reorder, + distribution_parameters=distribution_parameters, sfXB=sfXB, perm_is=perm_is, + distribution_name=distribution_name, permutation_name=permutation_name) + self.viewer.pushFormat(format=format) + # tmesh.topology_dm has already been redistributed. + sfXCtemp = tmesh.sfXB.compose(tmesh.sfBC) if tmesh.sfBC is not None else tmesh.sfXB + plex.labelsLoad(self.viewer, sfXCtemp) + self.viewer.popFormat() + # These labels are distribution dependent. + # We should be able to save/load labels selectively. + plex.removeLabel("pyop2_core") + plex.removeLabel("pyop2_owned") + plex.removeLabel("pyop2_ghost") + self._tmesh_cache[tmesh_key] = tmesh return tmesh @PETSc.Log.EventDecorator("LoadFunctionSpace") def _load_function_space(self, mesh, name): mesh.init() - mesh_key = self._generate_mesh_key(mesh.name, mesh.topology._did_reordering, - mesh.topology._distribution_parameters) + mesh_key = self._generate_mesh_key_from_names(mesh.name, + mesh.topology._distribution_name, + mesh.topology._permutation_name) V_key = mesh_key + (name, ) if V_key in self._function_spaces: return self._function_spaces[V_key] @@ -964,30 +1080,31 @@ def _load_function_space_topology(self, tmesh, element): tmesh.init() if element.family() == "Real": return impl.RealFunctionSpace(tmesh, element, "unused_name") - tmesh_key = self._generate_mesh_key(tmesh.name, tmesh._did_reordering, tmesh._distribution_parameters) + tmesh_key = self._generate_mesh_key_from_names(tmesh.name, + tmesh._distribution_name, + tmesh._permutation_name) sd_key = self._get_shared_data_key_for_checkpointing(tmesh, element) - if tmesh_key + sd_key in self._function_load_utils: - return impl.FunctionSpace(tmesh, element) - topology_dm = tmesh.topology_dm - dm = PETSc.DMShell().create(comm=topology_dm.comm) - dm.setName(self._get_dm_name_for_checkpointing(tmesh, element)) - dm.setPointSF(topology_dm.getPointSF()) - section = PETSc.Section().create(comm=topology_dm.comm) - section.setPermutation(tmesh._plex_renumbering) - dm.setSection(section) - base_tmesh = tmesh._base_mesh if isinstance(tmesh, ExtrudedMeshTopology) else tmesh - sfXC = base_tmesh.sfXC - topology_dm.setName(tmesh.name) - gsf, lsf = topology_dm.sectionLoad(self.viewer, dm, sfXC) - topology_dm.setName(base_tmesh.name) - nodes_per_entity, real_tensorproduct, block_size = sd_key - # Don't cache if the section has been expanded by block_size - if block_size == 1: - cached_section = get_global_numbering(tmesh, (nodes_per_entity, real_tensorproduct), global_numbering=dm.getSection()) - if dm.getSection() is not cached_section: - # The same section has already been cached. - dm.setSection(cached_section) - self._function_load_utils[tmesh_key + sd_key] = (dm, gsf, lsf) + if tmesh_key + sd_key not in self._function_load_utils: + topology_dm = tmesh.topology_dm + dm = PETSc.DMShell().create(comm=topology_dm.comm) + dm.setName(self._get_dm_name_for_checkpointing(tmesh, element)) + dm.setPointSF(topology_dm.getPointSF()) + section = PETSc.Section().create(comm=topology_dm.comm) + section.setPermutation(tmesh._plex_renumbering) + dm.setSection(section) + base_tmesh = tmesh._base_mesh if isinstance(tmesh, ExtrudedMeshTopology) else tmesh + sfXC = base_tmesh.sfXC + topology_dm.setName(tmesh.name) + gsf, lsf = topology_dm.sectionLoad(self.viewer, dm, sfXC) + topology_dm.setName(base_tmesh.name) + nodes_per_entity, real_tensorproduct, block_size = sd_key + # Don't cache if the section has been expanded by block_size + if block_size == 1: + cached_section = get_global_numbering(tmesh, (nodes_per_entity, real_tensorproduct), global_numbering=dm.getSection()) + if dm.getSection() is not cached_section: + # The same section has already been cached. + dm.setSection(cached_section) + self._function_load_utils[tmesh_key + sd_key] = (dm, gsf, lsf) return impl.FunctionSpace(tmesh, element) @PETSc.Log.EventDecorator("LoadFunction") @@ -996,7 +1113,7 @@ def load_function(self, mesh, name, idx=None): :arg mesh: the mesh on which the function is defined. :arg name: the name of the :class:`~.Function` to load. - :arg idx: optional timestepping index. A function can + :kwarg idx: optional timestepping index. A function can be loaded with idx only when it was saved with idx. :returns: the loaded :class:`~.Function`. """ @@ -1071,7 +1188,9 @@ def _load_function_topology(self, tmesh, element, tf_name, idx=None): with tf.dat.vec_wo as vec: vec.setName(tf_name) sd_key = self._get_shared_data_key_for_checkpointing(tmesh, element) - tmesh_key = self._generate_mesh_key(tmesh.name, tmesh._did_reordering, tmesh._distribution_parameters) + tmesh_key = self._generate_mesh_key_from_names(tmesh.name, + tmesh._distribution_name, + tmesh._permutation_name) dm, sf, _ = self._function_load_utils[tmesh_key + sd_key] base_tmesh_name = topology_dm.getName() topology_dm.setName(tmesh.name) @@ -1081,9 +1200,12 @@ def _load_function_topology(self, tmesh, element, tf_name, idx=None): self.viewer.popTimestepping() return tf - def _generate_mesh_key(self, mesh_name, reorder, distribution_parameters): + def _generate_mesh_key(self, mesh_name, distribution_name, permutation_name, reorder, distribution_parameters): dist_key = frozenset(distribution_parameters.items()) - return (self.filename, self.commkey, mesh_name, reorder, dist_key) + return (self.filename, self.commkey, mesh_name, distribution_name, permutation_name, reorder, dist_key) + + def _generate_mesh_key_from_names(self, mesh_name, distribution_name, permutation_name): + return (self.filename, self.commkey, mesh_name, distribution_name, permutation_name) def _generate_function_space_name(self, V): """Return a unique function space name.""" @@ -1208,6 +1330,18 @@ def _path_to_mixed_functions(self, mesh_name, V_name): def _path_to_mixed_function(self, mesh_name, V_name, function_name): return os.path.join(self._path_to_mixed_functions(mesh_name, V_name), function_name) + def _path_to_distributions(self, tmesh_name): + return os.path.join(self._path_to_topology(tmesh_name), "distributions") + + def _path_to_distribution(self, tmesh_name, distribution_name): + return os.path.join(self._path_to_distributions(tmesh_name), distribution_name) + + def _path_to_permutations(self, tmesh_name, distribution_name): + return os.path.join(self._path_to_distribution(tmesh_name, distribution_name), "permutations") + + def _path_to_permutation(self, tmesh_name, distribution_name, permutation_name): + return os.path.join(self._path_to_permutations(tmesh_name, distribution_name), permutation_name) + def _pickle(self, obj): return np.void(pickle.dumps(obj)) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 866d5ae54c..4fd847deea 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -63,6 +63,29 @@ def _generate_default_mesh_topology_name(name): return "_".join([name, "topology"]) +def _generate_default_mesh_topology_distribution_name(comm_size, dist_param): + """Generate the default mesh topology permutation name. + + :arg comm_size: the size of comm. + :arg dist_param: the distribution_parameter dict. + :returns: the default mesh topology distribution name. + """ + return "_".join(["firedrake", "default", + str(comm_size), + str(dist_param["partition"]).replace(' ', ''), + str(dist_param["partitioner_type"]), + "(" + dist_param["overlap_type"][0].name + "," + str(dist_param["overlap_type"][1]) + ")"]) + + +def _generate_default_mesh_topology_permutation_name(reorder): + """Generate the default mesh topology permutation name. + + :arg reorder: the flag indicating if the reordering happened or not. + :returns: the default mesh topology permutation name. + """ + return "_".join(["firedrake", "default", str(reorder)]) + + class DistributedMeshOverlapType(enum.Enum): """How should the mesh overlap be grown for distributed meshes? @@ -711,7 +734,7 @@ class MeshTopology(AbstractMeshTopology): """A representation of mesh topology implemented on a PETSc DMPlex.""" @PETSc.Log.EventDecorator("CreateMesh") - def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None): + def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None, perm_is=None): """Half-initialise a mesh topology. :arg plex: :class:`DMPlex` representing the mesh topology @@ -719,10 +742,14 @@ def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None): :arg reorder: whether to reorder the mesh (bool) :arg distribution_parameters: options controlling mesh distribution, see :func:`Mesh` for details. - :arg sfXB: :class:`PetscSF` that pushes forward the global point number + :kwarg sfXB: :class:`PetscSF` that pushes forward the global point number slab :math:`[0, NX)` to input (naive) plex (only significant when the mesh topology is loaded from file and only passed from inside :class:`~.CheckpointFile`). + :kwarg perm_is: :class:`IS` that is used as `_plex_renumbering`; only + makes sense if we know the exact parallel distribution of `plex` + at the time of mesh topology construction like when we load mesh + along with its distribution. If given, `reorder` param will be ignored. """ super().__init__(name) @@ -732,10 +759,13 @@ def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None): distribute = distribution_parameters.get("partition") if distribute is None: distribute = True + self._distribution_parameters["partition"] = distribute partitioner_type = distribution_parameters.get("partitioner_type") + self._distribution_parameters["partitioner_type"] = partitioner_type overlap_type, overlap = distribution_parameters.get("overlap_type", (DistributedMeshOverlapType.FACET, 1)) + self._distribution_parameters["overlap_type"] = (overlap_type, overlap) if overlap < 0: raise ValueError("Overlap depth must be >= 0") if overlap_type == DistributedMeshOverlapType.NONE: @@ -792,6 +822,7 @@ def add_overlap(): # refine this mesh in parallel. Later, when we actually use # it, we grow the halo. self.set_partitioner(distribute, partitioner_type) + self._distribution_parameters["partitioner_type"] = self.get_partitioner().getType() original_name = plex.getName() sfBC = plex.distribute(overlap=0) plex.setName(original_name) @@ -823,6 +854,9 @@ def add_overlap(): cell = ufl.Cell(_cells[tdim][nfacets]) self._ufl_mesh = ufl.Mesh(ufl.VectorElement("Lagrange", cell, 1, dim=cell.topological_dimension())) + self._distribution_name = distribution_name or _generate_default_mesh_topology_distribution_name(self.topology_dm.comm.size, self._distribution_parameters) + self._permutation_name = permutation_name or _generate_default_mesh_topology_permutation_name(reorder) + def callback(self): """Finish initialisation.""" del self._callback @@ -831,7 +865,18 @@ def callback(self): if self.sfXB is not None: self.sfXC = sfXB.compose(self.sfBC) if self.sfBC else self.sfXB dmcommon.complete_facet_labels(self.topology_dm) - + # TODO: Allow users to set distribution name if they want to save + # conceptually the same mesh but with different distributions, + # e.g., those generated by different partitioners. + # This currently does not make sense since those mesh instances + # of different distributions in general have different global + # point numbers (so they must be saved under different mesh names + # even though they are conceptually the same). + # The name set here almost uniquely identifies a distribution, but + # there is no gurantee that it really does or it continues to do so + # there are lots of parameters that can change distributions. + # Thus, when using CheckpointFile, it is recommended that the user set + # distribution_name explicitly. if reorder: with PETSc.Log.Event("Mesh: reorder"): old_to_new = self.topology_dm.getOrdering(PETSc.Mat.OrderingType.RCM).indices @@ -841,24 +886,23 @@ def callback(self): # No reordering reordering = None self._did_reordering = bool(reorder) - # Mark OP2 entities and derive the resulting Plex renumbering with PETSc.Log.Event("Mesh: numbering"): dmcommon.mark_entity_classes(self.topology_dm) self._entity_classes = dmcommon.get_entity_classes(self.topology_dm).astype(int) - self._plex_renumbering = dmcommon.plex_renumbering(self.topology_dm, - self._entity_classes, - reordering) - + if perm_is: + self._plex_renumbering = perm_is + else: + self._plex_renumbering = dmcommon.plex_renumbering(self.topology_dm, + self._entity_classes, + reordering) # Derive a cell numbering from the Plex renumbering entity_dofs = np.zeros(tdim+1, dtype=IntType) entity_dofs[-1] = 1 - self._cell_numbering = self.create_section(entity_dofs) entity_dofs[:] = 0 entity_dofs[0] = 1 self._vertex_numbering = self.create_section(entity_dofs) - entity_dofs[:] = 0 entity_dofs[-2] = 1 facet_numbering = self.create_section(entity_dofs) @@ -1067,6 +1111,11 @@ def set_partitioner(self, distribute, partitioner_type=None): # Command line option `-petscpartitioner_type ` overrides. partitioner.setFromOptions() + @PETSc.Log.EventDecorator() + def get_partitioner(self): + """Get partitioner actually used for (re)distributing underlying plex over comm.""" + return self.topology_dm.getPartitioner() + class ExtrudedMeshTopology(MeshTopology): """Representation of an extruded mesh topology.""" @@ -1262,6 +1311,14 @@ def _order_data_by_cell_index(self, column_list, cell_data): cell_list += list(range(col, col + (self.layers - 1))) return cell_data[cell_list] + @property + def _distribution_name(self): + return self._base_mesh._distribution_name + + @property + def _permutation_name(self): + return self._base_mesh._permutation_name + # TODO: Could this be merged with MeshTopology given that dmcommon.pyx # now covers DMSwarms and DMPlexes? diff --git a/tests/output/test_io_freeze_distribution_permutation.py b/tests/output/test_io_freeze_distribution_permutation.py new file mode 100644 index 0000000000..a1c71fb587 --- /dev/null +++ b/tests/output/test_io_freeze_distribution_permutation.py @@ -0,0 +1,81 @@ +import pytest +from firedrake import * +from pyop2.mpi import COMM_WORLD +import numpy as np +import os + +cwd = os.path.abspath(os.path.dirname(__file__)) +mesh_name = "m" +func_name = "f" + + +@pytest.mark.parallel(nprocs=7) +@pytest.mark.parametrize('case', ["interval", + "interval_small", + "interval_periodic", + "triangle", + "triangle_small", + "quadrilateral", + "triangle_periodic", + "quadrilateral_periodic", + "triangle_extrusion"]) +def test_io_freeze_dist_perm_base(case, tmpdir): + filename = os.path.join(str(tmpdir), "test_io_freeze_dist_perm_base_dump.h5") + filename = COMM_WORLD.bcast(filename, root=0) + if case == "interval": + mesh = UnitIntervalMesh(17, name=mesh_name) + V = FunctionSpace(mesh, "P", 5) + x, = SpatialCoordinate(mesh) + f = Function(V, name=func_name).interpolate(cos(x)) + elif case == "interval_small": + mesh = UnitIntervalMesh(1, name=mesh_name) + V = FunctionSpace(mesh, "P", 3) + x, = SpatialCoordinate(mesh) + f = Function(V, name=func_name).interpolate(cos(x)) + elif case == "interval_periodic": + mesh = PeriodicUnitIntervalMesh(7, name=mesh_name) + V = FunctionSpace(mesh, "P", 5) + x, = SpatialCoordinate(mesh) + f = Function(V, name=func_name).interpolate(cos(2 * pi * x)) + elif case == "triangle": + mesh = Mesh("./docs/notebooks/stokes-control.msh", name=mesh_name) + V = FunctionSpace(mesh, "BDMF", 3) + x, y = SpatialCoordinate(mesh) + f = Function(V, name=func_name).interpolate(as_vector([cos(x), 2 * sin(y)])) + elif case == "triangle_small": + mesh = UnitSquareMesh(2, 2, name=mesh_name) + V = FunctionSpace(mesh, "BDME", 2) + x, y = SpatialCoordinate(mesh) + f = Function(V, name=func_name).interpolate(as_vector([cos(x), 2 * sin(y)])) + elif case == "quadrilateral": + mesh = Mesh(os.path.join(cwd, "..", "meshes", "unitsquare_unstructured_quadrilaterals.msh"), name=mesh_name) + V = FunctionSpace(mesh, "RTCE", 3) + x, y = SpatialCoordinate(mesh) + f = Function(V, name=func_name).project(as_vector([cos(x), 2 * sin(y)])) + elif case == "triangle_periodic": + mesh = PeriodicUnitSquareMesh(20, 20, direction="both", name=mesh_name) + x, y = SpatialCoordinate(mesh) + V = FunctionSpace(mesh, "P", 4) + f = Function(V, name=func_name).interpolate(cos(2 * pi * x) + sin(4 * pi * y)) + elif case == "quadrilateral_periodic": + mesh = PeriodicUnitSquareMesh(20, 20, quadrilateral=True, direction="both", name=mesh_name) + x, y = SpatialCoordinate(mesh) + V = FunctionSpace(mesh, "CG", 4) + f = Function(V, name=func_name).interpolate(cos(2 * pi * x) + cos(8 * pi * y)) + elif case == "triangle_extrusion": + base_mesh = UnitSquareMesh(10, 10) + mesh = ExtrudedMesh(base_mesh, layers=4, name=mesh_name) + V = VectorFunctionSpace(mesh, "CG", 2) + x, y, z = SpatialCoordinate(mesh) + f = Function(V, name=func_name).interpolate(as_vector([cos(x), cos(y), cos(z)])) + else: + raise NotImplementedError("no such test") + ref = f.copy(deepcopy=True) + for i in range(4): + with CheckpointFile(filename, "w") as cf: + cf.save_mesh(mesh) + cf.save_function(f) + with CheckpointFile(filename, "r") as cf: + mesh = cf.load_mesh(mesh_name) + f = cf.load_function(mesh, func_name) + assert np.allclose(f.dat.data_ro, ref.dat.data_ro) From a698c7262faa4eb16ef5fb13ee4dff236ed71e31 Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Sat, 18 Jun 2022 23:45:57 +0100 Subject: [PATCH 2/3] io: allow users to set distribution_name and permutation_name --- firedrake/checkpointing.py | 4 +- firedrake/mesh.py | 20 ++- firedrake/utility_meshes.py | 328 ++++++++++++++++++++++++++++-------- 3 files changed, 276 insertions(+), 76 deletions(-) diff --git a/firedrake/checkpointing.py b/firedrake/checkpointing.py index 968adec713..dc93525200 100644 --- a/firedrake/checkpointing.py +++ b/firedrake/checkpointing.py @@ -539,10 +539,12 @@ def __exit__(self, *args): self.close() @PETSc.Log.EventDecorator("SaveMesh") - def save_mesh(self, mesh): + def save_mesh(self, mesh, distribution_name=None, permutation_name=None): r"""Save a mesh. :arg mesh: the mesh to save. + :kwarg distribution_name: the name under which distribution is saved; if `None`, auto-generated name will be used. + :kwarg permutation_name: the name under which permutation is saved; if `None`, auto-generated name will be used. """ mesh.init() # Handle extruded mesh diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 4fd847deea..6877d7426a 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -734,7 +734,7 @@ class MeshTopology(AbstractMeshTopology): """A representation of mesh topology implemented on a PETSc DMPlex.""" @PETSc.Log.EventDecorator("CreateMesh") - def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None, perm_is=None): + def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None, perm_is=None, distribution_name=None, permutation_name=None): """Half-initialise a mesh topology. :arg plex: :class:`DMPlex` representing the mesh topology @@ -750,6 +750,10 @@ def __init__(self, plex, name, reorder, distribution_parameters, sfXB=None, perm makes sense if we know the exact parallel distribution of `plex` at the time of mesh topology construction like when we load mesh along with its distribution. If given, `reorder` param will be ignored. + :kwarg distribution_name: name of the parallel distribution; + if `None`, automatically generated. + :kwarg permutation_name: name of the entity permutation (reordering); + if `None`, automatically generated. """ super().__init__(name) @@ -853,7 +857,7 @@ def add_overlap(): # corresponding UFL mesh. cell = ufl.Cell(_cells[tdim][nfacets]) self._ufl_mesh = ufl.Mesh(ufl.VectorElement("Lagrange", cell, 1, dim=cell.topological_dimension())) - + # Set/Generate names to be used when checkpointing. self._distribution_name = distribution_name or _generate_default_mesh_topology_distribution_name(self.topology_dm.comm.size, self._distribution_parameters) self._permutation_name = permutation_name or _generate_default_mesh_topology_permutation_name(reorder) @@ -1951,6 +1955,14 @@ def Mesh(meshfile, **kwargs): :class:`DistributedMeshOverlapType` instance, the second the number of levels of overlap. + :param distribution_name: the name of parallel distribution used + when checkpointing; if not given, the name is automatically + generated. + + :param permutation_name: the name of entity permutation (reordering) used + when checkpointing; if not given, the name is automatically + generated. + :param comm: the communicator to use when creating the mesh. If not supplied, then the mesh will be created on COMM_WORLD. Ignored if ``meshfile`` is a DMPlex object (in which case @@ -2027,7 +2039,9 @@ def Mesh(meshfile, **kwargs): plex.setName(_generate_default_mesh_topology_name(name)) # Create mesh topology topology = MeshTopology(plex, name=plex.getName(), reorder=reorder, - distribution_parameters=distribution_parameters) + distribution_parameters=distribution_parameters, + distribution_name=kwargs.get("distribution_name"), + permutation_name=kwargs.get("permutation_name")) return make_mesh_from_mesh_topology(topology, name) diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index 62a1343075..c264a9e975 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -34,7 +34,7 @@ @PETSc.Log.EventDecorator() -def IntervalMesh(ncells, length_or_left, right=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def IntervalMesh(ncells, length_or_left, right=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """ Generate a uniform mesh of an interval. @@ -47,6 +47,12 @@ def IntervalMesh(ncells, length_or_left, right=None, distribution_parameters=Non :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The left hand boundary point has boundary marker 1, while the right hand point has marker 2. @@ -81,11 +87,11 @@ def IntervalMesh(ncells, length_or_left, right=None, distribution_parameters=Non if vcoord[0] == coords[-1]: plex.setLabelValue(dmcommon.FACE_SETS_LABEL, v, 2) - return mesh.Mesh(plex, reorder=False, distribution_parameters=distribution_parameters, name=name) + return mesh.Mesh(plex, reorder=False, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def UnitIntervalMesh(ncells, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def UnitIntervalMesh(ncells, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """ Generate a uniform mesh of the interval [0,1]. @@ -93,16 +99,22 @@ def UnitIntervalMesh(ncells, distribution_parameters=None, comm=COMM_WORLD, name :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The left hand (:math:`x=0`) boundary point has boundary marker 1, while the right hand (:math:`x=1`) point has marker 2. """ - return IntervalMesh(ncells, length_or_left=1.0, distribution_parameters=distribution_parameters, comm=comm, name=name) + return IntervalMesh(ncells, length_or_left=1.0, distribution_parameters=distribution_parameters, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def PeriodicIntervalMesh(ncells, length, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def PeriodicIntervalMesh(ncells, length, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a periodic mesh of an interval. :arg ncells: The number of cells over the interval. @@ -110,13 +122,19 @@ def PeriodicIntervalMesh(ncells, length, distribution_parameters=None, comm=COMM :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ if ncells < 3: raise ValueError("1D periodic meshes with fewer than 3 \ cells are not currently supported") - m = CircleManifoldMesh(ncells, distribution_parameters=distribution_parameters, comm=comm, name=name) + m = CircleManifoldMesh(ncells, distribution_parameters=distribution_parameters, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) coord_fs = VectorFunctionSpace(m, FiniteElement('DG', interval, 1, variant="equispaced"), dim=1) old_coordinates = m.coordinates new_coordinates = Function(coord_fs, name=mesh._generate_default_mesh_coordinates_name(name)) @@ -149,23 +167,29 @@ def PeriodicIntervalMesh(ncells, length, distribution_parameters=None, comm=COMM "L": (cL, READ)}, is_loopy_kernel=True) - return mesh.Mesh(new_coordinates, name=name) + return mesh.Mesh(new_coordinates, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def PeriodicUnitIntervalMesh(ncells, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def PeriodicUnitIntervalMesh(ncells, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a periodic mesh of the unit interval :arg ncells: The number of cells in the interval. :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ - return PeriodicIntervalMesh(ncells, length=1.0, distribution_parameters=distribution_parameters, comm=comm, name=name) + return PeriodicIntervalMesh(ncells, length=1.0, distribution_parameters=distribution_parameters, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def OneElementThickMesh(ncells, Lx, Ly, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def OneElementThickMesh(ncells, Lx, Ly, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """ Generate a rectangular mesh in the domain with corners [0,0] and [Lx, Ly] with ncells, that is periodic in the x-direction. @@ -176,6 +200,12 @@ def OneElementThickMesh(ncells, Lx, Ly, distribution_parameters=None, comm=COMM_ :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ left = np.arange(ncells, dtype=np.int32) @@ -296,7 +326,7 @@ def OneElementThickMesh(ncells, Lx, Ly, distribution_parameters=None, comm=COMM_ Vc = VectorFunctionSpace(mesh1, fe_dg) fc = Function(Vc, name=mesh._generate_default_mesh_coordinates_name(name)).interpolate(mesh1.coordinates) - mash = mesh.Mesh(fc, name=name) + mash = mesh.Mesh(fc, name=name, distribution_name=distribution_name, permutation_name=permutation_name) topverts = Vc.cell_node_list[:, 1::2].flatten() mash.coordinates.dat.data_with_halos[topverts, 1] = Ly @@ -323,23 +353,29 @@ def OneElementThickMesh(ncells, Lx, Ly, distribution_parameters=None, comm=COMM_ @PETSc.Log.EventDecorator() -def UnitTriangleMesh(comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def UnitTriangleMesh(comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a mesh of the reference triangle :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ coords = [[0., 0.], [1., 0.], [0., 1.]] cells = [[0, 1, 2]] plex = mesh._from_cell_list(2, cells, coords, comm, mesh._generate_default_mesh_topology_name(name)) - return mesh.Mesh(plex, reorder=False, name=name) + return mesh.Mesh(plex, reorder=False, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() def RectangleMesh(nx, ny, Lx, Ly, quadrilateral=False, reorder=None, - diagonal="left", distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + diagonal="left", distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a rectangular mesh :arg nx: The number of cells in the x direction @@ -354,6 +390,12 @@ def RectangleMesh(nx, ny, Lx, Ly, quadrilateral=False, reorder=None, from bottom left to top right (``"right"``), or top left to bottom right (``"left"``), or put in both diagonals (``"crossed"``). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The boundary edges in this mesh are numbered as follows: @@ -375,14 +417,18 @@ def RectangleMesh(nx, ny, Lx, Ly, quadrilateral=False, reorder=None, diagonal=diagonal, distribution_parameters=distribution_parameters, comm=comm, - name=name) + name=name, + distribution_name=distribution_name, + permutation_name=permutation_name) def TensorRectangleMesh(xcoords, ycoords, quadrilateral=False, reorder=None, diagonal="left", distribution_parameters=None, comm=COMM_WORLD, - name=mesh.DEFAULT_MESH_NAME): + name=mesh.DEFAULT_MESH_NAME, + distribution_name=None, + permutation_name=None): """Generate a rectangular mesh :arg xcoords: mesh points for the x direction @@ -473,11 +519,11 @@ def TensorRectangleMesh(xcoords, ycoords, quadrilateral=False, if abs(face_coords[1] - y1) < ytol and abs(face_coords[3] - y1) < ytol: plex.setLabelValue(dmcommon.FACE_SETS_LABEL, face, 4) plex.removeLabel("boundary_faces") - return mesh.Mesh(plex, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + return mesh.Mesh(plex, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def SquareMesh(nx, ny, L, reorder=None, quadrilateral=False, diagonal="left", distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def SquareMesh(nx, ny, L, reorder=None, quadrilateral=False, diagonal="left", distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a square mesh :arg nx: The number of cells in the x direction @@ -488,6 +534,12 @@ def SquareMesh(nx, ny, L, reorder=None, quadrilateral=False, diagonal="left", di :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The boundary edges in this mesh are numbered as follows: @@ -501,11 +553,13 @@ def SquareMesh(nx, ny, L, reorder=None, quadrilateral=False, diagonal="left", di diagonal=diagonal, distribution_parameters=distribution_parameters, comm=comm, - name=name) + name=name, + distribution_name=distribution_name, + permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def UnitSquareMesh(nx, ny, reorder=None, diagonal="left", quadrilateral=False, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def UnitSquareMesh(nx, ny, reorder=None, diagonal="left", quadrilateral=False, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a unit square mesh :arg nx: The number of cells in the x direction @@ -515,6 +569,12 @@ def UnitSquareMesh(nx, ny, reorder=None, diagonal="left", quadrilateral=False, d :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The boundary edges in this mesh are numbered as follows: @@ -528,7 +588,9 @@ def UnitSquareMesh(nx, ny, reorder=None, diagonal="left", quadrilateral=False, d diagonal=diagonal, distribution_parameters=distribution_parameters, comm=comm, - name=name) + name=name, + distribution_name=distribution_name, + permutation_name=permutation_name) @PETSc.Log.EventDecorator() @@ -537,7 +599,9 @@ def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both", distribution_parameters=None, diagonal=None, comm=COMM_WORLD, - name=mesh.DEFAULT_MESH_NAME): + name=mesh.DEFAULT_MESH_NAME, + distribution_name=None, + permutation_name=None): """Generate a periodic rectangular mesh :arg nx: The number of cells in the x direction @@ -553,6 +617,12 @@ def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both", :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. If direction == "x" the boundary edges in this mesh are numbered as follows: @@ -566,7 +636,7 @@ def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both", """ if direction == "both" and ny == 1 and quadrilateral: - return OneElementThickMesh(nx, Lx, Ly, distribution_parameters=distribution_parameters, name=name) + return OneElementThickMesh(nx, Lx, Ly, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) if direction not in ("both", "x", "y"): raise ValueError("Cannot have a periodic mesh with periodicity '%s'" % direction) @@ -575,12 +645,12 @@ def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both", quadrilateral=quadrilateral, reorder=reorder, distribution_parameters=distribution_parameters, - diagonal=diagonal, comm=comm, name=name) + diagonal=diagonal, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) if nx < 3 or ny < 3: raise ValueError("2D periodic meshes with fewer than 3 \ cells in each direction are not currently supported") - m = TorusMesh(nx, ny, 1.0, 0.5, quadrilateral=quadrilateral, reorder=reorder, distribution_parameters=distribution_parameters, comm=comm, name=name) + m = TorusMesh(nx, ny, 1.0, 0.5, quadrilateral=quadrilateral, reorder=reorder, distribution_parameters=distribution_parameters, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) coord_family = 'DQ' if quadrilateral else 'DG' cell = 'quadrilateral' if quadrilateral else 'triangle' coord_fs = VectorFunctionSpace(m, FiniteElement(coord_family, cell, 1, variant="equispaced"), dim=2) @@ -629,12 +699,12 @@ def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both", "Ly": (cLy, READ)}, is_loopy_kernel=True) - return mesh.Mesh(new_coordinates, name=name) + return mesh.Mesh(new_coordinates, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() def PeriodicSquareMesh(nx, ny, L, direction="both", quadrilateral=False, reorder=None, - distribution_parameters=None, diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + distribution_parameters=None, diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a periodic square mesh :arg nx: The number of cells in the x direction @@ -649,6 +719,12 @@ def PeriodicSquareMesh(nx, ny, L, direction="both", quadrilateral=False, reorder :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. If direction == "x" the boundary edges in this mesh are numbered as follows: @@ -663,13 +739,13 @@ def PeriodicSquareMesh(nx, ny, L, direction="both", quadrilateral=False, reorder return PeriodicRectangleMesh(nx, ny, L, L, direction=direction, quadrilateral=quadrilateral, reorder=reorder, distribution_parameters=distribution_parameters, - diagonal=diagonal, comm=comm, name=name) + diagonal=diagonal, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() def PeriodicUnitSquareMesh(nx, ny, direction="both", reorder=None, quadrilateral=False, distribution_parameters=None, - diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a periodic unit square mesh :arg nx: The number of cells in the x direction @@ -683,6 +759,12 @@ def PeriodicUnitSquareMesh(nx, ny, direction="both", reorder=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. If direction == "x" the boundary edges in this mesh are numbered as follows: @@ -697,11 +779,11 @@ def PeriodicUnitSquareMesh(nx, ny, direction="both", reorder=None, return PeriodicSquareMesh(nx, ny, 1.0, direction=direction, reorder=reorder, quadrilateral=quadrilateral, distribution_parameters=distribution_parameters, - diagonal=diagonal, comm=comm, name=name) + diagonal=diagonal, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def CircleManifoldMesh(ncells, radius=1, degree=1, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def CircleManifoldMesh(ncells, radius=1, degree=1, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generated a 1D mesh of the circle, immersed in 2D. :arg ncells: number of cells the circle should be @@ -713,6 +795,12 @@ def CircleManifoldMesh(ncells, radius=1, degree=1, distribution_parameters=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ if ncells < 3: raise ValueError("CircleManifoldMesh must have at least three cells") @@ -725,19 +813,19 @@ def CircleManifoldMesh(ncells, radius=1, degree=1, distribution_parameters=None, plex = mesh._from_cell_list(1, cells, vertices, comm, mesh._generate_default_mesh_topology_name(name)) - m = mesh.Mesh(plex, dim=2, reorder=False, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=2, reorder=False, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) if degree > 1: new_coords = function.Function(functionspace.VectorFunctionSpace(m, "CG", degree)) new_coords.interpolate(ufl.SpatialCoordinate(m)) # "push out" to circle new_coords.dat.data[:] *= (radius / np.linalg.norm(new_coords.dat.data, axis=1)).reshape(-1, 1) - m = mesh.Mesh(new_coords, name=name) + m = mesh.Mesh(new_coords, name=name, distribution_name=distribution_name, permutation_name=permutation_name) m._radius = radius return m @PETSc.Log.EventDecorator() -def UnitDiskMesh(refinement_level=0, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def UnitDiskMesh(refinement_level=0, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a mesh of the unit disk in 2D :kwarg refinement_level: optional number of refinements (0 is a diamond) @@ -745,6 +833,12 @@ def UnitDiskMesh(refinement_level=0, reorder=None, distribution_parameters=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ vertices = np.array([[0, 0], [1, 0], @@ -787,12 +881,12 @@ def UnitDiskMesh(refinement_level=0, reorder=None, distribution_parameters=None, t = np.max(np.abs(x)) / norm x[:] *= t - m = mesh.Mesh(plex, dim=2, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=2, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) return m @PETSc.Log.EventDecorator() -def UnitBallMesh(refinement_level=0, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def UnitBallMesh(refinement_level=0, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a mesh of the unit ball in 3D :kwarg refinement_level: optional number of refinements (0 is an octahedron) @@ -800,6 +894,12 @@ def UnitBallMesh(refinement_level=0, reorder=None, distribution_parameters=None, :kwarg comm: Optional MPI communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ vertices = np.array([[0, 0, 0], [1, 0, 0], @@ -839,27 +939,33 @@ def UnitBallMesh(refinement_level=0, reorder=None, distribution_parameters=None, t = np.sum(np.abs(x)) / norm x[:] *= t - m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) return m @PETSc.Log.EventDecorator() -def UnitTetrahedronMesh(comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def UnitTetrahedronMesh(comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a mesh of the reference tetrahedron. :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ coords = [[0., 0., 0.], [1., 0., 0.], [0., 1., 0.], [0., 0., 1.]] cells = [[0, 1, 2, 3]] plex = mesh._from_cell_list(3, cells, coords, comm, mesh._generate_default_mesh_topology_name(name)) - return mesh.Mesh(plex, reorder=False, name=name) + return mesh.Mesh(plex, reorder=False, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def BoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameters=None, diagonal="default", comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def BoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameters=None, diagonal="default", comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a mesh of a 3D box. :arg nx: The number of cells in the x direction @@ -963,11 +1069,11 @@ def BoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameters=None, plex.setLabelValue(dmcommon.FACE_SETS_LABEL, face, 6) plex.removeLabel("boundary_faces") - return mesh.Mesh(plex, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + return mesh.Mesh(plex, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def CubeMesh(nx, ny, nz, L, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def CubeMesh(nx, ny, nz, L, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a mesh of a cube :arg nx: The number of cells in the x direction @@ -978,6 +1084,12 @@ def CubeMesh(nx, ny, nz, L, reorder=None, distribution_parameters=None, comm=COM :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The boundary surfaces are numbered as follows: @@ -989,11 +1101,11 @@ def CubeMesh(nx, ny, nz, L, reorder=None, distribution_parameters=None, comm=COM * 6: plane z == L """ return BoxMesh(nx, ny, nz, L, L, L, reorder=reorder, distribution_parameters=distribution_parameters, - comm=comm, name=name) + comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def UnitCubeMesh(nx, ny, nz, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def UnitCubeMesh(nx, ny, nz, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a mesh of a unit cube :arg nx: The number of cells in the x direction @@ -1003,6 +1115,12 @@ def UnitCubeMesh(nx, ny, nz, reorder=None, distribution_parameters=None, comm=CO :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The boundary surfaces are numbered as follows: @@ -1014,11 +1132,11 @@ def UnitCubeMesh(nx, ny, nz, reorder=None, distribution_parameters=None, comm=CO * 6: plane z == 1 """ return CubeMesh(nx, ny, nz, 1, reorder=reorder, distribution_parameters=distribution_parameters, - comm=comm, name=name) + comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() -def PeriodicBoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def PeriodicBoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a periodic mesh of a 3D box. :arg nx: The number of cells in the x direction @@ -1031,6 +1149,12 @@ def PeriodicBoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameter :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ for n in (nx, ny, nz): if n < 3: @@ -1061,7 +1185,7 @@ def PeriodicBoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameter cells = np.asarray(cells).swapaxes(0, 3).reshape(-1, 4) plex = mesh._from_cell_list(3, cells, coords, comm, mesh._generate_default_mesh_topology_name(name)) - m = mesh.Mesh(plex, reorder=reorder, distribution_parameters=distribution_parameters) + m = mesh.Mesh(plex, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) old_coordinates = m.coordinates new_coordinates = Function(VectorFunctionSpace(m, FiniteElement('DG', tetrahedron, 1, variant="equispaced")), name=mesh._generate_default_mesh_coordinates_name(name)) @@ -1111,12 +1235,12 @@ def PeriodicBoxMesh(nx, ny, nz, Lx, Ly, Lz, reorder=None, distribution_parameter "hy": (hy, READ), "hz": (hz, READ)}, is_loopy_kernel=True) - m1 = mesh.Mesh(new_coordinates, name=name) + m1 = mesh.Mesh(new_coordinates, name=name, distribution_name=distribution_name, permutation_name=permutation_name) return m1 @PETSc.Log.EventDecorator() -def PeriodicUnitCubeMesh(nx, ny, nz, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): +def PeriodicUnitCubeMesh(nx, ny, nz, reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a periodic mesh of a unit cube :arg nx: The number of cells in the x direction @@ -1126,13 +1250,19 @@ def PeriodicUnitCubeMesh(nx, ny, nz, reorder=None, distribution_parameters=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ - return PeriodicBoxMesh(nx, ny, nz, 1., 1., 1., reorder=reorder, distribution_parameters=distribution_parameters, comm=comm, name=name) + return PeriodicBoxMesh(nx, ny, nz, 1., 1., 1., reorder=reorder, distribution_parameters=distribution_parameters, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() def IcosahedralSphereMesh(radius, refinement_level=0, degree=1, reorder=None, - distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate an icosahedral approximation to the surface of the sphere. @@ -1152,6 +1282,12 @@ def IcosahedralSphereMesh(radius, refinement_level=0, degree=1, reorder=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ if refinement_level < 0 or refinement_level % 1: raise RuntimeError("Number of refinements must be a non-negative integer") @@ -1205,20 +1341,20 @@ def IcosahedralSphereMesh(radius, refinement_level=0, degree=1, reorder=None, coords = plex.getCoordinatesLocal().array.reshape(-1, 3) scale = (radius / np.linalg.norm(coords, axis=1)).reshape(-1, 1) coords *= scale - m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) if degree > 1: new_coords = function.Function(functionspace.VectorFunctionSpace(m, "CG", degree)) new_coords.interpolate(ufl.SpatialCoordinate(m)) # "push out" to sphere new_coords.dat.data[:] *= (radius / np.linalg.norm(new_coords.dat.data, axis=1)).reshape(-1, 1) - m = mesh.Mesh(new_coords, name=name) + m = mesh.Mesh(new_coords, name=name, distribution_name=distribution_name, permutation_name=permutation_name) m._radius = radius return m @PETSc.Log.EventDecorator() def UnitIcosahedralSphereMesh(refinement_level=0, degree=1, reorder=None, - distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate an icosahedral approximation to the unit sphere. :kwarg refinement_level: optional number of refinements (0 is an @@ -1229,9 +1365,15 @@ def UnitIcosahedralSphereMesh(refinement_level=0, degree=1, reorder=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ return IcosahedralSphereMesh(1.0, refinement_level=refinement_level, - degree=degree, reorder=reorder, comm=comm, name=name) + degree=degree, reorder=reorder, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) # mesh is mainly used as a utility, so it's unnecessary to annotate the construction @@ -1244,7 +1386,7 @@ def OctahedralSphereMesh(radius, refinement_level=0, degree=1, reorder=None, distribution_parameters=None, comm=COMM_WORLD, - name=mesh.DEFAULT_MESH_NAME): + name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate an octahedral approximation to the surface of the sphere. @@ -1262,6 +1404,12 @@ def OctahedralSphereMesh(radius, refinement_level=0, degree=1, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ if refinement_level < 0 or refinement_level % 1: raise ValueError("Number of refinements must be a non-negative integer") @@ -1302,10 +1450,10 @@ def OctahedralSphereMesh(radius, refinement_level=0, degree=1, plex.setName(mesh._generate_default_mesh_topology_name(name)) # build the initial mesh - m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) if degree > 1: # use it to build a higher-order mesh - m = mesh.Mesh(interpolate(ufl.SpatialCoordinate(m), VectorFunctionSpace(m, "CG", degree)), name=name) + m = mesh.Mesh(interpolate(ufl.SpatialCoordinate(m), VectorFunctionSpace(m, "CG", degree)), name=name, distribution_name=distribution_name, permutation_name=permutation_name) # remap to a cone x, y, z = ufl.SpatialCoordinate(m) @@ -1357,7 +1505,7 @@ def OctahedralSphereMesh(radius, refinement_level=0, degree=1, @PETSc.Log.EventDecorator() def UnitOctahedralSphereMesh(refinement_level=0, degree=1, hemisphere="both", z0=0.8, reorder=None, - distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate an octahedral approximation to the unit sphere. :kwarg refinement_level: optional number of refinements (0 is an @@ -1373,12 +1521,18 @@ def UnitOctahedralSphereMesh(refinement_level=0, degree=1, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ return OctahedralSphereMesh(1.0, refinement_level=refinement_level, degree=degree, hemisphere=hemisphere, z0=z0, reorder=reorder, - distribution_parameters=distribution_parameters, comm=comm, name=name) + distribution_parameters=distribution_parameters, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) def _cubedsphere_cells_and_coords(radius, refinement_level): @@ -1510,7 +1664,7 @@ def coordinates_on_panel(panel_num, X, Y, Z): @PETSc.Log.EventDecorator() def CubedSphereMesh(radius, refinement_level=0, degree=1, - reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + reorder=None, distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate an cubed approximation to the surface of the sphere. @@ -1522,6 +1676,12 @@ def CubedSphereMesh(radius, refinement_level=0, degree=1, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ if refinement_level < 0 or refinement_level % 1: raise RuntimeError("Number of refinements must be a non-negative integer") @@ -1533,21 +1693,21 @@ def CubedSphereMesh(radius, refinement_level=0, degree=1, plex = mesh._from_cell_list(2, cells, coords, comm, mesh._generate_default_mesh_topology_name(name)) - m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) if degree > 1: new_coords = function.Function(functionspace.VectorFunctionSpace(m, "Q", degree)) new_coords.interpolate(ufl.SpatialCoordinate(m)) # "push out" to sphere new_coords.dat.data[:] *= (radius / np.linalg.norm(new_coords.dat.data, axis=1)).reshape(-1, 1) - m = mesh.Mesh(new_coords) + m = mesh.Mesh(new_coords, distribution_name=distribution_name, permutation_name=permutation_name) m._radius = radius return m @PETSc.Log.EventDecorator() def UnitCubedSphereMesh(refinement_level=0, degree=1, reorder=None, - distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a cubed approximation to the unit sphere. :kwarg refinement_level: optional number of refinements (0 is a cube). @@ -1557,14 +1717,20 @@ def UnitCubedSphereMesh(refinement_level=0, degree=1, reorder=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ return CubedSphereMesh(1.0, refinement_level=refinement_level, - degree=degree, reorder=reorder, comm=comm, name=name) + degree=degree, reorder=reorder, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) @PETSc.Log.EventDecorator() def TorusMesh(nR, nr, R, r, quadrilateral=False, reorder=None, - distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + distribution_parameters=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generate a toroidal mesh :arg nR: The number of cells in the major direction (min 3) @@ -1576,6 +1742,12 @@ def TorusMesh(nR, nr, R, r, quadrilateral=False, reorder=None, :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. """ if nR < 3 or nr < 3: @@ -1606,14 +1778,14 @@ def TorusMesh(nR, nr, R, r, quadrilateral=False, reorder=None, plex = mesh._from_cell_list(2, cells, vertices, comm, mesh._generate_default_mesh_topology_name(name)) - m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) return m @PETSc.Log.EventDecorator() def CylinderMesh(nr, nl, radius=1, depth=1, longitudinal_direction="z", quadrilateral=False, reorder=None, - distribution_parameters=None, diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + distribution_parameters=None, diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generates a cylinder mesh. :arg nr: number of cells the cylinder circumference should be @@ -1631,6 +1803,12 @@ def CylinderMesh(nr, nl, radius=1, depth=1, longitudinal_direction="z", :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. The boundary edges in this mesh are numbered as follows: @@ -1738,13 +1916,13 @@ def CylinderMesh(nr, nl, radius=1, depth=1, longitudinal_direction="z", plex.setLabelValue(dmcommon.FACE_SETS_LABEL, face, 2) plex.removeLabel("boundary_faces") - m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name) + m = mesh.Mesh(plex, dim=3, reorder=reorder, distribution_parameters=distribution_parameters, name=name, distribution_name=distribution_name, permutation_name=permutation_name) return m @PETSc.Log.EventDecorator() def PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction="x", quadrilateral=False, - reorder=None, distribution_parameters=None, diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME): + reorder=None, distribution_parameters=None, diagonal=None, comm=COMM_WORLD, name=mesh.DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None): """Generates RectangleMesh that is periodic in the x or y direction. :arg nx: The number of cells in the x direction @@ -1759,6 +1937,12 @@ def PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction="x", quadrilateral= :kwarg comm: Optional communicator to build the mesh on (defaults to COMM_WORLD). :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. If direction == "x" the boundary edges in this mesh are numbered as follows: @@ -1786,7 +1970,7 @@ def PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction="x", quadrilateral= m = CylinderMesh(na, nb, 1.0, 1.0, longitudinal_direction="z", quadrilateral=quadrilateral, reorder=reorder, distribution_parameters=distribution_parameters, - diagonal=diagonal, comm=comm, name=name) + diagonal=diagonal, comm=comm, name=name, distribution_name=distribution_name, permutation_name=permutation_name) coord_family = 'DQ' if quadrilateral else 'DG' cell = 'quadrilateral' if quadrilateral else 'triangle' coord_fs = VectorFunctionSpace(m, FiniteElement(coord_family, cell, 1, variant="equispaced"), dim=2) @@ -1829,4 +2013,4 @@ def PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction="x", quadrilateral= [1, 0]]) new_coordinates.dat.data[:] = np.dot(new_coordinates.dat.data, operator.T) - return mesh.Mesh(new_coordinates, name=name) + return mesh.Mesh(new_coordinates, name=name, distribution_name=distribution_name, permutation_name=permutation_name) From 90ef3e9cf63abce8fe00872772917daeebdd596e Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Sun, 19 Jun 2022 01:32:22 +0100 Subject: [PATCH 3/3] io: update manual --- docs/source/checkpointing.rst | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/docs/source/checkpointing.rst b/docs/source/checkpointing.rst index c08f5716a0..63e5f89470 100644 --- a/docs/source/checkpointing.rst +++ b/docs/source/checkpointing.rst @@ -22,6 +22,29 @@ to pass through a single process. It also supports flexible checkpointing, where one can save meshes and :class:`~.Function` s on :math:`N` processes and later load them on :math:`P` processes. +If :math:`P == N`, the parallel distribution and entity permutation +(reordering) of the saved mesh is recovered on the loaded mesh by default. +The distribution and permutation data are by default stored under names +automatically generated by Firedrake. + +.. warning:: + + If the mesh has a non-standard distribution, e.g., generated by a + partitioner with some non-standard parameters, it is recommended that + the user set the distribution name explicitly when constructing a mesh; + see, e.g., :py:func:`~.Mesh`. + +.. warning:: + + If :math:`P != N` or `P == N` but `distribution_parameters` dict and/or + `reorder` parameter are passed when loading a mesh, the saved mesh and + the loaded mesh (and thus the saved function and the loaded function) + will in general be represented by different global numbering systems and + they are merely guaranteed to be analytically the same; as a consequence, + it is currently not allowed to save the once loaded mesh or function back + to the same file under the same name as this would cause conflict with + other data stored using incompatible global numbering system. + We plan to remove this restriction in the future. Saving ------ @@ -72,6 +95,14 @@ the HDF5 installation; "h5dump -n example.h5", for instance, shows: group /topologies/firedrake_mixed_meshes/meshA/firedrake_mixed_function_spaces/firedrake_function_space_meshA_CG2(None,None)_meshA_CG1(None,None)/firedrake_functions/g/0 group /topologies/firedrake_mixed_meshes/meshA/firedrake_mixed_function_spaces/firedrake_function_space_meshA_CG2(None,None)_meshA_CG1(None,None)/firedrake_functions/g/1 group /topologies/meshA_topology + group /topologies/meshA_topology/distributions + group /topologies/meshA_topology/distributions/firedrake_default_1_True_None_(FACET,1) + dataset /topologies/meshA_topology/distributions/firedrake_default_1_True_None_(FACET,1)/chart_sizes + dataset /topologies/meshA_topology/distributions/firedrake_default_1_True_None_(FACET,1)/global_point_numbers + dataset /topologies/meshA_topology/distributions/firedrake_default_1_True_None_(FACET,1)/owners + group /topologies/meshA_topology/distributions/firedrake_default_1_True_None_(FACET,1)/permutations + group /topologies/meshA_topology/distributions/firedrake_default_1_True_None_(FACET,1)/permutations/firedrake_default_True + dataset /topologies/meshA_topology/distributions/firedrake_default_1_True_None_(FACET,1)/permutations/firedrake_default_True/permutation group /topologies/meshA_topology/dms group /topologies/meshA_topology/dms/coordinateDM dataset /topologies/meshA_topology/dms/coordinateDM/order