warp-lang 1.7.2__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/examples/fem/example_elastic_shape_optimization.py

@@ -0,0 +1,387 @@
+# 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");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example: Shape Optimization of a 2D Elastic Cantilever Beam
+#
+# This example demonstrates shape optimization of a 2D elastic cantilever beam using
+# finite element analysis and gradient-based optimization.
+#
+# Problem Setup:
+# - The computational domain is a 2D beam (rectangular mesh), fixed on the left edge
+#   (Dirichlet boundary condition: zero displacement).
+# - A constant external load is applied to the right edge of the beam, causing it to deform.
+# - The beam is discretized using finite elements, and the displacement field is solved
+#   for the current geometry at each optimization step using a linear elasticity formulation.
+#
+# Shape Optimization Strategy:
+# - The goal is to optimize the shape of the beam to minimize the total squared norm of the
+#   stress field (e.g., compliance or strain energy) over the domain.
+# - The positions of the left and right boundary vertices are fixed throughout the optimization
+#   to maintain the beam's support and loading conditions.
+# - A volume constraint is enforced to preserve the total material volume, preventing trivial
+#   solutions that simply shrink the structure.
+# - An ad-hoc "quality" term is included in the loss function to penalize degenerate or inverted elements,
+#   helping to maintain mesh quality (as remeshing would be out of scope for this example).
+# - The optimization is performed by computing the gradient of the objective with respect to
+#   the nodal positions (shape derivatives) using the adjoint method or automatic differentiation.
+# - At each iteration, the nodal positions are updated in the direction of decreasing objective,
+#   subject to the volume constraint and boundary conditions.
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.optim import Adam
+
+
+@fem.integrand(kernel_options={"max_unroll": 1})
+def boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def classify_boundary_sides(
+    s: fem.Sample,
+    domain: fem.Domain,
+    left: wp.array(dtype=int),
+    right: wp.array(dtype=int),
+):
+    nor = fem.normal(domain, s)
+
+    if nor[0] < -0.5:
+        left[s.qp_index] = 1
+    elif nor[0] > 0.5:
+        right[s.qp_index] = 1
+
+
+@wp.func
+def hooke_stress(strain: wp.mat22, lame: wp.vec2):
+    """Hookean elasticity"""
+    return 2.0 * lame[1] * strain + lame[0] * wp.trace(strain) * wp.identity(n=2, dtype=float)
+
+
+@fem.integrand
+def stress_field(s: fem.Sample, u: fem.Field, lame: wp.vec2):
+    return hooke_stress(fem.D(u, s), lame)
+
+
+@fem.integrand
+def hooke_elasticity_form(s: fem.Sample, u: fem.Field, v: fem.Field, lame: wp.vec2):
+    return wp.ddot(fem.D(v, s), stress_field(s, u, lame))
+
+
+@fem.integrand
+def applied_load_form(s: fem.Sample, domain: fem.Domain, v: fem.Field, load: wp.vec2):
+    return wp.dot(v(s), load)
+
+
+@fem.integrand
+def loss_form(
+    s: fem.Sample, domain: fem.Domain, u: fem.Field, lame: wp.vec2, quality_threshold: float, quality_weight: float
+):
+    stress = stress_field(s, u, lame)
+    stress_norm_sq = wp.ddot(stress, stress)
+
+    # As we're not remeshing, add a "quality" term
+    # to avoid degenerate and inverted elements
+
+    F = fem.deformation_gradient(domain, s)
+    U, S, V = wp.svd2(F)
+
+    quality = wp.min(S) / wp.max(S) / quality_threshold
+    quality_pen = -wp.log(wp.max(quality, 0.0001)) * wp.min(0.0, quality - 1.0) * wp.min(0.0, quality - 1.0)
+
+    return stress_norm_sq + quality_pen * quality_weight
+
+
+@fem.integrand
+def volume_form():
+    return 1.0
+
+
+@wp.kernel
+def add_volume_loss(
+    loss: wp.array(dtype=wp.float32), vol: wp.array(dtype=wp.float32), target_vol: wp.float32, weight: wp.float32
+):
+    loss[0] += weight * (vol[0] - target_vol) * (vol[0] - target_vol)
+
+
+class Example:
+    def __init__(
+        self,
+        quiet=False,
+        degree=2,
+        resolution=25,
+        mesh="tri",
+        poisson_ratio=0.5,
+        load=(0.0, -1),
+        lr=1.0e-3,
+    ):
+        self._quiet = quiet
+
+        # Lame coefficients from Young modulus and Poisson ratio
+        self._lame = wp.vec2(1.0 / (1.0 + poisson_ratio) * np.array([poisson_ratio / (1.0 - poisson_ratio), 0.5]))
+        self._load = load
+
+        # procedural rectangular domain definition
+        bounds_lo = wp.vec2(0.0, 0.8)
+        bounds_hi = wp.vec2(1.0, 1.0)
+        self._initial_volume = (bounds_hi - bounds_lo)[0] * (bounds_hi - bounds_lo)[1]
+
+        if mesh == "tri":
+            # triangle mesh, optimize vertices directly
+            positions, tri_vidx = fem_example_utils.gen_trimesh(
+                res=wp.vec2i(resolution, resolution // 5), bounds_lo=bounds_lo, bounds_hi=bounds_hi
+            )
+            self._geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+            self._start_geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=wp.clone(positions))
+            self._vertex_positions = positions
+        elif mesh == "quad":
+            # quad mesh, optimize vertices directly
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(
+                res=wp.vec2i(resolution, resolution // 5), bounds_lo=bounds_lo, bounds_hi=bounds_hi
+            )
+            self._geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+            self._start_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=wp.clone(positions))
+            self._vertex_positions = positions
+        else:
+            # grid, optimize nodes of deformation field
+            self._start_geo = fem.Grid2D(
+                wp.vec2i(resolution, resolution // 5), bounds_lo=bounds_lo, bounds_hi=bounds_hi
+            )
+            vertex_displacement_space = fem.make_polynomial_space(self._start_geo, degree=1, dtype=wp.vec2)
+            vertex_position_field = fem.make_discrete_field(space=vertex_displacement_space)
+            vertex_position_field.dof_values = vertex_displacement_space.node_positions()
+            self._geo = vertex_position_field.make_deformed_geometry(relative=False)
+            self._vertex_positions = vertex_position_field.dof_values
+
+        # make sure positions are differentiable
+        self._vertex_positions.requires_grad = True
+
+        # Store initial node positions (for rendering)
+        self._u_space = fem.make_polynomial_space(self._geo, degree=degree, dtype=wp.vec2)
+        self._start_node_positions = self._u_space.node_positions()
+
+        # displacement field, make sure gradient is stored
+        self._u_field = fem.make_discrete_field(space=self._u_space)
+        self._u_field.dof_values.requires_grad = True
+
+        # Trial and test functions
+        self._u_test = fem.make_test(space=self._u_space)
+        self._u_trial = fem.make_trial(space=self._u_space)
+
+        # Identify left and right sides for boundary conditions
+        boundary = fem.BoundarySides(self._geo)
+
+        left_mask = wp.zeros(shape=boundary.element_count(), dtype=int)
+        right_mask = wp.zeros(shape=boundary.element_count(), dtype=int)
+
+        fem.interpolate(
+            classify_boundary_sides,
+            quadrature=fem.RegularQuadrature(boundary, order=0),
+            values={"left": left_mask, "right": right_mask},
+        )
+
+        self._left = fem.Subdomain(boundary, element_mask=left_mask)
+        self._right = fem.Subdomain(boundary, element_mask=right_mask)
+
+        # Build projectors for the left-side homogeneous Dirichlet condition
+        u_left_bd_test = fem.make_test(space=self._u_space, domain=self._left)
+        u_left_bd_trial = fem.make_trial(space=self._u_space, domain=self._left)
+        u_left_bd_matrix = fem.integrate(
+            boundary_projector_form,
+            fields={"u": u_left_bd_trial, "v": u_left_bd_test},
+            assembly="nodal",
+            output_dtype=float,
+        )
+        fem.normalize_dirichlet_projector(u_left_bd_matrix)
+        self._bd_projector = u_left_bd_matrix
+
+        # Fixed vertices (that the shape optimization should not move)
+        # Build projectors for the left and right subdomains and add them together
+        vertex_space = fem.make_polynomial_space(self._geo, degree=1, dtype=wp.vec2)
+        u_left_vertex_bd_test = fem.make_test(space=vertex_space, domain=self._left)
+        u_left_vertex_bd_trial = fem.make_trial(space=vertex_space, domain=self._left)
+        u_right_vertex_bd_test = fem.make_test(space=vertex_space, domain=self._right)
+        u_right_vertex_bd_trial = fem.make_trial(space=vertex_space, domain=self._right)
+        u_fixed_vertex_matrix = fem.integrate(
+            boundary_projector_form,
+            fields={"u": u_left_vertex_bd_trial, "v": u_left_vertex_bd_test},
+            assembly="nodal",
+            output_dtype=float,
+        ) + fem.integrate(
+            boundary_projector_form,
+            fields={"u": u_right_vertex_bd_trial, "v": u_right_vertex_bd_test},
+            assembly="nodal",
+            output_dtype=float,
+        )
+        fem.normalize_dirichlet_projector(u_fixed_vertex_matrix)
+        self._fixed_vertex_projector = u_fixed_vertex_matrix
+
+        self._u_right_test = fem.make_test(space=self._u_space, domain=self._right)
+
+        # initialize renderer
+        self.renderer = fem_example_utils.Plot()
+
+        # Initialize Adam optimizer
+        # Current implementation assumes scalar arrays, so cast our vec2 arrays to scalars
+        self._vertex_positions_scalar = wp.array(self._vertex_positions, dtype=wp.float32).flatten()
+        self._vertex_positions_scalar.grad = wp.array(self._vertex_positions.grad, dtype=wp.float32).flatten()
+        self.optimizer = Adam([self._vertex_positions_scalar], lr=lr)
+
+    def step(self):
+        # Forward step, record adjoint tape for forces
+        u = self._u_field.dof_values
+        u.zero_()
+
+        u_rhs = wp.empty(self._u_space.node_count(), dtype=wp.vec2f, requires_grad=True)
+
+        tape = wp.Tape()
+
+        with tape:
+            fem.integrate(
+                applied_load_form,
+                fields={"v": self._u_right_test},
+                values={"load": self._load},
+                output=u_rhs,
+            )
+            # the elastic force will be zero at the first iteration,
+            # but including it on the tape is necessary to compute the gradient of the force equilibrium
+            # using the implicit function theorem
+            # Note that this will be evaluated in the backward pass using the updated values for "_u_field"
+            fem.integrate(
+                hooke_elasticity_form,
+                fields={"u": self._u_field, "v": self._u_test},
+                values={"lame": -self._lame},
+                output=u_rhs,
+                add=True,
+            )
+
+        u_matrix = fem.integrate(
+            hooke_elasticity_form,
+            fields={"u": self._u_trial, "v": self._u_test},
+            values={"lame": self._lame},
+            output_dtype=float,
+        )
+        fem.project_linear_system(u_matrix, u_rhs, self._bd_projector, normalize_projector=False)
+
+        fem_example_utils.bsr_cg(u_matrix, b=u_rhs, x=u, quiet=self._quiet, tol=1e-6, max_iters=1000)
+
+        # Record adjoint of linear solve
+        # (For nonlinear elasticity, this should use the final hessian, as per implicit function theorem)
+        def solve_linear_system():
+            fem_example_utils.bsr_cg(u_matrix, b=u.grad, x=u_rhs.grad, quiet=self._quiet, tol=1e-6, max_iters=1000)
+            u_rhs.grad -= self._bd_projector @ u_rhs.grad
+            self._u_field.dof_values.grad.zero_()
+
+        tape.record_func(solve_linear_system, arrays=(u_rhs, u))
+
+        # Evaluate residual
+        # Integral of squared difference between simulated position and target positions
+        loss = wp.empty(shape=1, dtype=wp.float32, requires_grad=True)
+        vol = wp.empty(shape=1, dtype=wp.float32, requires_grad=True)
+
+        with tape:
+            fem.integrate(
+                loss_form,
+                fields={"u": self._u_field},
+                values={"lame": self._lame, "quality_threshold": 0.2, "quality_weight": 20.0},
+                domain=self._u_test.domain,
+                output=loss,
+            )
+            print(f"Loss: {loss}")
+
+            # Add penalization term enforcing constant volume
+            fem.integrate(
+                volume_form,
+                domain=self._u_test.domain,
+                output=vol,
+            )
+
+            vol_loss_weight = 100.0
+            wp.launch(
+                add_volume_loss,
+                dim=1,
+                inputs=(loss, vol, self._initial_volume, vol_loss_weight),
+            )
+
+        # perform backward step
+        tape.backward(loss=loss)
+
+        # enforce fixed vertices
+        self._vertex_positions.grad -= self._fixed_vertex_projector @ self._vertex_positions.grad
+
+        # update positions and reset tape
+        self.optimizer.step([self._vertex_positions_scalar.grad])
+        tape.zero()
+
+    def render(self):
+        # Render using fields defined on start geometry
+        # (renderer assumes geometry remains fixed for timesampled fields)
+        u_space = fem.make_polynomial_space(self._start_geo, degree=self._u_space.degree, dtype=wp.vec2)
+        u_field = fem.make_discrete_field(space=u_space)
+        rest_field = fem.make_discrete_field(space=u_space)
+
+        geo_displacement = self._u_space.node_positions() - self._start_node_positions
+        u_field.dof_values = self._u_field.dof_values + geo_displacement
+        rest_field.dof_values = geo_displacement
+
+        self.renderer.add_field("displacement", u_field)
+        self.renderer.add_field("rest", rest_field)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--device", type=str, default=None, help="Override the default Warp device.")
+    parser.add_argument("--resolution", type=int, default=10, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=1, help="Polynomial degree of shape functions.")
+    parser.add_argument("--mesh", choices=("tri", "quad", "grid"), default="tri", help="Mesh type")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+    parser.add_argument("--lr", type=float, default=1.0e-3, help="Learning rate.")
+    parser.add_argument("--num_iters", type=int, default=250, help="Number of iterations.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=True,
+            degree=args.degree,
+            resolution=args.resolution,
+            mesh=args.mesh,
+            poisson_ratio=0.95,
+            load=wp.vec2(0.0, -0.1),
+            lr=args.lr,
+        )
+
+        for _k in range(args.num_iters):
+            example.step()
+            example.render()
+
+        if not args.headless:
+            example.renderer.plot(options={"displacement": {"displacement": {}}, "rest": {"displacement": {}}})

warp/examples/fem/example_magnetostatics.py

@@ -24,9 +24,9 @@
 # Div. B = 0
 #
 # solved over field A such that B = Curl A,
-# and Direchlet homogeneous essential boundary conditions
+# and Dirichlet homogeneous essential boundary conditions
 #
-# This example also illustrates using an ImplictField to warp a grid mesh
+# This example also illustrates using an ImplicitField to warp a grid mesh
 # to a cylindrical domain
 ###########################################################################
 
@@ -199,7 +199,9 @@ class Example:
         boundary = fem.BoundarySides(sim_geo)
         u_bd = fem.make_trial(space=A_space, domain=boundary)
         v_bd = fem.make_test(space=A_space, domain=boundary)
-        dirichlet_bd_proj = fem.integrate(mass_form, fields={"u": u_bd, "v": v_bd}, nodal=True, output_dtype=float)
+        dirichlet_bd_proj = fem.integrate(
+            mass_form, fields={"u": u_bd, "v": v_bd}, assembly="nodal", output_dtype=float
+        )
         fem.project_linear_system(lhs, rhs, dirichlet_bd_proj)
 
         # solve using Conjugate Residual (numerically rhs may not be in image of lhs)

warp/examples/fem/example_mixed_elasticity.py

@@ -192,11 +192,11 @@ class Example:
             vertical_displacement_form,
             fields={"v": u_bd_test},
             values={"displacement": self._displacement},
-            nodal=True,
+            assembly="nodal",
             output_dtype=wp.vec2d,
         )
         u_bd_matrix = fem.integrate(
-            vertical_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True
+            vertical_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, assembly="nodal"
         )
 
         # Stress/velocity coupling
