warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/examples/fem/example_stokes_transfer.py

@@ -0,0 +1,262 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# 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 numpy as np
+
+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.utils import array_cast
+
+
+@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, quiet=False, resolution=50):
+        self._quiet = quiet
+
+        self.res = resolution
+        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 = fem_example_utils.Plot()
+
+    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
+        u_matrix += u_bd_matrix * self.bd_strength
+
+        gradient_matrix = div_matrix.transpose() @ inv_p_mass_matrix
+        u_matrix += gradient_matrix @ div_matrix / self.compliance
+
+        # scale u_rhs
+        array_axpy(u_rhs, u_rhs, alpha=0.0, beta=self.bd_strength)
+
+        # Solve for displacement
+        u_res = wp.zeros_like(u_rhs)
+        fem_example_utils.bsr_cg(u_matrix, x=u_res, b=u_rhs, quiet=self._quiet)
+
+        # Compute pressure from displacement
+        div_u = div_matrix @ u_res
+        p_res = -inv_p_mass_matrix @ div_u
+
+        # 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_field("pressure", self._p_field)
+        self.renderer.add_field("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__":
+    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(
+        "--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, resolution=args.resolution)
+        example.step()
+        example.render()
+
+        if not args.headless:
+            example.renderer.plot(options={"velocity": {"streamlines": {}}})

warp/examples/fem/example_streamlines.py

@@ -0,0 +1,352 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+###########################################################################
+# Example Streamlines
+#
+# Shows how to generate 3D streamlines by tracing through a velocity field
+# using the `warp.fem.lookup` operator.
+# Also illustrates using `warp.fem.Subdomain` to define subsets of elements.
+#
+###########################################################################
+
+import numpy as np
+
+import warp as wp
+import warp.examples.fem.utils as fem_example_utils
+import warp.fem as fem
+from warp.examples.fem.example_apic_fluid import divergence_form, solve_incompressibility
+
+
+@fem.integrand
+def classify_boundary_sides(
+    s: fem.Sample,
+    domain: fem.Domain,
+    outflow: wp.array(dtype=int),
+    freeslip: wp.array(dtype=int),
+    inflow: wp.array(dtype=int),
+):
+    x = fem.position(domain, s)
+    n = fem.normal(domain, s)
+
+    if n[0] < -0.5:
+        # left side
+        inflow[s.qp_index] = 1
+    elif n[0] > 0.5:
+        if x[1] > 0.33 or x[2] < 0.33:
+            # right side, top
+            freeslip[s.qp_index] = 1
+        else:
+            # right side, bottom
+            outflow[s.qp_index] = 1
+    else:
+        freeslip[s.qp_index] = 1
+
+
+@fem.integrand
+def inflow_velocity(
+    s: fem.Sample,
+    domain: fem.Domain,
+):
+    n = fem.normal(domain, s)
+    return -n
+
+
+@fem.integrand
+def noslip_projector_form(
+    s: fem.Sample,
+    u: fem.Field,
+    v: fem.Field,
+):
+    return wp.dot(u(s), v(s))
+
+
+@fem.integrand
+def freeslip_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    n = fem.normal(domain, s)
+    return wp.dot(u(s), n) * wp.dot(n, 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 spawn_streamlines(s: fem.Sample, domain: fem.Domain, jitter: float):
+    rng = wp.rand_init(s.qp_index)
+    random_offset = wp.vec3(wp.randf(rng), wp.randf(rng), wp.randf(rng)) - wp.vec3(0.5)
+
+    # remove jistter along normal
+    n = fem.normal(domain, s)
+    random_offset -= wp.dot(random_offset, n) * n
+
+    return domain(s) + jitter * random_offset
+
+
+@fem.integrand
+def gen_streamlines(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    spawn_points: wp.array(dtype=wp.vec3),
+    point_count: int,
+    dx: float,
+    pos: wp.array2d(dtype=wp.vec3),
+    speed: wp.array2d(dtype=float),
+):
+    idx = s.qp_index
+
+    p = spawn_points[idx]
+    s = fem.lookup(domain, p)
+    for k in range(point_count):
+        v = u(s)
+        pos[idx, k] = p
+        speed[idx, k] = wp.length(v)
+
+        flow_dir = wp.normalize(v)
+        adv_p = p + flow_dir * dx
+        adv_s = fem.lookup(domain, adv_p, s)
+
+        if adv_s.element_index != fem.NULL_ELEMENT_INDEX:
+            # if the lookup result position is different from adv_p,
+            # it means we have been projected back onto the domain;
+            # align back with flow and terminate streamline
+            new_p = domain(adv_s)
+            if wp.length_sq(new_p - adv_p) > 0.000001:
+                p = p + wp.dot(new_p - p, flow_dir) * flow_dir
+                s = fem.lookup(domain, p, s)
+                dx = 0.0
+            else:
+                s = adv_s
+                p = new_p
+
+
+class Example:
+    def __init__(self, quiet=False, degree=2, resolution=16, mesh="grid", headless: bool = False):
+        self._quiet = quiet
+        self._degree = degree
+
+        self._streamline_dx = 0.5 / resolution
+        self._streamline_point_count = 4 * resolution
+
+        res = wp.vec3i(resolution)
+
+        if mesh == "tet":
+            pos, tet_vtx_indices = fem_example_utils.gen_tetmesh(
+                res=res,
+            )
+            self._geo = fem.Tetmesh(tet_vtx_indices, pos, build_bvh=True)
+        elif mesh == "nano":
+            volume = fem_example_utils.gen_volume(
+                res=res,
+            )
+            self._geo = fem.Nanogrid(volume)
+        else:
+            self._geo = fem.Grid3D(
+                res=res,
+            )
+
+        # Mark sides with boundary conditions that should apply
+        boundary = fem.BoundarySides(self._geo)
+        inflow_mask = wp.zeros(shape=boundary.element_count(), dtype=int)
+        freeslip_mask = wp.zeros(shape=boundary.element_count(), dtype=int)
+        outflow_mask = wp.zeros(shape=boundary.element_count(), dtype=int)
+
+        fem.interpolate(
+            classify_boundary_sides,
+            quadrature=fem.RegularQuadrature(boundary, order=0),
+            values={"outflow": outflow_mask, "freeslip": freeslip_mask, "inflow": inflow_mask},
+        )
+
+        self._inflow = fem.Subdomain(boundary, element_mask=inflow_mask)
+        self._freeslip = fem.Subdomain(boundary, element_mask=freeslip_mask)
+        self._outflow = fem.Subdomain(boundary, element_mask=outflow_mask)
+
+        self.plot = fem_example_utils.Plot()
+
+        self.renderer = None
+        if not headless:
+            try:
+                self.renderer = wp.render.OpenGLRenderer(
+                    camera_pos=(2.0, 0.5, 3.0),
+                    camera_front=(-0.66, 0.0, -1.0),
+                    draw_axis=False,
+                )
+            except Exception as err:
+                wp.utils.warn(f"Could not initialize OpenGL renderer: {err}")
+                pass
+
+    def step(self):
+        self._generate_incompressible_flow()
+
+        # first generate spawn points for the streamlines
+        # we do this by regularly sampling the inflow boundary with a small amount of jitter
+        streamline_spawn = fem.RegularQuadrature(
+            domain=self._inflow, order=self._degree, family=fem.Polynomial.GAUSS_LEGENDRE
+        )
+        n_streamlines = streamline_spawn.total_point_count()
+        spawn_points = wp.empty(dtype=wp.vec3, shape=n_streamlines)
+
+        jitter_amount = self._streamline_dx / self._degree
+        fem.interpolate(
+            spawn_streamlines, dest=spawn_points, quadrature=streamline_spawn, values={"jitter": jitter_amount}
+        )
+
+        # now forward-trace the velocity field to generate the streamlines
+        # here we use a fixed number of points per streamline, otherwise we would need to
+        # do a first pass to count points, then array_scan the offsets, then a second pass
+        # to populate the per-point data
+
+        point_count = self._streamline_point_count
+        points = wp.empty(dtype=wp.vec3, shape=(n_streamlines, point_count))
+        speed = wp.empty(dtype=float, shape=(n_streamlines, point_count))
+
+        fem.interpolate(
+            gen_streamlines,
+            domain=fem.Cells(self._geo),
+            dim=n_streamlines,
+            fields={"u": self.velocity_field},
+            values={
+                "spawn_points": spawn_points,
+                "point_count": self._streamline_point_count,
+                "dx": self._streamline_dx,
+                "pos": points,
+                "speed": speed,
+            },
+        )
+
+        self._points = points
+        self._speed = speed
+
+    def render(self):
+        # self.renderer.add_field("solution", self.pressure_field)
+        self.plot.add_field("pressure", self.pressure_field)
+        # self.plot.add_field("velocity", self.velocity_field)
+
+        if self.renderer is not None:
+            streamline_count = self._points.shape[0]
+            point_count = self._streamline_point_count
+
+            vertices = self._points.flatten().numpy()
+
+            line_offsets = np.arange(streamline_count) * point_count
+            indices_beg = np.arange(point_count - 1)[np.newaxis, :] + line_offsets[:, np.newaxis]
+            indices_end = indices_beg + 1
+            indices = np.vstack((indices_beg.flatten(), indices_end.flatten())).T.flatten()
+
+            colors = self._speed.numpy()[:, :-1].flatten()
+            colors = [wp.render.bourke_color_map(0.0, 3.0, c) for c in colors]
+
+            self.renderer.begin_frame(0)
+            self.renderer.render_line_list("streamlines", vertices, indices)
+            self.renderer.render_line_list("streamlines", vertices, indices, colors)
+
+            self.renderer.paused = True
+            self.renderer.end_frame()
+
+    def _generate_incompressible_flow(self):
+        # Function spaces for velocity and pressure (RT1 / P0)
+        u_space = fem.make_polynomial_space(
+            geo=self._geo, element_basis=fem.ElementBasis.RAVIART_THOMAS, degree=1, dtype=wp.vec3
+        )
+        p_space = fem.make_polynomial_space(geo=self._geo, degree=0, dtype=float)
+
+        self.pressure_field = p_space.make_field()
+        self.velocity_field = u_space.make_field()
+
+        # Boundary condition projector and matrices
+        inflow_test = fem.make_test(u_space, domain=self._inflow)
+        inflow_trial = fem.make_trial(u_space, domain=self._inflow)
+        dirichlet_projector = fem.integrate(
+            noslip_projector_form, fields={"u": inflow_test, "v": inflow_trial}, nodal=True, output_dtype=float
+        )
+
+        freeslip_test = fem.make_test(u_space, domain=self._freeslip)
+        freeslip_trial = fem.make_trial(u_space, domain=self._freeslip)
+        dirichlet_projector += fem.integrate(
+            freeslip_projector_form,
+            fields={"u": freeslip_test, "v": freeslip_trial},
+            nodal=True,
+            output_dtype=float,
+        )
+        fem.normalize_dirichlet_projector(dirichlet_projector)
+
+        # Initialize velocity field with BC
+        fem.interpolate(inflow_velocity, dest=fem.make_restriction(self.velocity_field, domain=self._inflow))
+
+        # (Diagonal) mass matrix
+        rho_test = fem.make_test(u_space)
+        rho_trial = fem.make_trial(u_space)
+        inv_mass_matrix = fem.integrate(
+            mass_form, fields={"u": rho_trial, "v": rho_test}, nodal=True, output_dtype=float
+        )
+        fem_example_utils.invert_diagonal_bsr_matrix(inv_mass_matrix)
+
+        # Assemble divergence operator matrix
+        p_test = fem.make_test(p_space)
+        u_trial = fem.make_trial(u_space)
+        divergence_matrix = fem.integrate(
+            divergence_form,
+            fields={"u": u_trial, "psi": p_test},
+            output_dtype=float,
+        )
+
+        # Solve incompressibility
+        solve_incompressibility(
+            divergence_matrix,
+            dirichlet_projector,
+            inv_mass_matrix.values,
+            self.pressure_field.dof_values,
+            self.velocity_field.dof_values,
+            quiet=self._quiet,
+        )
+
+
+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=8, help="Grid resolution.")
+    parser.add_argument("--degree", type=int, default=2, help="Polynomial degree of shape functions.")
+    parser.add_argument("--mesh", choices=("grid", "tet", "hex", "nano"), default="grid", 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, headless=args.headless
+        )
+
+        example.step()
+        example.render()