Expression.h

// Copyright (C) 2020 Jack S. Hale
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <dolfinx/common/array2d.h>
#include <dolfinx/mesh/Mesh.h>
#include <functional>
#include <utility>
#include <vector>
#include <xtensor/xarray.hpp>
#include <xtl/xspan.hpp>
 
namespace dolfinx::fem
{
template <typename T>
class Constant;
 
 
template <typename T>
class Expression
{
public:
  Expression(
      const std::vector<std::shared_ptr<const fem::Function<T>>>& coefficients,
      const std::vector<std::shared_ptr<const fem::Constant<T>>>& constants,
      const std::shared_ptr<const mesh::Mesh>& mesh, const array2d<double>& X,
      const std::function<void(T*, const T*, const T*, const double*)> fn,
      const std::size_t value_size)
      : _coefficients(coefficients), _constants(constants), _mesh(mesh), _x(X),
        _fn(fn), _value_size(value_size)
  {
    // Do nothing
  }
 
  Expression(Expression&& form) = default;
 
  virtual ~Expression() = default;
 
  const std::vector<std::shared_ptr<const fem::Function<T>>>&
  coefficients() const
  {
    return _coefficients;
  }
 
  std::vector<int> coefficient_offsets() const
  {
    std::vector<int> n{0};
    for (const auto& c : _coefficients)
    {
      if (!c)
        throw std::runtime_error("Not all form coefficients have been set.");
      n.push_back(n.back() + c->function_space()->element()->space_dimension());
    }
    return n;
  }
 
  template <typename U>
  void eval(const xtl::span<const std::int32_t>& active_cells, U& values) const
  {
    static_assert(std::is_same<T, typename U::value_type>::value,
                  "Expression and array types must be the same");
 
    // Extract data from Expression
    assert(_mesh);
 
    // Prepare coefficients and constants
    const array2d<T> coeffs = pack_coefficients(*this);
    const std::vector<T> constant_data = pack_constants(*this);
 
    const auto& fn = this->get_tabulate_expression();
 
    // Prepare cell geometry
    const graph::AdjacencyList<std::int32_t>& x_dofmap
        = _mesh->geometry().dofmap();
    const fem::CoordinateElement& cmap = _mesh->geometry().cmap();
 
    // FIXME: Add proper interface for num coordinate dofs
    const std::size_t num_dofs_g = x_dofmap.num_links(0);
    const xt::xtensor<double, 2>& x_g = _mesh->geometry().x();
 
    // Create data structures used in evaluation
    std::vector<double> coordinate_dofs(3 * num_dofs_g);
 
    // Iterate over cells and 'assemble' into values
    std::vector<T> values_e(this->num_points() * this->value_size(), 0);
    for (std::size_t c = 0; c < active_cells.size(); ++c)
    {
      const std::int32_t cell = active_cells[c];
 
      auto x_dofs = x_dofmap.links(cell);
      for (std::size_t i = 0; i < x_dofs.size(); ++i)
      {
        std::copy_n(xt::row(x_g, x_dofs[i]).cbegin(), 3,
                    std::next(coordinate_dofs.begin(), 3 * i));
      }
 
      auto coeff_cell = coeffs.row(cell);
      std::fill(values_e.begin(), values_e.end(), 0.0);
      fn(values_e.data(), coeff_cell.data(), constant_data.data(),
         coordinate_dofs.data());
 
      for (std::size_t j = 0; j < values_e.size(); ++j)
        values(c, j) = values_e[j];
    }
  }
 
  const std::function<void(T*, const T*, const T*, const double*)>&
  get_tabulate_expression() const
  {
    return _fn;
  }
 
  const std::vector<std::shared_ptr<const fem::Constant<T>>>& constants() const
  {
    return _constants;
  }
 
  std::shared_ptr<const mesh::Mesh> mesh() const { return _mesh; }
 
  const array2d<double>& x() const { return _x; }
 
  const std::size_t value_size() const { return _value_size; }
 
  const std::size_t num_points() const { return _x.shape[0]; }
 
  using scalar_type = T;
 
private:
  // Coefficients associated with the Expression
  std::vector<std::shared_ptr<const fem::Function<T>>> _coefficients;
 
  // Constants associated with the Expression
  std::vector<std::shared_ptr<const fem::Constant<T>>> _constants;
 
  // Function to evaluate the Expression
  std::function<void(T*, const T*, const T*, const double*)> _fn;
 
  // Evaluation points on reference cell
  array2d<double> _x;
 
  // The mesh.
  std::shared_ptr<const mesh::Mesh> _mesh;
 
  // Evaluation size
  std::size_t _value_size;
};
} // namespace dolfinx::fem