@@ -207,7 +207,9 @@ class Example:
         gradient_matrix = fem.integrate(displacement_gradient_form, fields={"u": u_trial, "tau": tau_test}).transpose()
 
         # Compute inverse of the (block-diagonal) tau mass matrix
-        tau_inv_mass_matrix = fem.integrate(tensor_mass_form, fields={"sig": tau_trial, "tau": tau_test}, nodal=True)
+        tau_inv_mass_matrix = fem.integrate(
+            tensor_mass_form, fields={"sig": tau_trial, "tau": tau_test}, assembly="nodal"
+        )
         fem_example_utils.invert_diagonal_bsr_matrix(tau_inv_mass_matrix)
 
         # Newton iterations (without line-search for simplicity)

warp/examples/fem/example_navier_stokes.py

@@ -89,20 +89,22 @@ def div_form(
 
 
 class Example:
-    def __init__(self, quiet=False, degree=2, resolution=25, Re=1000.0, top_velocity=1.0, tri_mesh=False):
+    def __init__(self, quiet=False, degree=2, resolution=25, Re=1000.0, top_velocity=1.0, mesh: str = "grid"):
         self._quiet = quiet
 
-        res = resolution
         self.sim_dt = 1.0 / resolution
         self.current_frame = 0
 
         viscosity = top_velocity / Re
 
-        if tri_mesh:
-            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(res))
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
             geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions, build_bvh=True)
