warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/fem/integrate.py

@@ -19,6 +19,7 @@ import textwrap
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Set, Union
 
 import warp as wp
+import warp.fem.operator as operator
 from warp.codegen import get_annotations
 from warp.fem import cache
 from warp.fem.domain import GeometryDomain
@@ -35,7 +36,11 @@ from warp.fem.field import (
 )
 from warp.fem.field.virtual import make_bilinear_dispatch_kernel, make_linear_dispatch_kernel
 from warp.fem.linalg import array_axpy, basis_coefficient
-from warp.fem.operator import Integrand, Operator, at_node, integrand
+from warp.fem.operator import (
+    Integrand,
+    Operator,
+    integrand,
+)
 from warp.fem.quadrature import Quadrature, RegularQuadrature
 from warp.fem.types import (
     NULL_DOF_INDEX,
@@ -49,8 +54,9 @@ from warp.fem.types import (
     Sample,
     make_free_sample,
 )
+from warp.fem.utils import type_zero_element
 from warp.sparse import BsrMatrix, bsr_set_from_triplets, bsr_zeros
-from warp.types import type_length
+from warp.types import is_array, type_size
 from warp.utils import array_cast
 
 
@@ -111,6 +117,8 @@ class IntegrandVisitor(ast.NodeTransformer):
         def get_concrete_type(field: Union[FieldLike, Domain]):
             if isinstance(field, FieldLike):
                 return field.ElementEvalArg
+            elif isinstance(field, GeometryDomain):
+                return field.DomainArg
             return field.ElementArg
 
         return {
@@ -232,7 +240,7 @@ class IntegrandOperatorParser(IntegrandVisitor):
 
     @staticmethod
     def apply(
-        integrand: Integrand, field_args: Dict[str, FieldLike], operator_callback: Callable = None
+        integrand: Integrand, field_args: Dict[str, FieldLike], operator_callback: Optional[Callable] = None
     ) -> wp.Function:
         field_info = IntegrandVisitor._build_field_info(integrand, field_args)
         IntegrandOperatorParser(integrand, field_info, callback=operator_callback)._apply()
@@ -267,7 +275,11 @@ class IntegrandTransformer(IntegrandVisitor):
             setattr(field_info.concrete_type, pointer.key, pointer)
 
             # also insert callee as first argument
-            call.args = [ast.Name(id=callee, ctx=ast.Load())] + call.args
+            call.args = [ast.Name(id=callee, ctx=ast.Load()), *call.args]
+
+        # replace first argument with selected attribute
+        if operator.attr:
+            call.args[0] = ast.Attribute(value=call.args[0], attr=operator.attr)
 
     def _process_integrand_call(
         self, call: ast.Call, callee: Integrand, callee_field_args: Dict[str, IntegrandVisitor.FieldInfo]
@@ -456,6 +468,7 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
         fields_var_name: str = "fields",
         values_var_name: str = "values",
         domain_var_name: str = "domain_arg",
+        domain_index_var_name: str = "domain_index_arg",
         sample_var_name: str = "sample",
         field_wrappers_attr: str = "_field_wrappers",
     ):
@@ -470,6 +483,7 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
         self._fields_var_name = fields_var_name
         self._values_var_name = values_var_name
         self._domain_var_name = domain_var_name
+        self._domain_index_var_name = domain_index_var_name
         self._sample_var_name = sample_var_name
 
         self._field_wrappers_attr = field_wrappers_attr
@@ -485,8 +499,28 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
         for name, field in fields.items():
             if isinstance(field, FieldLike):
                 setattr(field_wrappers, name, field.ElementEvalArg)
+            elif isinstance(field, GeometryDomain):
+                setattr(field_wrappers, name, field.DomainArg)
         setattr(integrand_func, self._field_wrappers_attr, field_wrappers)
 
+    def _emit_field_wrapper_call(self, field_name, *data_arguments):
+        return ast.Call(
+            func=ast.Attribute(
+                value=ast.Attribute(
+                    value=ast.Name(id=self._func_name, ctx=ast.Load()),
+                    attr=self._field_wrappers_attr,
+                    ctx=ast.Load(),
+                ),
+                attr=field_name,
+                ctx=ast.Load(),
+            ),
+            args=[
+                ast.Name(id=self._domain_var_name, ctx=ast.Load()),
+                *data_arguments,
+            ],
+            keywords=[],
+        )
+
     def visit_Call(self, call: ast.Call):
         call = self.generic_visit(call)
 
@@ -498,33 +532,25 @@ class PassFieldArgsToIntegrand(ast.NodeTransformer):
             for arg in self._arg_names:
                 if arg == self._domain_name:
                     call.args.append(
-                        ast.Name(id=self._domain_var_name, ctx=ast.Load()),
+                        self._emit_field_wrapper_call(
+                            arg,
+                            ast.Name(id=self._domain_index_var_name, ctx=ast.Load()),
+                        )
                     )
+
                 elif arg == self._sample_name:
                     call.args.append(
                         ast.Name(id=self._sample_var_name, ctx=ast.Load()),
                     )
                 elif arg in self._field_args:
                     call.args.append(
-                        ast.Call(
-                            func=ast.Attribute(
-                                value=ast.Attribute(
-                                    value=ast.Name(id=self._func_name, ctx=ast.Load()),
-                                    attr=self._field_wrappers_attr,
-                                    ctx=ast.Load(),
-                                ),
+                        self._emit_field_wrapper_call(
+                            arg,
+                            ast.Attribute(
+                                value=ast.Name(id=self._fields_var_name, ctx=ast.Load()),
                                 attr=arg,
                                 ctx=ast.Load(),
                             ),
-                            args=[
-                                ast.Name(id=self._domain_var_name, ctx=ast.Load()),
-                                ast.Attribute(
-                                    value=ast.Name(id=self._fields_var_name, ctx=ast.Load()),
-                                    attr=arg,
-                                    ctx=ast.Load(),
-                                ),
-                            ],
-                            keywords=[],
                         )
                     )
                 elif arg in self._value_args:
@@ -704,7 +730,7 @@ def get_integrate_linear_nodal_kernel(
 
             coords = test.space.node_coords_in_element(
                 domain_arg,
-                _get_test_arg(),
+                _get_test_arg().space_arg,
                 element_index,
                 node_element_index.node_index_in_element,
             )
@@ -712,7 +738,7 @@ def get_integrate_linear_nodal_kernel(
             if coords[0] != OUTSIDE:
                 node_weight = test.space.node_quadrature_weight(
                     domain_arg,
-                    _get_test_arg(),
+                    _get_test_arg().space_arg,
                     element_index,
                     node_element_index.node_index_in_element,
                 )
@@ -913,7 +939,7 @@ def get_integrate_bilinear_nodal_kernel(
 
             coords = test.space.node_coords_in_element(
                 domain_arg,
-                _get_test_arg(),
+                _get_test_arg().space_arg,
                 element_index,
                 node_element_index.node_index_in_element,
             )
@@ -921,7 +947,7 @@ def get_integrate_bilinear_nodal_kernel(
             if coords[0] != OUTSIDE:
                 node_weight = test.space.node_quadrature_weight(
                     domain_arg,
-                    _get_test_arg(),
+                    _get_test_arg().space_arg,
                     element_index,
                     node_element_index.node_index_in_element,
                 )
@@ -1153,7 +1179,7 @@ def _launch_integrate_kernel(
     field_arg_values = FieldStruct()
     for k, v in fields.items():
         if not isinstance(v, GeometryDomain):
-            setattr(field_arg_values, k, v.eval_arg_value(device=device))
+            v.fill_eval_arg(getattr(field_arg_values, k), device=device)
 
     value_struct_values = cache.populate_argument_struct(ValueStruct, values, func_name=integrand.name)
 
@@ -1203,7 +1229,7 @@ def _launch_integrate_kernel(
             array_cast(in_array=accumulate_array, out_array=output)
         return output
 
-    test_arg = test.space_restriction.node_arg(device=device)
+    test_arg = test.space_restriction.node_arg_value(device=device)
     nodal = quadrature is None
 
     # Linear form
@@ -1211,9 +1237,9 @@ def _launch_integrate_kernel(
         # If an output array is provided with the correct type, accumulate directly into it
         # Otherwise, grab a temporary array
         if output is None:
-            if type_length(output_dtype) == test.node_dof_count:
+            if type_size(output_dtype) == test.node_dof_count:
                 output_shape = (test.space_partition.node_count(),)
-            elif type_length(output_dtype) == 1:
+            elif type_size(output_dtype) == 1:
                 output_shape = (test.space_partition.node_count(), test.node_dof_count)
             else:
                 raise RuntimeError(
@@ -1236,8 +1262,8 @@ def _launch_integrate_kernel(
                 raise RuntimeError(f"Output array must have at least {test.space_partition.node_count()} rows")
 
             output_dtype = output.dtype
-            if type_length(output_dtype) != test.node_dof_count:
-                if type_length(output_dtype) != 1:
+            if type_size(output_dtype) != test.node_dof_count:
+                if type_size(output_dtype) != 1:
                     raise RuntimeError(
                         f"Incompatible output type {wp.types.type_repr(output_dtype)}, must be scalar or vector of length {test.node_dof_count}"
                     )
@@ -1302,21 +1328,28 @@ def _launch_integrate_kernel(
                 device=device,
             )
 
-            dispatch_kernel = make_linear_dispatch_kernel(test, quadrature, accumulate_dtype)
-            wp.launch(
-                kernel=dispatch_kernel,
-                dim=(test.space_restriction.node_count(), test.node_dof_count),
-                inputs=[
-                    qp_arg,
-                    domain_elt_arg,
-                    domain_elt_index_arg,
-                    test_arg,
-                    test.global_field.eval_arg_value(device),
-                    local_result.array,
-                    output_view,
-                ],
-                device=device,
-            )
+            if test.TAYLOR_DOF_COUNT == 0:
+                wp.utils.warn(
+                    f"Test field is never evaluated in integrand '{integrand.name}', result will be zero",
+                    category=UserWarning,
+                    stacklevel=2,
+                )
+            else:
+                dispatch_kernel = make_linear_dispatch_kernel(test, quadrature, accumulate_dtype)
+                wp.launch(
+                    kernel=dispatch_kernel,
+                    dim=(test.space_restriction.node_count(), test.node_dof_count),
+                    inputs=[
+                        qp_arg,
+                        domain_elt_arg,
+                        domain_elt_index_arg,
+                        test_arg,
+                        test.space.space_arg_value(device),
+                        local_result.array,
+                        output_view,
+                    ],
+                    device=device,
+                )
 
             local_result.release()
 
@@ -1433,34 +1466,42 @@ def _launch_integrate_kernel(
             dtype=vec_array_dtype,
         )
 
-        dispatch_kernel = make_bilinear_dispatch_kernel(test, trial, quadrature, accumulate_dtype)
+        if test.TAYLOR_DOF_COUNT * trial.TAYLOR_DOF_COUNT == 0:
+            wp.utils.warn(
+                f"Test and/or trial fields are never evaluated in integrand '{integrand.name}', result will be zero",
+                category=UserWarning,
+                stacklevel=2,
+            )
+            triplet_rows.fill_(-1)
+        else:
+            dispatch_kernel = make_bilinear_dispatch_kernel(test, trial, quadrature, accumulate_dtype)
 
-        trial_partition_arg = trial.space_partition.partition_arg_value(device)
-        trial_topology_arg = trial.space_partition.space_topology.topo_arg_value(device)
-        wp.launch(
-            kernel=dispatch_kernel,
-            dim=(
-                test.space_restriction.node_count(),
-                test.node_dof_count,
-                trial.node_dof_count,
-                trial.space.topology.MAX_NODES_PER_ELEMENT,
-            ),
-            inputs=[
-                qp_arg,
-                domain_elt_arg,
-                domain_elt_index_arg,
-                test_arg,
-                test.global_field.eval_arg_value(device),
-                trial_partition_arg,
-                trial_topology_arg,
-                trial.global_field.eval_arg_value(device),
-                local_result_as_vec,
-                triplet_rows,
-                triplet_cols,
-                triplet_values,
-            ],
-            device=device,
-        )
+            trial_partition_arg = trial.space_partition.partition_arg_value(device)
+            trial_topology_arg = trial.space_partition.space_topology.topo_arg_value(device)
+            wp.launch(
+                kernel=dispatch_kernel,
+                dim=(
+                    test.space_restriction.node_count(),
+                    test.node_dof_count,
+                    trial.node_dof_count,
+                    trial.space.topology.MAX_NODES_PER_ELEMENT,
+                ),
+                inputs=[
+                    qp_arg,
+                    domain_elt_arg,
+                    domain_elt_index_arg,
+                    test_arg,
+                    test.space.space_arg_value(device),
+                    trial_partition_arg,
+                    trial_topology_arg,
+                    trial.space.space_arg_value(device),
+                    local_result_as_vec,
+                    triplet_rows,
+                    triplet_cols,
+                    triplet_values,
+                ],
+                device=device,
+            )
 
         local_result.release()
 
@@ -1529,21 +1570,30 @@ def _pick_assembly_strategy(
         if assembly not in ("generic", "nodal", "dispatch"):
             raise ValueError(f"Invalid assembly strategy'{assembly}'")
         return assembly
-    elif nodal:
-        return "nodal"
+    elif nodal is not None:
+        wp.utils.warn(
+            "'nodal' argument of `warp.fem.integrate` is deprecated and will be removed in a future version. Please use `assembly='nodal'` instead.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        if nodal:
+            return "nodal"
 
-    test_operators = operators.get(arguments.test_name, {})
-    trial_operators = operators.get(arguments.trial_name, {})
-    uses_at_node = at_node in test_operators or at_node in trial_operators
+    test_operators = operators.get(arguments.test_name, set())
+    trial_operators = operators.get(arguments.trial_name, set())
 
-    return "generic" if uses_at_node else "dispatch"
+    uses_virtual_node_operator = {operator.at_node, operator.node_count, operator.node_index} & (
+        test_operators | trial_operators
+    )
+
+    return "generic" if uses_virtual_node_operator else "dispatch"
 
 
 def integrate(
     integrand: Integrand,
     domain: Optional[GeometryDomain] = None,
     quadrature: Optional[Quadrature] = None,
-    nodal: bool = False,
+    nodal: Optional[bool] = None,
     fields: Optional[Dict[str, FieldLike]] = None,
     values: Optional[Dict[str, Any]] = None,
     accumulate_dtype: type = wp.float64,
@@ -1575,7 +1625,7 @@ def integrate(
         assembly: Specifies the strategy for assembling the integrated vector or matrix:
             - "nodal": For linear or bilinear forms, use the test function nodes as the quadrature points. Assumes Lagrange interpolation functions are used, and no differential or DG operator is evaluated on the test or trial functions.
             - "generic": Single-pass integration and shape-function evaluation. Makes no assumption about the integrand's content, but may lead to many redundant computations.
-            - "dispatch": For linear or bilinear forms, first evaluate the form at quadrature points then dispatch to nodes in a second pass. More efficient for integrands that are expensive to evaluate. Incompatible with `at_node` operator on test or trial functions.
+            - "dispatch": For linear or bilinear forms, first evaluate the form at quadrature points then dispatch to nodes in a second pass. More efficient for integrands that are expensive to evaluate. Incompatible with `at_node` and `node_index` operators on test or trial functions.
             - `None` (default): Automatically picks a suitable assembly strategy (either "generic" or "dispatch")
         add: If True and `output` is provided, add the integration result to `output` instead of replacing its content
         bsr_options: Additional options to be passed to the sparse matrix construction algorithm. See :func:`warp.sparse.bsr_set_from_triplets()`
@@ -1622,6 +1672,9 @@ def integrate(
 
     _find_integrand_operators(integrand, arguments.field_args)
 
+    if operator.lookup in integrand.operators.get(arguments.domain_name, []) and not domain.supports_lookup(device):
+        wp.utils.warn(f"{integrand.name}: using lookup() operator on a domain that does not support it")
+
     assembly = _pick_assembly_strategy(assembly, nodal, arguments=arguments, operators=integrand.operators)
     # print("assembly for ", integrand.name, ":", strategy)
 
@@ -1703,7 +1756,7 @@ def get_interpolate_to_field_function(
     ValueStruct: wp.codegen.Struct,
     dest: FieldRestriction,
 ):
-    value_type = dest.space.dtype
+    zero_value = type_zero_element(dest.space.dtype)
 
     def interpolate_to_field_fn(
         local_node_index: int,
@@ -1724,7 +1777,7 @@ def get_interpolate_to_field_function(
         # Volume-weighted average across elements
         # Superfluous if the interpolated function is continuous, but helpful for visualizing discontinuous spaces
 
-        val_sum = value_type(0.0)
+        val_sum = zero_value()
         vol_sum = float(0.0)
 
         for n in range(element_beg, element_end):
@@ -1969,6 +2022,7 @@ def get_interpolate_free_kernel(
     def interpolate_free_nonvalued_kernel_fn(
         dim: int,
         domain_arg: domain.ElementArg,
+        domain_index_arg: domain.ElementIndexArg,
         fields: FieldStruct,
         values: ValueStruct,
         result: wp.array(dtype=float),
@@ -1987,6 +2041,7 @@ def get_interpolate_free_kernel(
     def interpolate_free_kernel_fn(
         dim: int,
         domain_arg: domain.ElementArg,
+        domain_index_arg: domain.ElementIndexArg,
         fields: FieldStruct,
         values: ValueStruct,
         result: wp.array(dtype=value_type),
@@ -2143,12 +2198,12 @@ def _launch_interpolate_kernel(
     field_arg_values = FieldStruct()
     for k, v in fields.items():
         if not isinstance(v, GeometryDomain):
-            setattr(field_arg_values, k, v.eval_arg_value(device=device))
+            v.fill_eval_arg(getattr(field_arg_values, k), device=device)
 
     value_struct_values = cache.populate_argument_struct(ValueStruct, values, func_name=integrand.name)
 
     if isinstance(dest, FieldRestriction):
-        dest_node_arg = dest.space_restriction.node_arg(device=device)
+        dest_node_arg = dest.space_restriction.node_arg_value(device=device)
         dest_eval_arg = dest.field.eval_arg_value(device=device)
 
         wp.launch(
@@ -2167,33 +2222,49 @@ def _launch_interpolate_kernel(
         return
 
     if quadrature is None:
+        if dest is not None and (not is_array(dest) or dest.shape[0] != dim):
+            raise ValueError(f"dest must be a warp array with {dim} rows")
+
         wp.launch(
             kernel=kernel,
             dim=dim,
-            inputs=[dim, elt_arg, field_arg_values, value_struct_values, dest],
+            inputs=[dim, elt_arg, elt_index_arg, field_arg_values, value_struct_values, dest],
             device=device,
         )
         return
 
     qp_arg = quadrature.arg_value(device)
+    qp_eval_count = quadrature.evaluation_point_count()
+    qp_index_count = quadrature.total_point_count()
+
+    if qp_eval_count != qp_index_count:
+        wp.utils.warn(
+            f"Quadrature used for interpolation of {integrand.name} has different number of evaluation and indexed points, this may lead to incorrect results",
+            category=UserWarning,
+            stacklevel=2,
+        )
+
     qp_element_index_arg = quadrature.element_index_arg_value(device)
     if trial is None:
+        if dest is not None and (not is_array(dest) or dest.shape[0] != qp_index_count):
+            raise ValueError(f"dest must be a warp array with {qp_index_count} rows")
+
         wp.launch(
             kernel=kernel,
-            dim=quadrature.evaluation_point_count(),
+            dim=qp_eval_count,
             inputs=[qp_arg, qp_element_index_arg, elt_arg, elt_index_arg, field_arg_values, value_struct_values, dest],
             device=device,
         )
         return
 
-    nnz = quadrature.total_point_count() * trial.space.topology.MAX_NODES_PER_ELEMENT
+    nnz = qp_eval_count * trial.space.topology.MAX_NODES_PER_ELEMENT
 
-    if dest.nrow != quadrature.total_point_count() or dest.ncol != trial.space_partition.node_count():
+    if dest.nrow != qp_index_count or dest.ncol != trial.space_partition.node_count():
         raise RuntimeError(
-            f"'dest' matrix must have {quadrature.total_point_count()} rows and {trial.space_partition.node_count()} columns of blocks"
+            f"'dest' matrix must have {qp_index_count} rows and {trial.space_partition.node_count()} columns of blocks"
         )
     if dest.block_shape[1] != trial.node_dof_count:
-        raise f"'dest' matrix blocks must have {trial.node_dof_count} columns"
+        raise RuntimeError(f"'dest' matrix blocks must have {trial.node_dof_count} columns")
 
     triplet_rows_temp = cache.borrow_temporary(temporary_store, shape=(nnz,), dtype=int, device=device)
     triplet_cols_temp = cache.borrow_temporary(temporary_store, shape=(nnz,), dtype=int, device=device)
@@ -2243,7 +2314,7 @@ def interpolate(
     integrand: Union[Integrand, FieldLike],
     dest: Optional[Union[DiscreteField, FieldRestriction, wp.array]] = None,
     quadrature: Optional[Quadrature] = None,
-    dim: int = 0,
+    dim: Optional[int] = None,
     domain: Optional[Domain] = None,
     fields: Optional[Dict[str, FieldLike]] = None,
     values: Optional[Dict[str, Any]] = None,
@@ -2290,11 +2361,13 @@ def interpolate(
     arguments = _parse_integrand_arguments(integrand, fields)
     if arguments.test_name:
         raise ValueError(f"Test field '{arguments.test_name}' maybe not be used for interpolation")
-    if arguments.trial_name and (quadrature is None or not isinstance(dest, BsrMatrix)):
+    if arguments.trial_name and not isinstance(dest, BsrMatrix):
         raise ValueError(
-            f"Interpolation using trial field '{arguments.trial_name}' requires 'quadrature' to be provided and 'dest' to be a `warp.sparse.BsrMatrix`"
+            f"Interpolation using trial field '{arguments.trial_name}' requires 'dest' to be a `warp.sparse.BsrMatrix`"
         )
 
+    trial = arguments.field_args.get(arguments.trial_name, None)
+
     if isinstance(dest, DiscreteField):
         dest = make_restriction(dest, domain=domain)
 
@@ -2302,12 +2375,25 @@ def interpolate(
         domain = dest.domain
     elif quadrature is not None:
         domain = quadrature.domain
+    elif dim is None:
+        if trial is not None:
+            domain = trial.domain
+        elif domain is None:
+            raise ValueError(
+                "Unable to determine interpolation domain, provide an explicit field restriction or quadrature"
+            )
+
+        # Default to one sample per domain element
+        quadrature = RegularQuadrature(domain, order=0)
 
     if arguments.domain_name:
         arguments.field_args[arguments.domain_name] = domain
 
     _find_integrand_operators(integrand, arguments.field_args)
 
+    if operator.lookup in integrand.operators.get(arguments.domain_name, []) and not domain.supports_lookup(device):
+        wp.utils.warn(f"{integrand.name}: using lookup() operator on a domain that does not support it")
+
     kernel, FieldStruct, ValueStruct = _generate_interpolate_kernel(
         integrand=integrand,
         domain=domain,
@@ -2326,7 +2412,7 @@ def interpolate(
         dest=dest,
         quadrature=quadrature,
         dim=dim,
-        trial=fields.get(arguments.trial_name),
+        trial=trial,
         fields=arguments.field_args,
         values=values,
         temporary_store=temporary_store,