Spaces:
Sleeping
Sleeping
File size: 4,956 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 | """
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
@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 update(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.init()
wp.set_module_options({"enable_backward": False})
args = Example.parser.parse_args()
example = Example(args=args)
example.update()
example.render()
example.renderer.plot()
|