warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/examples/fem/example_stokes_transfer.py

@@ -0,0 +1,249 @@
+# 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 Stokes Transfer
+#
+# This example computes a 2D weakly-compressible Stokes flow around
+# a moving object, including:
+# - defining active cells from a mask, and restricting the computation domain to those
+# - utilizing the PicQuadrature to integrate over unstructured particles 
+###########################################################################
+
+import math
+import warp as wp
+import numpy as np
+
+import warp.fem as fem
+
+from warp.utils import array_cast
+from warp.fem.utils import array_axpy
+from warp.sparse import bsr_transposed, bsr_mm, bsr_axpy, bsr_mv
+
+# 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
+
+wp.init()
+
+
+@fem.integrand
+def vel_from_particles_form(s: fem.Sample, particle_vel: wp.array(dtype=wp.vec2), v: fem.Field):
+    vel = particle_vel[s.qp_index]
+    return wp.dot(vel, v(s))
+
+
+@fem.integrand
+def viscosity_form(s: fem.Sample, u: fem.Field, v: fem.Field, nu: float):
+    return nu * wp.ddot(fem.D(u, s), fem.D(v, s))
+
+
+@fem.integrand
+def mass_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def scalar_mass_form(
+    s: fem.Sample,
+    p: fem.Field,
+    q: fem.Field,
+):
+    return p(s) * q(s)
+
+
+@fem.integrand
+def div_form(
+    s: fem.Sample,
+    u: fem.Field,
+    q: fem.Field,
+):
+    return q(s) * fem.div(u, s)
+
+
+@fem.integrand
+def cell_activity(s: fem.Sample, domain: fem.Domain, c1: wp.vec2, c2: wp.vec2, radius: float):
+    pos = domain(s)
+    if wp.length(pos - c1) < radius:
+        return 0.0
+    if wp.length(pos - c2) < radius:
+        return 0.0
+    return 1.0
+
+
+@wp.kernel
+def inverse_array_kernel(m: wp.array(dtype=wp.float64)):
+    m[wp.tid()] = wp.float64(1.0) / m[wp.tid()]
+
+
+class Example:
+    def __init__(self, stage=None, quiet=False):
+        self._quiet = quiet
+
+        self.res = 50
+        self.cell_size = 1.0 / self.res
+
+        self.vel = 1.0
+        self.viscosity = 100.0
+        self.compliance = 0.01
+        self.bd_strength = 100000.0
+
+        geo = fem.Grid2D(res=wp.vec2i(self.res))
+
+        # Displacement boundary conditions are defined by two circles going in opposite directions
+        # Sample particles along those
+        circle_radius = 0.15
+        c1_center = wp.vec2(0.25, 0.5)
+        c2_center = wp.vec2(0.75, 0.5)
+        particles, particle_areas, particle_velocities = self._gen_particles(circle_radius, c1_center, c2_center)
+
+        # Disable cells that are interior to the circles
+        cell_space = fem.make_polynomial_space(geo, degree=0)
+        activity = cell_space.make_field()
+        fem.interpolate(
+            cell_activity,
+            dest=activity,
+            values={"c1": c1_center, "c2": c2_center, "radius": circle_radius - self.cell_size},
+        )
+
+        # Explicitly define the active geometry partition from those cells
+        self._active_partition = fem.ExplicitGeometryPartition(geo, wp.array(activity.dof_values.numpy(), dtype=int))
+        if not self._quiet:
+            print("Active cells:", self._active_partition.cell_count())
+
+        # Function spaces -- Q1 for vel, Q0 for pressure
+        u_space = fem.make_polynomial_space(geo, degree=1, dtype=wp.vec2)
+        p_space = fem.make_polynomial_space(geo, degree=0)
+
+        self._active_space_partition = fem.make_space_partition(
+            space=u_space, geometry_partition=self._active_partition
+        )
+        self._active_p_space_partition = fem.make_space_partition(
+            space=p_space, geometry_partition=self._active_partition
+        )
+
+        self._u_field = u_space.make_field()
+        self._p_field = p_space.make_field()
+
+        # Particle-based quadrature rule over active cells
+        domain = fem.Cells(geometry=self._active_partition)
+        self._pic_quadrature = fem.PicQuadrature(domain, particles, particle_areas)
+        self._particle_velocities = particle_velocities
+
+        self.renderer = Plot(stage)
+
+    def step(self):
+        u_space = self._u_field.space
+        p_space = self._p_field.space
+
+        # Weakly-enforced boundary condition on particles
+        u_test = fem.make_test(space=u_space, space_partition=self._active_space_partition)
+        u_trial = fem.make_trial(space=u_space, space_partition=self._active_space_partition)
+
+        u_rhs = fem.integrate(
+            vel_from_particles_form,
+            quadrature=self._pic_quadrature,
+            fields={"v": u_test},
+            values={"particle_vel": self._particle_velocities},
+            output_dtype=wp.vec2d,
+        )
+        u_bd_matrix = fem.integrate(mass_form, quadrature=self._pic_quadrature, fields={"u": u_trial, "v": u_test})
+
+        # Viscosity
+        u_visc_matrix = fem.integrate(
+            viscosity_form,
+            fields={"u": u_trial, "v": u_test},
+            values={"nu": self.viscosity},
+        )
+
+        # Pressure-velocity coupling
+        p_test = fem.make_test(space=p_space, space_partition=self._active_p_space_partition)
+        p_trial = fem.make_trial(space=p_space, space_partition=self._active_p_space_partition)
+
+        div_matrix = fem.integrate(div_form, fields={"u": u_trial, "q": p_test})
+        inv_p_mass_matrix = fem.integrate(scalar_mass_form, fields={"p": p_trial, "q": p_test})
+        wp.launch(
+            kernel=inverse_array_kernel,
+            dim=inv_p_mass_matrix.values.shape,
+            device=inv_p_mass_matrix.values.device,
+            inputs=[inv_p_mass_matrix.values],
+        )
+
+        # Assemble linear system
+        u_matrix = u_visc_matrix
+        bsr_axpy(u_bd_matrix, u_matrix, alpha=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)
+
+        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)
+
+        # 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)
+
+        # Copy to fields
+        u_nodes = wp.indexedarray(self._u_field.dof_values, indices=self._active_space_partition.space_node_indices())
+        p_nodes = wp.indexedarray(self._p_field.dof_values, indices=self._active_p_space_partition.space_node_indices())
+
+        array_cast(in_array=u_res, out_array=u_nodes)
+        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)
+
+    def _gen_particles(self, circle_radius, c1_center, c2_center):
+        """Generate some particles along two circles defining velocity boundary conditions"""
+
+        # Generate particles defining the transfer displacement
+        particles_per_circle = int(2.0 * math.pi * circle_radius * self.res)
+
+        angles = np.linspace(0, 2.0 * math.pi, particles_per_circle, endpoint=False)
+
+        n_particles = 2 * particles_per_circle
+        particles = np.empty((n_particles, 2), dtype=float)
+
+        particles[:particles_per_circle, 0] = c1_center[0] + circle_radius * np.cos(angles)
+        particles[:particles_per_circle, 1] = c1_center[1] + circle_radius * np.sin(angles)
+        particles[particles_per_circle:, 0] = c2_center[0] + circle_radius * np.cos(angles)
+        particles[particles_per_circle:, 1] = c2_center[1] + circle_radius * np.sin(angles)
+
+        particle_areas = np.ones(n_particles) * self.cell_size**2
+        particle_velocities = np.zeros_like(particles)
+        particle_velocities[:particles_per_circle, 0] = self.vel
+        particle_velocities[particles_per_circle:, 0] = -self.vel
+
+        particles = wp.array(particles, dtype=wp.vec2)
+        particle_areas = wp.array(particle_areas, dtype=float)
+        particle_velocities = wp.array(particle_velocities, dtype=wp.vec2)
+
+        return particles, particle_areas, particle_velocities
+
+
+if __name__ == "__main__":
+    wp.set_module_options({"enable_backward": False})
+
+    example = Example()
+    example.step()
+    example.render()
+
+    example.renderer.plot(streamlines=["velocity"])

