Update utils

physics_driven_ml/utils/__init__.py

@@ -0,0 +1 @@
+from physics_driven_ml.utils.utils import ModelConfig, get_logger    # noqa: F401

physics_driven_ml/utils/graph/adjacency_dofs.pyx

@@ -0,0 +1,116 @@
+# Which one is faster: getAdjacency for every entity or use hashmap ?
+
+# cython: language_level=3
+
+from firedrake.petsc import PETSc
+from firedrake.utils import IntType
+
+import numpy as np
+
+cimport cython
+cimport numpy as np
+cimport petsc4py.PETSc as PETSc
+
+include "petschdr.pxi"
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def build_dof_adjacency(ndofs, num_entities, dmplex, section):
+    cdef:
+        PetscInt i, j, k, dof, size_entity_adj, max_adjacent_dofs
+        np.ndarray[PetscInt, ndim=1, mode="c"] number_adj_entities
+        np.ndarray[PetscInt, ndim=1, mode="c"] adj_entity_offset
+        np.ndarray[PetscInt, ndim=1, mode="c"] entity_per_dof_count
+        np.ndarray[PetscInt, ndim=1, mode="c"] dof_to_entity_offset
+        np.ndarray[PetscInt, ndim=1, mode="c"] adj_entities
+        np.ndarray[PetscInt, ndim=1, mode="c"] dofs_to_entities
+        np.ndarray[PetscInt, ndim=1, mode="c"] dofs_idx
+        np.ndarray[PetscInt, ndim=2, mode="c"] adjacency_dofs
+
+    # Number of adjacent entities of each entity
+    number_adj_entities = np.zeros(num_entities + 1, dtype=IntType)
+    # Offset of each entity in `adj_entities`
+    adj_entity_offset = np.empty(num_entities + 1, dtype=IntType)
+    # Number of entities associated with each dof
+    entity_per_dof_count = np.zeros(ndofs + 1, dtype=IntType)
+    # Offset of each dof in `dofs_to_entities`
+    dof_to_entity_offset = np.empty(ndofs + 1, dtype=IntType)
+    # Mapping from dof to entity
+    dofs_to_entities = np.empty(ndofs, dtype=IntType)
+    # Helper array to populate `dofs_to_entities`
+    dofs_idx = np.zeros(ndofs, dtype=IntType)
+    # Adjacency list of degrees of freedom
+    # -> 100 is the upper bound on the number of adjacent dofs
+    adjacency_dofs = np.empty((100 * ndofs, 2), dtype=IntType)
+
+
+    # Preprocessing loop: populate `number_adj_entities` and `entity_per_dof_count`
+    size_entity_adj = 0
+    for i in range(num_entities):
+        # Number of adjacent entities of entity i
+        number_adj_entities[i + 1] = dmplex.getAdjacency(i).size
+        size_entity_adj += number_adj_entities[i + 1]
+        # Number of dofs associated with entity i
+        dof_offset = section.getOffset(i) + 1
+        for n in range(section.getDof(i)):
+            entity_per_dof_count[dof_offset + n] += 1
+
+    # Compute offsets
+    dof_to_entity_offset = np.cumsum(entity_per_dof_count, dtype=IntType)
+    adj_entity_offset = np.cumsum(number_adj_entities, dtype=IntType)
+    # Free memory of both: `entity_per_dof_count` and `number_adj_entities`
+    del entity_per_dof_count, number_adj_entities
+
+    # Adjacency list of entities
+    adj_entities = np.empty(size_entity_adj, dtype=IntType)
+
+    # Populate `adj_entities`, `dofs_to_entities`, and `dofs_idx`
+    for i in range(num_entities):
+        # Build adjacency list of entities
+        # adj_entities:
+        #                 entity i     
+        #  | --------- | --------- | --------- |
+        #              ^           ^
+        #              |           |
+        #              a           b
+        # with `a = adj_entity_offset[i]` and  `b = adj_entity_offset[i + 1]`
+        start_ent, end_ent = adj_entity_offset[i], adj_entity_offset[i+1]
+        adj_i = dmplex.getAdjacency(i)
+        for j in range(start_ent, end_ent):
+            adj_entities[j] = adj_i[j - start_ent]
+        # Build mapping from dofs to entities
+        # dofs_to_entities:
+        #                 dof j
+        # | --------- | --------- | --------- |
+        #             ^           ^
+        #             |           |
+        #             a           b
+        # with `a = dof_to_entity_offset[j]` and  `b = dof_to_entity_offset[j + 1]`
+        start_dof = section.getOffset(i)
+        end_dof = start_dof + section.getDof(i)
+        for dof in range(start_dof, end_dof):
+            dof_offset = dof_to_entity_offset[dof]
+            dofs_to_entities[dof_offset + dofs_idx[dof]] = i
+            dofs_idx[dof] += 1
+
+    # Populate `adjacency_dofs`
+    max_adjacent_dofs = 0
+    for i in range(ndofs):
+        # Loop over entities related to `dof`
+        for j in range(dof_to_entity_offset[i], dof_to_entity_offset[i + 1]):
+            e = dofs_to_entities[j]
+            # Loop over adjacent entities of `e`
+            start_ent, end_ent = adj_entity_offset[e], adj_entity_offset[e+1]
+            for k in range(start_ent, end_ent):
+                ek = adj_entities[k]
+                # Loop over dofs associated with `ek`
+                start_dof = section.getOffset(ek)
+                end_dof = start_dof + section.getDof(ek)
+                for dof in range(start_dof, end_dof):
+                    # Add edge to adjacency list if not already present
+                    if i >= dof:
+                        adjacency_dofs[max_adjacent_dofs, 0] = dof
+                        adjacency_dofs[max_adjacent_dofs, 1] = i
+                        max_adjacent_dofs += 1
+
+    return adjacency_dofs[:max_adjacent_dofs, :]

