warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/fem/field/field.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,659 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, ClassVar, Dict, Optional, Set
+# isort: skip_file
 
-import warp as wp
-from warp.fem import cache
-from warp.fem.domain import GeometryDomain, Sides
-from warp.fem.geometry import DeformedGeometry, Geometry
-from warp.fem.operator import Operator, integrand
-from warp.fem.space import FunctionSpace, SpacePartition
-from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind, Sample
-from warp.fem.utils import type_zero_element
+# TODO: Remove after cleaning up the public API.
 
+from warp._src.fem.field import field as _field
 
-class FieldLike:
-    """Base class for integrable fields"""
 
-    EvalArg: wp.codegen.Struct
-    """Structure containing field-level arguments passed to device functions for field evaluation"""
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
 
-    ElementEvalArg: wp.codegen.Struct
-    """Structure combining geometry-level and field-level arguments passed to device functions for field evaluation"""
-
-    def eval_arg_value(self, device) -> "EvalArg":  # noqa: F821
-        """Value of the field-level arguments to be passed to device functions"""
-        raise NotImplementedError
-
-    def fill_eval_arg(self, arg: "FieldLike.EvalArg", device):
-        """Fill the field-level arguments to be passed to device functions"""
-        raise NotImplementedError
-
-    @property
-    def degree(self) -> int:
-        """Polynomial degree of the field, used to estimate necessary quadrature order"""
-        raise NotImplementedError
-
-    @property
-    def name(self) -> str:
-        raise NotImplementedError
-
-    @property
-    def __str__(self) -> str:
-        return self.name
-
-    def gradient_valid(self) -> bool:
-        """Whether the gradient operator is implemented for this field."""
-        return False
-
-    def divergence_valid(self) -> bool:
-        """Whether the divergence operator is implemented for this field."""
-        return False
-
-    @staticmethod
-    def eval_inner(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the inner field value at a sample point"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_grad_inner(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the inner field gradient at a sample point"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_div_inner(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the inner field divergence at a sample point"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_outer(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the outer field value at a sample point"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_grad_outer(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the outer field gradient at a sample point"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_div_outer(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the outer field divergence at a sample point"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_degree(args: "ElementEvalArg"):  # noqa: F821
-        """Polynomial degree of the field is applicable, or hint for determination of interpolation order"""
-        raise NotImplementedError
-
-    def notify_operator_usage(self, ops: Set[Operator]):
-        """Makes the Domain aware that the operators `ops` will be applied"""
-        pass
-
-
-class GeometryField(FieldLike):
-    """Base class for fields defined over a geometry"""
-
-    @property
-    def geometry(self) -> Geometry:
-        """Geometry over which the field is expressed"""
-        raise NotImplementedError
-
-    @property
-    def element_kind(self) -> ElementKind:
-        """Kind of element over which the field is expressed"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_reference_grad_inner(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the inner field gradient with respect to reference element coordinates at a sample point"""
-        raise NotImplementedError
-
-    @staticmethod
-    def eval_reference_grad_outer(args: "ElementEvalArg", s: Sample):  # noqa: F821
-        """Device function evaluating the outer field gradient with respect to reference element coordinates at a sample point"""
-        raise NotImplementedError
-
-    def trace(self) -> FieldLike:
-        """Trace of this field over lower-dimensional elements"""
-        raise NotImplementedError
-
-    def make_deformed_geometry(self, relative=True) -> Geometry:
-        """Returns a deformed version of the underlying geometry, with positions displaced according to this field's values.
-
-        Args:
-            relative: If ``True``, the field is interpreted as a relative displacement over the original geometry.
-              If ``False``, the field values are interpreted as absolute positions.
-
-        """
-        return DeformedGeometry(self, relative=relative)
-
-    @property
-    def gradient_dtype(self):
-        """Return type of the (world space) gradient operator. Assumes self.gradient_valid()"""
-        if wp.types.type_is_matrix(self.dtype):
-            return None
-
-        if wp.types.type_is_vector(self.dtype):
-            return cache.cached_mat_type(
-                shape=(wp.types.type_size(self.dtype), self.geometry.dimension),
-                dtype=wp.types.type_scalar_type(self.dtype),
-            )
-        if wp.types.type_is_quaternion(self.dtype):
-            return cache.cached_mat_type(
-                shape=(4, self.geometry.dimension),
-                dtype=wp.types.type_scalar_type(self.dtype),
-            )
-        return cache.cached_vec_type(length=self.geometry.dimension, dtype=wp.types.type_scalar_type(self.dtype))
-
-    @property
-    def reference_gradient_dtype(self):
-        """Return type of the reference space gradient operator. Assumes self.gradient_valid()"""
-        if wp.types.type_is_matrix(self.dtype):
-            return None
-
-        if wp.types.type_is_vector(self.dtype):
-            return cache.cached_mat_type(
-                shape=(wp.types.type_size(self.dtype), self.geometry.cell_dimension),
-                dtype=wp.types.type_scalar_type(self.dtype),
-            )
-        if wp.types.type_is_quaternion(self.dtype):
-            return cache.cached_mat_type(
-                shape=(4, self.geometry.cell_dimension),
-                dtype=wp.types.type_scalar_type(self.dtype),
-            )
-        return cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=wp.types.type_scalar_type(self.dtype))
-
-    @property
-    def divergence_dtype(self):
-        """Return type of the divergence operator. Assumes self.divergence_valid()"""
-        if wp.types.type_is_vector(self.dtype):
-            return wp.types.type_scalar_type(self.dtype)
-        if wp.types.type_is_matrix(self.dtype):
-            return cache.cached_vec_type(length=self.dtype._shape_[1], dtype=wp.types.type_scalar_type(self.dtype))
-        return None
-
-
-class SpaceField(GeometryField):
-    """Base class for fields defined over a function space"""
-
-    def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
-        self._space = space
-        self._space_partition = space_partition
-
-        self.gradient_valid = self.space.gradient_valid
-        self.divergence_valid = self.space.divergence_valid
-
-    @property
-    def geometry(self) -> Geometry:
-        return self._space.geometry
-
-    @property
-    def element_kind(self) -> ElementKind:
-        return self._space.element_kind
-
-    @property
-    def space(self) -> FunctionSpace:
-        return self._space
-
-    @property
-    def space_partition(self) -> SpacePartition:
-        return self._space_partition
-
-    @property
-    def degree(self) -> int:
-        return self.space.degree
-
-    @property
-    def dtype(self) -> type:
-        return self.space.dtype
-
-    @property
-    def dof_dtype(self) -> type:
-        return self.space.dof_dtype
-
-    def _make_eval_degree(self):
-        ORDER = self.space.ORDER
-
-        @cache.dynamic_func(suffix=self.name)
-        def degree(args: self.ElementEvalArg):
-            return ORDER
-
-        return degree
-
-
-class DiscreteField(SpaceField):
-    """Explicitly-valued field defined over a partition of a discrete function space"""
-
-    @property
-    def dof_values(self) -> wp.array:
-        """Array of degrees of freedom values"""
-        raise NotImplementedError
-
-    @dof_values.setter
-    def dof_values(self, values: wp.array):
-        """Sets degrees of freedom values from an array"""
-        raise NotImplementedError
-
-    @staticmethod
-    def set_node_value(args: "FieldLike.EvalArg", node_index: int, value: Any):
-        """Device function setting the value at given node"""
-        raise NotImplementedError
-
-    @property
-    def name(self) -> str:
-        return f"{self.__class__.__qualname__}_{self.space.name}_{self.space_partition.name}"
-
-
-class ImplicitField(GeometryField):
-    """Field defined from an arbitrary function over a domain.
-    Does not support autodiff yet, so if gradient/divergence evaluation is required corresponding functions must be provided.
-
-    Args:
-        domain: Domain over which the field is defined
-        func: Warp function evaluating the field at a given position. Must accept at least one argument, with the first argument being the evaluation position (``wp.vec2`` or ``wp.vec3``).
-        values: Optional dictionary of additional argument values to be passed to the evaluation function.
-        grad_func: Optional gradient evaluation function; must take same arguments as `func`
-        div_func: Optional divergence evaluation function; must take same arguments as `func`
-        degree: Optional hint for automatic determination of quadrature orders when integrating this field
-    """
-
-    _dynamic_attribute_constructors: ClassVar = {
-        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
-        "eval_degree": lambda obj: obj._make_eval_degree(),
-        "eval_inner": lambda obj: obj._make_eval_func(obj._func),
-        "eval_grad_inner": lambda obj: obj._make_eval_func(obj._grad_func),
-        "eval_div_inner": lambda obj: obj._make_eval_func(obj._div_func),
-        "eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
-        "eval_outer": lambda obj: obj.eval_inner,
-        "eval_grad_outer": lambda obj: obj.eval_grad_inner,
-        "eval_div_outer": lambda obj: obj.eval_div_inner,
-        "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
-    }
-
-    def __init__(
-        self,
-        domain: GeometryDomain,
-        func: wp.Function,
-        values: Optional[Dict[str, Any]] = None,
-        grad_func: Optional[wp.Function] = None,
-        div_func: Optional[wp.Function] = None,
-        degree=0,
-    ):
-        self.domain = domain
-        self._degree = degree
-
-        if not isinstance(func, wp.Function):
-            raise ValueError("Implicit field function must be a warp Function (decorated with `wp.func`)")
-
-        self._func = func
-        self._grad_func = grad_func
-        self._div_func = div_func
-
-        argspec = integrand(func.func).argspec
-        arg_types = argspec.annotations
-
-        pos_arg_type = arg_types.pop(argspec.args[0]) if arg_types else None
-        if not pos_arg_type or not wp.types.types_equal(
-            pos_arg_type, wp.vec(length=domain.geometry.dimension, dtype=float), match_generic=True
-        ):
-            raise ValueError(
-                f"Implicit field function '{func.func.__name__}' must accept a position as its first argument"
-            )
-
-        self.EvalArg = cache.get_argument_struct(arg_types)
-        self.values = values
-
-        cache.setup_dynamic_attributes(self)
-
-    @property
-    def values(self):
-        return self._func_arg
-
-    @values.setter
-    def values(self, v):
-        self._values = v
-        self._func_arg = cache.populate_argument_struct(self.EvalArg, v, self._func.func.__name__)
-
-    @property
-    def geometry(self) -> Geometry:
-        return self.domain.geometry
-
-    @property
-    def element_kind(self) -> ElementKind:
-        return self.domain.element_kind
-
-    def eval_arg_value(self, device):
-        return self._func_arg
-
-    def fill_eval_arg(self, arg, device):
-        cache.populate_argument_struct(self.EvalArg, self._values, self._func.func.__name__, arg)
-
-    @property
-    def degree(self) -> int:
-        return self._degree
-
-    @property
-    def name(self) -> str:
-        return f"Implicit_{self.domain.name}_{self.degree}_{self.EvalArg.key}"
-
-    def gradient_valid(self) -> bool:
-        return self._grad_func is not None
-
-    def divergence_valid(self) -> bool:
-        return self._div_func is not None
-
-    def _make_eval_func(self, func):
-        if func is None:
-            return None
-
-        @cache.dynamic_func(
-            suffix=f"{self.name}_{func.key}",
-            code_transformers=[cache.ExpandStarredArgumentStruct({"args.eval_arg": self.EvalArg})],
-        )
-        def eval_inner(args: self.ElementEvalArg, s: Sample):
-            pos = self.domain.element_position(args.elt_arg, s)
-            return func(pos, *args.eval_arg)
-
-        return eval_inner
-
-    def _make_eval_reference_grad(self):
-        if self.eval_grad_inner is None:
-            return None
-
-        @cache.dynamic_func(suffix=f"{self.eval_grad_inner.key}")
-        def eval_reference_grad_inner(args: self.ElementEvalArg, s: Sample):
-            return self.eval_grad_inner(args, s) * self.domain.element_deformation_gradient(args.elt_arg, s)
-
-        return eval_reference_grad_inner
-
-    def _make_element_eval_arg(self):
-        @cache.dynamic_struct(suffix=self.name)
-        class ElementEvalArg:
-            elt_arg: self.domain.ElementArg
-            eval_arg: self.EvalArg
-
-        return ElementEvalArg
-
-    def _make_eval_degree(self):
-        ORDER = wp.constant(self._degree)
-
-        @cache.dynamic_func(suffix=self.name)
-        def degree(args: self.ElementEvalArg):
-            return ORDER
-
-        return degree
-
-    def trace(self):
-        if self.element_kind == ElementKind.SIDE:
-            raise RuntimeError("Trace only available for field defined on cell elements")
-
-        return ImplicitField(
-            domain=Sides(self.domain.geometry_partition),
-            func=self._func,
-            values={name: getattr(self.values, name) for name in self.EvalArg.vars},
-            grad_func=self._grad_func,
-            div_func=self._div_func,
-            degree=self._degree,
-        )
-
-
-class UniformField(GeometryField):
-    """Field defined as a constant value over a domain.
-
-    Args:
-        domain: Domain over which the field is defined
-        value: Uniform value over the domain
-    """
-
-    _dynamic_attribute_constructors: ClassVar = {
-        "EvalArg": lambda obj: obj._make_eval_arg(),
-        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
-        "eval_degree": lambda obj: obj._make_eval_degree(),
-        "eval_inner": lambda obj: obj._make_eval_inner(),
-        "eval_grad_inner": lambda obj: obj._make_eval_zero(obj.gradient_dtype),
-        "eval_div_inner": lambda obj: obj._make_eval_zero(obj.divergence_dtype),
-        "eval_reference_grad_inner": lambda obj: obj._make_eval_zero(obj.reference_gradient_dtype),
-        "eval_outer": lambda obj: obj.eval_inner,
-        "eval_grad_outer": lambda obj: obj.eval_grad_inner,
-        "eval_div_outer": lambda obj: obj.eval_div_inner,
-        "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
-    }
-
-    def __init__(self, domain: GeometryDomain, value: Any):
-        self.domain = domain
-
-        if not wp.types.is_value(value):
-            raise ValueError("value must be a Warp scalar, vector or matrix")
-
-        self.dtype = wp.types.type_to_warp(type(value))
-        self._value = self.dtype(value)
-
-        cache.setup_dynamic_attributes(self)
-
-    @property
-    def value(self):
-        return self._value
-
-    @value.setter
-    def value(self, v):
-        value_type = wp.types.type_to_warp(type(v))
-        assert wp.types.types_equal(value_type, self.dtype)
-        self._value = self.dtype(v)
-
-    @property
-    def geometry(self) -> Geometry:
-        return self.domain.geometry
-
-    @property
-    def element_kind(self) -> ElementKind:
-        return self.domain.element_kind
-
-    def eval_arg_value(self, device):
-        arg = self.EvalArg()
-        arg.value = self.value
-        return arg
-
-    def fill_eval_arg(self, arg, device):
-        arg.value = self.value
-
-    @property
-    def degree(self) -> int:
-        return 0
-
-    def gradient_valid(self) -> bool:
-        return self.gradient_dtype is not None
-
-    def divergence_valid(self) -> bool:
-        return self.divergence_dtype is not None
-
-    @property
-    def name(self) -> str:
-        return f"Uniform{self.domain.name}_{wp.types.get_type_code(self.dtype)}"
-
-    def _make_eval_inner(self):
-        @cache.dynamic_func(suffix=self.name)
-        def eval_inner(args: self.ElementEvalArg, s: Sample):
-            return args.eval_arg.value
-
-        return eval_inner
-
-    def _make_eval_zero(self, dtype):
-        if dtype is None:
-            return None
-
-        zero_element = type_zero_element(dtype)
-
-        @cache.dynamic_func(suffix=f"{self.name}_{wp.types.get_type_code(dtype)}")
-        def eval_zero(args: self.ElementEvalArg, s: Sample):
-            return zero_element()
-
-        return eval_zero
-
-    def _make_eval_arg(self):
-        @cache.dynamic_struct(suffix=self.name)
-        class EvalArg:
-            value: self.dtype
-
-        return EvalArg
-
-    def _make_element_eval_arg(self):
-        @cache.dynamic_struct(suffix=self.name)
-        class ElementEvalArg:
-            elt_arg: self.domain.ElementArg
-            eval_arg: self.EvalArg
-
-        return ElementEvalArg
-
-    def _make_eval_degree(self):
-        @cache.dynamic_func(suffix=self.name)
-        def degree(args: self.ElementEvalArg):
-            return 0
-
-        return degree
-
-    def trace(self):
-        if self.element_kind == ElementKind.SIDE:
-            raise RuntimeError("Trace only available for field defined on cell elements")
-
-        return UniformField(domain=Sides(self.domain.geometry_partition), value=self.value)
-
-
-class NonconformingField(GeometryField):
-    """Field defined as the map of a DiscreteField over a non-conforming geometry.
-
-    Args:
-        domain: The new domain over which the nonconforming field will be evaluated
-        field: Nonconforming discrete field
-        background: Uniform value or domain-conforming field determining the value outside of the geometry of definition of `field`
-    """
-
-    _LOOKUP_EPS = wp.constant(1.0e-6)
-
-    _dynamic_attribute_constructors: ClassVar = {
-        "EvalArg": lambda obj: obj._make_eval_arg(),
-        "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
-        "eval_degree": lambda obj: obj._make_eval_degree(),
-        "eval_inner": lambda obj: obj._make_nonconforming_eval("eval_inner"),
-        "eval_grad_inner": lambda obj: obj._make_nonconforming_eval("eval_grad_inner"),
-        "eval_div_inner": lambda obj: obj._make_nonconforming_eval("eval_div_inner"),
-        "eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
-        "eval_outer": lambda obj: obj.eval_inner,
-        "eval_grad_outer": lambda obj: obj.eval_grad_inner,
-        "eval_div_outer": lambda obj: obj.eval_div_inner,
-        "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
-    }
-
-    def __init__(self, domain: GeometryDomain, field: DiscreteField, background: Any = 0.0):
-        self.domain = domain
-
-        self.field = field
-        self.dtype = field.dtype
-
-        if not isinstance(background, GeometryField):
-            background = UniformField(domain, self.dtype(background))
-        elif background.geometry != domain.geometry or background.element_kind != domain.element_kind:
-            raise ValueError("Background field must be conforming to the domain")
-        self.background = background
-
-        cache.setup_dynamic_attributes(self)
-
-    @property
-    def geometry(self) -> Geometry:
-        return self.domain.geometry
-
-    @property
-    def element_kind(self) -> ElementKind:
-        return self.domain.element_kind
-
-    def eval_arg_value(self, device):
-        arg = self.EvalArg()
-        self.fill_eval_arg(arg, device)
-        return arg
-
-    def fill_eval_arg(self, arg, device):
-        self.field.fill_eval_arg(arg.field_cell_eval_arg.eval_arg, device)
-        self.field.geometry.fill_cell_arg(arg.field_cell_eval_arg.elt_arg, device)
-        self.background.fill_eval_arg(arg.background_arg, device)
-
-    @property
-    def degree(self) -> int:
-        return self.field.degree
-
-    def gradient_valid(self) -> bool:
-        return self.field.gradient_valid() and self.background.gradient_valid()
-
-    def divergence_valid(self) -> bool:
-        return self.field.divergence_valid() and self.background.divergence_valid()
-
-    @property
-    def name(self) -> str:
-        return f"{self.domain.name}_{self.field.name}_{self.background.name}"
-
-    def _make_nonconforming_eval(self, eval_func_name):
-        field_eval = getattr(self.field, eval_func_name)
-        bg_eval = getattr(self.background, eval_func_name)
-
-        if field_eval is None or bg_eval is None:
-            return None
-
-        cell_lookup = self.field.geometry.cell_lookup
-
-        @cache.dynamic_func(suffix=f"{eval_func_name}_{self.name}")
-        def eval_nc(args: self.ElementEvalArg, s: Sample):
-            pos = self.domain.element_position(args.elt_arg, s)
-            cell_arg = args.eval_arg.field_cell_eval_arg.elt_arg
-            nonconforming_s = cell_lookup(cell_arg, pos, NonconformingField._LOOKUP_EPS)
-            if nonconforming_s.element_index != NULL_ELEMENT_INDEX:
-                if (
-                    wp.length_sq(pos - self.field.geometry.cell_position(cell_arg, nonconforming_s))
-                    <= NonconformingField._LOOKUP_EPS
-                ):
-                    return field_eval(
-                        self.field.ElementEvalArg(cell_arg, args.eval_arg.field_cell_eval_arg.eval_arg), nonconforming_s
-                    )
-
-            return bg_eval(self.background.ElementEvalArg(args.elt_arg, args.eval_arg.background_arg), s)
-
-        return eval_nc
-
-    def _make_eval_reference_grad(self):
-        if self.eval_grad_inner is None:
-            return None
-
-        @cache.dynamic_func(suffix=f"{self.eval_grad_inner.key}")
-        def eval_reference_grad_inner(args: self.ElementEvalArg, s: Sample):
-            return self.eval_grad_inner(args, s) * self.domain.element_deformation_gradient(args.elt_arg, s)
-
-        return eval_reference_grad_inner
-
-    def _make_eval_arg(self):
-        @cache.dynamic_struct(suffix=self.name)
-        class EvalArg:
-            field_cell_eval_arg: self.field.ElementEvalArg
-            background_arg: self.background.EvalArg
-
-        return EvalArg
-
-    def _make_element_eval_arg(self):
-        @cache.dynamic_struct(suffix=self.name)
-        class ElementEvalArg:
-            elt_arg: self.domain.ElementArg
-            eval_arg: self.EvalArg
-
-        return ElementEvalArg
-
-    def _make_eval_degree(self):
-        @cache.dynamic_func(suffix=self.name)
-        def degree(args: self.ElementEvalArg):
-            return self.field.eval_degree(args.eval_arg.field_cell_eval_arg)
-
-        return degree
-
-    def trace(self):
-        if self.element_kind == ElementKind.SIDE:
-            raise RuntimeError("Trace only available for field defined on cell elements")
-
-        return NonconformingField(
-            domain=Sides(self.domain.geometry_partition), field=self.field, background=self.background.trace()
-        )
+    return get_deprecated_api(_field, "wp.fem.field", name)

warp/fem/geometry/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,20 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .adaptive_nanogrid import AdaptiveNanogrid
-from .deformed_geometry import DeformedGeometry
-from .element import Element
-from .geometry import Geometry
-from .grid_2d import Grid2D
-from .grid_3d import Grid3D
-from .hexmesh import Hexmesh
-from .nanogrid import Nanogrid
-from .partition import (
-    ExplicitGeometryPartition,
-    GeometryPartition,
-    LinearGeometryPartition,
-    WholeGeometryPartition,
-)
-from .quadmesh import Quadmesh, Quadmesh2D, Quadmesh3D
-from .tetmesh import Tetmesh
-from .trimesh import Trimesh, Trimesh2D, Trimesh3D
+# isort: skip_file
+
+from . import closest_point as closest_point
+
+from warp._src.fem.geometry import DeformedGeometry as DeformedGeometry
+from warp._src.fem.geometry import WholeGeometryPartition as WholeGeometryPartition