Spaces:
Sleeping
Sleeping
File size: 4,971 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 | """
This example solves a 3d diffusion problem:
nu Div u = 1
with homogeneous Neumann conditions on horizontal sides and homogeneous Dirichlet boundary conditions other sides.
"""
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 .example_diffusion import diffusion_form, linear_form
from .bsr_utils import bsr_cg
from .mesh_utils import gen_tetmesh
from .plot_utils import Plot
except ImportError:
from example_diffusion import diffusion_form, linear_form
from bsr_utils import bsr_cg
from mesh_utils import gen_tetmesh, gen_hexmesh
from plot_utils import Plot
@fem.integrand
def vert_boundary_projector_form(
s: fem.Sample,
domain: fem.Domain,
u: fem.Field,
v: fem.Field,
):
# Non-zero mass on vertical sides only
w = 1.0 - wp.abs(fem.normal(domain, s)[1])
return w * u(s) * v(s)
class Example:
parser = argparse.ArgumentParser()
parser.add_argument("--resolution", type=int, default=10)
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("--boundary_compliance", type=float, default=0, help="Dirichlet boundary condition compliance")
parser.add_argument("--mesh", choices=("grid", "tet", "hex"), 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 = wp.vec3i(args.resolution, args.resolution // 2, args.resolution * 2)
if args.mesh == "tet":
pos, tet_vtx_indices = gen_tetmesh(
res=res,
bounds_lo=wp.vec3(0.0, 0.0, 0.0),
bounds_hi=wp.vec3(1.0, 0.5, 2.0),
)
self._geo = fem.Tetmesh(tet_vtx_indices, pos)
elif args.mesh == "hex":
pos, hex_vtx_indices = gen_hexmesh(
res=res,
bounds_lo=wp.vec3(0.0, 0.0, 0.0),
bounds_hi=wp.vec3(1.0, 0.5, 2.0),
)
self._geo = fem.Hexmesh(hex_vtx_indices, pos)
else:
self._geo = fem.Grid3D(
res=res,
bounds_lo=wp.vec3(0.0, 0.0, 0.0),
bounds_hi=wp.vec3(1.0, 0.5, 2.0),
)
# Domain and function spaces
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: fem.DiscreteField = self._scalar_space.make_field()
self.renderer = Plot(stage)
def update(self):
args = self._args
geo = self._geo
domain = fem.Cells(geometry=geo)
# Right-hand-side
test = fem.make_test(space=self._scalar_space, domain=domain)
rhs = fem.integrate(linear_form, fields={"v": test})
# Weakly-imposed boundary conditions on Y sides
with wp.ScopedTimer("Integrate"):
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(vert_boundary_projector_form, fields={"u": bd_trial, "v": bd_test}, nodal=True)
# 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})
if args.boundary_compliance == 0.0:
# Hard BC: project linear system
bd_rhs = wp.zeros_like(rhs)
fem.project_linear_system(matrix, rhs, bd_matrix, bd_rhs)
else:
# Weak BC: add toegether diffusion and boundary condition matrices
boundary_strength = 1.0 / args.boundary_compliance
bsr_axpy(x=bd_matrix, y=matrix, alpha=boundary_strength, beta=1)
with wp.ScopedTimer("CG solve"):
x = wp.zeros_like(rhs)
bsr_cg(matrix, b=rhs, x=x, quiet=self._quiet)
self._scalar_field.dof_values = x
def render(self):
self.renderer.add_volume("solution", self._scalar_field)
if __name__ == "__main__":
wp.init()
wp.set_module_options({"enable_backward": False})
args = Example.parser.parse_args()
example = Example(args=args)
example.update()
example.render()
example.renderer.plot()
|