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

warp/examples/fem/example_darcy_ls_optimization.py

@@ -0,0 +1,489 @@
+# 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 2D Darcy Flow using Level Set Method
+#
+# This example demonstrates the use of a level set-based shape optimization technique
+# to maximize the total Darcy flow (inflow) through a 2D square domain. The domain
+# contains a material region whose shape is implicitly represented by a level set function.
+#
+# Physical Setup:
+# - The computational domain is a unit square, discretized using either a structured grid
+#   or a triangular mesh.
+# - The material region within the domain is defined by the zero level set of a scalar field.
+# - The permeability of the domain is a smooth function of the level set, with high permeability
+#   inside the material region and low permeability outside.
+# - Boundary conditions are set such that one side of the domain has a prescribed inflow pressure
+#   and the opposite side has a prescribed outflow (Dirichlet) pressure, driving Darcy flow
+#   through the material region.
+#
+# Optimization Goal:
+# - The objective is to optimize the shape of the material region (by evolving the level set)
+#   to maximize the total inflow (Darcy flux) across the inflow boundary, subject to a constraint
+#   on the total volume of the material region.
+#
+# Numerical Approach:
+# - The pressure field is solved using the finite element method (FEM) for the current material
+#   configuration.
+# - The level set function is updated using the adjoint method: the gradient of the objective
+#   with respect to the level set is computed via automatic differentiation, and the level set
+#   is advected in the direction of increasing inflow.
+# - The optimization is performed iteratively, with each iteration consisting of a forward
+#   pressure solve, loss evaluation, backward (adjoint) computation, and level set update.
+# - The code supports both continuous and discontinuous Galerkin formulations for the level set
+#   advection step.
+#
+# Visualization:
+# - The script provides visualization of the velocity field and the evolving material region.
+#
+###########################################################################
+
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+
+@fem.integrand
+def classify_boundary_sides(
+    s: fem.Sample,
+    domain: fem.Domain,
+    dirichlet: wp.array(dtype=int),
+    inflow: wp.array(dtype=int),
+):
+    """Assign boundary sides to inflow or Dirichlet subdomains based on normal direction"""
+    nor = fem.normal(domain, s)
+
+    if nor[0] < -0.5:
+        inflow[s.qp_index] = 1
+        dirichlet[s.qp_index] = 1
+    elif nor[0] > 0.5:
+        dirichlet[s.qp_index] = 1
+
+
+@wp.func
+def initial_level_set(x: wp.vec2, radius: float):
+    """Initial level set function for the material region -- three circles"""
+
+    return (
+        wp.min(
+            wp.vec3(
+                wp.length(x - wp.vec2(0.667, 0.5)),
+                wp.length(x - wp.vec2(0.333, 0.333)),
+                wp.length(x - wp.vec2(0.333, 0.667)),
+            )
+        )
+        - radius
+    )
+
+
+@fem.integrand
+def identity_form(
+    s: fem.Sample,
+    p: fem.Field,
+    q: fem.Field,
+):
+    return p(s) * q(s)
+
+
+@fem.integrand
+def material_fraction(s: fem.Sample, level_set: fem.Field, smoothing: float):
+    """Sigmoid approximation of the level set interior"""
+    return 1.0 / (1.0 + wp.exp(-level_set(s) / smoothing))
+
+
+@fem.integrand
+def permeability(s: fem.Sample, level_set: fem.Field, smoothing: float):
+    """Define permeability as strictly proportional to material fraction (arbitrary choice)"""
+    return material_fraction(s, level_set, smoothing)
+
+
+@fem.integrand
+def velocity_field(s: fem.Sample, level_set: fem.Field, p: fem.Field, smoothing: float):
+    """Velocity field based on permeability and pressure gradient according to Darcy's law"""
+    return -permeability(s, level_set, smoothing) * fem.grad(p, s)
+
+
+@fem.integrand
+def diffusion_form(s: fem.Sample, level_set: fem.Field, p: fem.Field, q: fem.Field, smoothing: float, scale: float):
+    """Inhomogeneous diffusion form"""
+    return scale * wp.dot(velocity_field(s, level_set, p, smoothing), fem.grad(q, s))
+
+
+@fem.integrand
+def inflow_velocity(s: fem.Sample, domain: fem.Domain, level_set: fem.Field, p: fem.Field, smoothing: float):
+    return wp.dot(velocity_field(s, level_set, p, smoothing), fem.normal(domain, s))
+
+
+@fem.integrand
+def volume_form(s: fem.Sample, level_set: fem.Field, smoothing: float):
+    return material_fraction(s, level_set, smoothing)
+
+
+@wp.kernel
+def combine_losses(
+    loss: wp.array(dtype=wp.float32),
+    vol: wp.array(dtype=wp.float32),
+    target_vol: wp.float32,
+    vol_weight: wp.float32,
+):
+    loss[0] += vol_weight * (vol[0] - target_vol) * (vol[0] - target_vol)
+
+
+@fem.integrand
+def advected_level_set_semi_lagrangian(
+    s: fem.Sample, domain: fem.Domain, level_set: fem.Field, velocity: fem.Field, dt: float
+):
+    x_prev = domain(s) - velocity(s) * dt
+    s_prev = fem.lookup(domain, x_prev, guess=s)
+    return level_set(s_prev)
+
+
+# Discontinuous Galerkin variant of level set advection
+
+
+@fem.integrand
+def level_set_transport_form(s: fem.Sample, level_set: fem.Field, psi: fem.Field, velocity: fem.Field, dt: float):
+    return dt * wp.dot(fem.grad(level_set, s), velocity(s)) * psi(s)
+
+
+@fem.integrand
+def level_set_transport_form_upwind(
+    s: fem.Sample, domain: fem.Domain, level_set: fem.Field, psi: fem.Field, velocity: fem.Field, dt: float
+):
+    vel = dt * velocity(s)
+    vel_n = wp.dot(vel, fem.normal(domain, s))
+    return fem.jump(level_set, s) * (-fem.average(psi, s) * vel_n + 0.5 * fem.jump(psi, s) * wp.abs(vel_n))
+
+
+@fem.integrand
+def advected_level_set_upwind(
+    s: fem.Sample, domain: fem.Domain, level_set: fem.Field, transport_integrals: wp.array(dtype=float)
+):
+    return level_set(s) - transport_integrals[s.qp_index] / (fem.measure(domain, s) * s.qp_weight)
+
+
+class Example:
+    def __init__(
+        self, quiet=False, degree=2, resolution=25, mesh: str = "grid", dt: float = 1.0, discontinuous: bool = False
+    ):
+        self._quiet = quiet
+
+        self._smoothing = 0.5 / resolution  # smoothing for level set interface approximation as sigmoid
+        self._dt = dt  # level set advection time step (~gradient step size)
+        self._discontinuous = discontinuous
+
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution, resolution))
+            self._geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions, build_bvh=True)
+        else:
+            self._geo = fem.Grid2D(res=wp.vec2i(resolution, resolution))
+
+        # Pressure, level set, and level set velocity spaces
+        self._p_space = fem.make_polynomial_space(self._geo, degree=degree, dtype=float)
+        self._ls_space = fem.make_polynomial_space(self._geo, degree=degree, dtype=float, discontinuous=discontinuous)
+        self._v_space = fem.make_polynomial_space(self._geo, degree=degree, dtype=wp.vec2)
+
+        # pressure field
+        self._p_field = fem.make_discrete_field(space=self._p_space)
+        self._p_field.dof_values.requires_grad = True
+
+        # level set field
+        self._level_set_field = fem.make_discrete_field(space=self._ls_space)
+        self._level_set_field.dof_values.requires_grad = True
+
+        # level set advection velocity field
+        self._level_set_velocity_field = fem.make_discrete_field(space=self._v_space)
+        self._level_set_velocity_field.dof_values.requires_grad = True
+
+        fem.interpolate(
+            fem.ImplicitField(fem.Cells(self._geo), initial_level_set, values={"radius": 0.125}),
+            dest=self._level_set_field,
+        )
+
+        # recording initial volume, we want to preserve that when doing optimization
+        self._target_vol = fem.integrate(
+            volume_form,
+            domain=fem.Cells(self._geo),
+            fields={"level_set": self._level_set_field},
+            values={"smoothing": self._smoothing},
+        )
+
+        # Trial and test functions for pressure solve
+        self._p_test = fem.make_test(space=self._p_space)
+        self._p_trial = fem.make_trial(space=self._p_space)
+
+        # For discontinuous level set advection
+        if self._discontinuous:
+            self._ls_test = fem.make_test(space=self._ls_space)
+            self._ls_sides_test = fem.make_test(space=self._ls_space, domain=fem.Sides(self._geo))
+
+        # Identify inflow and outflow sides
+        boundary = fem.BoundarySides(self._geo)
+
+        inflow_mask = wp.zeros(shape=boundary.element_count(), dtype=int)
+        dirichlet_mask = wp.zeros(shape=boundary.element_count(), dtype=int)
+
+        fem.interpolate(
+            classify_boundary_sides,
+            quadrature=fem.RegularQuadrature(boundary, order=0),
+            values={"inflow": inflow_mask, "dirichlet": dirichlet_mask},
+        )
+
+        self._inflow = fem.Subdomain(boundary, element_mask=inflow_mask)
+        self._dirichlet = fem.Subdomain(boundary, element_mask=dirichlet_mask)
+
+        # Build projector for the inflow and outflow homogeneous Dirichlet condition
+        p_dirichlet_bd_test = fem.make_test(space=self._p_space, domain=self._dirichlet)
+        p_dirichlet_bd_trial = fem.make_trial(space=self._p_space, domain=self._dirichlet)
+        p_dirichlet_bd_matrix = fem.integrate(
+            identity_form,
+            fields={"p": p_dirichlet_bd_trial, "q": p_dirichlet_bd_test},
+            assembly="nodal",
+            output_dtype=float,
+        )
+
+        # Inflow prescribed pressure
+        p_inflow_bd_test = fem.make_test(space=self._p_space, domain=self._inflow)
+        p_inflow_bd_value = fem.integrate(
+            identity_form,
+            fields={"p": fem.UniformField(self._inflow, 1.0), "q": p_inflow_bd_test},
+            assembly="nodal",
+            output_dtype=float,
+        )
+        fem.normalize_dirichlet_projector(p_dirichlet_bd_matrix, p_inflow_bd_value)
+
+        self._bd_projector = p_dirichlet_bd_matrix
+        self._bd_prescribed_value = p_inflow_bd_value
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        p = self._p_field.dof_values
+        p.zero_()
+        v = self._level_set_velocity_field.dof_values
+        v.zero_()
+
+        # Advected level set field, used in adjoint computations
+        advected_level_set = fem.make_discrete_field(space=self._ls_space)
+        advected_level_set.dof_values.assign(self._level_set_field.dof_values)
+        advected_level_set.dof_values.requires_grad = True
+        advected_level_set_restriction = fem.make_restriction(advected_level_set, domain=self._p_test.domain)
+
+        # Forward step, record adjoint tape for forces
+        p_rhs = wp.empty(self._p_space.node_count(), dtype=wp.float32, requires_grad=True)
+
+        tape = wp.Tape()
+        with tape:
+            # Dummy advection step, so backward pass can compute adjoint w.r.t advection velocity
+            self.advect_level_set(
+                level_set_in=self._level_set_field,
+                level_set_out=advected_level_set_restriction,
+                velocity=self._level_set_velocity_field,
+                dt=1.0,
+            )
+
+            # Left-hand-side of implicit solve (zero if p=0, but required for adjoint computation through implicit function theorem)
+            fem.integrate(
+                diffusion_form,
+                fields={
+                    "level_set": advected_level_set,
+                    "p": self._p_field,
+                    "q": self._p_test,
+                },
+                values={"smoothing": self._smoothing, "scale": -1.0},
+                output=p_rhs,
+            )
+
+        # Diffusion matrix (inhomogeneous Poisson)
+        p_matrix = fem.integrate(
+            diffusion_form,
+            fields={
+                "level_set": advected_level_set,
+                "p": self._p_trial,
+                "q": self._p_test,
+            },
+            values={"smoothing": self._smoothing, "scale": 1.0},
+            output_dtype=float,
+        )
+
+        # Project to enforce Dirichlet boundary conditions then solve linear system
+        fem.project_linear_system(
+            p_matrix, p_rhs, self._bd_projector, self._bd_prescribed_value, normalize_projector=False
+        )
+
+        fem_example_utils.bsr_cg(p_matrix, b=p_rhs, x=p, quiet=self._quiet, tol=1e-6, max_iters=1000)
+
+        # Record adjoint of linear solve
+        def solve_linear_system():
+            fem_example_utils.bsr_cg(p_matrix, b=p.grad, x=p_rhs.grad, quiet=self._quiet, tol=1e-6, max_iters=1000)
+            p_rhs.grad -= self._bd_projector @ p_rhs.grad
+
+        tape.record_func(solve_linear_system, arrays=(p_rhs, p))
+
+        # Evaluate losses
+        loss = wp.empty(shape=1, dtype=wp.float32, requires_grad=True)
+        vol = wp.empty(shape=1, dtype=wp.float32, requires_grad=True)
+
+        with tape:
+            # Main objective: inflow flux
+            fem.integrate(
+                inflow_velocity,
+                fields={"level_set": advected_level_set.trace(), "p": self._p_field.trace()},
+                values={"smoothing": self._smoothing},
+                domain=self._inflow,
+                output=loss,
+            )
+
+            # Add penalization term enforcing constant volume
+            fem.integrate(
+                volume_form,
+                fields={"level_set": advected_level_set},
+                values={"smoothing": self._smoothing},
+                domain=self._p_test.domain,
+                output=vol,
+            )
+
+            print("Total inflow", loss, "Volume", vol)
+
+            vol_loss_weight = 1000.0
+            wp.launch(
+                combine_losses,
+                dim=1,
+                inputs=(loss, vol, self._target_vol, vol_loss_weight),
+            )
+
+        # perform backward step
+        tape.backward(loss=loss)
+
+        # Advect level set with velocity field adjoint
+        v.assign(v.grad)
+        self.advect_level_set(
+            level_set_in=advected_level_set,
+            level_set_out=self._level_set_field,
+            velocity=self._level_set_velocity_field,
+            dt=-self._dt,
+        )
+
+        # Zero-out gradients used in tape
+        tape.zero()
+
+    def advect_level_set(self, level_set_in: fem.Field, level_set_out: fem.Field, velocity: fem.Field, dt: float):
+        if self._discontinuous:
+            # Discontinuous Galerkin version with (explicit) upwind transport:
+            # Integrate in-cell transport + side flux
+            transport_integrals = wp.empty(
+                shape=self._ls_space.node_count(),
+                dtype=float,
+                requires_grad=True,
+            )
+            fem.integrate(
+                level_set_transport_form,
+                fields={"psi": self._ls_test, "level_set": level_set_in, "velocity": velocity},
+                values={"dt": dt},
+                output=transport_integrals,
+            )
+
+            fem.integrate(
+                level_set_transport_form_upwind,
+                fields={"level_set": level_set_in.trace(), "psi": self._ls_sides_test, "velocity": velocity.trace()},
+                values={"dt": dt},
+                output=transport_integrals,
+                add=True,
+            )
+
+            # Divide by mass matrix and write back to advected field
+            out_field = level_set_out if isinstance(level_set_out, fem.DiscreteField) else level_set_out.field
+            fem.interpolate(
+                advected_level_set_upwind,
+                fields={"level_set": level_set_in, "velocity": velocity},
+                values={"transport_integrals": transport_integrals},
+                quadrature=fem.NodalQuadrature(self._p_test.domain, self._ls_space),
+                dest=out_field.dof_values,
+            )
+        else:
+            # Continuous Galerkin version with semi-Lagrangian transport:
+            fem.interpolate(
+                advected_level_set_semi_lagrangian,
+                fields={"level_set": level_set_in, "velocity": velocity},
+                values={"dt": dt},
+                dest=level_set_out,
+            )
+
+    def render(self):
+        # velocity field
+        u_space = fem.make_polynomial_space(self._geo, degree=self._p_space.degree, dtype=wp.vec2)
+        u_field = fem.make_discrete_field(space=u_space)
+        fem.interpolate(
+            velocity_field,
+            fields={"level_set": self._level_set_field, "p": self._p_field},
+            values={"smoothing": self._smoothing},
+            dest=u_field,
+        )
+
+        # material fraction field
+        mat_field = fem.make_discrete_field(space=self._ls_space)
+        fem.interpolate(
+            material_fraction,
+            fields={"level_set": self._level_set_field},
+            values={"smoothing": self._smoothing},
+            dest=mat_field,
+        )
+
+        self.renderer.add_field("velocity", u_field)
+        self.renderer.add_field("material", mat_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=100, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=1, help="Polynomial degree of shape functions.")
+    parser.add_argument("--discontinuous", action="store_true", help="Use discontinuous level set advection.")
+    parser.add_argument("--mesh", type=str, default="grid", help="Mesh type.")
+    parser.add_argument("--num_iters", type=int, default=100, help="Number of iterations.")
+    parser.add_argument("--dt", type=float, default=0.05, help="Level set update time step.")
+    parser.add_argument(
+        "--headless",
+        action="store_true",
+        help="Run in headless mode, suppressing the opening of any graphical windows.",
+    )
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=True,
+            degree=args.degree,
+            resolution=args.resolution,
+            discontinuous=args.discontinuous,
+            mesh=args.mesh,
+            dt=args.dt,
+        )
+
+        for _iter in range(args.num_iters):
+            example.step()
+            example.render()
+
+        if not args.headless:
+            example.renderer.plot(
+                options={
+                    "velocity": {"arrows": {"glyph_scale": 0.1}},
+                    "material": {"contours": {"levels": [0.0, 0.5, 1.0001]}},
+                },
+            )

