warp-lang 1.2.2__py3-none-macosx_10_13_universal2.whl → 1.3.1__py3-none-macosx_10_13_universal2.whl

warp/examples/fem/example_nonconforming_contact.py

@@ -0,0 +1,290 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+###########################################################################
+# Example Nonconforming Contact
+#
+# This example demonstrates using nonconforming fields (warp.fem.NonconformingField)
+# to solve a no-slip contact problem between two elastic bodies discretized separately.
+#
+# Div[ E: D(u) ] = g  over each body
+# u_top = u_bottom    along the contact surface (top body bottom boundary)
+# u_bottom = 0        along the bottom boundary of the bottom body
+#
+# with E the rank-4 elasticity tensor
+#
+# Below we use a simple staggered scheme for solving bodies iteratively,
+# but more robust methods could be considered (e.g. Augmented Lagrangian)
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+
+@wp.func
+def compute_stress(tau: wp.mat22, E: wp.mat33):
+    """Strain to stress computation (using Voigt notation to drop tensor order)"""
+    tau_sym = wp.vec3(tau[0, 0], tau[1, 1], tau[0, 1] + tau[1, 0])
+    sig_sym = E * tau_sym
+    return wp.mat22(sig_sym[0], 0.5 * sig_sym[2], 0.5 * sig_sym[2], sig_sym[1])
+
+
+@fem.integrand
+def stress_form(s: fem.Sample, u: fem.Field, tau: fem.Field, E: wp.mat33):
+    """Stress inside body:  (E : D(u)) : tau"""
+    return wp.ddot(tau(s), compute_stress(fem.D(u, s), E))
+
+
+@fem.integrand
+def boundary_stress_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    tau: fem.Field,
+):
+    """Stress on boundary: u' tau n"""
+    return wp.dot(tau(s) * fem.normal(domain, s), u(s))
+
+
+@fem.integrand
+def symmetric_grad_form(
+    s: fem.Sample,
+    u: fem.Field,
+    tau: fem.Field,
+):
+    """Symmetric part of gradient of displacement: D(u) : tau"""
+    return wp.ddot(tau(s), fem.D(u, s))
+
+
+@fem.integrand
+def gravity_form(
+    s: fem.Sample,
+    v: fem.Field,
+    gravity: float,
+):
+    return -gravity * v(s)[1]
+
+
+@fem.integrand
+def bottom_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # non zero on bottom boundary only
+    nor = fem.normal(domain, s)
+    return wp.dot(u(s), v(s)) * wp.max(0.0, -nor[1])
+
+
+@fem.integrand
+def tensor_mass_form(
+    s: fem.Sample,
+    sig: fem.Field,
+    tau: fem.Field,
+):
+    return wp.ddot(tau(s), sig(s))
+
+
+class Example:
+    def __init__(
+        self,
+        degree=2,
+        resolution=16,
+        young_modulus=1.0,
+        poisson_ratio=0.5,
+        nonconforming_stresses=False,
+    ):
+        self._geo1 = fem.Grid2D(bounds_hi=wp.vec2(1.0, 0.5), res=wp.vec2i(resolution))
+        self._geo2 = fem.Grid2D(bounds_lo=(0.33, 0.5), bounds_hi=(0.67, 0.5 + 0.33), res=wp.vec2i(resolution))
+
+        # Strain-stress matrix
+        young = young_modulus
+        poisson = poisson_ratio
+        self._elasticity_mat = wp.mat33(
+            young
+            / (1.0 - poisson * poisson)
+            * np.array(
+                [
+                    [1.0, poisson, 0.0],
+                    [poisson, 1.0, 0.0],
+                    [0.0, 0.0, (2.0 * (1.0 + poisson)) * (1.0 - poisson * poisson)],
+                ]
+            )
+        )
+
+        # Displacement spaces and fields -- S_k
+        self._u1_space = fem.make_polynomial_space(
+            self._geo1, degree=degree, dtype=wp.vec2, element_basis=fem.ElementBasis.SERENDIPITY
+        )
+        self._u2_space = fem.make_polynomial_space(
+            self._geo2, degree=degree, dtype=wp.vec2, element_basis=fem.ElementBasis.SERENDIPITY
+        )
+        self._u1_field = self._u1_space.make_field()
+        self._u2_field = self._u2_space.make_field()
+
+        # Stress spaces and fields -- Q_{k-1}d
+        # Store stress degrees of freedom as symmetric tensors (3 dof) rather than full 2x2 matrices
+        self._tau1_space = fem.make_polynomial_space(
+            self._geo1,
+            degree=degree - 1,
+            discontinuous=True,
+            element_basis=fem.ElementBasis.LAGRANGE,
+            dof_mapper=fem.SymmetricTensorMapper(wp.mat22),
+        )
+        self._tau2_space = fem.make_polynomial_space(
+            self._geo2,
+            degree=degree - 1,
+            discontinuous=True,
+            element_basis=fem.ElementBasis.LAGRANGE,
+            dof_mapper=fem.SymmetricTensorMapper(wp.mat22),
+        )
+
+        self._sig1_field = self._tau1_space.make_field()
+        self._sig2_field = self._tau2_space.make_field()
+        self._sig2_field_new = self._tau2_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        # Solve for the two bodies separately
+        # Body (top) is 25x more dense and 5x stiffer than top body
+        # Body 1 affects body 2 through bottom displacement dirichlet BC
+        # Body 2 affects body 1 through applied strain on top
+        self.solve_solid(self._u1_field, self._sig1_field, self._u2_field, self._sig2_field, gravity=1.0, stiffness=1.0)
+        self.solve_solid(
+            self._u2_field, self._sig2_field_new, self._u1_field, self._sig1_field, gravity=25.0, stiffness=5.0
+        )
+
+        # Damped update of coupling stress (for stability)
+        alpha = 0.1
+        fem.utils.array_axpy(
+            x=self._sig2_field_new.dof_values, y=self._sig2_field.dof_values, alpha=alpha, beta=1.0 - alpha
+        )
+
+    def solve_solid(
+        self,
+        u_field,
+        stress_field,
+        other_u_field,
+        other_stress_field,
+        gravity: float,
+        stiffness: float,
+    ):
+        u_space = u_field.space
+        stress_space = stress_field.space
+        geo = u_field.space.geometry
+
+        domain = fem.Cells(geometry=geo)
+        boundary = fem.BoundarySides(geometry=geo)
+
+        u_test = fem.make_test(space=u_space, domain=domain)
+        u_trial = fem.make_trial(space=u_space, domain=domain)
+        tau_test = fem.make_test(space=stress_space, domain=domain)
+        tau_trial = fem.make_trial(space=stress_space, domain=domain)
+
+        u_bd_test = fem.make_test(space=u_space, domain=boundary)
+        u_bd_trial = fem.make_trial(space=u_space, domain=boundary)
+
+        # Assemble stiffness matrix
+        # (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)
+        fem_example_utils.invert_diagonal_bsr_matrix(tau_inv_mass_matrix)
+
+        stress_matrix = tau_inv_mass_matrix @ fem.integrate(
+            stress_form, fields={"u": u_trial, "tau": tau_test}, values={"E": self._elasticity_mat * stiffness}
+        )
+        stiffness_matrix = sym_grad_matrix.transpose() @ stress_matrix
+
+        # Right-hand-side
+        u_rhs = fem.integrate(gravity_form, fields={"v": u_test}, values={"gravity": gravity}, output_dtype=wp.vec2d)
+
+        # Add boundary stress from other solid field
+        other_stress_field = fem.field.field.NonconformingField(boundary, other_stress_field)
+        fem.utils.array_axpy(
+            y=u_rhs,
+            x=fem.integrate(
+                boundary_stress_form, fields={"u": u_bd_test, "tau": other_stress_field}, output_dtype=wp.vec2d
+            ),
+        )
+
+        # Enforce boundary conditions
+        u_bd_matrix = fem.integrate(
+            bottom_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True
+        )
+
+        # 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
+        )
+
+        fem.project_linear_system(stiffness_matrix, u_rhs, u_bd_matrix, u_bd_rhs)
+
+        # solve
+        x = wp.zeros_like(u_rhs)
+        wp.utils.array_cast(in_array=u_field.dof_values, out_array=x)
+        fem_example_utils.bsr_cg(stiffness_matrix, b=u_rhs, x=x, tol=1.0e-6, quiet=True)
+
+        # Extract result
+        stress = stress_matrix @ x
+        wp.utils.array_cast(in_array=x, out_array=u_field.dof_values)
+        wp.utils.array_cast(in_array=stress, out_array=stress_field.dof_values)
+
+    def render(self):
+        self.renderer.add_field("u1", self._u1_field)
+        self.renderer.add_field("u2", self._u2_field)
+        self.renderer.add_field("sig1", self._sig1_field)
+        self.renderer.add_field("sig2", self._sig2_field)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    wp.set_module_options({"enable_backward": False})
+
+    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=32, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--young_modulus", type=float, default=10.0)
+    parser.add_argument("--poisson_ratio", type=float, default=0.9)
+    parser.add_argument("--num_steps", type=int, default=50)
+    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(
+            degree=args.degree,
+            resolution=args.resolution,
+            young_modulus=args.young_modulus,
+            poisson_ratio=args.poisson_ratio,
+        )
+
+        for i in range(args.num_steps):
+            print("Step", i)
+            example.step()
+            example.render()
+
+        if not args.headless:
+            example.renderer.plot(
+                {
+                    "rows": 2,
+                    "u1": {"displacement": {}},
+                    "u2": {"displacement": {}},
+                },
+            )

