From cf8035b12aa69ca59628b587fef7534e0751e7d5 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Sun, 29 Aug 2021 23:36:31 -0500
Subject: [PATCH] Memoize normals instead of recomputing them, tweak other
metric memoizations
---
grudge/geometry/metrics.py | 97 +++++++++++++++++++++-----------------
1 file changed, 53 insertions(+), 44 deletions(-)
diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py
index fe0030dd..919bdf46 100644
--- a/grudge/geometry/metrics.py
+++ b/grudge/geometry/metrics.py
@@ -58,10 +58,12 @@ THE SOFTWARE.
import numpy as np
-from grudge import DiscretizationCollection
from arraycontext import thaw, freeze, ArrayContext
from meshmode.dof_array import DOFArray
+from grudge import DiscretizationCollection
+import grudge.dof_desc as dof_desc
+
from grudge.dof_desc import (
DD_VOLUME, DOFDesc, DISCR_TAG_BASE
)
@@ -72,6 +74,10 @@ from pytools.obj_array import make_obj_array
from pytools import memoize_in
+from arraycontext import register_multivector_as_array_container
+register_multivector_as_array_container()
+
+
# {{{ Metric computations
def forward_metric_nth_derivative(
@@ -388,9 +394,12 @@ def inverse_surface_metric_derivative(
dim = dcoll.dim
ambient_dim = dcoll.ambient_dim
+ if dd is None:
+ dd = DD_VOLUME
+ dd = dof_desc.as_dofdesc(dd)
+
@memoize_in(dcoll, (inverse_surface_metric_derivative, dd,
- "inv_metric_deriv_rst%s_xyz%s_adim%s_gdim%s"
- % (rst_axis, xyz_axis, ambient_dim, dim)))
+ rst_axis, xyz_axis))
def _inv_surf_metric_deriv():
if ambient_dim == dim:
imd = inverse_metric_derivative(
@@ -404,6 +413,7 @@ def inverse_surface_metric_derivative(
) for d in range(dim)
)
return freeze(imd, actx)
+
return thaw(_inv_surf_metric_deriv(), actx)
@@ -492,11 +502,7 @@ def area_element(
if dd is None:
dd = DD_VOLUME
- dim = dcoll.discr_from_dd(dd).dim
-
- @memoize_in(dcoll, (area_element, dd,
- "area_elements_adim%s_gdim%s"
- % (dcoll.ambient_dim, dim)))
+ @memoize_in(dcoll, (area_element, dd, actx.supports_nonscalar_broadcasting))
def _area_elements():
return freeze(actx.np.sqrt(
pseudoscalar(actx, dcoll, dd=dd).norm_squared()), actx)
@@ -516,7 +522,6 @@ def rel_mv_normal(
:arg dd: a :class:`~grudge.dof_desc.DOFDesc`, or a value convertible to one.
"""
- import grudge.dof_desc as dof_desc
dd = dof_desc.as_dofdesc(dd)
@@ -543,46 +548,50 @@ def mv_normal(
:returns: a :class:`~pymbolic.geometric_algebra.MultiVector`
containing the unit normals.
"""
- import grudge.dof_desc as dof_desc
-
dd = dof_desc.as_dofdesc(dd)
- dim = dcoll.discr_from_dd(dd).dim
- ambient_dim = dcoll.ambient_dim
+ @memoize_in(dcoll, (mv_normal, dd))
+ def _normal():
+ dim = dcoll.discr_from_dd(dd).dim
+ ambient_dim = dcoll.ambient_dim
- if dim == ambient_dim:
- raise ValueError("may only request normals on domains whose topological "
- f"dimension ({dim}) differs from "
- f"their ambient dimension ({ambient_dim})")
-
- if dim == ambient_dim - 1:
- return rel_mv_normal(actx, dcoll, dd=dd)
-
- # NOTE: In the case of (d - 2)-dimensional curves, we don't really have
- # enough information on the face to decide what an "exterior face normal"
- # is (e.g the "normal" to a 1D curve in 3D space is actually a
- # "normal plane")
- #
- # The trick done here is that we take the surface normal, move it to the
- # face and then take a cross product with the face tangent to get the
- # correct exterior face normal vector.
- assert dim == ambient_dim - 2
-
- from grudge.op import project
- import grudge.dof_desc as dof_desc
-
- volm_normal = MultiVector(
- project(dcoll, dof_desc.DD_VOLUME, dd,
- rel_mv_normal(
- actx, dcoll,
- dd=dof_desc.DD_VOLUME
- ).as_vector(dtype=object))
- )
- pder = pseudoscalar(actx, dcoll, dd=dd)
+ if dim == ambient_dim:
+ raise ValueError("may only request normals on domains whose topological "
+ f"dimension ({dim}) differs from "
+ f"their ambient dimension ({ambient_dim})")
+
+ if dim == ambient_dim - 1:
+ result = rel_mv_normal(actx, dcoll, dd=dd)
+ else:
+ # NOTE: In the case of (d - 2)-dimensional curves, we don't really have
+ # enough information on the face to decide what an "exterior face normal"
+ # is (e.g the "normal" to a 1D curve in 3D space is actually a
+ # "normal plane")
+ #
+ # The trick done here is that we take the surface normal, move it to the
+ # face and then take a cross product with the face tangent to get the
+ # correct exterior face normal vector.
+ assert dim == ambient_dim - 2
+
+ from grudge.op import project
+
+ volm_normal = MultiVector(
+ project(dcoll, dof_desc.DD_VOLUME, dd,
+ rel_mv_normal(
+ actx, dcoll,
+ dd=dof_desc.DD_VOLUME
+ ).as_vector(dtype=object))
+ )
+ pder = pseudoscalar(actx, dcoll, dd=dd)
+
+ mv = -(volm_normal ^ pder) << volm_normal.I.inv()
+
+ result = mv / actx.np.sqrt(mv.norm_squared())
- mv = -(volm_normal ^ pder) << volm_normal.I.inv()
+ return freeze(result, actx)
- return mv / actx.np.sqrt(mv.norm_squared())
+ n = _normal()
+ return thaw(n, actx)
def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd):
--
GitLab