warp-lang 1.2.1__py3-none-manylinux2014_aarch64.whl → 1.3.0__py3-none-manylinux2014_aarch64.whl

warp/examples/fem/example_deformed_geometry.py

@@ -8,7 +8,7 @@
 ###########################################################################
 # Example Deformed Geometry
 #
-# This example solves a 2d diffusion problem:
+# This example solves a 2d diffusion problem on a deformed (curved) mesh:
 #
 # nu Div u = 1
 #
@@ -17,20 +17,9 @@
 ###########################################################################
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-
-# 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 .example_diffusion import diffusion_form, linear_form
-    from .mesh_utils import gen_quadmesh, gen_trimesh
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import bsr_cg
-    from example_diffusion import diffusion_form, linear_form
-    from mesh_utils import gen_quadmesh, gen_trimesh
-    from plot_utils import Plot
+from warp.examples.fem.example_diffusion import diffusion_form, linear_form
 
 
 @fem.integrand
@@ -46,7 +35,7 @@ def deformation_field_expr(
     r = x[1] + 0.5
     t = 0.5 * 3.1416 * x[0]
 
-    return r * wp.vec2(wp.sin(t), wp.cos(t)) - x
+    return r * wp.vec2(wp.sin(t), wp.cos(t))
 
 
 @fem.integrand
@@ -80,10 +69,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))
             base_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))
             base_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
         else:
             base_geo = fem.Grid2D(res=wp.vec2i(resolution))
@@ -93,7 +82,7 @@ class Example:
         deformation_field = deformation_space.make_field()
         fem.interpolate(deformation_field_expr, dest=deformation_field)
 
-        self._geo = deformation_field.make_deformed_geometry()
+        self._geo = deformation_field.make_deformed_geometry(relative=False)
 
         # Scalar function space on deformed geometry
         element_basis = fem.ElementBasis.SERENDIPITY if serendipity else None
@@ -102,7 +91,7 @@ class Example:
         # Scalar field over our function space
         self._scalar_field = self._scalar_space.make_field()
 
-        self.renderer = Plot()
+        self.renderer = fem_example_utils.Plot()
 
     def step(self):
         geo = self._geo
@@ -128,13 +117,13 @@ class Example:
 
         # Solve linear system using Conjugate Gradient
         x = wp.zeros_like(rhs)
-        bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet, tol=1.0e-6)
+        fem_example_utils.bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet, tol=1.0e-6)
 
         # Assign system result to our discrete field
         self._scalar_field.dof_values = x
 
     def render(self):
-        self.renderer.add_surface("solution", self._scalar_field)
+        self.renderer.add_field("solution", self._scalar_field)
 
 
 if __name__ == "__main__":
@@ -148,7 +137,7 @@ if __name__ == "__main__":
     parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
     parser.add_argument("--serendipity", action="store_true", default=False, help="Use Serendipity basis functions.")
     parser.add_argument("--viscosity", type=float, default=2.0, help="Fluid viscosity parameter.")
-    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="tri", help="Mesh type")
     parser.add_argument(
         "--headless",
         action="store_true",

warp/examples/fem/example_diffusion.py

@@ -17,20 +17,9 @@
 ###########################################################################
 
 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
-
-# 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 .mesh_utils import gen_quadmesh, gen_trimesh
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import bsr_cg
-    from mesh_utils import gen_quadmesh, gen_trimesh
-    from plot_utils import Plot
 
 
 @fem.integrand
@@ -92,10 +81,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))
             self._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))
             self._geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
         else:
             self._geo = fem.Grid2D(res=wp.vec2i(resolution))
@@ -107,7 +96,7 @@ class Example:
         # Scalar field over our function space
         self._scalar_field = self._scalar_space.make_field()
 
-        self.renderer = Plot()
+        self.renderer = fem_example_utils.Plot()
 
     def step(self):
         geo = self._geo
@@ -140,18 +129,18 @@ class Example:
         else:
             # Weak BC: add together diffusion and boundary condition matrices
             boundary_strength = 1.0 / self._boundary_compliance
