#include <ceed.h>
#include <ceed/backend.h>
#include <math.h>
#include <stdbool.h>
#include <string.h>
#include "ceed-ref.h"
static int CeedBasisApply_Ref(CeedBasis basis, CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector U, CeedVector V) {
Ceed ceed;
bool is_tensor_basis;
CeedInt dim, num_comp, q_comp, num_nodes, num_qpts;
const CeedInt add = (t_mode == CEED_TRANSPOSE);
const CeedScalar *u;
CeedScalar *v;
CeedTensorContract contract;
CeedBasis_Ref *impl;
CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
CeedCallBackend(CeedBasisGetData(basis, &impl));
CeedCallBackend(CeedBasisGetDimension(basis, &dim));
CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
CeedCallBackend(CeedBasisGetNumQuadratureComponents(basis, eval_mode, &q_comp));
CeedCallBackend(CeedBasisGetNumNodes(basis, &num_nodes));
CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis, &num_qpts));
CeedCallBackend(CeedBasisGetTensorContract(basis, &contract));
if (U != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(U, CEED_MEM_HOST, &u));
else CeedCheck(eval_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
CeedCallBackend(CeedVectorGetArrayWrite(V, CEED_MEM_HOST, &v));
if (t_mode == CEED_TRANSPOSE) {
const CeedInt v_size = num_elem * num_comp * num_nodes;
for (CeedInt i = 0; i < v_size; i++) v[i] = (CeedScalar)0.0;
}
CeedCallBackend(CeedBasisIsTensor(basis, &is_tensor_basis));
if (is_tensor_basis) {
CeedInt P_1d, Q_1d;
CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
switch (eval_mode) {
case CEED_EVAL_INTERP: {
if (impl->has_collo_interp) {
memcpy(v, u, num_elem * num_comp * num_nodes * sizeof(u[0]));
} else {
CeedInt P = P_1d, Q = Q_1d;
if (t_mode == CEED_TRANSPOSE) {
P = Q_1d;
Q = P_1d;
}
CeedInt pre = num_comp * CeedIntPow(P, dim - 1), post = num_elem;
CeedScalar tmp[2][num_elem * num_comp * Q * CeedIntPow(P > Q ? P : Q, dim - 1)];
const CeedScalar *interp_1d;
CeedCallBackend(CeedBasisGetInterp1D(basis, &interp_1d));
for (CeedInt d = 0; d < dim; d++) {
CeedCallBackend(CeedTensorContractApply(contract, pre, P, post, Q, interp_1d, t_mode, add && (d == dim - 1), d == 0 ? u : tmp[d % 2],
d == dim - 1 ? v : tmp[(d + 1) % 2]));
pre /= P;
post *= Q;
}
}
} break;
case CEED_EVAL_GRAD: {
CeedInt P = P_1d, Q = Q_1d;
if (t_mode == CEED_TRANSPOSE) {
P = Q_1d;
Q = Q_1d;
}
CeedInt pre = num_comp * CeedIntPow(P, dim - 1), post = num_elem;
const CeedScalar *interp_1d;
CeedCallBackend(CeedBasisGetInterp1D(basis, &interp_1d));
if (impl->collo_grad_1d) {
CeedScalar tmp[2][num_elem * num_comp * Q * CeedIntPow(P > Q ? P : Q, dim - 1)];
CeedScalar interp[num_elem * num_comp * Q * CeedIntPow(P > Q ? P : Q, dim - 1)];
for (CeedInt d = 0; d < dim; d++) {
CeedCallBackend(CeedTensorContractApply(contract, pre, P, post, Q, (t_mode == CEED_NOTRANSPOSE ? interp_1d : impl->collo_grad_1d), t_mode,
add && (d > 0),
(t_mode == CEED_NOTRANSPOSE ? (d == 0 ? u : tmp[d % 2]) : u + d * num_qpts * num_comp * num_elem),
(t_mode == CEED_NOTRANSPOSE ? (d == dim - 1 ? interp : tmp[(d + 1) % 2]) : interp)));
pre /= P;
post *= Q;
}
P = Q_1d, Q = Q_1d;
if (t_mode == CEED_TRANSPOSE) {
P = Q_1d;
Q = P_1d;
}
pre = num_comp * CeedIntPow(P, dim - 1), post = num_elem;
for (CeedInt d = 0; d < dim; d++) {
CeedCallBackend(CeedTensorContractApply(
contract, pre, P, post, Q, (t_mode == CEED_NOTRANSPOSE ? impl->collo_grad_1d : interp_1d), t_mode, add && (d == dim - 1),
(t_mode == CEED_NOTRANSPOSE ? interp : (d == 0 ? interp : tmp[d % 2])),
(t_mode == CEED_NOTRANSPOSE ? v + d * num_qpts * num_comp * num_elem : (d == dim - 1 ? v : tmp[(d + 1) % 2]))));
pre /= P;
post *= Q;
}
} else if (impl->has_collo_interp) { const CeedScalar *grad_1d;
CeedCallBackend(CeedBasisGetGrad1D(basis, &grad_1d));
CeedInt pre = num_comp * CeedIntPow(P, dim - 1), post = num_elem;
for (CeedInt d = 0; d < dim; d++) {
CeedCallBackend(CeedTensorContractApply(contract, pre, P, post, Q, grad_1d, t_mode, add && (d > 0),
t_mode == CEED_NOTRANSPOSE ? u : u + d * num_comp * num_qpts * num_elem,
t_mode == CEED_TRANSPOSE ? v : v + d * num_comp * num_qpts * num_elem));
pre /= P;
post *= Q;
}
} else { const CeedScalar *grad_1d;
CeedCallBackend(CeedBasisGetGrad1D(basis, &grad_1d));
if (t_mode == CEED_TRANSPOSE) {
P = Q_1d;
Q = P_1d;
}
CeedScalar tmp[2][num_elem * num_comp * Q * CeedIntPow(P > Q ? P : Q, dim - 1)];
for (CeedInt p = 0; p < dim; p++) {
CeedInt pre = num_comp * CeedIntPow(P, dim - 1), post = num_elem;
for (CeedInt d = 0; d < dim; d++) {
CeedCallBackend(CeedTensorContractApply(
contract, pre, P, post, Q, (p == d) ? grad_1d : interp_1d, t_mode, add && (d == dim - 1),
(d == 0 ? (t_mode == CEED_NOTRANSPOSE ? u : u + p * num_comp * num_qpts * num_elem) : tmp[d % 2]),
(d == dim - 1 ? (t_mode == CEED_TRANSPOSE ? v : v + p * num_comp * num_qpts * num_elem) : tmp[(d + 1) % 2])));
pre /= P;
post *= Q;
}
}
}
} break;
case CEED_EVAL_WEIGHT: {
CeedInt Q = Q_1d;
const CeedScalar *q_weight_1d;
CeedCheck(t_mode == CEED_NOTRANSPOSE, ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT incompatible with CEED_TRANSPOSE");
CeedCallBackend(CeedBasisGetQWeights(basis, &q_weight_1d));
for (CeedInt d = 0; d < dim; d++) {
CeedInt pre = CeedIntPow(Q, dim - d - 1), post = CeedIntPow(Q, d);
for (CeedInt i = 0; i < pre; i++) {
for (CeedInt j = 0; j < Q; j++) {
for (CeedInt k = 0; k < post; k++) {
const CeedScalar w = q_weight_1d[j] * (d == 0 ? 1 : v[((i * Q + j) * post + k) * num_elem]);
for (CeedInt e = 0; e < num_elem; e++) v[((i * Q + j) * post + k) * num_elem + e] = w;
}
}
}
}
} break;
case CEED_EVAL_DIV:
return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_DIV not supported");
case CEED_EVAL_CURL:
return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_CURL not supported");
case CEED_EVAL_NONE:
return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
}
} else {
CeedInt P = num_nodes, Q = num_qpts;
switch (eval_mode) {
case CEED_EVAL_INTERP: {
const CeedScalar *interp;
CeedCallBackend(CeedBasisGetInterp(basis, &interp));
CeedCallBackend(CeedTensorContractStridedApply(contract, num_comp, P, num_elem, q_comp, Q, interp, t_mode, add, u, v));
} break;
case CEED_EVAL_GRAD: {
const CeedScalar *grad;
CeedCallBackend(CeedBasisGetGrad(basis, &grad));
CeedCallBackend(CeedTensorContractStridedApply(contract, num_comp, P, num_elem, q_comp, Q, grad, t_mode, add, u, v));
} break;
case CEED_EVAL_DIV: {
const CeedScalar *div;
CeedCallBackend(CeedBasisGetDiv(basis, &div));
CeedCallBackend(CeedTensorContractStridedApply(contract, num_comp, P, num_elem, q_comp, Q, div, t_mode, add, u, v));
} break;
case CEED_EVAL_CURL: {
const CeedScalar *curl;
CeedCallBackend(CeedBasisGetCurl(basis, &curl));
CeedCallBackend(CeedTensorContractStridedApply(contract, num_comp, P, num_elem, q_comp, Q, curl, t_mode, add, u, v));
} break;
case CEED_EVAL_WEIGHT: {
const CeedScalar *q_weight;
CeedCheck(t_mode == CEED_NOTRANSPOSE, ceed, CEED_ERROR_BACKEND, "CEED_EVAL_WEIGHT incompatible with CEED_TRANSPOSE");
CeedCallBackend(CeedBasisGetQWeights(basis, &q_weight));
for (CeedInt i = 0; i < num_qpts; i++) {
for (CeedInt e = 0; e < num_elem; e++) v[i * num_elem + e] = q_weight[i];
}
} break;
case CEED_EVAL_NONE:
return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
}
}
if (U != CEED_VECTOR_NONE) {
CeedCallBackend(CeedVectorRestoreArrayRead(U, &u));
}
CeedCallBackend(CeedVectorRestoreArray(V, &v));
return CEED_ERROR_SUCCESS;
}
static int CeedBasisDestroyTensor_Ref(CeedBasis basis) {
CeedBasis_Ref *impl;
CeedCallBackend(CeedBasisGetData(basis, &impl));
CeedCallBackend(CeedFree(&impl->collo_grad_1d));
CeedCallBackend(CeedFree(&impl));
return CEED_ERROR_SUCCESS;
}
int CeedBasisCreateTensorH1_Ref(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
Ceed ceed, ceed_parent;
CeedBasis_Ref *impl;
CeedTensorContract contract;
CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
CeedCallBackend(CeedGetParent(ceed, &ceed_parent));
CeedCallBackend(CeedCalloc(1, &impl));
if (Q_1d == P_1d) {
bool has_collocated = true;
for (CeedInt i = 0; i < P_1d; i++) {
has_collocated = has_collocated && (fabs(interp_1d[i + P_1d * i] - 1.0) < 1e-14);
for (CeedInt j = 0; j < P_1d; j++) {
if (j != i) has_collocated = has_collocated && (fabs(interp_1d[j + P_1d * i]) < 1e-14);
}
}
impl->has_collo_interp = has_collocated;
}
if (Q_1d >= P_1d && !impl->has_collo_interp) {
CeedCallBackend(CeedMalloc(Q_1d * Q_1d, &impl->collo_grad_1d));
CeedCallBackend(CeedBasisGetCollocatedGrad(basis, impl->collo_grad_1d));
}
CeedCallBackend(CeedBasisSetData(basis, impl));
CeedCallBackend(CeedTensorContractCreate(ceed_parent, &contract));
CeedCallBackend(CeedBasisSetTensorContract(basis, contract));
CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApply_Ref));
CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyTensor_Ref));
return CEED_ERROR_SUCCESS;
}
int CeedBasisCreateH1_Ref(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp, const CeedScalar *grad,
const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
Ceed ceed, ceed_parent;
CeedTensorContract contract;
CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
CeedCallBackend(CeedGetParent(ceed, &ceed_parent));
CeedCallBackend(CeedTensorContractCreate(ceed_parent, &contract));
CeedCallBackend(CeedBasisSetTensorContract(basis, contract));
CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApply_Ref));
return CEED_ERROR_SUCCESS;
}
int CeedBasisCreateHdiv_Ref(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp, const CeedScalar *div,
const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
Ceed ceed, ceed_parent;
CeedTensorContract contract;
CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
CeedCallBackend(CeedGetParent(ceed, &ceed_parent));
CeedCallBackend(CeedTensorContractCreate(ceed_parent, &contract));
CeedCallBackend(CeedBasisSetTensorContract(basis, contract));
CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApply_Ref));
return CEED_ERROR_SUCCESS;
}
int CeedBasisCreateHcurl_Ref(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp,
const CeedScalar *curl, const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
Ceed ceed, ceed_parent;
CeedTensorContract contract;
CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
CeedCallBackend(CeedGetParent(ceed, &ceed_parent));
CeedCallBackend(CeedTensorContractCreate(ceed_parent, &contract));
CeedCallBackend(CeedBasisSetTensorContract(basis, contract));
CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApply_Ref));
return CEED_ERROR_SUCCESS;
}