warp/examples/fem/example_stokes.py

@@ -17,24 +17,10 @@
 ###########################################################################
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-import warp.sparse as sparse
 from warp.fem.utils import array_axpy
 
-try:
-    from .bsr_utils import SaddleSystem, bsr_solve_saddle
-    from .mesh_utils import gen_quadmesh, gen_trimesh
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import SaddleSystem, bsr_solve_saddle
-    from mesh_utils import gen_quadmesh, gen_trimesh
-    from plot_utils import Plot
-
-
-@fem.integrand
-def constant_form(val: wp.vec2):
-    return val
-
 
 @fem.integrand
 def viscosity_form(s: fem.Sample, u: fem.Field, v: fem.Field, nu: float):
@@ -90,10 +76,10 @@ class Example:
 
         # Grid or triangle mesh geometry
         if mesh == "tri":
-            positions, tri_vidx = gen_trimesh(res=wp.vec2i(resolution))
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
             geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
         elif mesh == "quad":
-            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(resolution))
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
             geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
         else:
             geo = fem.Grid2D(res=wp.vec2i(resolution))
@@ -111,13 +97,9 @@ class Example:
         self._u_field = u_space.make_field()
         self._p_field = p_space.make_field()
 
-        # Interpolate initial condition on boundary (for example purposes)
-        self._bd_field = u_space.make_field()
-        f_boundary = fem.make_restriction(self._bd_field, domain=fem.BoundarySides(geo))
-        top_velocity = wp.vec2(top_velocity, 0.0)
-        fem.interpolate(constant_form, dest=f_boundary, values={"val": top_velocity})
+        self._bd_field = fem.UniformField(domain=fem.BoundarySides(geo), value=wp.vec2(top_velocity, 0.0))
 