physics_driven_ml/utils/graph/mesh_to_graph.py

@@ -0,0 +1,13 @@
+from adjacency_dofs import build_dof_adjacency
+
+
+def build_adjacency(mesh, V):
+    dmplex = mesh.topology_dm
+    section = V.dm.getLocalSection()
+    # Get number of dofs
+    ndofs = V.dof_count
+    # Get number of entities of each dimension
+    dimension = mesh.geometric_dimension()
+    # Total number of mesh entities
+    num_entities = sum(mesh.num_entities(d) for d in range(dimension + 1))
+    return build_dof_adjacency(ndofs, num_entities, dmplex, section)

physics_driven_ml/utils/graph/mpi-compat.h

@@ -0,0 +1,19 @@
+/* Author:  Lisandro Dalcin   */
+/* Contact: [email protected] */
+
+#ifndef MPI_COMPAT_H
+#define MPI_COMPAT_H
+
+#include <mpi.h>
+
+#if (MPI_VERSION < 3) && !defined(PyMPI_HAVE_MPI_Message)
+typedef void *PyMPI_MPI_Message;
+#define MPI_Message PyMPI_MPI_Message
+#endif
+
+#if (MPI_VERSION < 4) && !defined(PyMPI_HAVE_MPI_Session)
+typedef void *PyMPI_MPI_Session;
+#define MPI_Session PyMPI_MPI_Session
+#endif
+
+#endif/*MPI_COMPAT_H*/

physics_driven_ml/utils/graph/petschdr.pxi