+        elif mesh == "quad":
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions, build_bvh=True)
         else:
-            geo = fem.Grid2D(res=wp.vec2i(res))
+            geo = fem.Grid2D(res=wp.vec2i(resolution))
 
         domain = fem.Cells(geometry=geo)
         boundary = fem.BoundarySides(geo)
@@ -130,7 +132,7 @@ class Example:
         # build projector for velocity left- and right-hand-sides
         u_bd_test = fem.make_test(space=u_space, domain=boundary)
         u_bd_trial = fem.make_trial(space=u_space, domain=boundary)
-        u_bd_projector = fem.integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True)
+        u_bd_projector = fem.integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test}, assembly="nodal")
 
         # Define an implicit field for our boundary condition value and integrate
         u_bd_field = fem.ImplicitField(
@@ -139,7 +141,7 @@ class Example:
         u_bd_value = fem.integrate(
             mass_form,
             fields={"u": u_bd_field, "v": u_bd_test},
-            nodal=True,
+            assembly="nodal",
             output_dtype=wp.vec2d,
         )
 
@@ -225,7 +227,7 @@ if __name__ == "__main__":
         help="Horizontal velocity boundary condition at the top of the domain.",
     )
     parser.add_argument("--Re", type=float, default=1000.0, help="Reynolds number.")