-        self.renderer = Plot()
+        self.renderer = fem_example_utils.Plot()
 
     def step(self):
         u_space = self._u_field.space
@@ -140,9 +122,7 @@ class Example:
         # Weak velocity boundary conditions
         u_bd_test = fem.make_test(space=u_space, domain=boundary)
         u_bd_trial = fem.make_trial(space=u_space, domain=boundary)
-        u_rhs = fem.integrate(
-            top_mass_form, fields={"u": self._bd_field.trace(), "v": u_bd_test}, output_dtype=wp.vec2d
-        )
+        u_rhs = fem.integrate(top_mass_form, fields={"u": self._bd_field, "v": u_bd_test}, output_dtype=wp.vec2d)
         u_bd_matrix = fem.integrate(mass_form, fields={"u": u_bd_trial, "v": u_bd_test})
 
         # Pressure-velocity coupling
@@ -151,23 +131,28 @@ class Example:
 
         # Define and solve the saddle-point system
         u_matrix = u_visc_matrix
-        sparse.bsr_axpy(x=u_bd_matrix, y=u_matrix, alpha=self.boundary_strength, beta=1.0)
+        u_matrix += self.boundary_strength * u_bd_matrix
         array_axpy(x=u_rhs, y=u_rhs, alpha=0.0, beta=self.boundary_strength)
 
         p_rhs = wp.zeros(p_space.node_count(), dtype=wp.float64)
         x_u = wp.zeros_like(u_rhs)
         x_p = wp.zeros_like(p_rhs)
 
