Add support for CEED_EVAL_GRADIENT and CEED_TRANSPOSE in CeedBasisApplyAtPoints (#1388)

Author: zatkins-dev

interface/ceed-basis.c

@@ -1513,8 +1513,6 @@ int CeedBasisApplyAtPoints(CeedBasis basis, CeedInt num_points, CeedTransposeMod
 
   // Default implementation
   CeedCheck(basis->is_tensor_basis, basis->ceed, CEED_ERROR_UNSUPPORTED, "Evaluation at arbitrary points only supported for tensor product bases");
-  CeedCheck(eval_mode == CEED_EVAL_INTERP || t_mode == CEED_NOTRANSPOSE, basis->ceed, CEED_ERROR_UNSUPPORTED, "%s evaluation only supported for %s",
-            CeedEvalModes[eval_mode], CeedTransposeModes[CEED_NOTRANSPOSE]);
   if (!basis->basis_chebyshev) {
     // Build matrix mapping from quadrature point values to Chebyshev coefficients
     CeedScalar       *tau, *C, *I, *chebyshev_coeffs_1d;
@@ -1662,22 +1660,52 @@ int CeedBasisApplyAtPoints(CeedBasis basis, CeedInt num_points, CeedTransposeMod
       CeedCall(CeedVectorGetArrayWrite(basis->vec_chebyshev, CEED_MEM_HOST, &chebyshev_coeffs));
       CeedCall(CeedVectorGetArrayRead(x_ref, CEED_MEM_HOST, &x_array_read));
       CeedCall(CeedVectorGetArrayRead(u, CEED_MEM_HOST, &u_array));
-      {
-        CeedScalar tmp[2][num_comp * CeedIntPow(Q_1d, dim)], chebyshev_x[Q_1d];
-
-        // ---- Values at point
-        for (CeedInt p = 0; p < num_points; p++) {
-          CeedInt pre = num_comp * 1, post = 1;
-
-          for (CeedInt d = 0; d < dim; d++) {
-            // ------ Tensor contract with current Chebyshev polynomial values
-            CeedCall(CeedChebyshevPolynomialsAtPoint(x_array_read[p * dim + d], Q_1d, chebyshev_x));
-            CeedCall(CeedTensorContractApply(basis->contract, pre, 1, post, Q_1d, chebyshev_x, t_mode, p > 0 && d == (dim - 1),
-                                             d == 0 ? &u_array[p * num_comp] : tmp[d % 2], d == (dim - 1) ? chebyshev_coeffs : tmp[(d + 1) % 2]));
-            pre /= 1;
-            post *= Q_1d;
+
+      switch (eval_mode) {
+        case CEED_EVAL_INTERP: {
+          CeedScalar tmp[2][num_comp * CeedIntPow(Q_1d, dim)], chebyshev_x[Q_1d];
+
+          // ---- Values at point
+          for (CeedInt p = 0; p < num_points; p++) {
+            CeedInt pre = num_comp * 1, post = 1;
+
+            for (CeedInt d = 0; d < dim; d++) {
+              // ------ Tensor contract with current Chebyshev polynomial values
+              CeedCall(CeedChebyshevPolynomialsAtPoint(x_array_read[p * dim + d], Q_1d, chebyshev_x));
+              CeedCall(CeedTensorContractApply(basis->contract, pre, 1, post, Q_1d, chebyshev_x, t_mode, p > 0 && d == (dim - 1),
+                                               d == 0 ? &u_array[p * num_comp] : tmp[d % 2], d == (dim - 1) ? chebyshev_coeffs : tmp[(d + 1) % 2]));
+              pre /= 1;
+              post *= Q_1d;
+            }
+          }
+          break;
+        }
+        case CEED_EVAL_GRAD: {
+          CeedScalar tmp[2][num_comp * CeedIntPow(Q_1d, dim)], chebyshev_x[Q_1d];
+
+          // ---- Values at point
+          for (CeedInt p = 0; p < num_points; p++) {
+            // Dim**2 contractions, apply grad when pass == dim
+            for (CeedInt pass = 0; pass < dim; pass++) {
+              CeedInt pre = num_comp * 1, post = 1;
+
+              for (CeedInt d = 0; d < dim; d++) {
+                // ------ Tensor contract with current Chebyshev polynomial values
+                if (pass == d) CeedCall(CeedChebyshevDerivativeAtPoint(x_array_read[p * dim + d], Q_1d, chebyshev_x));
+                else CeedCall(CeedChebyshevPolynomialsAtPoint(x_array_read[p * dim + d], Q_1d, chebyshev_x));
+                CeedCall(CeedTensorContractApply(
+                    basis->contract, pre, 1, post, Q_1d, chebyshev_x, t_mode, (p > 0 || (p == 0 && pass > 0)) && d == (dim - 1),
+                    d == 0 ? &u_array[p * num_comp * dim + pass] : tmp[d % 2], d == (dim - 1) ? chebyshev_coeffs : tmp[(d + 1) % 2]));
+                pre /= 1;
+                post *= Q_1d;
+              }
+            }
           }
+          break;
         }
+        default:
+          // Nothing to do, excluded above
+          break;
       }
       CeedCall(CeedVectorRestoreArray(basis->vec_chebyshev, &chebyshev_coeffs));
       CeedCall(CeedVectorRestoreArrayRead(x_ref, &x_array_read));

tests/t357-basis.c

@@ -0,0 +1,114 @@
+/// @file
+/// Test gradient transpose in multiple dimensions at arbitrary points
+/// \test Test gradient transpose in multiple dimensions at arbitrary points
+#include <ceed.h>
+#include <math.h>
+#include <stdio.h>
+
+static CeedScalar Eval(CeedInt dim, const CeedScalar x[]) {
+  CeedScalar result = tanh(x[0] + 0.1);
+  if (dim > 1) result += atan(x[1] + 0.2);
+  if (dim > 2) result += exp(-(x[2] + 0.3) * (x[2] + 0.3));
+  return result;
+}
+
+static CeedScalar GetTolerance(CeedScalarType scalar_type, int dim) {
+  CeedScalar tol;
+  if (scalar_type == CEED_SCALAR_FP32) {
+    if (dim == 3) tol = 0.005;
+    else tol = 1.e-4;
+  } else {
+    tol = 1.e-11;
+  }
+  return tol;
+}
+
+int main(int argc, char **argv) {
+  Ceed ceed;
+
+  CeedInit(argv[1], &ceed);
+
+  for (CeedInt dim = 1; dim <= 3; dim++) {
+    CeedVector    x, x_nodes, x_points, u, u_points, v, ones;
+    CeedBasis     basis_x, basis_u;
+    const CeedInt p = 9, q = 9, num_points = 4, x_dim = CeedIntPow(2, dim), p_dim = CeedIntPow(p, dim);
+    CeedScalar    sum_1 = 0, sum_2 = 0;
+
+    CeedVectorCreate(ceed, x_dim * dim, &x);
+    CeedVectorCreate(ceed, p_dim * dim, &x_nodes);
+    CeedVectorCreate(ceed, num_points * dim, &x_points);
+    CeedVectorCreate(ceed, p_dim, &u);
+    CeedVectorCreate(ceed, num_points * dim, &u_points);
+    CeedVectorCreate(ceed, p_dim, &v);
+    CeedVectorCreate(ceed, num_points * dim, &ones);
+
+    CeedVectorSetValue(ones, 1);
+    CeedVectorSetValue(v, 0);
+
+    // Get nodal coordinates
+    CeedBasisCreateTensorH1Lagrange(ceed, dim, dim, 2, p, CEED_GAUSS_LOBATTO, &basis_x);
+    {
+      CeedScalar x_array[x_dim * dim];
+
+      for (CeedInt d = 0; d < dim; d++) {
+        for (CeedInt i = 0; i < x_dim; i++) x_array[d * x_dim + i] = (i % CeedIntPow(2, d + 1)) / CeedIntPow(2, d) ? 1 : -1;
+      }
+      CeedVectorSetArray(x, CEED_MEM_HOST, CEED_COPY_VALUES, x_array);
+    }
+    CeedBasisApply(basis_x, 1, CEED_NOTRANSPOSE, CEED_EVAL_INTERP, x, x_nodes);
+
+    // Set values of u at nodes
+    {
+      const CeedScalar *x_array;
+      CeedScalar        u_array[p_dim];
+
+      CeedVectorGetArrayRead(x_nodes, CEED_MEM_HOST, &x_array);
+      for (CeedInt i = 0; i < p_dim; i++) {
+        CeedScalar coord[dim];
+
+        for (CeedInt d = 0; d < dim; d++) coord[d] = x_array[d * p_dim + i];
+        u_array[i] = Eval(dim, coord);
+      }
+      CeedVectorRestoreArrayRead(x_nodes, &x_array);
+      CeedVectorSetArray(u, CEED_MEM_HOST, CEED_COPY_VALUES, (CeedScalar *)&u_array);
+    }
+
+    // Interpolate to arbitrary points
+    CeedBasisCreateTensorH1Lagrange(ceed, dim, 1, p, q, CEED_GAUSS, &basis_u);
+    {
+      CeedScalar x_array[12] = {-0.33, -0.65, 0.16, 0.99, -0.65, 0.16, 0.99, -0.33, 0.16, 0.99, -0.33, -0.65};
+
+      CeedVectorSetArray(x_points, CEED_MEM_HOST, CEED_COPY_VALUES, x_array);
+    }
+
+    // Calculate G u at arbitrary points, G' * 1 at dofs
+    CeedBasisApplyAtPoints(basis_u, num_points, CEED_NOTRANSPOSE, CEED_EVAL_GRAD, x_points, u, u_points);
+    CeedBasisApplyAtPoints(basis_u, num_points, CEED_TRANSPOSE, CEED_EVAL_GRAD, x_points, ones, v);
+    {
+      const CeedScalar *u_array, *v_array, *u_points_array;
+
+      CeedVectorGetArrayRead(u, CEED_MEM_HOST, &u_array);
+      CeedVectorGetArrayRead(v, CEED_MEM_HOST, &v_array);
+      CeedVectorGetArrayRead(u_points, CEED_MEM_HOST, &u_points_array);
+      for (CeedInt i = 0; i < p_dim; i++) sum_1 += v_array[i] * u_array[i];
+      for (CeedInt i = 0; i < num_points * dim; i++) sum_2 += u_points_array[i];
+      CeedVectorRestoreArrayRead(u, &u_array);
+      CeedVectorRestoreArrayRead(v, &v_array);
+      CeedVectorRestoreArrayRead(u_points, &u_points_array);
+    }
+    CeedScalar tol = GetTolerance(CEED_SCALAR_TYPE, dim);
+    if (fabs(sum_1 - sum_2) > tol) printf("[%" CeedInt_FMT "] %f != %f\n", dim, sum_1, sum_2);
+
+    CeedVectorDestroy(&x);
+    CeedVectorDestroy(&x_nodes);
+    CeedVectorDestroy(&x_points);
+    CeedVectorDestroy(&u);
+    CeedVectorDestroy(&u_points);
+    CeedVectorDestroy(&ones);
+    CeedVectorDestroy(&v);
+    CeedBasisDestroy(&basis_x);
+    CeedBasisDestroy(&basis_u);
+  }
+  CeedDestroy(&ceed);
+  return 0;
+}