warp/examples/fem/example_deformed_geometry.py

@@ -119,7 +119,7 @@ class Example:
         bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
         bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
 
-        bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+        bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, assembly="nodal")
 
         fem.project_linear_system(matrix, rhs, bd_matrix)
 

warp/examples/fem/example_diffusion.py

@@ -125,9 +125,9 @@ class Example:
         bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
         bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
 
-        bd_matrix = fem.integrate(y_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+        bd_matrix = fem.integrate(y_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, assembly="nodal")
         bd_rhs = fem.integrate(
-            y_boundary_value_form, fields={"v": bd_test}, values={"val": self._boundary_value}, nodal=True
+            y_boundary_value_form, fields={"v": bd_test}, values={"val": self._boundary_value}, assembly="nodal"
         )
 
         # Assemble linear system

warp/examples/fem/example_diffusion_3d.py

@@ -157,7 +157,7 @@ class Example:
             boundary_projector_form = (
                 vertical_boundary_projector_form if self._geo.cell_dimension == 3 else y_boundary_projector_form
             )
-            bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+            bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, assembly="nodal")
 
             # Diffusion form
             trial = fem.make_trial(space=self._scalar_space, domain=domain)

warp/examples/fem/example_distortion_energy.py

@@ -158,7 +158,7 @@ class Example:
         u_bd_test = fem.make_test(space=self._u_space, domain=boundary)
         u_bd_trial = fem.make_trial(space=self._u_space, domain=boundary)
         u_bd_matrix = fem.integrate(
-            boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True, output_dtype=float
+            boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, assembly="nodal", output_dtype=float
         )
         fem.normalize_dirichlet_projector(u_bd_matrix)