-    parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh.")
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type.")
     parser.add_argument(
         "--headless",
         action="store_true",
@@ -242,7 +244,7 @@ if __name__ == "__main__":
             resolution=args.resolution,
             Re=args.Re,
             top_velocity=args.top_velocity,
-            tri_mesh=args.tri_mesh,
+            mesh=args.mesh,
         )
 
         for k in range(args.num_frames):

warp/examples/fem/example_nonconforming_contact.py

@@ -205,7 +205,9 @@ class Example:
         # (Note: this is constant per body, this could be precomputed)
         sym_grad_matrix = fem.integrate(symmetric_grad_form, fields={"u": u_trial, "tau": tau_test})
 
-        tau_inv_mass_matrix = fem.integrate(tensor_mass_form, fields={"sig": tau_trial, "tau": tau_test}, nodal=True)
+        tau_inv_mass_matrix = fem.integrate(
+            tensor_mass_form, fields={"sig": tau_trial, "tau": tau_test}, assembly="nodal"
+        )
         fem_example_utils.invert_diagonal_bsr_matrix(tau_inv_mass_matrix)
 
         stress_matrix = tau_inv_mass_matrix @ fem.integrate(
@@ -227,13 +229,13 @@ class Example:
 
         # Enforce boundary conditions
         u_bd_matrix = fem.integrate(
-            bottom_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True
+            bottom_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, assembly="nodal"
         )
 
         # read displacement from other body set create bottom boundary Dirichlet BC
         other_u_field = fem.field.field.NonconformingField(boundary, other_u_field)
         u_bd_rhs = fem.integrate(
-            bottom_boundary_projector_form, fields={"u": other_u_field, "v": u_bd_test}, nodal=True
+            bottom_boundary_projector_form, fields={"u": other_u_field, "v": u_bd_test}, assembly="nodal"
         )
 
         fem.project_linear_system(stiffness_matrix, u_rhs, u_bd_matrix, u_bd_rhs)

warp/examples/fem/example_streamlines.py

@@ -158,6 +158,11 @@ class Example:
                 res=res,
             )
             self._geo = fem.Tetmesh(tet_vtx_indices, pos, build_bvh=True)
