From 81d08cb0ac4c2f37589c46bf9d7354259118e079 Mon Sep 17 00:00:00 2001
From: Thomas Gibson <gibsonthomas1120@hotmail.com>
Date: Wed, 28 Apr 2021 21:49:24 -0500
Subject: [PATCH] Implement local gradient operator
---
grudge/geometry/metrics.py | 13 +++++++---
grudge/op.py | 50 ++++++++++++++++++++++++++++++++++++++
2 files changed, 59 insertions(+), 4 deletions(-)
diff --git a/grudge/geometry/metrics.py b/grudge/geometry/metrics.py
index 17d07a06..4e7b07ff 100644
--- a/grudge/geometry/metrics.py
+++ b/grudge/geometry/metrics.py
@@ -168,8 +168,12 @@ def inverse_metric_derivative(actx, dcoll, rst_axis, xyz_axis, dd):
for rst_axis in range(dim)
).inv()
- return (outprod_with_unit(xyz_axis, rst_axis)
- * volume_pseudoscalar_inv).as_scalar()
+ # NOTE: this always produces a DOFArray of shape (1,).
+ # the thing we want is actually inside
+ result, = (outprod_with_unit(xyz_axis, rst_axis)
+ * volume_pseudoscalar_inv).as_scalar()
+
+ return result
@memoize_on_first_arg
@@ -186,7 +190,6 @@ def inverse_surface_metric_derivative(actx, dcoll, rst_axis, xyz_axis, dd=None):
actx, dcoll, d, rst_axis, dd=dd
) for d in range(dim)
)
-
return imd
@@ -236,7 +239,9 @@ def parametrization_derivative(actx, dcoll, dd):
)
-def pseudoscalar(actx, dcoll, dd):
+def pseudoscalar(actx, dcoll, dd=None):
+ if dd is None:
+ dd = DD_VOLUME
return parametrization_derivative(
actx, dcoll, dd
diff --git a/grudge/op.py b/grudge/op.py
index 8acf3c83..bce19ed1 100644
--- a/grudge/op.py
+++ b/grudge/op.py
@@ -161,6 +161,56 @@ def local_grad(dcoll, vec):
return _bound_grad(dcoll)(u=vec)
+def reference_derivative_matrices(actx, element_group):
+ @keyed_memoize_in(
+ actx, reference_derivative_matrices,
+ lambda grp: grp.discretization_key())
+ def get_ref_derivative_mats(grp):
+ from meshmode.discretization.poly_element import diff_matrices
+ return tuple(
+ actx.freeze(actx.from_numpy(dfmat))
+ for dfmat in diff_matrices(grp)
+ )
+
+ return get_ref_derivative_mats(element_group)
+
+
+def _compute_local_gradient(dcoll, vec, xyz_axis):
+ from grudge.geometry import \
+ inverse_surface_metric_derivative
+
+ discr = dcoll.discr_from_dd(dof_desc.DD_VOLUME)
+ actx = vec.array_context
+
+ inverse_jac_t = [
+ inverse_surface_metric_derivative(
+ actx, dcoll, rst_axis, xyz_axis
+ ) for rst_axis in range(dcoll.dim)
+ ]
+
+ data = []
+ for grp, vec_i in zip(discr.groups, vec):
+ data.append(sum(
+ actx.einsum(
+ "ij,ej,ej->ei",
+ reference_derivative_matrices(actx, grp)[rst_axis],
+ vec_i,
+ inverse_jac_t[rst_axis],
+ arg_names=("ref_diff_mat", "vec", "inv_jac_t"),
+ tagged=(HasElementwiseMatvecTag(),)
+ ) for rst_axis in range(dcoll.dim))
+ )
+ return DOFArray(actx, data=tuple(data))
+
+
+def local_gradient(dcoll, vec):
+ from pytools.obj_array import make_obj_array
+ return make_obj_array(
+ [_compute_local_gradient(dcoll, vec, xyz_axis)
+ for xyz_axis in range(dcoll.dim)]
+ )
+
+
@memoize_on_first_arg
def _bound_d_dx(dcoll, xyz_axis):
return bind(dcoll, sym.nabla(dcoll.dim)[xyz_axis] * sym.Variable("u"),
--
GitLab