Skip to content

Commit

Permalink
include stretch factor test on torus
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexfikl
alexfikl committed Dec 12, 2021
1 parent 791d3c6 commit ce151e8
Showing 1 changed file with 88 additions and 34 deletions.
122 changes: 88 additions & 34 deletions test/test_symbolic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -337,50 +337,87 @@ def test_node_reduction(actx_factory):

# {{{ test_stretch_factor

def make_simplex_mesh(order):
def make_stretch_mesh(order, cls):
vertices = np.array([
[-1, -1, 0], [1, -1, 0], [-1, 1, 0], [1, 1, 0], [3, -1, 0], [3, 1, 0],
], dtype=np.float64).T
vertex_indices = np.array([
(0, 1, 2), (1, 3, 2),
(1, 4, 3), (4, 5, 3),
], dtype=np.int32)

from meshmode.mesh import SimplexElementGroup
import meshmode.mesh as mm
if issubclass(cls, mm.SimplexElementGroup):
vertex_indices = np.array([
(0, 1, 2), (1, 3, 2),
# (1, 4, 3), (4, 5, 3),
], dtype=np.int32)
elif issubclass(cls, mm.TensorProductElementGroup):
vertex_indices = np.array([
(0, 1, 2, 3),
# (1, 4, 3, 5)
], dtype=np.int32)
else:
raise ValueError

from meshmode.mesh.generation import make_group_from_vertices
grp = make_group_from_vertices(
vertices, vertex_indices, order,
unit_nodes=None,
group_cls=SimplexElementGroup)
group_cls=cls)

from meshmode.mesh import Mesh
return Mesh(vertices, [grp], is_conforming=True)
return mm.Mesh(vertices, [grp], is_conforming=True)


def make_square_mesh(order):
vertices = np.array([
[-1, -1, 0], [1, -1, 0], [-1, 1, 0], [1, 1, 0],
[3, -1, 0], [3, 1, 0],
], dtype=np.float64).T
vertex_indices = np.array([
(0, 1, 2, 3), (1, 4, 3, 5)
], dtype=np.int32)
def make_torus_mesh(order, cls, a=2.0, b=1.0, n_major=12, n_minor=6):
u, v = np.mgrid[0:2*np.pi:2*np.pi/n_major, 0:2*np.pi:2*np.pi/n_minor]

x = np.cos(u) * (a + b*np.cos(v))
y = np.sin(u) * (a + b*np.cos(v))
z = b * np.sin(v)
del u, v

vertices = (
np.vstack((x[np.newaxis], y[np.newaxis], z[np.newaxis]))
.transpose(0, 2, 1).copy().reshape(3, -1))

def idx(i, j):
return (i % n_major) + (j % n_minor) * n_major

import meshmode.mesh as mm
# i, j = 0, 0
i, j = 0, n_minor // 3
if issubclass(cls, mm.SimplexElementGroup):
vertex_indices = [
(idx(i, j), idx(i+1, j), idx(i, j+1)),
(idx(i+1, j), idx(i+1, j+1), idx(i, j+1)),
]
elif issubclass(cls, mm.TensorProductElementGroup):
vertex_indices = [(idx(i, j), idx(i+1, j), idx(i, j+1), idx(i+1, j+1))]
else:
raise TypeError(f"unsupported 'group_cls': {cls}")

from meshmode.mesh import TensorProductElementGroup
from meshmode.mesh.generation import make_group_from_vertices
vertex_indices = np.array(vertex_indices, dtype=np.int32)
grp = make_group_from_vertices(
vertices, vertex_indices, order,
unit_nodes=None,
group_cls=TensorProductElementGroup)
group_cls=cls)

# NOTE: project the nodes back to the torus surface
u = np.arctan2(grp.nodes[1], grp.nodes[0])
v = np.arctan2(
grp.nodes[2],
grp.nodes[0] * np.cos(u) + grp.nodes[1] * np.sin(u) - a)

from meshmode.mesh import Mesh
return Mesh(vertices, [grp], is_conforming=True)
nodes = np.empty_like(grp.nodes)
nodes[0] = np.cos(u) * (a + b*np.cos(v))
nodes[1] = np.sin(u) * (a + b*np.cos(v))
nodes[2] = b * np.sin(v)

