Add flag to disable computation of doflocs (to save memory)

kinnala · Mar 19, 2024 · 213bcfd · 213bcfd
1 parent 65b4a17
commit 213bcfd
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 16 deletions.
diff --git a/skfem/assembly/basis/abstract_basis.py b/skfem/assembly/basis/abstract_basis.py
@@ -47,7 +47,8 @@ def __init__(self,
                  intorder: Optional[int] = None,
                  quadrature: Optional[Tuple[ndarray, ndarray]] = None,
                  refdom: Type[Refdom] = Refdom,
-                 dofs: Optional[Dofs] = None):
+                 dofs: Optional[Dofs] = None,
+                 disable_doflocs: bool = False):
 
         if mesh.refdom != elem.refdom:
             raise ValueError("Incompatible Mesh and Element.")
@@ -56,17 +57,18 @@ def __init__(self,
         self.dofs = Dofs(mesh, elem) if dofs is None else dofs
 
         # global degree-of-freedom location
-        try:
-            doflocs = self.mapping.F(elem.doflocs.T)
-            self.doflocs = np.zeros((doflocs.shape[0], self.N))
-
-            # match mapped dofs and global dof numbering
-            for itr in range(doflocs.shape[0]):
-                for jtr in range(self.dofs.element_dofs.shape[0]):
-                    self.doflocs[itr, self.dofs.element_dofs[jtr]] =\
-                        doflocs[itr, :, jtr]
-        except Exception:
-            logger.warning("Unable to calculate global DOF locations.")
+        if not disable_doflocs:
+            try:
+                doflocs = self.mapping.F(elem.doflocs.T)
+                self.doflocs = np.zeros((doflocs.shape[0], self.N))
+
+                # match mapped dofs and global dof numbering
+                for itr in range(doflocs.shape[0]):
+                    for jtr in range(self.dofs.element_dofs.shape[0]):
+                        self.doflocs[itr, self.dofs.element_dofs[jtr]] =\
+                            doflocs[itr, :, jtr]
+            except Exception:
+                logger.warning("Unable to calculate global DOF locations.")
 
         self.mesh = mesh
         self.elem = elem

diff --git a/skfem/assembly/basis/cell_basis.py b/skfem/assembly/basis/cell_basis.py
@@ -46,7 +46,8 @@ def __init__(self,
                  intorder: Optional[int] = None,
                  elements: Optional[Any] = None,
                  quadrature: Optional[Tuple[ndarray, ndarray]] = None,
-                 dofs: Optional[Dofs] = None):
+                 dofs: Optional[Dofs] = None,
+                 disable_doflocs: bool = False):
         """Combine :class:`~skfem.mesh.Mesh` and
         :class:`~skfem.element.Element` into a set of precomputed global basis
         functions.
@@ -70,6 +71,10 @@ def __init__(self,
             Optional tuple of quadrature points and weights.
         dofs
             Optional :class:`~skfem.assembly.Dofs` object.
+        disable_doflocs
+            If `True`, the computation of global DOF locations is
+            disabled.  This may save memory on large meshes if DOF
+            locations are not required.
 
         """
         logger.info("Initializing {}({}, {})".format(type(self).__name__,
@@ -83,6 +88,7 @@ def __init__(self,
             quadrature,
             mesh.refdom,
             dofs,
+            disable_doflocs,
         )
 
         if elements is None:

diff --git a/skfem/assembly/basis/facet_basis.py b/skfem/assembly/basis/facet_basis.py
@@ -29,7 +29,8 @@ def __init__(self,
                  quadrature: Optional[Tuple[ndarray, ndarray]] = None,
                  facets: Optional[Any] = None,
                  dofs: Optional[Dofs] = None,
-                 side: int = 0):
+                 side: int = 0,
+                 disable_doflocs: bool = False):
         """Precomputed global basis on boundary facets.
 
         Parameters
@@ -51,6 +52,10 @@ def __init__(self,
             Optional subset of facet indices.
         dofs
             Optional :class:`~skfem.assembly.Dofs` object.
+        disable_doflocs
+            If `True`, the computation of global DOF locations is
+            disabled.  This may save memory on large meshes if DOF
+            locations are not required.
 
         """
         typestr = ("{}({}, {})".format(type(self).__name__,
@@ -65,6 +70,7 @@ def __init__(self,
             quadrature,
             mesh.brefdom,
             dofs,
+            disable_doflocs,
         )
 
         # by default use boundary facets

diff --git a/skfem/assembly/basis/interior_facet_basis.py b/skfem/assembly/basis/interior_facet_basis.py
@@ -25,7 +25,8 @@ def __init__(self,
                  quadrature: Optional[Tuple[ndarray, ndarray]] = None,
                  facets: Optional[Any] = None,
                  dofs: Optional[Dofs] = None,
-                 side: int = 0):
+                 side: int = 0,
+                 disable_doflocs: bool = False):
         """Precomputed global basis on interior facets."""
 
         if facets is None:
@@ -42,4 +43,5 @@ def __init__(self,
             facets=facets,
             dofs=dofs,
             side=side,
+            disable_doflocs=disable_doflocs,
         )
diff --git a/skfem/assembly/form/coo_data.py b/skfem/assembly/form/coo_data.py
@@ -55,7 +55,6 @@ def tolocal(self, basis=None):
         if self.local_shape is None:
             raise NotImplementedError("Cannot build local matrices if "
                                       "local_shape is not specified.")
-        assert len(self.local_shape) == 2
 
         local = np.moveaxis(self.data.reshape(self.local_shape + (-1,),
                                               order='C'), -1, 0)