Spaces:
Sleeping
Sleeping
File size: 4,880 Bytes
66c9c8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """
This example simulates a convection-diffusion PDE using semi-Lagrangian advection
D phi / dt - nu d2 phi / dx^2 = 0
"""
import argparse
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
return 0.0
@wp.func
def velocity(pos: wp.vec2, ang_vel: float):
center = wp.vec2(0.5, 0.5)
offset = pos - center
return wp.vec2(offset[1], -offset[0]) * ang_vel
@fem.integrand
def inertia_form(s: fem.Sample, phi: fem.Field, psi: fem.Field, dt: float):
return phi(s) * psi(s) / dt
@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)
# semi-Lagrangian advection; evaluate phi upstream
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(fem.lookup(domain, conv_pos, s))
return conv_phi * psi(s) / dt
@fem.integrand
def diffusion_form(
s: fem.Sample,
u: fem.Field,
v: fem.Field,
):
return wp.dot(
fem.grad(u, s),
fem.grad(v, s),
)
@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)
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=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")
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
# right-hand-side -- advected inertia
rhs = fem.integrate(
transported_inertia_form,
fields={"phi": self._phi_field, "psi": self._test},
values={"ang_vel": self._args.ang_vel, "dt": self.sim_dt},
output_dtype=float,
)
# Solve linear system
bsr_cg(self._matrix, x=self._phi_field.dof_values, b=rhs, quiet=self._quiet, tol=1.0e-12)
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.init()
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.update()
example.render()
example.renderer.plot()
|