warp-lang 1.1.0__py3-none-macosx_10_13_universal2.whl → 1.2.1__py3-none-macosx_10_13_universal2.whl

warp/examples/fem/example_burgers.py

@@ -0,0 +1,262 @@
+# 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 Burgers
+#
+# This example simulates an inviscid non-conservative Burgers PDE using
+# Discontinuous Galerkin with minmod slope limiter
+#
+# d u /dt + (u . grad) u = 0
+#
+###########################################################################
+
+import warp as wp
+import warp.fem as fem
+import warp.sparse as sp
+
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .bsr_utils import invert_diagonal_bsr_mass_matrix
+    from .plot_utils import Plot
+except ImportError:
+    from bsr_utils import invert_diagonal_bsr_mass_matrix
+    from plot_utils import Plot
+
+
+@fem.integrand
+def vel_mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(v(s), u(s))
+
+
+@fem.integrand
+def upwind_transport_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, w: fem.Field):
+    # Upwinding transport with discontinuous convection velocity,
+    # using jump(w v) = jump(w) avg(v) + avg(w) jump(w)
+
+    nor = fem.normal(domain, s)
+    w_avg_n = wp.dot(fem.average(w, s), nor)
+    w_jump_n = wp.dot(fem.jump(w, s), nor)
+
+    x = domain(s)
+    v_avg = fem.average(v, s)
+    if x[0] <= 0.0 or x[0] >= 1.0:  # out
+        # if x[0] >= 1.0:  # out
+        v_jump = v(s)
+    else:
+        v_jump = fem.jump(v, s)
+
+    u_avg = fem.average(u, s)
+    u_jump = fem.jump(u, s)
+
+    return wp.dot(u_avg, v_jump * w_avg_n + v_avg * w_jump_n) + 0.5 * wp.dot(v_jump, u_jump) * (
+        wp.abs(w_avg_n) + 0.5 * wp.abs(w_jump_n)
+    )
+
+
+@fem.integrand
+def cell_transport_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, v: fem.Field, w: fem.Field):
+    # ((w . grad) u) . v = v^T (grad u) w =  grad(u) : (v w^T)
+    # with integration by parts
+    # u . Div (w v^T) = u^T grad(v) w  + u^T v div (w)
+
+    return -wp.dot(fem.div(w, s) * v(s) + fem.grad(v, s) * w(s), u(s))
+
+
+@fem.integrand
+def initial_condition(s: fem.Sample, domain: fem.Domain):
+    x = domain(s)[0] * 2.0
+    wave = wp.sin(x * wp.pi)
+    return wp.vec2(wp.select(x <= 1.0, 0.0, wave), 0.0)
+
+
+@fem.integrand
+def velocity_norm(s: fem.Sample, u: fem.Field):
+    return wp.length(u(s))
+
+
+@wp.func
+def minmod(a: float, b: float):
+    sa = wp.sign(a)
+    sb = wp.sign(b)
+    return wp.select(sa == sb, 0.0, sa * wp.min(wp.abs(a), wp.abs(b)))
+
+
+@fem.integrand
+def slope_limiter(domain: fem.Domain, s: fem.Sample, u: fem.Field, dx: wp.vec2):
+    # Minmod slope limiter against P0 discretization (evaluation at cell centers)
+    # Assumes regular grid topology
+
+    center_coords = fem.Coords(0.5, 0.5, 0.0)
+    cell_center = fem.types.make_free_sample(s.element_index, center_coords)
+    center_pos = domain(cell_center)
+
+    u_center = u(cell_center)
+
+    delta_coords = s.element_coords - center_coords
+
+    neighbour_xp = fem.lookup(domain, center_pos + wp.vec2(dx[0], 0.0))
+    neighbour_yp = fem.lookup(domain, center_pos + wp.vec2(0.0, dx[1]))
+    neighbour_xm = fem.lookup(domain, center_pos - wp.vec2(dx[0], 0.0))
+    neighbour_ym = fem.lookup(domain, center_pos - wp.vec2(0.0, dx[1]))
+
+    u_nxp = u(neighbour_xp)
+    u_nyp = u(neighbour_yp)
+    u_nxm = u(neighbour_xm)
+    u_nym = u(neighbour_ym)
+
+    gx = minmod(u_nxp[0] - u_center[0], u_center[0] - u_nxm[0]) * delta_coords[0]
+    gy = minmod(u_nyp[1] - u_center[1], u_center[1] - u_nym[1]) * delta_coords[1]
+
+    delta_u = u(s) - u_center
+    return u_center + wp.vec2(minmod(gx, delta_u[0]), minmod(gy, delta_u[1]))
+
+
+class Example:
+    def __init__(self, quiet=False, resolution=50, degree=1):
+        self._quiet = quiet
+
+        res = resolution
+        self.sim_dt = 1.0 / res
+        self.current_frame = 0
+
+        geo = fem.Grid2D(res=wp.vec2i(resolution))
+
+        domain = fem.Cells(geometry=geo)
+        sides = fem.Sides(geo)
+
+        basis_space = fem.make_polynomial_basis_space(geo, degree=degree, discontinuous=True)
+        vector_space = fem.make_collocated_function_space(basis_space, dtype=wp.vec2)
+        scalar_space = fem.make_collocated_function_space(basis_space, dtype=float)
+
+        # Test function for ou vector space
+        self._test = fem.make_test(space=vector_space, domain=domain)
+        # Test function for integration on sides
+        self._side_test = fem.make_test(space=vector_space, domain=sides)
+
+        # Inertia matrix
+        # For simplicity, use nodal integration so that inertia matrix is diagonal
+        trial = fem.make_trial(space=vector_space, domain=domain)
+        matrix_inertia = fem.integrate(
+            vel_mass_form, fields={"u": trial, "v": self._test}, output_dtype=wp.float32, nodal=True
+        )
+        self._inv_mass_matrix = sp.bsr_copy(matrix_inertia)
+        invert_diagonal_bsr_mass_matrix(self._inv_mass_matrix)
+
+        # Initial condition
+        self.velocity_field = vector_space.make_field()
+        fem.interpolate(initial_condition, dest=self.velocity_field)
+
+        # Velocity nor field -- for visualization purposes
+        self.velocity_norm_field = scalar_space.make_field()
+        fem.interpolate(velocity_norm, dest=self.velocity_norm_field, fields={"u": self.velocity_field})
+
+        self.renderer = Plot()
+        self.renderer.add_surface("u_norm", self.velocity_norm_field)
+
+    def _velocity_delta(self, trial_velocity):
+        # Integration on sides
+        rhs = fem.integrate(
+            upwind_transport_form,
+            fields={"u": trial_velocity.trace(), "v": self._side_test, "w": trial_velocity.trace()},
+            output_dtype=wp.vec2,
+        )
+
+        if self.velocity_field.space.degree > 0:
+            # Integration on cells (if not piecewise-constant)
+            fem.utils.array_axpy(
+                x=fem.integrate(
+                    cell_transport_form,
+                    fields={"u": trial_velocity, "v": self._test, "w": trial_velocity},
+                    output_dtype=wp.vec2,
+                    quadrature=fem.RegularQuadrature(
+                        order=3, domain=self._test.domain, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE
+                    ),
+                ),
+                y=rhs,
+                alpha=1.0,
+                beta=1.0,
+            )
+        return sp.bsr_mv(self._inv_mass_matrix, rhs)
+
+    def step(self):
+        self.current_frame += 1
+
+        # Third-order Strong Stability Preserving Runge-Kutta (SSPRK3)
+
+        k1 = self._velocity_delta(self.velocity_field)
+
+        # tmp = v0 - dt * k1
+        tmp = self.velocity_field.space.make_field()
+        fem.utils.array_axpy(y=tmp.dof_values, x=self.velocity_field.dof_values, alpha=1.0, beta=0.0)
+        fem.utils.array_axpy(y=tmp.dof_values, x=k1, alpha=-self.sim_dt, beta=1.0)
+        k2 = self._velocity_delta(tmp)
+
+        # tmp = v0 - dt * (0.25 * k1 + 0.25 * k2)
+        fem.utils.array_axpy(y=tmp.dof_values, x=k1, alpha=0.75 * self.sim_dt, beta=1.0)
+        fem.utils.array_axpy(y=tmp.dof_values, x=k2, alpha=-0.25 * self.sim_dt, beta=1.0)
+        k3 = self._velocity_delta(tmp)
+
+        # v = v0 - dt * (1/6 * k1 + 1/6 * k2 + 2/3 * k3)
+        fem.utils.array_axpy(y=self.velocity_field.dof_values, x=k1, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
+        fem.utils.array_axpy(y=self.velocity_field.dof_values, x=k2, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
+        fem.utils.array_axpy(y=self.velocity_field.dof_values, x=k3, alpha=-2.0 / 3.0 * self.sim_dt, beta=1.0)
+
+        # Apply slope limiter
+        if self.velocity_field.space.degree > 0:
+            res = self.velocity_field.space.geometry.res
+            dx = wp.vec2(1.0 / res[0], 1.0 / res[1])
+            fem.interpolate(slope_limiter, dest=tmp, fields={"u": self.velocity_field}, values={"dx": dx})
+            wp.copy(src=tmp.dof_values, dest=self.velocity_field.dof_values)
+
+        # Update velocity norm (for visualization)
+        fem.interpolate(velocity_norm, dest=self.velocity_norm_field, fields={"u": self.velocity_field})
+
+    def render(self):
+        self.renderer.begin_frame(time=self.current_frame * self.sim_dt)
+        self.renderer.add_surface("u_norm", self.velocity_norm_field)
+        self.renderer.end_frame()
+
+
+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=50, help="Grid resolution.")
+    parser.add_argument("--num_frames", type=int, default=250, help="Total number of frames.")
+    parser.add_argument("--degree", choices=(0, 1), type=int, default=1, help="Discretization order.")
+    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")
+
+    args = parser.parse_known_args()[0]
+
+    with wp.ScopedDevice(args.device):
+        example = Example(
+            quiet=args.quiet,
+            resolution=args.resolution,
+            degree=args.degree,
+        )
+
+        for k in range(args.num_frames):
+            print(f"Frame {k}:")
+            example.step()
+            example.render()
+
+        if not args.headless:
+            example.renderer.plot()

warp/examples/fem/example_convection_diffusion.py

@@ -28,8 +28,6 @@ except ImportError:
     from mesh_utils import gen_trimesh
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def initial_condition(domain: fem.Domain, s: fem.Sample):

warp/examples/fem/example_convection_diffusion_dg.py

@@ -41,8 +41,6 @@ except ImportError:
     from mesh_utils import gen_quadmesh, gen_trimesh
     from plot_utils import Plot
 
-wp.init()
-
 
 # Standard transport term, on cells' interior
 @fem.integrand

warp/examples/fem/example_deformed_geometry.py

@@ -32,8 +32,6 @@ except ImportError:
     from mesh_utils import gen_quadmesh, gen_trimesh
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def deformation_field_expr(

warp/examples/fem/example_diffusion.py

@@ -32,8 +32,6 @@ except ImportError:
     from mesh_utils import gen_quadmesh, gen_trimesh
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def linear_form(

warp/examples/fem/example_diffusion_3d.py

@@ -25,7 +25,7 @@ from warp.sparse import bsr_axpy
 try:
     from .bsr_utils import bsr_cg
     from .example_diffusion import diffusion_form, linear_form
-    from .mesh_utils import gen_tetmesh
+    from .mesh_utils import gen_hexmesh, gen_tetmesh
     from .plot_utils import Plot
 except ImportError:
     from bsr_utils import bsr_cg
@@ -33,8 +33,6 @@ except ImportError:
     from mesh_utils import gen_hexmesh, gen_tetmesh
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def vert_boundary_projector_form(
@@ -80,6 +78,9 @@ class Example:
                 bounds_hi=wp.vec3(1.0, 0.5, 2.0),
             )
             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)
+            self._geo = fem.Nanogrid(volume)
         else:
             self._geo = fem.Grid3D(
                 res=res,
@@ -149,7 +150,7 @@ 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"), default="grid", help="Mesh type.")
+    parser.add_argument("--mesh", choices=("grid", "tet", "hex", "nano"), default="grid", help="Mesh type.")
     parser.add_argument(
         "--headless",
         action="store_true",

warp/examples/fem/example_diffusion_mgpu.py

@@ -30,8 +30,6 @@ except ImportError:
     from example_diffusion import diffusion_form, linear_form
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def mass_form(

warp/examples/fem/example_mixed_elasticity.py

@@ -32,8 +32,6 @@ except ImportError:
     from mesh_utils import gen_quadmesh, gen_trimesh
     from plot_utils import Plot
 
-wp.init()
-
 
 @wp.func
 def compute_stress(tau: wp.mat22, E: wp.mat33):

warp/examples/fem/example_navier_stokes.py

@@ -31,8 +31,6 @@ except ImportError:
     from mesh_utils import gen_trimesh
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def u_boundary_value(s: fem.Sample, domain: fem.Domain, v: fem.Field, top_vel: float):

warp/examples/fem/example_stokes.py

@@ -30,8 +30,6 @@ except ImportError:
     from mesh_utils import gen_quadmesh, gen_trimesh
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def constant_form(val: wp.vec2):

warp/examples/fem/example_stokes_transfer.py

@@ -33,8 +33,6 @@ except ImportError:
     from bsr_utils import bsr_cg
     from plot_utils import Plot
 
-wp.init()
-
 
 @fem.integrand
 def vel_from_particles_form(s: fem.Sample, particle_vel: wp.array(dtype=wp.vec2), v: fem.Field):

warp/examples/optim/example_bounce.py

@@ -23,8 +23,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 @wp.kernel
 def loss_kernel(pos: wp.array(dtype=wp.vec3), target: wp.vec3, loss: wp.array(dtype=float)):

warp/examples/optim/example_cloth_throw.py

@@ -23,8 +23,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 @wp.kernel
 def com_kernel(positions: wp.array(dtype=wp.vec3), n: int, com: wp.array(dtype=wp.vec3)):

warp/examples/optim/example_diffray.py

@@ -23,8 +23,6 @@ import warp as wp
 import warp.examples
 from warp.optim import SGD
 
-wp.init()
-
 
 class RenderMode:
     """Rendering modes

warp/examples/optim/example_drone.py

@@ -28,8 +28,6 @@ from warp.sim.collide import box_sdf, capsule_sdf, cone_sdf, cylinder_sdf, mesh_
 
 DEFAULT_DRONE_PATH = os.path.join(warp.examples.get_asset_directory(), "crazyflie.usd")  # Path to input drone asset
 
-wp.init()
-
 
 @wp.struct
 class Propeller:

warp/examples/optim/example_inverse_kinematics.py

@@ -19,8 +19,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 TARGET = wp.constant(wp.vec3(2.0, 1.0, 0.0))
 
 

warp/examples/optim/example_inverse_kinematics_torch.py

@@ -21,8 +21,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class ForwardKinematics(torch.autograd.Function):
     @staticmethod

warp/examples/optim/example_spring_cage.py

@@ -20,8 +20,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 @wp.kernel
 def compute_loss_kernel(

warp/examples/optim/example_trajectory.py

@@ -20,8 +20,6 @@ import warp.sim
 import warp.sim.render
 from warp.optim import Adam
 
-wp.init()
-
 
 @wp.kernel
 def loss_l2(

warp/examples/optim/example_walker.py

@@ -29,8 +29,6 @@ import warp.optim
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 @wp.kernel
 def loss_kernel(com: wp.array(dtype=wp.vec3), loss: wp.array(dtype=float)):

warp/examples/sim/example_cartpole.py

@@ -24,8 +24,6 @@ import warp.examples
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_cartpole.usd", num_envs=8):

warp/examples/sim/example_cloth.py

@@ -25,8 +25,6 @@ import warp.examples
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class IntegratorType(Enum):
     EULER = "euler"

warp/examples/sim/example_granular.py

@@ -17,8 +17,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_granular.usd"):

warp/examples/sim/example_granular_collision_sdf.py

@@ -25,8 +25,6 @@ import warp.examples
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_granular_collision_sdf.usd"):

warp/examples/sim/example_jacobian_ik.py

@@ -26,8 +26,6 @@ import warp.examples
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 @wp.kernel
 def compute_endeffector_position(

warp/examples/sim/example_particle_chain.py

@@ -19,8 +19,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_particle_chain.usd"):

warp/examples/sim/example_quadruped.py

@@ -24,8 +24,6 @@ import warp.examples
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 # Taken from env/environment.py
 def compute_env_offsets(num_envs, env_offset=(5.0, 0.0, 5.0), up_axis="Y"):

warp/examples/sim/example_rigid_chain.py

@@ -20,8 +20,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_rigid_chain.usd"):

warp/examples/sim/example_rigid_contact.py

@@ -24,8 +24,6 @@ import warp.examples
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_rigid_contact.usd"):

warp/examples/sim/example_rigid_force.py

@@ -17,8 +17,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_rigid_force.usd", use_opengl=False):

warp/examples/sim/example_rigid_gyroscopic.py

@@ -17,8 +17,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_rigid_gyroscopic.usd"):

warp/examples/sim/example_rigid_soft_contact.py

@@ -17,8 +17,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 class Example:
     def __init__(self, stage_path="example_rigid_soft_contact.usd"):

warp/examples/sim/example_soft_body.py

@@ -18,8 +18,6 @@ import warp as wp
 import warp.sim
 import warp.sim.render
 
-wp.init()
-
 
 @wp.kernel
 def twist_points(

warp/fem/__init__.py

@@ -10,6 +10,7 @@ from .geometry import (
     Grid3D,
     Hexmesh,
     LinearGeometryPartition,
+    Nanogrid,
     Quadmesh2D,
     Tetmesh,
     Trimesh2D,

warp/fem/cache.py

@@ -228,7 +228,7 @@ class Temporary:
 
     def _view_as(self, shape, dtype) -> "Temporary":
         def _view_reshaped_truncated(array):
-            return wp.types.array(
+            view = wp.types.array(
                 ptr=array.ptr,
                 dtype=dtype,
                 shape=shape,
@@ -238,6 +238,8 @@ class Temporary:
                 copy=False,
                 grad=None if array.grad is None else _view_reshaped_truncated(array.grad),
             )
+            view._ref = array
+            return view
 
         self._array_view = _view_reshaped_truncated(self._raw_array)
         return self

warp/fem/geometry/__init__.py

@@ -4,6 +4,7 @@ from .geometry import Geometry
 from .grid_2d import Grid2D
 from .grid_3d import Grid3D
 from .hexmesh import Hexmesh
+from .nanogrid import Nanogrid
 from .partition import (
     ExplicitGeometryPartition,
     GeometryPartition,

warp/fem/geometry/element.py

@@ -14,6 +14,10 @@ class Element:
         """Returns a quadrature of a given order for a prototypical element"""
         raise NotImplementedError
 
+    def center(self) -> Tuple[float]:
+        coords, _ = self.instantiate_quadrature(order=0, family=None)
+        return coords[0]
+
 
 def _point_count_from_order(order: int, family: Polynomial):
     if family == Polynomial.GAUSS_LEGENDRE:

warp/fem/geometry/grid_3d.py

@@ -263,10 +263,6 @@ class Grid3D(Geometry):
     def cell_normal(args: CellArg, s: Sample):
         return wp.vec3(0.0)
 
-    @wp.func
-    def cell_transform_reference_gradient(args: CellArg, cell_index: ElementIndex, coords: Coords, ref_grad: wp.vec3):
-        return wp.cw_div(ref_grad, args.cell_size)
-
     @cached_arg_value
     def side_arg_value(self, device) -> SideArg:
         args = self.SideArg()