warp-lang 1.4.1__py3-none-manylinux2014_aarch64.whl → 1.5.0__py3-none-manylinux2014_aarch64.whl

warp/examples/fem/example_apic_fluid.py

@@ -12,6 +12,8 @@
 # grid and the PicQuadrature class.
 ###########################################################################
 
+from typing import Any
+
 import numpy as np
 
 import warp as wp
@@ -123,7 +125,7 @@ def scalar_vector_multiply(
 @wp.kernel
 def scale_transposed_divergence_mat(
     tr_divergence_mat_offsets: wp.array(dtype=int),
-    tr_divergence_mat_values: wp.array(dtype=wp.mat(shape=(3, 1), dtype=float)),
+    tr_divergence_mat_values: wp.array(dtype=Any),
     inv_fraction_int: wp.array(dtype=float),
 ):
     # In-place scaling of gradient operator rows with inverse mass
@@ -203,7 +205,6 @@ class Example:
         particle_grid_offset = wp.vec3(self.radius, self.radius, self.radius)
 
         # Initialize warp.sim model, spawn particles
-        np.random.seed(0)
         builder = wp.sim.ModelBuilder()
         builder.add_particle_grid(
             dim_x=particle_grid_res[0],

warp/examples/fem/example_convection_diffusion_dg.py

@@ -88,7 +88,7 @@ class Example:
             geo,
             discontinuous=True,
             degree=degree,
-            family=fem.Polynomial.GAUSS_LEGENDRE,
+            family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE,
         )
 
         # Assemble transport, diffusion and inertia matrices
@@ -110,23 +110,22 @@ class Example:
 
         side_test = fem.make_test(space=scalar_space, domain=sides)
         side_trial = fem.make_trial(space=scalar_space, domain=sides)
-
-        matrix_transport += fem.integrate(
+        fem.integrate(
             upwind_transport_form,
             fields={"phi": side_trial, "psi": side_test},
             values={"ang_vel": ang_vel},
+            output=matrix_transport,
+            add=True,
         )
 
         matrix_diffusion = fem.integrate(
             diffusion_form,
             fields={"u": trial, "v": self._test},
         )
-
         matrix_diffusion += fem.integrate(
             sip_diffusion_form,
             fields={"phi": side_trial, "psi": side_test},
         )
-
         self._matrix = matrix_inertia + matrix_transport + viscosity * matrix_diffusion
 
         # Initial condition

warp/examples/fem/example_deformed_geometry.py

@@ -28,7 +28,7 @@ def deformation_field_expr(
     domain: fem.Domain,
 ):
     """
-    Deformation field mapping the unique square to a circular band
+    Deformation field mapping the unit square to a circular band
     """
     x = domain(s)
 

warp/examples/fem/example_diffusion_3d.py

@@ -16,6 +16,8 @@
 # and homogeneous Dirichlet boundary conditions other sides.
 ###########################################################################
 
+import numpy as np
+
 import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
@@ -24,7 +26,7 @@ from warp.sparse import bsr_axpy
 
 
 @fem.integrand
-def vert_boundary_projector_form(
+def vertical_boundary_projector_form(
     s: fem.Sample,
     domain: fem.Domain,
     u: fem.Field,
@@ -36,6 +38,18 @@ def vert_boundary_projector_form(
     return w * u(s) * v(s)
 
 
+@fem.integrand
+def y_boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # Constrain Y edges
+    tangent = fem.deformation_gradient(domain, s)
+    return wp.abs(tangent[1]) * u(s) * v(s)
+
+
 class Example:
     def __init__(
         self,
@@ -77,6 +91,28 @@ class Example:
                 bounds_hi=bounds_hi,
             )
             self._geo = fem.Nanogrid(volume)
+        elif mesh == "tri":
+            pos, quad_vtx_indices = fem_example_utils.gen_trimesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            pos = pos.numpy()
+            pos_z = np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            pos = wp.array(pos, dtype=wp.vec3)
+            self._geo = fem.Trimesh3D(quad_vtx_indices, pos)
+        elif mesh == "quad":
+            pos, quad_vtx_indices = fem_example_utils.gen_quadmesh(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
+            pos = pos.numpy()
+            pos_z = np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            pos = wp.array(pos, dtype=wp.vec3)
+            self._geo = fem.Quadmesh3D(quad_vtx_indices, pos)
         else:
             self._geo = fem.Grid3D(
                 res=res,
@@ -108,7 +144,12 @@ 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(vert_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+
+            # Pick boundary conditions depending on whether our geometry is a 3d surface or a volume
+            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)
 
             # Diffusion form
             trial = fem.make_trial(space=self._scalar_space, domain=domain)
@@ -146,7 +187,9 @@ if __name__ == "__main__":
     parser.add_argument(
         "--boundary_compliance", type=float, default=0.0, help="Dirichlet boundary condition compliance."
     )
-    parser.add_argument("--mesh", choices=("grid", "tet", "hex", "nano", "anano"), default="grid", help="Mesh type.")
+    parser.add_argument(
+        "--mesh", choices=("grid", "tet", "hex", "nano", "anano", "tri", "quad"), default="grid", help="Mesh type."
+    )
     parser.add_argument(
         "--headless",
         action="store_true",
@@ -171,4 +214,4 @@ if __name__ == "__main__":
         example.render()
 
         if not args.headless:
-            example.renderer.plot(backend="matplotlib")
+            example.renderer.plot()

warp/examples/fem/example_distortion_energy.py

@@ -0,0 +1,220 @@
+# Copyright (c) 2022 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 Distortion Energy
+#
+# This example illustrates using a Newton loop to minimize distortion of a
+# 3D surface (u,v) parameterization under a Symmetric Dirichlet energy,
+#
+# E(F) = 1/2 |F|^2 + |F^{-1}|^2
+#
+# with F := dx/du
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+
+@fem.integrand
+def distortion_gradient_form(s: fem.Sample, u_cur: fem.Field, v: fem.Field):
+    # Symmetric Dirichlet energy gradient (linear form)
+    # E = 1/2 (F:F + F^-T:F^-T)
+
+    F = fem.grad(u_cur, s)
+
+    F_inv_sq = wp.inverse(F * wp.transpose(F))
+    F_inv = F_inv_sq * F
+
+    dE_dF = F - F_inv_sq * F_inv
+
+    return wp.ddot(fem.grad(v, s), dE_dF)
+
+
+@fem.integrand
+def distortion_hessian_form(s: fem.Sample, u_cur: fem.Field, u: fem.Field, v: fem.Field):
+    # Symmetric Dirichlet energy approximate hessian (bilinear form)
+
+    # F:F term
+    H = wp.ddot(fem.grad(v, s), fem.grad(u, s))
+
+    # F^-T:F^-T term
+    F = fem.grad(u_cur, s)
+    F_inv_sq = wp.inverse(F * wp.transpose(F))
+
+    # Gauss--Newton (ignore F^-2 derivative)
+    H += wp.ddot(F_inv_sq * fem.grad(v, s), F_inv_sq * F_inv_sq * fem.grad(u, s))
+
+    return H
+
+
+@fem.integrand
+def initial_guess(
+    s: fem.Sample,
+    domain: fem.Domain,
+):
+    # initialization for UV parameter
+    x = domain(s)
+    return wp.vec2(x[0], x[1])
+
+
+@fem.integrand
+def boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    # Fix a single point
+    # (underconstrained, solution up to a rotation in UV space)
+    w = wp.select(s.qp_index == 0, 0.0, 1.0)
+    return w * wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def checkerboard(s: fem.Sample, domain: fem.Domain, u: fem.Field):
+    # checkerboard pattern for parameter visualization
+    u_s = u(s)
+    return wp.sign(wp.cos(16.0 * u_s[0]) * wp.sin(16.0 * u_s[1]))
+
+
+class Example:
+    def __init__(
+        self,
+        quiet=False,
+        degree=2,
+        resolution=25,
+        mesh="grid",
+        nonconforming_stresses=False,
+    ):
+        self._quiet = quiet
+
+        def deform_along_z(positions, z_scale=1.0):
+            pos = positions.numpy()
+            pos_z = z_scale * np.cos(3.0 * pos[:, 0]) * np.sin(4.0 * pos[:, 1])
+            pos = np.hstack((pos, np.expand_dims(pos_z, axis=1)))
+            return wp.array(pos, dtype=wp.vec3)
+
+        # Grid or mesh geometry
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
+            self._uv_geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=wp.zeros_like(positions))
+
+            positions = deform_along_z(positions)
+            self._geo = fem.Trimesh3D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif mesh == "quad":
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            self._uv_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=wp.zeros_like(positions))
+
+            positions = deform_along_z(positions)
+            self._geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            self._uv_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=wp.zeros_like(positions))
+
+            undef_positions = deform_along_z(positions, z_scale=0.0)
+            flat_geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=undef_positions)
+
+            deformation_field = fem.make_discrete_field(fem.make_polynomial_space(flat_geo, dtype=wp.vec3))
+            deformation_field.dof_values = deform_along_z(positions)
+
+            self._geo = deformation_field.make_deformed_geometry(relative=False)
+
+        # parameter space
+        self._u_space = fem.make_polynomial_space(self._geo, degree=degree, dtype=wp.vec2)
+        self._u_field = self._u_space.make_field()
+        self._du_field = self._u_space.make_field()
+        fem.interpolate(initial_guess, dest=self._u_field)
+
+        # scalar parameter visualization function
+        viz_space = fem.make_polynomial_space(self._geo, degree=3, dtype=float)
+        self.viz_field = viz_space.make_field()
+        # For visualization of uv in 2D space
+        uv_space = fem.make_polynomial_space(self._uv_geo, degree=degree, dtype=wp.vec2)
+        self._uv_field = uv_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        boundary = fem.BoundarySides(self._geo)
+        domain = fem.Cells(geometry=self._geo)
+
+        # Parameter boundary conditions
+        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
+        )
+        fem.normalize_dirichlet_projector(u_bd_matrix)
+
+        u_test = fem.make_test(space=self._u_space, domain=domain)
+        u_trial = fem.make_trial(space=self._u_space, domain=domain)
+
+        # Newton iterations (without line-search for simplicity)
+        for _newton_iteration in range(10):
+            u_matrix = fem.integrate(
+                distortion_hessian_form, fields={"u_cur": self._u_field, "u": u_trial, "v": u_test}, output_dtype=float
+            )
+
+            u_rhs = fem.integrate(
+                distortion_gradient_form, fields={"u_cur": self._u_field, "v": u_test}, output_dtype=wp.vec2
+            )
+
+            fem.project_linear_system(u_matrix, u_rhs, u_bd_matrix, normalize_projector=False)
+
+            # Solve for uv increment
+            du = self._du_field.dof_values
+            du.zero_()
+            fem_example_utils.bsr_cg(u_matrix, b=u_rhs, x=du, quiet=self._quiet)
+
+            # Accumulate to UV field
+            fem.utils.array_axpy(x=du, y=self._u_field.dof_values, alpha=-1.0, beta=1.0)
+
+    def render(self):
+        # Visualization
+        fem.interpolate(checkerboard, fields={"u": self._u_field}, dest=self.viz_field)
+
+        self._uv_field.dof_values = wp.clone(self._u_field.dof_values)
+
+        self.renderer.add_field("pattern", self.viz_field)
+        self.renderer.add_field("uv", self._uv_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=25, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=1, help="Polynomial degree of shape functions.")
+    parser.add_argument("--mesh", choices=("tri", "quad", "deformed"), 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("--quiet", action="store_true", help="Suppresses the printing out of iteration residuals.")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=args.quiet,
+            degree=args.degree,
+            resolution=args.resolution,
+            mesh=args.mesh,
+        )
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot(options={"uv": {"displacement": {}}})