-        bsr_solve_saddle(
-            SaddleSystem(A=u_matrix, B=div_matrix), x_u=x_u, x_p=x_p, b_u=u_rhs, b_p=p_rhs, quiet=self._quiet
+        fem_example_utils.bsr_solve_saddle(
+            fem_example_utils.SaddleSystem(A=u_matrix, B=div_matrix),
+            x_u=x_u,
+            x_p=x_p,
+            b_u=u_rhs,
+            b_p=p_rhs,
+            quiet=self._quiet,
         )
 
         wp.utils.array_cast(in_array=x_u, out_array=self._u_field.dof_values)
         wp.utils.array_cast(in_array=x_p, out_array=self._p_field.dof_values)
 
     def render(self):
-        self.renderer.add_surface("pressure", self._p_field)
-        self.renderer.add_surface_vector("velocity", self._u_field)
+        self.renderer.add_field("pressure", self._p_field)
+        self.renderer.add_field("velocity", self._u_field)
 
 
 if __name__ == "__main__":
@@ -215,4 +200,4 @@ if __name__ == "__main__":
         example.render()
 
         if not args.headless:
-            example.renderer.plot()
+            example.renderer.plot(options={"velocity": {"streamlines": {}}, "pressure": {"contours": {}}})

warp/examples/fem/example_stokes_transfer.py

@@ -19,20 +19,11 @@ import math
 import numpy as np
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
 from warp.fem.utils import array_axpy
-from warp.sparse import bsr_axpy, bsr_mm, bsr_mv, bsr_transposed
 from warp.utils import array_cast
 
-# Import example utilities
-# Make sure that works both when imported as module and run as standalone file
-try:
-    from .bsr_utils import bsr_cg
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import bsr_cg
-    from plot_utils import Plot
-
 
 @fem.integrand
 def vel_from_particles_form(s: fem.Sample, particle_vel: wp.array(dtype=wp.vec2), v: fem.Field):
@@ -141,7 +132,7 @@ class Example:
         self._pic_quadrature = fem.PicQuadrature(domain, particles, particle_areas)
         self._particle_velocities = particle_velocities
 
-        self.renderer = Plot()
+        self.renderer = fem_example_utils.Plot()
 
     def step(self):
         u_space = self._u_field.space
@@ -182,21 +173,21 @@ class Example:
 
         # Assemble linear system
         u_matrix = u_visc_matrix
-        bsr_axpy(u_bd_matrix, u_matrix, alpha=self.bd_strength)
+        u_matrix += u_bd_matrix * self.bd_strength
 
-        div_matrix_t = bsr_transposed(div_matrix)
-        gradient_matrix = bsr_mm(div_matrix_t, inv_p_mass_matrix)
-        bsr_mm(gradient_matrix, div_matrix, u_matrix, alpha=1.0 / self.compliance, beta=1.0)
+        gradient_matrix = div_matrix.transpose() @ inv_p_mass_matrix
+        u_matrix += gradient_matrix @ div_matrix / self.compliance
 
+        # scale u_rhs
         array_axpy(u_rhs, u_rhs, alpha=0.0, beta=self.bd_strength)
 
         # Solve for displacement
         u_res = wp.zeros_like(u_rhs)
-        bsr_cg(u_matrix, x=u_res, b=u_rhs, quiet=self._quiet)
+        fem_example_utils.bsr_cg(u_matrix, x=u_res, b=u_rhs, quiet=self._quiet)
 
         # Compute pressure from displacement
-        div_u = bsr_mv(A=div_matrix, x=u_res)
-        p_res = bsr_mv(A=inv_p_mass_matrix, x=div_u, alpha=-1)
+        div_u = div_matrix @ u_res
+        p_res = -inv_p_mass_matrix @ div_u
 
         # Copy to fields
         u_nodes = wp.indexedarray(self._u_field.dof_values, indices=self._active_space_partition.space_node_indices())
@@ -206,8 +197,8 @@ class Example:
         array_cast(in_array=p_res, out_array=p_nodes)
 
     def render(self):
-        self.renderer.add_surface("pressure", self._p_field)
-        self.renderer.add_surface_vector("velocity", self._u_field)
+        self.renderer.add_field("pressure", self._p_field)
+        self.renderer.add_field("velocity", self._u_field)
 
     def _gen_particles(self, circle_radius, c1_center, c2_center):
         """Generate some particles along two circles defining velocity boundary conditions"""
@@ -260,4 +251,4 @@ if __name__ == "__main__":
         example.render()
 
         if not args.headless:
-            example.renderer.plot(streamlines=["velocity"])
+            example.renderer.plot(options={"velocity": {"streamlines": {}}})