Expression.h

// Copyright (C) 2020-2026 Jack S. Hale, Michal Habera, Garth N. Wells and
// Jørgen S. Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "Constant.h"
#include "Function.h"
#include <algorithm>
#include <array>
#include <concepts>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/EntityMap.h>
#include <dolfinx/mesh/Mesh.h>
#include <functional>
#include <memory>
#include <span>
#include <utility>
#include <vector>
 
namespace dolfinx::fem
{
template <dolfinx::scalar T, std::floating_point U>
class Function;

template <dolfinx::scalar T, std::floating_point U = dolfinx::scalar_value_t<T>>
class Expression
{
public:
  using scalar_type = T;

  using geometry_type = U;

  Expression(
      const std::vector<std::shared_ptr<
          const Function<scalar_type, geometry_type>>>& coefficients,
      const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
          constants,
      std::span<const geometry_type> X, std::array<std::size_t, 2> Xshape,
      std::function<void(scalar_type*, const scalar_type*, const scalar_type*,
                         const geometry_type*, const int*, const uint8_t*,
                         void*)>
          fn,
      const std::vector<std::size_t>& value_shape,
      const std::vector<std::reference_wrapper<const dolfinx::mesh::EntityMap>>&
          entity_maps,
 
      std::uint64_t coordinate_element_hash,
      std::shared_ptr<const FunctionSpace<geometry_type>> argument_space
      = nullptr)
      : _argument_space(argument_space), _coefficients(coefficients),
        _constants(constants), _fn(fn), _value_shape(value_shape),
        _x_ref(std::vector<geometry_type>(X.begin(), X.end()), Xshape),
        _entity_maps(entity_maps),
        _coordinate_element_hash(coordinate_element_hash) {};

  Expression(Expression&& e) = default;

  virtual ~Expression() = default;

  std::shared_ptr<const FunctionSpace<geometry_type>> argument_space() const
  {
    return _argument_space;
  };

  const std::vector<
      std::shared_ptr<const Function<scalar_type, geometry_type>>>&
  coefficients() const
  {
    return _coefficients;
  }

  const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
  constants() const
  {
    return _constants;
  }

  std::vector<int> coefficient_offsets() const
  {
    std::vector<int> n{0};
    for (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;
  }

  const std::function<void(scalar_type*, const scalar_type*, const scalar_type*,
                           const geometry_type*, const int*, const uint8_t*,
                           void*)>&
  kernel() const
  {
    return _fn;
  }

  int value_size() const
  {
    return std::reduce(_value_shape.begin(), _value_shape.end(), 1,
                       std::multiplies{});
  }

  const std::vector<std::size_t>& value_shape() const { return _value_shape; }

  std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>> X() const
  {
    return _x_ref;
  }

  const std::vector<std::reference_wrapper<const dolfinx::mesh::EntityMap>>&
  entity_maps() const
  {
    return _entity_maps;
  }

  std::uint64_t coordinate_element_hash() const
  {
    return _coordinate_element_hash;
  }
 
private:
  // Function space for Argument
  std::shared_ptr<const FunctionSpace<geometry_type>> _argument_space;
 
  // Coefficients associated with the Expression
  std::vector<std::shared_ptr<const Function<scalar_type, geometry_type>>>
      _coefficients;
 
  // Constants associated with the Expression
  std::vector<std::shared_ptr<const Constant<scalar_type>>> _constants;
 
  // Function to evaluate the Expression
  std::function<void(scalar_type*, const scalar_type*, const scalar_type*,
                     const geometry_type*, const int*, const uint8_t*, void*)>
      _fn;
 
  // Shape of the evaluated Expression
  std::vector<std::size_t> _value_shape;
 
  // Evaluation points on reference cell
  std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>> _x_ref;
 
  // Set of bidirectional maps of mesh entities (cells, facets, ridges or peaks)
  // between meshes and submeshes of any codimension
  std::vector<std::reference_wrapper<const dolfinx::mesh::EntityMap>>
      _entity_maps;
 
  // Hash for coordinate element used to create the expression
  std::uint64_t _coordinate_element_hash;
};
} // namespace dolfinx::fem