return mm.Mesh(vertices, [grp.copy(nodes=nodes)], is_conforming=True)


def make_simplex_stretch_factors(ambient_dim):
from pytential.symbolic.primitives import \
_equilateral_parametrization_derivative_matrix
equi_pder = _equilateral_parametrization_derivative_matrix(ambient_dim)
# equi_pder = sym.parametrization_derivative_matrix(ambient_dim, ambient_dim - 1)
equi_form1 = sym.cse(equi_pder.T @ equi_pder, "pd_mat_jtj")

from pytential.symbolic.primitives import _small_mat_eigenvalues
Expand All @@ -406,16 +443,29 @@ def test_stretch_factor(actx_factory, order, visualize=False):
actx = actx_factory()

def wobble(x):
rx, ry = 2, 0.5
theta = np.pi / 4

result = np.empty_like(x)
result[0] = x[0]
result[1] = x[1]
result[2] = np.sin(x[1]) * np.sin(x[0])
result[0] = rx * (np.cos(theta) * x[0] - np.sin(theta) * x[1])
result[1] = np.sin(ry * (np.sin(theta) * x[0] + np.cos(theta) * x[1]))
result[2] = x[2]
# result[2] = np.sin(x[1]) * np.sin(x[0])

return result

from meshmode.mesh.processing import map_mesh
square_mesh = map_mesh(make_square_mesh(order), wobble)
simplex_mesh = map_mesh(make_simplex_mesh(order), wobble)
from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup
if True:
square_mesh = make_torus_mesh(order, TensorProductElementGroup)
simplex_mesh = make_torus_mesh(order, SimplexElementGroup)
else:
from meshmode.mesh.processing import map_mesh
square_mesh = map_mesh(
make_stretch_mesh(order, TensorProductElementGroup),
wobble)
simplex_mesh = map_mesh(
make_stretch_mesh(order, SimplexElementGroup),
wobble)

from meshmode.discretization import Discretization
import meshmode.discretization.poly_element as mpoly
Expand All @@ -429,6 +479,14 @@ def wobble(x):
q0, q1 = bind(square_discr,
make_quad_stretch_factors(square_discr.ambient_dim))(actx)

# s0 = (1 + (np.sqrt(3) / 4) / 4) * s0
# s1 = (1 + (np.sqrt(3) / 4) / 4) * s1

print(actx.to_numpy(actx.np.min(s0))[()], actx.to_numpy(actx.np.max(s0))[()])
print(actx.to_numpy(actx.np.min(q0))[()], actx.to_numpy(actx.np.max(q0))[()])
print(actx.to_numpy(actx.np.min(s1))[()], actx.to_numpy(actx.np.max(s1))[()])
print(actx.to_numpy(actx.np.min(q1))[()], actx.to_numpy(actx.np.max(q1))[()])

if not visualize:
return

Expand All @@ -450,13 +508,10 @@ def wobble(x):

xi = simplex_discr.groups[0].unit_nodes
for name, sv in zip(["s0", "s1"], [s0, s1]):
sv0 = actx.to_numpy(sv[0])[0].ravel()
sv1 = actx.to_numpy(sv[0])[0].ravel()
print(sv0, sv1)
sv = actx.to_numpy(sv[0])[0].ravel()

ax = fig.gca()
im = ax.tricontourf(xi[0], xi[1], sv0, levels=32)
im = ax.tricontourf(-xi[0], -xi[1], sv1, levels=32)
im = ax.tricontourf(xi[0], xi[1], sv, levels=32)
fig.colorbar(im, ax=ax)
fig.savefig(f"stretch_factor_simplex_{order:02d}_{name}")
fig.clf()
Expand All @@ -468,7 +523,6 @@ def wobble(x):
xi = square_discr.groups[0].unit_nodes
for name, sv in zip(["q0", "q1"], [q0, q1]):
sv = actx.to_numpy(sv[0])[0]
print(sv)

ax = fig.gca()
im = ax.tricontourf(xi[0], xi[1], sv, levels=32)
Expand Down

0 comments on commit ce151e8

Please sign in to comment.