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

warp/examples/fem/example_convection_diffusion_dg.py

@@ -0,0 +1,209 @@
+# 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 Convection Diffusion DG
+#
+# This example simulates a convection-diffusion PDE using Discontinuous
+# Galerkin with upwind transport and Symmetric Interior Penalty
+#
+# D phi / dt - nu d2 phi / dx^2 = 0
+###########################################################################
+
+import argparse
+
+import warp as wp
+import warp.fem as fem
+
+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_trimesh, gen_quadmesh
+    from .plot_utils import Plot
+    from .example_convection_diffusion import (
+        initial_condition,
+        velocity,
+        inertia_form,
+        diffusion_form,
+    )
+except ImportError:
+    from bsr_utils import bsr_cg
+    from mesh_utils import gen_trimesh, gen_quadmesh
+    from plot_utils import Plot
+    from example_convection_diffusion import (
+        initial_condition,
+        velocity,
+        inertia_form,
+        diffusion_form,
+    )
+
+wp.init()
+
+
+# Standard transport term, on cells' interior
+@fem.integrand
+def transport_form(s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi: fem.Field, ang_vel: float):
+    pos = domain(s)
+    vel = velocity(pos, ang_vel)
+
+    return psi(s) * wp.dot(fem.grad(phi, s), vel)
+
+
+# Upwind flux, on cell sides
+@fem.integrand
+def upwind_transport_form(s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi: fem.Field, ang_vel: float):
+    pos = domain(s)
+    vel = velocity(pos, ang_vel)
+    vel_n = wp.dot(vel, fem.normal(domain, s))
+
+    return fem.jump(phi, s) * (-fem.average(psi, s) * vel_n + 0.5 * fem.jump(psi, s) * wp.abs(vel_n))
+
+
+# Symmetric-Interior-Penalty diffusion term (See Pietro Ern 2012)
+@fem.integrand
+def sip_diffusion_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    psi: fem.Field,
+    phi: fem.Field,
+):
+    nor = fem.normal(domain, s)
+    penalty = fem.measure_ratio(domain, s) * float(fem.degree(psi) * fem.degree(phi))
+
+    return penalty * fem.jump(phi, s) * fem.jump(psi, s) - (
+        wp.dot(fem.grad_average(phi, s), nor) * fem.jump(psi, s)
+        + wp.dot(fem.grad_average(psi, s), nor) * fem.jump(phi, s)
+    )
+
+
+class Example:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resolution", type=int, default=50)
+    parser.add_argument("--degree", type=int, default=2)
+    parser.add_argument("--num_frames", type=int, default=100)
+    parser.add_argument("--viscosity", type=float, default=0.001)
+    parser.add_argument("--ang_vel", type=float, default=1.0)
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+
+    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.mesh == "tri":
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(args.resolution))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif args.mesh == "quad":
+            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(args.resolution))
+            geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(args.resolution))
+
+        domain = fem.Cells(geometry=geo)
+        sides = fem.Sides(geo)
+        scalar_space = fem.make_polynomial_space(
+            geo,
+            discontinuous=True,
+            degree=args.degree,
+            family=fem.Polynomial.GAUSS_LEGENDRE,
+        )
+
+        # Assemble transport, diffusion and inertia matrices
+
+        self._test = fem.make_test(space=scalar_space, domain=domain)
+        trial = fem.make_trial(space=scalar_space, domain=domain)
+
+        matrix_inertia = fem.integrate(
+            inertia_form,
+            fields={"phi": trial, "psi": self._test},
+            values={"dt": self.sim_dt},
+        )
+
+        matrix_transport = fem.integrate(
+            transport_form,
+            fields={"phi": trial, "psi": self._test},
+            values={"ang_vel": args.ang_vel},
+        )
+
+        side_test = fem.make_test(space=scalar_space, domain=sides)
+        side_trial = fem.make_trial(space=scalar_space, domain=sides)
+
+        bsr_axpy(
+            fem.integrate(
+                upwind_transport_form,
+                fields={"phi": side_trial, "psi": side_test},
+                values={"ang_vel": args.ang_vel},
+            ),
+            y=matrix_transport,
+        )
+
+        matrix_diffusion = fem.integrate(
+            diffusion_form,
+            fields={"u": trial, "v": self._test},
+        )
+        bsr_axpy(
+            fem.integrate(
+                sip_diffusion_form,
+                fields={"phi": side_trial, "psi": side_test},
+            ),
+            y=matrix_diffusion,
+        )
+
+        self._matrix = matrix_inertia
+        bsr_axpy(x=matrix_transport, y=self._matrix)
+        bsr_axpy(x=matrix_diffusion, y=self._matrix, alpha=args.viscosity)
+
+        # Initial condition
+        self._phi_field = scalar_space.make_field()
+        fem.interpolate(initial_condition, dest=self._phi_field)
+
+        self.renderer = Plot(stage)
+        self.renderer.add_surface("phi", self._phi_field)
+
+
+    def step(self):
+        self.current_frame += 1
+
+        rhs = fem.integrate(
+            inertia_form,
+            fields={"phi": self._phi_field, "psi": self._test},
+            values={"dt": self.sim_dt},
+        )
+
+        phi = wp.zeros_like(rhs)
+        bsr_cg(self._matrix, b=rhs, x=phi, method='bicgstab', quiet=self._quiet)
+
+        wp.utils.array_cast(in_array=phi, out_array=self._phi_field.dof_values)
+
+    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()
+
+
+if __name__ == "__main__":
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    for k in range(args.num_frames):
+        print(f"Frame {k}:")
+        example.step()
+        example.render()
+
+    example.renderer.plot()

