Expression.h

// Copyright (C) 2020-2025 Jack S. Hale, Michal Habera and Garth N. Wells
//
// 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/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,
             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)
 
  {
    for (auto& c : _coefficients)
    {
      assert(c);
      if (c->function_space()->mesh()
          != _coefficients.front()->function_space()->mesh())
      {
        throw std::runtime_error("Coefficients not all defined on same mesh.");
      }
    }
  }

  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;
  }
 
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;
};
} // namespace dolfinx::fem