warp/examples/fem/mesh_utils.py

@@ -0,0 +1,109 @@
+import numpy as np
+import warp as wp
+
+from warp.fem.utils import grid_to_tets, grid_to_tris, grid_to_quads, grid_to_hexes
+
+
+def gen_trimesh(res, bounds_lo: wp.vec2 = wp.vec2(0.0), bounds_hi: wp.vec2 = wp.vec2(1.0)):
+    """Constructs a triangular mesh by diving each cell of a dense 2D grid into two triangles
+
+    Args:
+        res: Resolution of the grid along each dimension
+        bounds_lo: Position of the lower bound of the axis-aligned grid
+        bounds_up: Position of the upper bound of the axis-aligned grid
+
+    Returns:
+        Tuple of ndarrays: (Vertex positions, Triangle vertex indices)
+    """
+
+    Nx = res[0]
+    Ny = res[1]
+
+    x = np.linspace(bounds_lo[0], bounds_hi[0], Nx + 1)
+    y = np.linspace(bounds_lo[1], bounds_hi[1], Ny + 1)
+
+    positions = np.transpose(np.meshgrid(x, y, indexing="ij"), axes=(1, 2, 0)).reshape(-1, 2)
+
+    vidx = grid_to_tris(Nx, Ny)
+
+    return wp.array(positions, dtype=wp.vec2), wp.array(vidx, dtype=int)
+
+
+def gen_tetmesh(res, bounds_lo: wp.vec3 = wp.vec3(0.0), bounds_hi: wp.vec3 = wp.vec3(1.0)):
+    """Constructs a tetrahedral mesh by diving each cell of a dense 3D grid into five tetrahedrons
+
+    Args:
+        res: Resolution of the grid along each dimension
+        bounds_lo: Position of the lower bound of the axis-aligned grid
+        bounds_up: Position of the upper bound of the axis-aligned grid
+
+    Returns:
+        Tuple of ndarrays: (Vertex positions, Tetrahedron vertex indices)
+    """
+
+    Nx = res[0]
+    Ny = res[1]
+    Nz = res[2]
+
+    x = np.linspace(bounds_lo[0], bounds_hi[0], Nx + 1)
+    y = np.linspace(bounds_lo[1], bounds_hi[1], Ny + 1)
+    z = np.linspace(bounds_lo[2], bounds_hi[2], Nz + 1)
+
+    positions = np.transpose(np.meshgrid(x, y, z, indexing="ij"), axes=(1, 2, 3, 0)).reshape(-1, 3)
+
+    vidx = grid_to_tets(Nx, Ny, Nz)
+
+    return wp.array(positions, dtype=wp.vec3), wp.array(vidx, dtype=int)
+
+
+def gen_quadmesh(res, bounds_lo: wp.vec2 = wp.vec2(0.0), bounds_hi: wp.vec2 = wp.vec2(1.0)):
+    """Constructs a quadrilateral mesh from a dense 2D grid
+
+    Args:
+        res: Resolution of the grid along each dimension
+        bounds_lo: Position of the lower bound of the axis-aligned grid
+        bounds_up: Position of the upper bound of the axis-aligned grid
+
+    Returns:
+        Tuple of ndarrays: (Vertex positions, Triangle vertex indices)
+    """
+
+    Nx = res[0]
+    Ny = res[1]
+
+    x = np.linspace(bounds_lo[0], bounds_hi[0], Nx + 1)
+    y = np.linspace(bounds_lo[1], bounds_hi[1], Ny + 1)
+
+    positions = np.transpose(np.meshgrid(x, y, indexing="ij"), axes=(1, 2, 0)).reshape(-1, 2)
+
+    vidx = grid_to_quads(Nx, Ny)
+
+    return wp.array(positions, dtype=wp.vec2), wp.array(vidx, dtype=int)
+
+
+def gen_hexmesh(res, bounds_lo: wp.vec3 = wp.vec3(0.0), bounds_hi: wp.vec3 = wp.vec3(1.0)):
+    """Constructs a quadrilateral mesh from a dense 2D grid
+
+    Args:
+        res: Resolution of the grid along each dimension
+        bounds_lo: Position of the lower bound of the axis-aligned grid
+        bounds_up: Position of the upper bound of the axis-aligned grid
+
+    Returns:
+        Tuple of ndarrays: (Vertex positions, Triangle vertex indices)
+    """
+
+    Nx = res[0]
+    Ny = res[1]
+    Nz = res[2]
+
+    x = np.linspace(bounds_lo[0], bounds_hi[0], Nx + 1)
+    y = np.linspace(bounds_lo[1], bounds_hi[1], Ny + 1)
+    z = np.linspace(bounds_lo[1], bounds_hi[1], Nz + 1)
+
+    positions = np.transpose(np.meshgrid(x, y, z, indexing="ij"), axes=(1, 2, 3, 0)).reshape(-1, 3)
+
+    vidx = grid_to_hexes(Nx, Ny, Nz)
+
+    return wp.array(positions, dtype=wp.vec3), wp.array(vidx, dtype=int)
+

