warp-lang 1.0.0b2__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

examples/fem/example_apic_fluid.py

@@ -1,26 +1,23 @@
 import os
-import math
-from typing import Any
 
 import warp as wp
 import numpy as np
 
-import warp.sim.render
 from warp.sim import Model, State
+import warp.sim.render
+
+import warp.fem as fem
+
+from warp.fem import integrand, lookup, normal, grad, at_node, div
+from warp.fem import Field, Sample, Domain
 
-from warp.fem.geometry import Grid3D
-from warp.fem.domain import Cells, BoundarySides
-from warp.fem.space import make_polynomial_space
-from warp.fem.quadrature import PicQuadrature
-from warp.fem.field import make_test, make_trial
-from warp.fem.types import vec3i, Field, Sample, Domain
-from warp.fem.integrate import integrate
-from warp.fem.operator import integrand, lookup, normal, grad, at_node, div
-from warp.fem.dirichlet import normalize_dirichlet_projector
+from warp.sparse import bsr_mv, bsr_copy, bsr_mm, bsr_transposed, bsr_zeros, BsrMatrix
 
-from warp.sparse import bsr_mv, bsr_copy, bsr_mm, bsr_transposed, BsrMatrix
 
-from bsr_utils import bsr_cg
+try:
+    from .bsr_utils import bsr_cg
+except ImportError:
+    from bsr_utils import bsr_cg
 
 
 @integrand