@@ -0,0 +1,185 @@
+cimport petsc4py.PETSc as PETSc
+cimport mpi4py.MPI as MPI
+cimport numpy as np
+
+cdef extern from "mpi-compat.h":
+    pass
+
+#IF COMPLEX:
+#    ctypedef np.complex128_t PetscScalar
+#ELSE:
+ctypedef double PetscScalar
+
+cdef extern from "petsc.h":
+    ctypedef long PetscInt
+    ctypedef double PetscReal
+    ctypedef enum PetscBool:
+        PETSC_TRUE, PETSC_FALSE
+    ctypedef enum PetscCopyMode:
+        PETSC_COPY_VALUES,
+        PETSC_OWN_POINTER,
+        PETSC_USE_POINTER
+    ctypedef enum PetscDataType:
+        PETSC_INT,
+        PETSC_REAL,
+        PETSC_SCALAR,
+        #PETSC_COMPLEX,
+        PETSC_DATATYPE_UNKNOWN
+
+cdef extern from "petscsys.h" nogil:
+    int PetscMalloc1(PetscInt,void*)
+    int PetscMalloc2(PetscInt,void*,PetscInt,void*)
+    int PetscFree(void*)
+    int PetscFree2(void*,void*)
+    int PetscSortIntWithArray(PetscInt,PetscInt[],PetscInt[])
+
+cdef extern from "petscdmplex.h" nogil:
+    int DMPlexGetHeightStratum(PETSc.PetscDM,PetscInt,PetscInt*,PetscInt*)
+    int DMPlexGetDepthStratum(PETSc.PetscDM,PetscInt,PetscInt*,PetscInt*)
+    int DMPlexGetPointHeight(PETSc.PetscDM,PetscInt,PetscInt*)
+    int DMPlexGetPointDepth(PETSc.PetscDM,PetscInt,PetscInt*)
+
+    int DMPlexGetChart(PETSc.PetscDM,PetscInt*,PetscInt*)
+    int DMPlexGetConeSize(PETSc.PetscDM,PetscInt,PetscInt*)
+    int DMPlexGetCone(PETSc.PetscDM,PetscInt,PetscInt*[])
+    int DMPlexGetConeOrientation(PETSc.PetscDM,PetscInt,PetscInt*[])
+    int DMPlexGetSupportSize(PETSc.PetscDM,PetscInt,PetscInt*)
+    int DMPlexGetSupport(PETSc.PetscDM,PetscInt,PetscInt*[])
+    int DMPlexGetMaxSizes(PETSc.PetscDM,PetscInt*,PetscInt*)
+
+    int DMPlexGetTransitiveClosure(PETSc.PetscDM,PetscInt,PetscBool,PetscInt *,PetscInt *[])
+    int DMPlexRestoreTransitiveClosure(PETSc.PetscDM,PetscInt,PetscBool,PetscInt *,PetscInt *[])
+    int DMPlexDistributeData(PETSc.PetscDM,PETSc.PetscSF,PETSc.PetscSection,MPI.MPI_Datatype,void*,PETSc.PetscSection,void**)
+    int DMPlexSetAdjacencyUser(PETSc.PetscDM,int(*)(PETSc.PetscDM,PetscInt,PetscInt*,PetscInt[],void*),void*)
+    int DMPlexCreatePointNumbering(PETSc.PetscDM,PETSc.PetscIS*)
+    int DMPlexLabelComplete(PETSc.PetscDM, PETSc.PetscDMLabel)
+
+cdef extern from "petscdmlabel.h" nogil:
+    struct _n_DMLabel
+    ctypedef _n_DMLabel* DMLabel "DMLabel"
+    int DMLabelCreateIndex(DMLabel, PetscInt, PetscInt)
+    int DMLabelDestroyIndex(DMLabel)
+    int DMLabelDestroy(DMLabel*)
+    int DMLabelHasPoint(DMLabel, PetscInt, PetscBool*)
+    int DMLabelSetValue(DMLabel, PetscInt, PetscInt)
+    int DMLabelGetValue(DMLabel, PetscInt, PetscInt*)
+    int DMLabelClearValue(DMLabel, PetscInt, PetscInt)
+    int DMLabelGetStratumSize(DMLabel, PetscInt, PetscInt*)
+    int DMLabelGetStratumIS(DMLabel, PetscInt, PETSc.PetscIS*)
+
+cdef extern from "petscdm.h" nogil:
+    int DMCreateLabel(PETSc.PetscDM,char[])
+    int DMGetLabel(PETSc.PetscDM,char[],DMLabel*)
+    int DMGetPointSF(PETSc.PetscDM,PETSc.PetscSF*)
+
+cdef extern from "petscdmswarm.h" nogil:
+    int DMSwarmGetLocalSize(PETSc.PetscDM,PetscInt*)
+    int DMSwarmGetCellDM(PETSc.PetscDM, PETSc.PetscDM*)
+    int DMSwarmGetField(PETSc.PetscDM,const char[],PetscInt*,PetscDataType*,void**)
+    int DMSwarmRestoreField(PETSc.PetscDM,const char[],PetscInt*,PetscDataType*,void**)
+
+cdef extern from "petscvec.h" nogil:
+    int VecGetArray(PETSc.PetscVec,PetscScalar**)
+    int VecRestoreArray(PETSc.PetscVec,PetscScalar**)
+    int VecGetArrayRead(PETSc.PetscVec,const PetscScalar**)
+    int VecRestoreArrayRead(PETSc.PetscVec,const PetscScalar**)
+
+cdef extern from "petscis.h" nogil:
+    int PetscSectionGetOffset(PETSc.PetscSection,PetscInt,PetscInt*)
+    int PetscSectionGetDof(PETSc.PetscSection,PetscInt,PetscInt*)
+    int PetscSectionSetDof(PETSc.PetscSection,PetscInt,PetscInt)
+    int PetscSectionSetFieldDof(PETSc.PetscSection,PetscInt,PetscInt,PetscInt)
+    int PetscSectionGetFieldDof(PETSc.PetscSection,PetscInt,PetscInt,PetscInt*)
+    int PetscSectionGetMaxDof(PETSc.PetscSection,PetscInt*)
+    int PetscSectionSetPermutation(PETSc.PetscSection,PETSc.PetscIS)
+    int ISGetIndices(PETSc.PetscIS,PetscInt*[])
+    int ISGetSize(PETSc.PetscIS,PetscInt*)
+    int ISRestoreIndices(PETSc.PetscIS,PetscInt*[])
+    int ISGeneralSetIndices(PETSc.PetscIS,PetscInt,PetscInt[],PetscCopyMode)
+    int ISLocalToGlobalMappingCreateIS(PETSc.PetscIS,PETSc.PetscLGMap*)
+    int ISLocalToGlobalMappingGetSize(PETSc.PetscLGMap,PetscInt*)
+    int ISLocalToGlobalMappingGetBlockIndices(PETSc.PetscLGMap, const PetscInt**)
+    int ISLocalToGlobalMappingRestoreBlockIndices(PETSc.PetscLGMap, const PetscInt**)
+    int ISDestroy(PETSc.PetscIS*)
+
+cdef extern from "petscsf.h" nogil:
+    struct PetscSFNode_:
+        PetscInt rank
+        PetscInt index
+    ctypedef PetscSFNode_ PetscSFNode "PetscSFNode"
+
+    int PetscSFGetGraph(PETSc.PetscSF,PetscInt*,PetscInt*,PetscInt**,PetscSFNode**)
+    int PetscSFSetGraph(PETSc.PetscSF,PetscInt,PetscInt,PetscInt*,PetscCopyMode,PetscSFNode*,PetscCopyMode)
+    int PetscSFBcastBegin(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,)
+    int PetscSFBcastEnd(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*)
+    int PetscSFReduceBegin(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,MPI.MPI_Op)
+    int PetscSFReduceEnd(PETSc.PetscSF,MPI.MPI_Datatype,const void*, void*,MPI.MPI_Op)
+
+ctypedef int (*PetscPCPatchComputeFunction)(PETSc.PetscPC,
+                                            PetscInt,
+                                            PETSc.PetscVec,
+                                            PETSc.PetscVec,
+                                            PETSc.PetscIS,
+                                            PetscInt,
+                                            const PetscInt*,
+                                            const PetscInt*,
+                                            void*)
+ctypedef int (*PetscPCPatchComputeOperator)(PETSc.PetscPC,
+                                            PetscInt,
+                                            PETSc.PetscVec,
+                                            PETSc.PetscMat,
+                                            PETSc.PetscIS,
+                                            PetscInt,
+                                            const PetscInt*,
+                                            const PetscInt*,
+                                            void*)
+cdef extern from "petscsnes.h" nogil:
+   int SNESPatchSetComputeFunction(PETSc.PetscSNES, PetscPCPatchComputeFunction, void *)
+   int SNESPatchSetComputeOperator(PETSc.PetscSNES, PetscPCPatchComputeOperator, void *)
+
+cdef extern from "petscpc.h" nogil:
+   int PCPatchSetComputeFunction(PETSc.PetscPC, PetscPCPatchComputeFunction, void *)
+   int PCPatchSetComputeFunctionInteriorFacets(PETSc.PetscPC, PetscPCPatchComputeFunction, void *)
+   int PCPatchSetComputeOperator(PETSc.PetscPC, PetscPCPatchComputeOperator, void *)
+   int PCPatchSetComputeOperatorInteriorFacets(PETSc.PetscPC, PetscPCPatchComputeOperator, void *)
+
+cdef extern from "petscbt.h" nogil:
+    ctypedef char * PetscBT
+    int PetscBTCreate(PetscInt,PetscBT*)
+    int PetscBTDestroy(PetscBT*)
+    char PetscBTLookup(PetscBT,PetscInt)
+    int PetscBTSet(PetscBT,PetscInt)
+
+cdef extern from "petscmat.h" nogil:
+    int MatSetValuesLocal(PETSc.PetscMat, PetscInt, const PetscInt[], PetscInt, const PetscInt[],
+                          const PetscScalar[], PetscInt)
+    int MatAssemblyBegin(PETSc.PetscMat, PetscInt)
+    int MatAssemblyEnd(PETSc.PetscMat, PetscInt)
+    PetscInt MAT_FINAL_ASSEMBLY = 0
+
+cdef extern from * nogil:
+    int PetscObjectTypeCompare(PETSc.PetscObject, char[], PetscBool*)
+
+# --- Error handling taken from petsc4py (src/PETSc.pyx) -------------
+
+cdef extern from *:
+    void PyErr_SetObject(object, object)
+    void *PyExc_RuntimeError
+
+cdef object PetscError = <object>PyExc_RuntimeError
+
+cdef inline int SETERR(int ierr) with gil:
+    if (<void*>PetscError) != NULL:
+        PyErr_SetObject(PetscError, <long>ierr)
+    else:
+        PyErr_SetObject(<object>PyExc_RuntimeError, <long>ierr)
+    return ierr
+
+cdef inline int CHKERR(int ierr) nogil except -1:
+    if ierr == 0:
+        return 0 # no error
+    else:
+        SETERR(ierr)
+        return -1
+
+# --------------------------------------------------------------------