warp/examples/fem/plot_utils.py

@@ -0,0 +1,287 @@
+from typing import Set
+
+import numpy as np
+
+from warp.fem import DiscreteField
+
+
+def plot_grid_surface(field, axes=None):
+    import matplotlib.pyplot as plt
+    from matplotlib import cm
+
+    if axes is None:
+        fig, axes = plt.subplots(subplot_kw={"projection": "3d"})
+
+    node_positions = field.space.node_grid()
+
+    # Make data.
+    X = node_positions[0]
+    Y = node_positions[1]
+    Z = field.dof_values.numpy().reshape(X.shape)
+
+    # Plot the surface.
+    return axes.plot_surface(X, Y, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False)
+
+
+def plot_tri_surface(field, axes=None):
+    import matplotlib.pyplot as plt
+    from matplotlib import cm
+    from matplotlib.tri.triangulation import Triangulation
+
+    if axes is None:
+        fig, axes = plt.subplots(subplot_kw={"projection": "3d"})
+
+    node_positions = field.space.node_positions().numpy()
+
+    triangulation = Triangulation(
+        x=node_positions[:, 0], y=node_positions[:, 1], triangles=field.space.node_triangulation()
+    )
+
+    Z = field.dof_values.numpy()
+
+    # Plot the surface.
+    return axes.plot_trisurf(triangulation, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False)
+
+
+def plot_scatter_surface(field, axes=None):
+    import matplotlib.pyplot as plt
+    from matplotlib import cm
+
+    if axes is None:
+        fig, axes = plt.subplots(subplot_kw={"projection": "3d"})
+
+    X, Y = field.space.node_positions().numpy().T
+
+    # Make data.
+    Z = field.dof_values.numpy().reshape(X.shape)
+
+    # Plot the surface.
+    return axes.scatter(X, Y, Z, c=Z, cmap=cm.coolwarm)
+
+
+def plot_surface(field, axes=None):
+    if hasattr(field.space, "node_grid"):
+        return plot_grid_surface(field, axes)
+    elif hasattr(field.space, "node_triangulation"):
+        return plot_tri_surface(field, axes)
+    else:
+        return plot_scatter_surface(field, axes)
+
+
+def plot_grid_color(field, axes=None):
+    import matplotlib.pyplot as plt
+    from matplotlib import cm
+
+    if axes is None:
+        fig, axes = plt.subplots()
+
+    node_positions = field.space.node_grid()
+
+    # Make data.
+    X = node_positions[0]
+    Y = node_positions[1]
+    Z = field.dof_values.numpy().reshape(X.shape)
+
+    # Plot the surface.
+    return axes.pcolormesh(X, Y, Z, cmap=cm.coolwarm)
+
+
+def plot_velocities(field, axes=None):
+    import matplotlib.pyplot as plt
+
+    if axes is None:
+        fig, axes = plt.subplots()
+
+    node_positions = field.space.node_positions().numpy()
+
+    # Make data.
+    X = node_positions[:, 0]
+    Y = node_positions[:, 1]
+
+    vel = field.dof_values.numpy()
+    u = np.ascontiguousarray(vel[:, 0])
+    v = np.ascontiguousarray(vel[:, 1])
+
+    u = u.reshape(X.shape)
+    v = v.reshape(X.shape)
+
+    return axes.quiver(X, Y, u, v)
+
+
+def plot_grid_streamlines(field, axes=None):
+    import matplotlib.pyplot as plt
+
+    if axes is None:
+        fig, axes = plt.subplots()
+
+    node_positions = field.space.node_grid()
+
+    # Make data.
+    X = node_positions[0][:, 0]
+    Y = node_positions[1][0, :]
+
+    vel = field.dof_values.numpy()
+    u = np.ascontiguousarray(vel[:, 0])
+    v = np.ascontiguousarray(vel[:, 1])
+
+    u = np.transpose(u.reshape(node_positions[0].shape))
+    v = np.transpose(v.reshape(node_positions[0].shape))
+
+    splot = axes.streamplot(X, Y, u, v, density=2)
+    splot.axes = axes
+    return splot
+
+
+def plot_3d_scatter(field, axes=None):
+    import matplotlib.pyplot as plt
+    from matplotlib import cm
+
+    if axes is None:
+        fig, axes = plt.subplots(subplot_kw={"projection": "3d"})
+
+    X, Y, Z = field.space.node_positions().numpy().T
+
+    # Make data.
+    f = field.dof_values.numpy().reshape(X.shape)
+
+    # Plot the surface.
+    return axes.scatter(X, Y, Z, c=f, cmap=cm.coolwarm)
+
+
+def plot_3d_velocities(field, axes=None):
+    import matplotlib.pyplot as plt
+
+    if axes is None:
+        fig, axes = plt.subplots(subplot_kw={"projection": "3d"})
+
+    X, Y, Z = field.space.node_positions().numpy().T
+
+    vel = field.dof_values.numpy()
+    u = np.ascontiguousarray(vel[:, 0])
+    v = np.ascontiguousarray(vel[:, 1])
+    w = np.ascontiguousarray(vel[:, 2])
+
+    u = u.reshape(X.shape)
+    v = v.reshape(X.shape)
+    w = w.reshape(X.shape)
+
+    return axes.quiver(X, Y, Z, u, v, w, length=1.0 / X.shape[0], normalize=False)
+
+
+class Plot:
+    def __init__(self, stage=None, default_point_radius=0.01):
+        self.default_point_radius = default_point_radius
+
+        self._surfaces = {}
+        self._surface_vectors = {}
+        self._volumes = {}
+
+        self._usd_renderer = None
+        if stage is not None:
+            try:
+                from warp.render import UsdRenderer
+
+                self._usd_renderer = UsdRenderer(stage)
+            except Exception as err:
+                print(f"Could not initialize UsdRenderer for stage '{stage}': {err}.")
+
+    def begin_frame(self, time):
+        if self._usd_renderer is not None:
+            self._usd_renderer.begin_frame(time=time)
+
+    def end_frame(self):
+        if self._usd_renderer is not None:
+            self._usd_renderer.end_frame()
+
+    def add_surface(self, name: str, field: DiscreteField):
+        if self._usd_renderer is not None:
+            points_2d = field.space.node_positions().numpy()
+            values = field.dof_values.numpy()
+            points_3d = np.hstack((points_2d, values.reshape(-1, 1)))
+
+            if hasattr(field.space, "node_triangulation"):
+                indices = field.space.node_triangulation()
+                self._usd_renderer.render_mesh(name, points=points_3d, indices=indices)
+            else:
+                self._usd_renderer.render_points(name, points=points_3d, radius=self.default_point_radius)
+
+        if name not in self._surfaces:
+            field_clone = field.space.make_field(space_partition=field.space_partition)
+            self._surfaces[name] = (field_clone, [])
+
+        self._surfaces[name][1].append(field.dof_values.numpy())
+
+    def add_surface_vector(self, name: str, field: DiscreteField):
+        if self._usd_renderer is not None:
+            points_2d = field.space.node_positions().numpy()
+            values = field.dof_values.numpy()
+            points_3d = np.hstack((points_2d + values, np.zeros_like(points_2d[:, 0]).reshape(-1, 1)))
+
+            if hasattr(field.space, "node_triangulation"):
+                indices = field.space.node_triangulation()
+                self._usd_renderer.render_mesh(name, points=points_3d, indices=indices)
+            else:
+                self._usd_renderer.render_points(name, points=points_3d, radius=self.default_point_radius)
+
+        if name not in self._surface_vectors:
+            field_clone = field.space.make_field(space_partition=field.space_partition)
+            self._surface_vectors[name] = (field_clone, [])
+
+        self._surface_vectors[name][1].append(field.dof_values.numpy())
+
+    def add_volume(self, name: str, field: DiscreteField):
+        if self._usd_renderer is not None:
+            points_3d = field.space.node_positions().numpy()
+            values = field.dof_values.numpy()
+
+            self._usd_renderer.render_points(name, points_3d, radius=values)
+
+        if name not in self._volumes:
+            field_clone = field.space.make_field(space_partition=field.space_partition)
+            self._volumes[name] = (field_clone, [])
+
+        self._volumes[name][1].append(field.dof_values.numpy())
+
+    def plot(self, streamlines: Set[str] = []):
+        return self._plot_matplotlib(streamlines)
+
+    def _plot_matplotlib(self, streamlines: Set[str] = []):
+        import matplotlib.pyplot as plt
+        import matplotlib.animation as animation
+
+        def make_animation(ax, field, values, plot_func, num_frames: int):
+            def animate(i):
+                ax.clear()
+                field.dof_values = values[i]
+                return plot_func(field, axes=ax)
+
+            return animation.FuncAnimation(
+                ax.figure,
+                animate,
+                interval=30,
+                blit=False,
+                frames=len(values),
+            )
+
+        for name, (field, values) in self._surfaces.items():
+            field.dof_values = values[0]
+            ax = plot_surface(field).axes
+
+            if len(values) > 1:
+                anim = make_animation(ax, field, values, plot_func=plot_surface, num_frames=len(values))
+
+        for name, (field, values) in self._surface_vectors.items():
+            field.dof_values = values[0]
+            if name in streamlines and hasattr(field.space, "node_grid"):
+                ax = plot_grid_streamlines(field).axes
+            else:
+                ax = plot_velocities(field).axes
+
+                if len(values) > 1:
+                    anim = make_animation(ax, field, values, plot_func=plot_velocities, num_frames=len(values))
+
+        for name, (field, values) in self._volumes.items():
+            field.dof_values = values[0]
+            ax = plot_3d_scatter(field).axes
+
+        plt.show()