warp/examples/fem/example_convection_diffusion_dg0.py

@@ -0,0 +1,194 @@
+# 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 Convection Diffusion DG0
+#
+# This example simulates a convection-diffusion PDE using 
+# FVM with upwind transport
+#
+# D phi / dt + nu Div f = 0
+# f = grad phi
+###########################################################################
+
+import argparse
+
+import warp as wp
+import warp.fem as fem
+
+from warp.sparse import bsr_mm, bsr_axpy, bsr_transposed
+
+wp.init()
+
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .bsr_utils import bsr_cg, invert_diagonal_bsr_mass_matrix
+    from .plot_utils import Plot
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+    from .example_convection_diffusion import initial_condition, velocity, inertia_form
+except ImportError:
+    from bsr_utils import bsr_cg, invert_diagonal_bsr_mass_matrix
+    from plot_utils import Plot
+    from mesh_utils import gen_trimesh, gen_quadmesh
+    from example_convection_diffusion import initial_condition, velocity, inertia_form
+
+
+@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 half_diffusion_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    psi: fem.Field,
+    u: fem.Field,
+):
+    return fem.jump(psi, s) * wp.dot(fem.average(u, s), fem.normal(domain, s))
+
+
+@fem.integrand
+def upwind_transport_form(s: fem.Sample, domain: fem.Domain, phi: fem.Field, psi: fem.Field, ang_vel: float):
+    pos = domain(s)
+
+    vel = velocity(pos, ang_vel)
+
+    vel_n = wp.dot(vel, fem.normal(domain, s))
+
+    return fem.jump(psi, s) * (fem.average(phi, s) * vel_n + 0.5 * fem.jump(phi, s) * wp.abs(vel_n))
+
+
+class Example:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resolution", type=int, default=50)
+    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("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+
+    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.mesh == "tri":
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(args.resolution))
+            geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif args.mesh == "quad":
+            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(args.resolution))
+            geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            geo = fem.Grid2D(res=wp.vec2i(args.resolution))
+
+        domain = fem.Cells(geometry=geo)
+        sides = fem.Sides(geo)
+        scalar_space = fem.make_polynomial_space(geo, degree=0)
+
+        # Inertia matrix
+        self._test = fem.make_test(space=scalar_space, domain=domain)
+        trial = fem.make_trial(space=scalar_space, domain=domain)
+        matrix_inertia = fem.integrate(
+            inertia_form,
+            fields={"phi": trial, "psi": self._test},
+            values={"dt": self.sim_dt},
+        )
+
+        # Upwind transport term
+        side_test = fem.make_test(space=scalar_space, domain=sides)
+        side_trial = fem.make_trial(space=scalar_space, domain=sides)
+        matrix_transport = fem.integrate(
+            upwind_transport_form,
+            fields={"phi": side_trial, "psi": side_test},
+            values={"ang_vel": args.ang_vel},
+        )
+
+        # Diffusion bilinear form
+        # Since we have piecewise constant element, we cannot use the classical diffusion form
+        # Instead we assemble the matrix B M^-1 B^T, with B associated to the form psi div(u)
+        # and the diagonal matrix M to the velocity mass form u.v
+
+        velocity_space = fem.make_polynomial_space(geo, degree=0, dtype=wp.vec2)
+        side_trial_vel = fem.make_trial(space=velocity_space, domain=sides)
+        matrix_half_diffusion = fem.integrate(
+            half_diffusion_form,
+            fields={"psi": side_test, "u": side_trial_vel},
+        )
+
+        # Diagonal velocity mass matrix
+        test_vel = fem.make_test(space=velocity_space, domain=domain)
+        trial_vel = fem.make_trial(space=velocity_space, domain=domain)
+        inv_vel_mass_matrix = fem.integrate(
+            vel_mass_form, domain=domain, fields={"u": trial_vel, "v": test_vel}, nodal=True
+        )
+        invert_diagonal_bsr_mass_matrix(inv_vel_mass_matrix)
+
+        # Assemble system matrix
+
+        self._matrix = matrix_inertia
+        # matrix += matrix_transport
+        bsr_axpy(x=matrix_transport, y=self._matrix)
+        # matrix += nu * B M^-1 B^T
+        bsr_mm(
+            x=bsr_mm(matrix_half_diffusion, inv_vel_mass_matrix),
+            y=bsr_transposed(matrix_half_diffusion),
+            z=self._matrix,
+            alpha=args.viscosity,
+            beta=1.0,
+        )
+
+        # Initial condition
+        self._phi_field = scalar_space.make_field()
+        fem.interpolate(initial_condition, dest=self._phi_field)
+
+        self.renderer = Plot(stage)
+        self.renderer.add_surface("phi", self._phi_field)
+
+    def step(self):
+        self.current_frame += 1
+
+        rhs = fem.integrate(
+            inertia_form,
+            fields={"phi": self._phi_field, "psi": self._test},
+            values={"dt": self.sim_dt},
+        )
+
+        phi = wp.zeros_like(rhs)
+        bsr_cg(self._matrix, b=rhs, x=phi, method='bicgstab', quiet=self._quiet)
+
+        wp.utils.array_cast(in_array=phi, out_array=self._phi_field.dof_values)
+
+    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()
+
+
+if __name__ == "__main__":
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    for k in range(args.num_frames):
+        print(f"Frame {k}:")
+        example.step()
+        example.render()
+
+    example.renderer.plot()