+        elif mesh == "hex":
+            pos, hex_vtx_indices = fem_example_utils.gen_hexmesh(
+                res=res,
+            )
+            self._geo = fem.Hexmesh(hex_vtx_indices, pos, assume_parallelepiped_cells=True, build_bvh=True)
         elif mesh == "nano":
             volume = fem_example_utils.gen_volume(
                 res=res,
@@ -280,7 +285,7 @@ class Example:
         inflow_test = fem.make_test(u_space, domain=self._inflow)
         inflow_trial = fem.make_trial(u_space, domain=self._inflow)
         dirichlet_projector = fem.integrate(
-            noslip_projector_form, fields={"u": inflow_test, "v": inflow_trial}, nodal=True, output_dtype=float
+            noslip_projector_form, fields={"u": inflow_test, "v": inflow_trial}, assembly="nodal", output_dtype=float
         )
 
         freeslip_test = fem.make_test(u_space, domain=self._freeslip)
@@ -288,7 +293,7 @@ class Example:
         dirichlet_projector += fem.integrate(
             freeslip_projector_form,
             fields={"u": freeslip_test, "v": freeslip_trial},
-            nodal=True,
+            assembly="nodal",
             output_dtype=float,
         )
         fem.normalize_dirichlet_projector(dirichlet_projector)
@@ -300,7 +305,7 @@ class Example:
         rho_test = fem.make_test(u_space)
         rho_trial = fem.make_trial(u_space)
         inv_mass_matrix = fem.integrate(
-            mass_form, fields={"u": rho_trial, "v": rho_test}, nodal=True, output_dtype=float
+            mass_form, fields={"u": rho_trial, "v": rho_test}, assembly="nodal", output_dtype=float
         )
         fem_example_utils.invert_diagonal_bsr_matrix(inv_mass_matrix)
 

warp/examples/fem/utils.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
@@ -23,15 +24,15 @@ from warp.optim.linear import LinearOperator, aslinearoperator, preconditioner
 from warp.sparse import BsrMatrix, bsr_get_diag, bsr_mv, bsr_transposed
 
 __all__ = [
+    "Plot",
+    "SaddleSystem",
+    "bsr_cg",
+    "bsr_solve_saddle",
     "gen_hexmesh",
     "gen_quadmesh",
     "gen_tetmesh",
     "gen_trimesh",
-    "bsr_cg",
-    "bsr_solve_saddle",
-    "SaddleSystem",
     "invert_diagonal_bsr_matrix",
-    "Plot",
 ]
 
 # matrix inversion routines contain nested loops,
@@ -576,7 +577,7 @@ class Plot:
         else:
             self._usd_renderer.render_points(name, points, radius=self.default_point_radius)
 
-    def plot(self, options: Dict[str, Any] = None, backend: str = "auto"):
+    def plot(self, options: Optional[Dict[str, Any]] = None, backend: str = "auto"):
         if options is None:
             options = {}
 
@@ -611,7 +612,7 @@ class Plot:
                 offsets = np.cumsum(counts)
                 ranges = np.array([offsets - counts, offsets]).T
                 faces = np.concatenate(
-                    [[count] + list(indices[beg:end]) for (count, (beg, end)) in zip(counts, ranges)]
+                    [[count, *list(indices[beg:end])] for (count, (beg, end)) in zip(counts, ranges)]
                 )
                 ref_geom = pyvista.PolyData(vertices, faces)
             else:
@@ -859,8 +860,8 @@ class Plot:
         value_range = field_options.get("clim", None)
         if value_range is None:
             value_range = (
-                min((np.min(_value_or_magnitude(v)) for v in values)),
-                max((np.max(_value_or_magnitude(v)) for v in values)),
+                min(np.min(_value_or_magnitude(v)) for v in values),
+                max(np.max(_value_or_magnitude(v)) for v in values),
             )
 
         return value_range

warp/examples/interop/example_jax_ffi_callback.py

@@ -45,11 +45,11 @@ def example1():
     # the Python function to call
     def print_args(inputs, outputs, attrs, ctx):
         def buffer_to_string(b):
-            return str(b.dtype) + str(list(b.shape)) + " @%x" % b.data
+            return f"{b.dtype}{list(b.shape)} @{b.data:x}"
 
         print("Inputs:     ", ", ".join([buffer_to_string(b) for b in inputs]))
         print("Outputs:    ", ", ".join([buffer_to_string(b) for b in outputs]))
-        print("Attributes: ", "".join(["\n  %s: %s" % (k, str(v)) for k, v in attrs.items()]))
+        print("Attributes: ", "".join([f"\n  {k}: {str(v)}" for k, v in attrs.items()]))  # noqa: RUF010
 
     # register callback
     register_ffi_callback("print_args", print_args)

warp/examples/optim/example_drone.py

@@ -708,7 +708,7 @@ class Example:
         self.tape.zero()
 
     def step(self):
-        if self.frame % int((self.num_frames / len(self.targets))) == 0:
+        if self.frame % int(self.num_frames / len(self.targets)) == 0:
             if self.verbose:
                 print(f"Choosing new flight target: {self.target_idx + 1}")
 

warp/examples/sim/example_cloth.py

@@ -219,7 +219,7 @@ if __name__ == "__main__":
             example.render()
 
         frame_times = example.profiler["step"]
-        print("\nAverage frame sim time: {:.2f} ms".format(sum(frame_times) / len(frame_times)))
+        print(f"\nAverage frame sim time: {sum(frame_times) / len(frame_times):.2f} ms")
 
         if example.renderer:
             example.renderer.save()