@@ -67,8 +64,6 @@ def update_particles(
     pos_proj = domain(lookup(domain, pos_adv))
     pos[s.qp_index] = pos_proj
 
-    return 0.0
-
 
 @integrand
 def velocity_boundary_projector_form(s: Sample, domain: Domain, u: Field, v: Field):
@@ -122,12 +117,7 @@ def scale_transposed_divergence_mat(
         tr_divergence_mat_values[b] = tr_divergence_mat_values[b] * inv_fraction_int[u_i]
 
 
-def solve_incompressibility(
-    divergence_mat: BsrMatrix,
-    inv_volume,
-    pressure,
-    velocity,
-):
+def solve_incompressibility(divergence_mat: BsrMatrix, inv_volume, pressure, velocity, quiet: bool = False):
     """Solve for divergence-free velocity delta:
 
     delta_velocity = inv_volume * transpose(divergence_mat) * pressure
@@ -151,24 +141,26 @@ def solve_incompressibility(
 
     rhs = wp.zeros_like(pressure)
     bsr_mv(A=divergence_mat, x=velocity, y=rhs, alpha=-1.0, beta=0.0)
-    bsr_cg(schur, b=rhs, x=pressure)
+    bsr_cg(schur, b=rhs, x=pressure, quiet=quiet)
 
     # Apply pressure to velocity
     bsr_mv(A=transposed_divergence_mat, x=pressure, y=velocity, alpha=1.0, beta=1.0)
 
 
 class Example:
-    def __init__(self, stage):
+    def __init__(self, stage, num_frames=1000, res=[32, 64, 16], quiet=False):
         self.frame_dt = 1.0 / 60
-        self.frame_count = 1000
+        self.num_frames = num_frames
+        self.current_frame = 0
 
         self.sim_substeps = 1
         self.sim_dt = self.frame_dt / self.sim_substeps
-        self.sim_steps = self.frame_count * self.sim_substeps
-        self.sim_time = 0.0
+        self.sim_steps = self.num_frames * self.sim_substeps
+
+        self._quiet = quiet
 
         # grid dimensions and particle emission
-        grid_res = np.array([32, 64, 16], dtype=int)
+        grid_res = np.array(res, dtype=int)
         particle_fill_frac = np.array([0.5, 0.5, 1.0])
         grid_lo = wp.vec3(0.0)
         grid_hi = wp.vec3(50, 100, 25)
@@ -177,10 +169,10 @@ class Example:
         grid_cell_volume = np.prod(grid_cell_size)
 
         PARTICLES_PER_CELL_DIM = 3
-        self.radius = np.max(grid_cell_size) / (2 * PARTICLES_PER_CELL_DIM)
+        self.radius = float(np.max(grid_cell_size) / (2 * PARTICLES_PER_CELL_DIM))
 
         particle_grid_res = np.array(particle_fill_frac * grid_res * PARTICLES_PER_CELL_DIM, dtype=int)
-        particle_grid_offset = self.radius * np.ones(3)
+        particle_grid_offset = wp.vec3(self.radius, self.radius, self.radius)
 
         np.random.seed(0)
         builder = wp.sim.ModelBuilder()
@@ -191,20 +183,20 @@ class Example:
             cell_x=self.radius * 2.0,
             cell_y=self.radius * 2.0,
             cell_z=self.radius * 2.0,
-            pos=(0.0, 0.0, 0.0) + particle_grid_offset,
+            pos=wp.vec3(0.0, 0.0, 0.0) + particle_grid_offset,
             rot=wp.quat_identity(),
-            vel=(0.0, 0.0, 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
             mass=grid_cell_volume / PARTICLES_PER_CELL_DIM**3,
             jitter=self.radius * 1.0,
             radius_mean=self.radius,
         )
 
-        self.grid = Grid3D(vec3i(grid_res), grid_lo, grid_hi)
+        self.grid = fem.Grid3D(wp.vec3i(grid_res), grid_lo, grid_hi)
 
         # Function spaces
-        self.velocity_space = make_polynomial_space(self.grid, dtype=wp.vec3, degree=1)
-        self.fraction_space = make_polynomial_space(self.grid, dtype=float, degree=1)
-        self.strain_space = make_polynomial_space(
+        self.velocity_space = fem.make_polynomial_space(self.grid, dtype=wp.vec3, degree=1)
+        self.fraction_space = fem.make_polynomial_space(self.grid, dtype=float, degree=1)
+        self.strain_space = fem.make_polynomial_space(
             self.grid,
             dtype=float,
             degree=0,
@@ -214,29 +206,39 @@ class Example:
         self.velocity_field = self.velocity_space.make_field()
 
         # Test and trial functions
-        self.domain = Cells(self.grid)
-        self.velocity_test = make_test(self.velocity_space, domain=self.domain)
-        self.velocity_trial = make_trial(self.velocity_space, domain=self.domain)
-        self.fraction_test = make_test(self.fraction_space, domain=self.domain)
-        self.strain_test = make_test(self.strain_space, domain=self.domain)
-        self.strain_trial = make_trial(self.strain_space, domain=self.domain)
+        self.domain = fem.Cells(self.grid)
+        self.velocity_test = fem.make_test(self.velocity_space, domain=self.domain)
+        self.velocity_trial = fem.make_trial(self.velocity_space, domain=self.domain)
+        self.fraction_test = fem.make_test(self.fraction_space, domain=self.domain)
+        self.strain_test = fem.make_test(self.strain_space, domain=self.domain)
+        self.strain_trial = fem.make_trial(self.strain_space, domain=self.domain)
 
         # Enforcing the Dirichlet boundary condition the hard way;
         # build projector for velocity left- and right-hand-sides
-        boundary = BoundarySides(self.grid)
-        u_bd_test = make_test(space=self.velocity_space, domain=boundary)
-        u_bd_trial = make_trial(space=self.velocity_space, domain=boundary)
-        u_bd_projector = integrate(
+        boundary = fem.BoundarySides(self.grid)
+        u_bd_test = fem.make_test(space=self.velocity_space, domain=boundary)
+        u_bd_trial = fem.make_trial(space=self.velocity_space, domain=boundary)
+        u_bd_projector = fem.integrate(
             velocity_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, nodal=True, output_dtype=float
         )
 
-        normalize_dirichlet_projector(u_bd_projector)
+        fem.normalize_dirichlet_projector(u_bd_projector)
         self.vel_bd_projector = u_bd_projector
 
+        # Storage for temporary variables
+        self.temporary_store = fem.TemporaryStore()
+
+        self._divergence_matrix = bsr_zeros(
+            self.strain_space.node_count(),
+            self.velocity_space.node_count(),
+            block_type=wp.mat(shape=(1, 3), dtype=float),
+        )
+
         # Warp.sim model
         self.model: Model = builder.finalize()
 
-        print("Particle count:", self.model.particle_count)
+        if not self._quiet:
+            print("Particle count:", self.model.particle_count)
 
         self.state_0: State = self.model.state()
         self.state_0.particle_qd_grad = wp.zeros(shape=(self.model.particle_count), dtype=wp.mat33)
@@ -244,28 +246,41 @@ class Example:
         self.state_1: State = self.model.state()
         self.state_1.particle_qd_grad = wp.zeros(shape=(self.model.particle_count), dtype=wp.mat33)
 
-        self.renderer = wp.sim.render.SimRenderer(self.model, stage, scaling=20.0)
+        self.renderer = warp.sim.render.SimRenderer(self.model, stage, scaling=20.0)
 
-    def update(self, frame_index):
-        with wp.ScopedTimer(f"simulate frame {frame_index}", active=True):
+    def update(self):
+        fem.set_default_temporary_store(self.temporary_store)
+
+        self.current_frame = self.current_frame + 1
+        with wp.ScopedTimer(f"simulate frame {self.current_frame}", active=True):
             for s in range(self.sim_substeps):
                 # Bin particles to grid cells
-                pic = PicQuadrature(
-                    domain=Cells(self.grid), positions=self.state_0.particle_q, measures=self.model.particle_mass
+                pic = fem.PicQuadrature(
+                    domain=fem.Cells(self.grid), positions=self.state_0.particle_q, measures=self.model.particle_mass
                 )
 
+                # Borrow some temporary arrays for storing integration results
+                inv_volume_temporary = fem.borrow_temporary(
+                    self.temporary_store, shape=(self.fraction_space.node_count()), dtype=float
+                )
+                velocity_int_temporary = fem.borrow_temporary(
+                    self.temporary_store, shape=(self.velocity_space.node_count()), dtype=wp.vec3
+                )
+                inv_volume = inv_volume_temporary.array
+                velocity_int = velocity_int_temporary.array
+
                 # Inverse volume fraction
-                fraction_test = make_test(space=self.fraction_space, domain=pic.domain)
-                inv_volume = integrate(
+                fem.integrate(
                     integrate_fraction,
                     quadrature=pic,
-                    fields={"phi": fraction_test},
+                    fields={"phi": self.fraction_test},
                     accumulate_dtype=float,
+                    output=inv_volume,
                 )
                 wp.launch(kernel=invert_volume_kernel, dim=inv_volume.shape, inputs=[inv_volume])
 
                 # Velocity right-hand side
-                velocity_int = integrate(
+                fem.integrate(
                     integrate_velocity,
                     quadrature=pic,
                     fields={"u": self.velocity_test},
@@ -276,7 +291,7 @@ class Example:
                         "gravity": self.model.gravity,
                     },
                     accumulate_dtype=float,
-                    output_dtype=wp.vec3,
+                    output=velocity_int,
                 )
 
                 # Compute constraint-free velocity
@@ -292,28 +307,29 @@ class Example:
                 bsr_mv(A=self.vel_bd_projector, x=velocity_int, y=self.velocity_field.dof_values, alpha=-1.0, beta=1.0)
 
                 # Divergence matrix
-                divergence_mat = integrate(
+                fem.integrate(
                     divergence_form,
                     quadrature=pic,
                     fields={"u": self.velocity_trial, "psi": self.strain_test},
                     accumulate_dtype=float,
-                    output_dtype=float,
+                    output=self._divergence_matrix,
                 )
 
                 # Project matrix to enforce boundary conditions
-                divergence_mat_tmp = bsr_copy(divergence_mat)
-                bsr_mm(alpha=-1.0, x=divergence_mat_tmp, y=self.vel_bd_projector, z=divergence_mat, beta=1.0)
+                divergence_mat_tmp = bsr_copy(self._divergence_matrix)
+                bsr_mm(alpha=-1.0, x=divergence_mat_tmp, y=self.vel_bd_projector, z=self._divergence_matrix, beta=1.0)
 
                 # Solve unilateral incompressibility
                 solve_incompressibility(
-                    divergence_mat,
+                    self._divergence_matrix,
                     inv_volume,
                     self.pressure_field.dof_values,
                     self.velocity_field.dof_values,
+                    quiet=self._quiet,
                 )
 
                 # (A)PIC advection
-                integrate(
+                fem.interpolate(
                     update_particles,
                     quadrature=pic,
                     values={
@@ -328,17 +344,17 @@ class Example:
 
                 # swap states
                 (self.state_0, self.state_1) = (self.state_1, self.state_0)
+        
+        fem.set_default_temporary_store(None)
 
     def render(self, is_live=False):
         with wp.ScopedTimer("render", active=True):
-            time = 0.0 if is_live else self.sim_time
+            time = self.current_frame * self.frame_dt
 
             self.renderer.begin_frame(time)
             self.renderer.render(self.state_0)
             self.renderer.end_frame()
 
-        self.sim_time += self.frame_dt
-
 
 if __name__ == "__main__":
     wp.set_module_options({"enable_backward": False})
@@ -348,8 +364,8 @@ if __name__ == "__main__":
 
     example = Example(stage_path)
 
-    for i in range(example.frame_count):
-        example.update(i)
+    for i in range(example.num_frames):
+        example.update()
         example.render()
 
     example.renderer.save()

examples/fem/example_convection_diffusion.py

@@ -7,28 +7,25 @@ This example simulates a convection-diffusion PDE using semi-Lagrangian advectio
 
 import argparse
 
-import warp as wp
-
-from warp.fem.types import *
-from warp.fem.geometry import Grid2D, Trimesh2D
-from warp.fem.field import make_test, make_trial
-from warp.fem.space import make_polynomial_space
-from warp.fem.domain import Cells
-from warp.fem.integrate import integrate, interpolate
-from warp.fem.operator import grad, integrand, lookup
-
-from bsr_utils import bsr_to_scipy
-from plot_utils import plot_surface
-from mesh_utils import gen_trimesh
-
-from scipy.sparse.linalg import factorized
-
-import matplotlib.pyplot as plt
-import matplotlib.animation as animation
-
 
-@integrand
-def initial_condition(domain: Domain, s: Sample):
+import warp as wp
+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 .mesh_utils import gen_trimesh
+    from .plot_utils import Plot
+except ImportError:
+    from bsr_utils import bsr_cg
+    from mesh_utils import gen_trimesh
+    from plot_utils import Plot
+
+
+@fem.integrand
+def initial_condition(domain: fem.Domain, s: fem.Sample):
+    """Initial condition: 1.0 in ]0.6, 0.4[ x ]0.2, 0.8[, 0.0 elsewhere"""
     pos = domain(s)
     if pos[0] > 0.4 and pos[0] < 0.6 and pos[1] > 0.2 and pos[1] < 0.8:
         return 1.0
@@ -42,13 +39,15 @@ def velocity(pos: wp.vec2, ang_vel: float):
     return wp.vec2(offset[1], -offset[0]) * ang_vel
 
 
-@integrand
-def inertia_form(s: Sample, phi: Field, psi: Field, dt: float):
+@fem.integrand
+def inertia_form(s: fem.Sample, phi: fem.Field, psi: fem.Field, dt: float):
     return phi(s) * psi(s) / dt
 
 
-@integrand
-def transported_inertia_form(s: Sample, domain: Domain, phi: Field, psi: Field, ang_vel: float, dt: float):
+@fem.integrand
+def transported_inertia_form(
+    s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi: fem.Field, ang_vel: float, dt: float
+):
     pos = domain(s)
     vel = velocity(pos, ang_vel)
 
@@ -56,100 +55,105 @@ def transported_inertia_form(s: Sample, domain: Domain, phi: Field, psi: Field,
     conv_pos = pos - vel * dt
     # lookup opertor constructs a Sample from a world position.
     # the optional last argument provides a initial guess for the lookup
-    conv_phi = phi(lookup(domain, conv_pos, s))
+    conv_phi = phi(fem.lookup(domain, conv_pos, s))
 
     return conv_phi * psi(s) / dt
 
 
-@integrand
+@fem.integrand
 def diffusion_form(
-    s: Sample,
-    u: Field,
-    v: Field,
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
 ):
     return wp.dot(
-        grad(u, s),
-        grad(v, s),
+        fem.grad(u, s),
+        fem.grad(v, s),
     )
 
 
-@integrand
-def diffusion_and_inertia_form(s: Sample, phi: Field, psi: Field, dt: float, nu: float):
+@fem.integrand
+def diffusion_and_inertia_form(s: fem.Sample, phi: fem.Field, psi: fem.Field, dt: float, nu: float):
     return inertia_form(s, phi, psi, dt) + nu * diffusion_form(s, phi, psi)
 
 
-if __name__ == "__main__":
-    wp.init()
-    wp.set_module_options({"enable_backward": False})
-
+class Example:
     parser = argparse.ArgumentParser()
     parser.add_argument("--resolution", type=int, default=50)
     parser.add_argument("--degree", type=int, default=2)
-    parser.add_argument("--n_frames", type=int, default=250)
+    parser.add_argument("--num_frames", type=int, default=250)
     parser.add_argument("--viscosity", type=float, default=0.001)
     parser.add_argument("--ang_vel", type=float, default=1.0)
     parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh")
-    args = parser.parse_args()
-
-    res = args.resolution
-    dt = 1.0 / (args.ang_vel * res)
-
-    if args.tri_mesh:
-        positions, tri_vidx = gen_trimesh(res=vec2i(res))
-        geo = Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
-    else:
-        geo = Grid2D(res=vec2i(res))
-
-    domain = Cells(geometry=geo)
-    scalar_space = make_polynomial_space(geo, degree=args.degree)
-    quadrature = None
-
-    # Initial condition
-    phi0 = scalar_space.make_field()
-    interpolate(initial_condition, dest=phi0)
-
-    # Assemble and factorize diffusion and inertia matrix
-    test = make_test(space=scalar_space, domain=domain)
-    trial = make_trial(space=scalar_space, domain=domain)
-    matrix = integrate(
-        diffusion_and_inertia_form,
-        quadrature=quadrature,
-        fields={"phi": trial, "psi": test},
-        values={"nu": args.viscosity, "dt": dt},
-    )
-    matrix_solve = factorized(bsr_to_scipy(matrix))
 
-    results = [phi0.dof_values.numpy()]
-    phik = phi0
+    def __init__(self, stage=None, quiet=False, args=None, **kwargs):
+        if args is None:
+            # Read args from kwargs, add default arg values from parser
+            args = argparse.Namespace(**kwargs)
+            args = Example.parser.parse_args(args=[], namespace=args)
+        self._args = args
+        self._quiet = quiet
+
+        res = args.resolution
+        self.sim_dt = 1.0 / (args.ang_vel * res)
+        self.current_frame = 0
+
+        if args.tri_mesh:
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(res))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(res))
+
+        domain = fem.Cells(geometry=geo)
+        scalar_space = fem.make_polynomial_space(geo, degree=args.degree)
+
+        # Initial condition
+        self._phi_field = scalar_space.make_field()
+        fem.interpolate(initial_condition, dest=self._phi_field)
+
+        # Assemble diffusion and inertia matrix
+        self._test = fem.make_test(space=scalar_space, domain=domain)
+        self._trial = fem.make_trial(space=scalar_space, domain=domain)
+        self._matrix = fem.integrate(
+            diffusion_and_inertia_form,
+            fields={"phi": self._trial, "psi": self._test},
+            values={"nu": args.viscosity, "dt": self.sim_dt},
+            output_dtype=float,
+        )
+
+        self.renderer = Plot(stage)
+        self.renderer.add_surface("phi", self._phi_field)
+
+    def update(self):
+        self.current_frame += 1
 
-    for k in range(args.n_frames):
         # right-hand-side -- advected inertia
-        rhs = integrate(
+        rhs = fem.integrate(
             transported_inertia_form,
-            quadrature=quadrature,
-            fields={"phi": phik, "psi": test},
-            values={"ang_vel": args.ang_vel, "dt": dt},
+            fields={"phi": self._phi_field, "psi": self._test},
+            values={"ang_vel": self._args.ang_vel, "dt": self.sim_dt},
+            output_dtype=float,
         )
 
-        # Solve using Scipy
-        x = matrix_solve(rhs.numpy().flatten())
+        # Solve linear system
+        bsr_cg(self._matrix, x=self._phi_field.dof_values, b=rhs, quiet=self._quiet, tol=1.0e-12)
 
-        phik.dof_values = x
-        results.append(x)
+    def render(self):
+        self.renderer.begin_frame(time = self.current_frame * self.sim_dt)
+        self.renderer.add_surface("phi", self._phi_field)
+        self.renderer.end_frame()
 
-    colormesh = plot_surface(phi0)
-    ax = colormesh.axes
 
-    def animate(i):
-        ax.clear()
-        phik.dof_values = results[i]
-        return plot_surface(phik, axes=ax)
+if __name__ == "__main__":
+    wp.init()
+    wp.set_module_options({"enable_backward": False})
 
-    anim = animation.FuncAnimation(
-        ax.figure,
-        animate,
-        interval=30,
-        blit=False,
-        frames=len(results),
-    )
-    plt.show()
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    for k in range(args.num_frames):
+        print(f"Frame {k}:")
+        example.update()
+        example.render()
+
+    example.renderer.plot()