warp/examples/fem/example_deformed_geometry.py

@@ -0,0 +1,159 @@
+# 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 Deformed Geometry
+#
+# This example solves a 2d diffusion problem:
+# 
+# nu Div u = 1
+#
+# with Dirichlet boundary conditions on vertical edges and 
+# homogeneous Neumann on horizontal edges.
+###########################################################################
+
+import argparse
+
+import warp as wp
+import warp.fem as fem
+
+from warp.sparse import bsr_axpy
+from warp.fem.utils import array_axpy
+
+# Import example utilities
+# Make sure that works both when imported as module and run as standalone file
+try:
+    from .example_diffusion import linear_form, diffusion_form
+    from .bsr_utils import bsr_cg
+    from .mesh_utils import gen_trimesh, gen_quadmesh
+    from .plot_utils import Plot
+except ImportError:
+    from example_diffusion import linear_form, diffusion_form 
+    from bsr_utils import bsr_cg
+    from mesh_utils import gen_trimesh, gen_quadmesh
+    from plot_utils import Plot
+
+wp.init()
+
+
+@fem.integrand
+def deformation_field_expr(
+    s: fem.Sample,
+    domain: fem.Domain,
+):
+    """
+    Deformation field mapping the unique square to a circular band
+    """
+    x = domain(s)
+
+    r = x[1] + 0.5
+    t = 0.5 * 3.1416 * x[0]
+
+    return r * wp.vec2(wp.sin(t), wp.cos(t)) - x
+
+
+@fem.integrand
+def boundary_projector_form(
+    s: fem.Sample,
+    domain: fem.Domain,
+    u: fem.Field,
+    v: fem.Field,
+):
+    """
+    Bilinear boundary condition projector form, non-zero on radial edges
+    """
+    nor = fem.normal(domain, s)
+    active = wp.select(nor[0] < -0.9999 or nor[1] < -0.9999, 0.0, 1.0)
+    return active * u(s) * v(s)
+
+
+class Example:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resolution", type=int, default=50)
+    parser.add_argument("--degree", type=int, default=2)
+    parser.add_argument("--serendipity", action="store_true", default=False)
+    parser.add_argument("--viscosity", type=float, default=2.0)
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type")
+
+    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
+
+        # Grid or triangle mesh geometry
+        if args.mesh == "tri":
+            positions, tri_vidx = gen_trimesh(res=wp.vec2i(args.resolution))
+            base_geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+        elif args.mesh == "quad":
+            positions, quad_vidx = gen_quadmesh(res=wp.vec2i(args.resolution))
+            base_geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+        else:
+            base_geo = fem.Grid2D(res=wp.vec2i(args.resolution))
+
+        # Construct deformation field on base geometry
+        deformation_space = fem.make_polynomial_space(base_geo, degree=args.degree, dtype=wp.vec2)
+        deformation_field = deformation_space.make_field()
+        fem.interpolate(deformation_field_expr, dest=deformation_field)
+
+        self._geo = deformation_field.make_deformed_geometry()
+
+        # Scalar function space on deformed geometry
+        element_basis = fem.ElementBasis.SERENDIPITY if args.serendipity else None
+        self._scalar_space = fem.make_polynomial_space(self._geo, degree=args.degree, element_basis=element_basis)
+
+        # Scalar field over our function space
+        self._scalar_field = self._scalar_space.make_field()
+
+        self.renderer = Plot(stage)
+
+    def step(self):
+        args = self._args
+        geo = self._geo
+
+        domain = fem.Cells(geometry=geo)
+
+        # Right-hand-side (forcing term)
+        test = fem.make_test(space=self._scalar_space, domain=domain)
+        rhs = fem.integrate(linear_form, fields={"v": test})
+
+        # Diffusion form
+        trial = fem.make_trial(space=self._scalar_space, domain=domain)
+        matrix = fem.integrate(diffusion_form, fields={"u": trial, "v": test}, values={"nu": args.viscosity})
+
+        # Weakly-imposed boundary conditions on all sides
+        boundary = fem.BoundarySides(geo)
+        bd_test = fem.make_test(space=self._scalar_space, domain=boundary)
+        bd_trial = fem.make_trial(space=self._scalar_space, domain=boundary)
+
+        bd_matrix = fem.integrate(boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
+
+        fem.project_linear_system(matrix, rhs, bd_matrix)
+
+        # 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)
+
+        # Assign system result to our discrete field
+        self._scalar_field.dof_values = x
+
+    def render(self):
+        self.renderer.add_surface("solution", self._scalar_field)
+
+
+if __name__ == "__main__":
+    wp.set_module_options({"enable_backward": False})
+
+    args = Example.parser.parse_args()
+
+    example = Example(args=args)
+    example.step()
+    example.render()
+
+    example.renderer.plot()