warp-lang 1.9.1__py3-none-win_amd64.whl → 1.10.0__py3-none-win_amd64.whl

warp/{fem → _src/fem}/space/basis_function_space.py

@@ -13,20 +13,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Any, ClassVar, Optional
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.geometry import Geometry
-from warp.fem.linalg import generalized_inner, generalized_outer
-from warp.fem.types import NULL_QP_INDEX, Coords, ElementIndex, make_free_sample
-from warp.fem.utils import type_basis_element
+from warp._src.fem import cache
+from warp._src.fem.geometry import Geometry
+from warp._src.fem.linalg import generalized_inner, generalized_outer
+from warp._src.fem.types import NULL_QP_INDEX, Coords, ElementIndex, make_free_sample
+from warp._src.fem.utils import type_basis_element
 
 from .basis_space import BasisSpace
 from .dof_mapper import DofMapper, IdentityMapper
 from .function_space import FunctionSpace
 from .partition import SpacePartition, make_space_partition
 
+_wp_module_name_ = "warp.fem.space.basis_function_space"
+
 
 class CollocatedFunctionSpace(FunctionSpace):
     """Function space where values are collocated at nodes"""
@@ -81,10 +84,14 @@ class CollocatedFunctionSpace(FunctionSpace):
         if hasattr(basis, "vtk_cells"):
             self.vtk_cells = basis.vtk_cells
 
-    @property
+    @cached_property
     def name(self):
         return f"{self._basis.name}_{self.dof_mapper}".replace(".", "_")
 
+    @property
+    def basis(self) -> BasisSpace:
+        return self._basis
+
     def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
         return self._basis.node_positions(out=out)
 
@@ -92,7 +99,7 @@ class CollocatedFunctionSpace(FunctionSpace):
         self,
         space_partition: Optional[SpacePartition] = None,
     ) -> "wp.fem.field.NodalField":
-        from warp.fem.field import NodalField
+        from warp._src.fem.field import NodalField
 
         if space_partition is None:
             space_partition = make_space_partition(space_topology=self.topology)
@@ -185,7 +192,7 @@ class CollocatedFunctionSpaceTrace(CollocatedFunctionSpace):
         self._space = space
         super().__init__(space._basis.trace(), space.dtype, space.dof_mapper)
 
-    @property
+    @cached_property
     def name(self):
         return f"{self._space.name}_Trace"
 
@@ -237,6 +244,10 @@ class VectorValuedFunctionSpace(FunctionSpace):
     def name(self):
         return self._basis.name
 
+    @property
+    def basis(self) -> BasisSpace:
+        return self._basis
+
     def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
         return self._basis.node_positions(out=out)
 
@@ -244,7 +255,7 @@ class VectorValuedFunctionSpace(FunctionSpace):
         self,
         space_partition: Optional[SpacePartition] = None,
     ) -> "wp.fem.field.NodalField":
-        from warp.fem.field import NodalField
+        from warp._src.fem.field import NodalField
 
         if space_partition is None:
             space_partition = make_space_partition(space_topology=self.topology)
@@ -367,7 +378,7 @@ class CovariantFunctionSpaceTrace(VectorValuedFunctionSpace):
 
         cache.setup_dynamic_attributes(self, cls=__class__)
 
-    @property
+    @cached_property
     def name(self):
         return f"{self._space.name}_Trace"
 
@@ -448,7 +459,7 @@ class ContravariantFunctionSpaceTrace(VectorValuedFunctionSpace):
 
         cache.setup_dynamic_attributes(self, cls=__class__)
 
-    @property
+    @cached_property
     def name(self):
         return f"{self._space.name}_Trace"