-            bsr_axpy(x=bd_matrix, y=matrix, alpha=boundary_strength, beta=1)
+            matrix += bd_matrix * boundary_strength
             array_axpy(x=bd_rhs, y=rhs, alpha=boundary_strength, beta=1)
 
         # Solve linear system using Conjugate Gradient
         x = wp.zeros_like(rhs)
-        bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
+        fem_example_utils.bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
 
         # Assign system result to our discrete field
         self._scalar_field.dof_values = x
 
     def render(self):
-        self.renderer.add_surface("solution", self._scalar_field)
+        self.renderer.add_field("solution", self._scalar_field)
 
 
 if __name__ == "__main__":

warp/examples/fem/example_diffusion_3d.py

@@ -17,22 +17,11 @@
 ###########################################################################
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
+from warp.examples.fem.example_diffusion import diffusion_form, linear_form
 from warp.sparse import bsr_axpy
 
-# 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 .example_diffusion import diffusion_form, linear_form
-    from .mesh_utils import gen_hexmesh, gen_tetmesh
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import bsr_cg
-    from example_diffusion import diffusion_form, linear_form
-    from mesh_utils import gen_hexmesh, gen_tetmesh
-    from plot_utils import Plot
-
 
 @fem.integrand
 def vert_boundary_projector_form(
@@ -41,8 +30,9 @@ def vert_boundary_projector_form(
     u: fem.Field,
     v: fem.Field,
 ):
-    # Non-zero mass on vertical sides only
-    w = 1.0 - wp.abs(fem.normal(domain, s)[1])
+    # Constrain XY and YZ faces
+    nor = fem.normal(domain, s)
+    w = 1.0 - wp.abs(nor[1])
     return w * u(s) * v(s)
 
 
@@ -62,30 +52,36 @@ class Example:
         self._viscosity = viscosity
         self._boundary_compliance = boundary_compliance
 
-        res = wp.vec3i(resolution, resolution // 2, resolution * 2)
+        res = wp.vec3i(resolution, max(1, resolution // 2), resolution * 2)
+        bounds_lo = wp.vec3(0.0, 0.0, 0.0)
+        bounds_hi = wp.vec3(1.0, 0.5, 2.0)
 
         if mesh == "tet":
-            pos, tet_vtx_indices = gen_tetmesh(
+            pos, tet_vtx_indices = fem_example_utils.gen_tetmesh(
                 res=res,
-                bounds_lo=wp.vec3(0.0, 0.0, 0.0),
-                bounds_hi=wp.vec3(1.0, 0.5, 2.0),
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
             )
             self._geo = fem.Tetmesh(tet_vtx_indices, pos)
         elif mesh == "hex":
-            pos, hex_vtx_indices = gen_hexmesh(
+            pos, hex_vtx_indices = fem_example_utils.gen_hexmesh(
                 res=res,
-                bounds_lo=wp.vec3(0.0, 0.0, 0.0),
-                bounds_hi=wp.vec3(1.0, 0.5, 2.0),
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
             )
             self._geo = fem.Hexmesh(hex_vtx_indices, pos)
         elif mesh == "nano":
-            volume = wp.Volume.allocate(min=[0, 0, 0], max=[1.0, 0.5, 2.0], voxel_size=1.0 / res[0], bg_value=None)
+            volume = fem_example_utils.gen_volume(
+                res=res,
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
+            )
             self._geo = fem.Nanogrid(volume)
         else:
             self._geo = fem.Grid3D(
                 res=res,
-                bounds_lo=wp.vec3(0.0, 0.0, 0.0),
-                bounds_hi=wp.vec3(1.0, 0.5, 2.0),
+                bounds_lo=bounds_lo,
+                bounds_hi=bounds_hi,
             )
 
         # Domain and function spaces
@@ -95,7 +91,7 @@ class Example:
         # Scalar field over our function space
         self._scalar_field: fem.DiscreteField = self._scalar_space.make_field()
 
-        self.renderer = Plot()
+        self.renderer = fem_example_utils.Plot()
 
     def step(self):
         geo = self._geo
@@ -129,11 +125,11 @@ class Example:
 
         with wp.ScopedTimer("CG solve"):
             x = wp.zeros_like(rhs)
-            bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
+            fem_example_utils.bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
             self._scalar_field.dof_values = x
 
     def render(self):
-        self.renderer.add_volume("solution", self._scalar_field)
+        self.renderer.add_field("solution", self._scalar_field)
 
 
 if __name__ == "__main__":

warp/examples/fem/example_diffusion_mgpu.py

@@ -15,21 +15,12 @@
 from typing import Tuple
 
 import warp as wp
+import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
+from warp.examples.fem.example_diffusion import diffusion_form, linear_form
 from warp.sparse import bsr_axpy, bsr_mv
 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 .example_diffusion import diffusion_form, linear_form
-    from .plot_utils import Plot
-except ImportError:
-    from bsr_utils import bsr_cg
-    from example_diffusion import diffusion_form, linear_form
-    from plot_utils import Plot
-
 
 @fem.integrand
 def mass_form(
@@ -115,7 +106,7 @@ class Example:
             self._scalar_space = fem.make_polynomial_space(self._geo, degree=3)
             self._scalar_field = self._scalar_space.make_field()
 
-        self.renderer = Plot()
+        self.renderer = fem_example_utils.Plot()
 
     def step(self):
         devices = wp.get_cuda_devices()
@@ -161,12 +152,14 @@ class Example:
         A.rank_data = (matrices, rhs_vecs, res_vecs, indices)
 
         with wp.ScopedDevice(main_device):
-            bsr_cg(A, x=global_res, b=glob_rhs, use_diag_precond=False, quiet=self._quiet, mv_routine=A.mv_routine)
+            fem_example_utils.bsr_cg(
+                A, x=global_res, b=glob_rhs, use_diag_precond=False, quiet=self._quiet, mv_routine=A.mv_routine
+            )
 
         array_cast(in_array=global_res, out_array=self._scalar_field.dof_values)
 
     def render(self):
-        self.renderer.add_surface("solution", self._scalar_field)
+        self.renderer.add_field("solution", self._scalar_field)
 
     def _assemble_local_system(self, geo_partition: fem.GeometryPartition):
         scalar_space = self._scalar_space

warp/examples/fem/example_magnetostatics.py

@@ -0,0 +1,190 @@
+# 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 Magnetostatics
+#
+# This example demonstrates solving an in-plane magnetostatics problem
+# using a curl-curl formulation
+#
+# 1/mu Curl B + j = 0
+# Div. B = 0
+#
+# solved over field A such that B = Curl A, A = (0, 0, a_z),
+# and a_z = 0 on the domain boundary
+#
+# This example also illustrates using an ImplictField to warp a square mesh
+# to a circular domain
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+
+# Vacuum and copper magnetic permeabilities
+MU_0 = wp.constant(np.pi * 4.0e-7)
+MU_C = wp.constant(1.25e-6)
+
+
+@wp.func
+def square_to_disk(x: wp.vec2):
+    # mapping from unit square to unit disk
+    return wp.normalize(x) * wp.max(wp.abs(x))
+
+
+@wp.func
+def square_to_disk_grad(x: wp.vec2):
+    # gradient of mapping from unit square to unit disk
+    if x == wp.vec2(0.0):
+        return wp.mat22(0.0)
+
+    d = wp.normalize(x)
+    d_grad = (wp.identity(n=2, dtype=float) - wp.outer(d, d)) / wp.length(x)
+
+    ax = wp.abs(x)
+    xinf = wp.max(ax)
+
+    xinf_grad = wp.select(ax[0] > ax[1], wp.vec2(0.0, wp.sign(x[1])), wp.vec(wp.sign(x[0]), 0.0))
+
+    return d_grad * xinf + wp.outer(d, xinf_grad)
+
+
+@wp.func
+def permeability_field(pos: wp.vec2, coil_height: float, coil_internal_radius: float, coil_external_radius: float):
+    # space-varying permeability
+
+    x = wp.abs(pos[0])
+    y = wp.abs(pos[1])
+    return wp.select(
+        y < coil_height and x > coil_internal_radius and x < coil_external_radius,
+        MU_0,
+        MU_C,
+    )
+
+
+@wp.func
+def current_field(
+    pos: wp.vec2, coil_height: float, coil_internal_radius: float, coil_external_radius: float, current: float
+):
+    # space-varying current direction along z axis (0, +1 or -1)
+    x = wp.abs(pos[0])
+    y = wp.abs(pos[1])
+    return (
+        wp.select(
+            y < coil_height and x > coil_internal_radius and x < coil_external_radius,
+            0.0,
+            wp.sign(pos[0]),
+        )
+        * current
+    )
+
+
+@fem.integrand
+def curl_z(u: fem.Field, s: fem.Sample):
+    # projection of curl((0, 0, u)) over z axis
+    du = fem.grad(u, s)
+    return wp.vec2(du[1], -du[0])
+
+
+@fem.integrand
+def curl_curl_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, mu: fem.Field):
+    return wp.dot(curl_z(u, s), curl_z(v, s)) / mu(s)
+
+
+@fem.integrand
+def mass_form(s: fem.Sample, domain: fem.Domain, v: fem.Field, u: fem.Field):
+    return u(s) * v(s)
+
+
+class Example:
+    def __init__(self, quiet=False, degree=2, resolution=32, domain_radius=2.0, current=1.0e6):
+        # We mesh the unit disk by first meshing the unit square, then building a deformed geometry
+        # from an implicit mapping field
+        square_geo = fem.Grid2D(
+            bounds_lo=wp.vec2(-domain_radius, -domain_radius),
+            bounds_hi=wp.vec2(domain_radius, domain_radius),
+            res=wp.vec2i(resolution, resolution),
+        )
+
+        def_field = fem.ImplicitField(domain=fem.Cells(square_geo), func=square_to_disk, grad_func=square_to_disk_grad)
+        disk_geo = def_field.make_deformed_geometry(relative=False)
+
+        coil_config = {"coil_height": 1.0, "coil_internal_radius": 0.1, "coil_external_radius": 0.3}
+
+        domain = fem.Cells(disk_geo)
+        self._permeability_field = fem.ImplicitField(domain, func=permeability_field, values=coil_config)
+        self._current_field = fem.ImplicitField(domain, func=current_field, values=dict(current=current, **coil_config))
+
+        z_space = fem.make_polynomial_space(disk_geo, degree=degree, element_basis=fem.ElementBasis.LAGRANGE)
+        xy_space = fem.make_polynomial_space(
+            disk_geo, degree=degree, element_basis=fem.ElementBasis.LAGRANGE, dtype=wp.vec2
+        )
+
+        self.A_field = z_space.make_field()
+        self.B_field = xy_space.make_field()
+
+        self.renderer = fem_example_utils.Plot()
+
+    def step(self):
+        z_space = self.A_field.space
+        disk_geo = z_space.geometry
+
+        u = fem.make_trial(space=z_space)
+        v = fem.make_test(space=z_space)
+        lhs = fem.integrate(curl_curl_form, fields={"u": u, "v": v, "mu": self._permeability_field})
+        rhs = fem.integrate(mass_form, fields={"v": v, "u": self._current_field})
+
+        # Dirichlet BC
+        boundary = fem.BoundarySides(disk_geo)
+        u_bd = fem.make_trial(space=z_space, domain=boundary)
+        v_bd = fem.make_test(space=z_space, domain=boundary)
+        dirichlet_bd_proj = fem.integrate(mass_form, fields={"u": u_bd, "v": v_bd}, nodal=True)
+        fem.project_linear_system(lhs, rhs, dirichlet_bd_proj)
+
+        x = wp.zeros_like(rhs)
+        fem_example_utils.bsr_cg(lhs, b=rhs, x=x, tol=1.0e-8, quiet=False)
+
+        # make sure result is exactly zero outisde of circle
+        wp.sparse.bsr_mv(dirichlet_bd_proj, x=x, y=x, alpha=-1.0, beta=1.0)
+        wp.utils.array_cast(in_array=x, out_array=self.A_field.dof_values)
+
+        # compute B as curl(A)
+        fem.interpolate(curl_z, dest=self.B_field, fields={"u": self.A_field})
+
+    def render(self):
+        self.renderer.add_field("A", self.A_field)
+        self.renderer.add_field("B", self.B_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("--radius", type=float, default=2.0, help="Radius of simulation domain.")
+    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, domain_radius=args.radius)
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot({"A": {"contours": {"levels": 30}}, "B": {"streamlines": {"density": 1.0}}})