Merge branch 'development' of github.com:underworldcode/underworld3 into development

Author: julesghub

src/underworld3/adaptivity.py

@@ -268,27 +268,30 @@ def mesh2mesh_swarm(mesh0, mesh1, swarm0, swarmVarList, proxy=True, verbose=Fals
     ## some of these. So let's add them now
 
     cell = mesh1.get_closest_local_cells(global_unallocated_coords)
-    found1 = np.where(cell >= 0)[0]
-    not_found1 = np.where(cell == -1)[0]
+    # found1 = np.where(cell >= 0)[0]
+    # not_found1 = np.where(cell == -1)[0]s
+    cell_arr = np.atleast_1d(cell)              
+    found1 = np.where(cell_arr >= 0)[0]
+    not_found1 = np.where(cell_arr == -1)[0]
 
     n_found1 = found1.shape[0]
     n_not_found1 = not_found1.shape[0]
 
-    psize = swarm1.dm.getLocalSize()
-    adds = n_found1 + 1
+    if n_found1 > 0:
+        psize = swarm1.dm.getLocalSize()        
+        adds = n_found1     # swarm is not blank, so only need to use N for addNPoints
 
-    swarm1.dm.addNPoints(adds)
+        swarm1.dm.addNPoints(adds)
 
-    cellid = swarm1.dm.getField("DMSwarm_cellid")
-    coords = swarm1.dm.getField("DMSwarmPIC_coor").reshape((-1, swarm1.dim))
+        cellid = swarm1.dm.getField("DMSwarm_cellid")
+        coords = swarm1.dm.getField("DMSwarmPIC_coor").reshape((-1, swarm1.dim))
 
-    coords[psize + 1 :, :] = global_unallocated_coords[found1, :]
+        coords[psize + 1 :, :] = global_unallocated_coords[found1, :]
 
-    if n_found1 > 0:
         cellid[psize + 1 :] = cell[found1]  ## gathered points
 
-    swarm1.dm.restoreField("DMSwarm_cellid")
-    swarm1.dm.restoreField("DMSwarmPIC_coor")
+        swarm1.dm.restoreField("DMSwarm_cellid")
+        swarm1.dm.restoreField("DMSwarmPIC_coor")
 
     # print(
     #     f"{uw.mpi.rank}/i: {swarm1.dm.getLocalSize()} / {swarm1.dm.getSize()} cf {swarm0.dm.getSize()}",

src/underworld3/cython/petsc_generic_snes_solvers.pyx

@@ -391,7 +391,124 @@ class SolverBaseClass(uw_object):
             print(f"{name}.constitutive_model = uw.constitutive_models...")
 
         return
+    
+    def get_dof_partition(self, 
+                          section_type: str, 
+                          filename: Optional[str | None] = None, 
+                          outputPath: Optional[str] = ""):
+        """
+        Obtains how the degrees of freedom (DOF) are distributed/divided among the processors and saves them in an h5 file.  
+        Parameters
+        ----------
+        section_type:
+            Can be: "local" which includes DOFs from ghost points or "global" which differentiates DOFs from ghost points by having negative values. 
+        filename:
+            Output file name. If None, will print out results; if set to a string, the final output file will be <filename>_<section_type>.u.h5.
+        outputPath:
+            Path of directory where data is saved. If left empty it will save the data in the current working directory.
+        """
+
+        self.validate_solver()  # mainly check if self.u is properly set 
+
+        u_id = self.Unknowns.u.field_id
+        fname = None if filename is None else f"{filename}_{section_type}.u.h5"
+
+        self._get_dof_partition_by_field_id(section_type    = section_type, 
+                                            field_id        = u_id, 
+                                            filename        = fname,
+                                            outputPath      = outputPath)
+        
+        return
+
+
+    def _get_dof_partition_by_field_id(self,
+                                       section_type: str,
+                                       field_id: int,  
+                                       filename: Optional[str | None] = None, 
+                                       outputPath: Optional[str] = ""):
+        """
+        Private version of get_dof_partition with field_id as an additional parameter. 
+        Parameters
+        ----------
+        section_type:
+            Can be: "local" which includes DOFs from ghost points or "global" which differentiates DOFs from ghost points by having negative values. 
+        field_id:
+            The field id
+        filename:
+            Output file name. If None, will print out results; if set to a string, resulting h5 file has the following keys: field_id, rank, dof.
+        outputPath:
+            Path of directory where data is saved. If left empty it will save the data in the current working directory.
+        """
+
+        import os
+        import h5py
+        import numpy as np
 
+        # check if section type is valid
+        if section_type not in ['local', 'global']:
+            raise("'section_type' unknown. Value must be either 'local' or 'global'")
+        
+        # check if path exists
+        if os.path.exists(os.path.abspath(outputPath)):  # easier to debug abs
+            pass
+        else:
+            raise RuntimeError(f"{os.path.abspath(outputPath)} does not exist")
+
+        # check if we have write access
+        if os.access(os.path.abspath(outputPath), os.W_OK):
+            pass
+        else:
+            raise RuntimeError(f"No write access to {os.path.abspath(outputPath)}")
+
+        
+        # get all points in the DAG of this partition
+        if section_type == "local":
+            section = self.mesh.dm.getLocalSection()
+        elif section_type == "global":
+            section = self.mesh.dm.getGlobalSection()
+          
+        # NOTE: negative DOFs mean that these are ghost ones and owned by a different process
+        
+        ptStart, ptEnd = section.getChart() # will give all DOFs including ghosts
+
+        fdofs = [section.getFieldDof(pt, field_id) for pt in range(ptStart, ptEnd)]
+        fdofs = np.array(fdofs)
+        pos_dof_data = np.array([field_id, uw.mpi.rank, fdofs[fdofs > 0].sum()])
+        
+        if section_type == "global":
+            neg_dof_data = np.array([field_id, uw.mpi.rank, fdofs[fdofs < 0].sum()])
+        
+        comm = uw.mpi.comm
+
+        # Gather the arrays on rank 0
+        gath_pos_dof_data = comm.gather(pos_dof_data, root = 0)
+        if section_type == "global":
+            gath_neg_dof_data = comm.gather(neg_dof_data, root = 0)
+
+        # pack data and save to a dataframe for formatted opening
+        if uw.mpi.rank == 0:
+            gath_dof_data = np.vstack(gath_pos_dof_data)
+            if section_type == "global":
+                gath_dof_data = np.vstack([gath_pos_dof_data, gath_neg_dof_data])
+
+            if filename is None: # print out
+                print(f"Section type: {section_type}")
+                print(f"| Field ID      | Rank           | # DOFs        |")
+                print(f"| ---------------------------------------------- |")
+                for i in range(gath_dof_data.shape[0]):
+                    print(
+                        f"| {gath_dof_data[i, 0]:<15}|{gath_dof_data[i, 1]:<15}|{gath_dof_data[i, 2]:<15}|"
+                    )
+                print(f"| ---------------------------------------------- |")
+                print("\n", flush = True)
+
+            else: # save
+                with h5py.File(f"{outputPath}/{filename}", "w") as f: 
+                    f.create_dataset("field_id", data = gath_dof_data[:, 0])
+                    f.create_dataset("rank", data = gath_dof_data[:, 1])
+                    f.create_dataset("dof", data = gath_dof_data[:, 2])
+        
+        return
 
 ## Specific to dimensionality
 
@@ -2077,9 +2194,45 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
             raise RuntimeError("Constitutive Model is required")
 
         return
-
-
-
+    
+    def get_dof_partition(self, 
+                          section_type: str, 
+                          filename: Optional[str | None] = None, 
+                          outputPath: Optional[str] = ""):
+        """
+        Obtains how the degrees of freedom (DOF) are distributed/divided among the processors and saves them in an h5 file.  
+        Parameters
+        ----------
+        section_type:
+            Can be: "local" which includes DOFs from ghost points or "global" which differentiates DOFs from ghost points by having negative values. 
+        filename:
+            Output file name. If None, will print out results; if set to a string, the output files will be <filename>_<section_type>.u.h5 and <filename>_<section_type>.p.h5.
+        outputPath:
+            Path of directory where data is saved. If left empty it will save the data in the current working directory.
+        """
+        # NOTE: supposed to inherit get_dof_partition from SolverBaseClass 
+        # NOTE: _get_dof_partition_by_field_id is defined in SolverBaseClass
+
+        self.validate_solver()
+
+        u_id = self.Unknowns.u.field_id
+        fname = None if filename is None else f"{filename}_{section_type}.u.h5"
+
+        self._get_dof_partition_by_field_id(section_type    = section_type, 
+                                            field_id        = u_id, 
+                                            filename        = fname,
+                                            outputPath      = outputPath)
+        
+        p_id = self.Unknowns.p.field_id
+        fname = None if filename is None else f"{filename}_{section_type}.p.h5"
+
+        self._get_dof_partition_by_field_id(section_type    = section_type, 
+                                            field_id        = p_id, 
+                                            filename        = fname,
+                                            outputPath      = outputPath)
+        
+        return
+    
     @timing.routine_timer_decorator
     def _setup_pointwise_functions(self, verbose=False, debug=False, debug_name=None):
         import sympy

tests/test_0002_swarm.py

@@ -41,3 +41,22 @@ def test_create_swarmvariable(setup_data):
     elements = swarm.mesh._centroids.shape[0]
     var.save("var.h5")
     assert shape == (elements * 6, 2)
+
+def test_addNPoints(setup_data):
+    
+    from underworld3 import swarm
+
+    swarm2 = setup_data
+    var = swarm.SwarmVariable(name = "test", swarm = swarm2, size = 1)
+    swarm2.dm.finalizeFieldRegister()
+
+    swarm2.dm.addNPoints(10) # since swarm is initially empty, will add (10 - 1) points
+    with swarm2.access():
+        npts = swarm2.data.shape[0]
+    assert npts == 9
+
+    swarm2.dm.addNPoints(1) # already has particles, so will add 1 point
+    with swarm2.access():
+        npts = swarm2.data.shape[0]
+    assert npts == 10
+