Question: reusing `ScalarFunction` creation code

[Victorlouisdg]

For my use case (deformable simulation), I find that I'm often writing very similar ScalarFunctions. For example, when calculating two energies for a range of vertices, only a single call to an elementEnergy function template differs. Below is a minimal example:

#include <Eigen/Core>
#include <TinyAD/ScalarFunction.hh>
#include <iostream>

template<typename T> T elementEnergy0(Eigen::Vector3<T> x) {
  return x.sum();
}

template<typename T> T elementEnergy1(Eigen::Vector3<T> x) {
  return x.prod();
}

TinyAD::ScalarFunction<3, double, Eigen::Index> createScalarFunction(int vertexCount) {
  auto scalarFunction = TinyAD::scalar_function<3>(TinyAD::range(vertexCount));

  scalarFunction.add_elements<1>(TinyAD::range(vertexCount), [&](auto &element) -> TINYAD_SCALAR_TYPE(element) {
    using ScalarT = TINYAD_SCALAR_TYPE(element);
    int index = element.handle;
    Eigen::Vector3<ScalarT> x = element.variables(index);
    ScalarT energy = elementEnergy0(x);  // I want this function call to be variable.
    return energy;
  });

  return scalarFunction;
}

int main() {
  Eigen::RowVector3d v0(0.0, 0.0, 0.0);
  Eigen::RowVector3d v1(1.0, 1.0, 1.0);
  Eigen::Matrix<double, 2, 3> vertexPositions;
  vertexPositions << v0, v1;
  int vertexCount = vertexPositions.rows();

  Eigen::VectorXd x = vertexPositions.reshaped<Eigen::RowMajor>(vertexPositions.rows() * 3, 1);

  auto scalarFunction0 = createScalarFunction(vertexCount /*, elementEnergy0 */);
  auto e0 = scalarFunction0.eval(x);
  std::cout << e0 << std::endl;

  // I want this function to use elementEnergy1:
  auto scalarFunction1 = createScalarFunction(vertexCount /*, elementEnergy1 */);
}

The problem is, I can't find a way to pass function templates into the createScalarFunction without assigning types. However, TinyAD "needs" these function templates to remain "templated", such that they can be called with both a passive type e.g. double and an active type e.g. TinyAD::Scalar.

So to conclude my question is: how can I create a function that returns a ScalarFunction, that internally uses a variable function template?

PS: Thanks for the great library, I'm really loving it. The sparse Hessians are exactly what I've been searching for!

[Victorlouisdg]

I've found a solution using generic lambdas I'm pretty happy with. Here it is in case anyone is trying something similar:

#include <Eigen/Core>
#include <TinyAD/ScalarFunction.hh>
#include <iostream>

template<typename T> T elementEnergy0(Eigen::Vector3<T> x) {
  return x.sum();
}

template<typename T> T elementEnergy1(Eigen::Vector3<T> x) {
  return x.prod();
}

template<typename F>
TinyAD::ScalarFunction<3, double, Eigen::Index> createScalarFunction(int vertexCount, F &&elementEnergy) {
  auto scalarFunction = TinyAD::scalar_function<3>(TinyAD::range(vertexCount));

  scalarFunction.add_elements<1>(TinyAD::range(vertexCount), [&](auto &element) -> TINYAD_SCALAR_TYPE(element) {
    using ScalarT = TINYAD_SCALAR_TYPE(element);
    int index = element.handle;
    Eigen::Vector3<ScalarT> x = element.variables(index);
    ScalarT energy = elementEnergy(x);
    return energy;
  });

  return scalarFunction;
}

int main() {
  Eigen::RowVector3d v0(0.0, 0.0, 0.0);
  Eigen::RowVector3d v1(1.0, 1.0, 1.0);
  Eigen::Matrix<double, 2, 3> vertexPositions;
  vertexPositions << v0, v1;
  int vertexCount = vertexPositions.rows();

  Eigen::VectorXd x = vertexPositions.reshaped<Eigen::RowMajor>(vertexPositions.rows() * 3, 1);

  auto scalarFunction0 = createScalarFunction(vertexCount, [](auto &&t) { return elementEnergy0(t); });
  auto scalarFunction1 = createScalarFunction(vertexCount, [](auto &&t) { return elementEnergy1(t); });

  auto e0 = scalarFunction0.eval(x);
  std::cout << e0 << std::endl;

  auto e1 = scalarFunction1.eval(x);
  std::cout << e1 << std::endl;
}

If there is a simpler way to achieve this functionality, please let me know! For me, this issue can be closed now.

[patr-schm]

Hi, I'm happy to see that you're using TinyAD. Thanks for sharing your answer!

If you actually want to create a separate ScalarFunction object for each variant, your solution is the best I can think of.

Not sure if this applies to your use case, but I usually just use branching inside the ScalarFunction to handle different variants of an energy. This might have a slight run time overhead compared to your version, but I often find this to be the most convenient option.

#include <TinyAD/ScalarFunction.hh>

void my_algorithm(
        Eigen::Matrix<double, 2, 3>& vertexPositions,
        const bool useEnergy0)
{
    int vertexCount = vertexPositions.rows();
    Eigen::VectorXd x = vertexPositions.reshaped<Eigen::RowMajor>(vertexPositions.rows() * 3, 1);

    auto func = TinyAD::scalar_function<3>(TinyAD::range(vertexCount));
    func.add_elements<1>(TinyAD::range(vertexCount), [&] (auto& element) -> TINYAD_SCALAR_TYPE(element)
    {
        using ScalarT = TINYAD_SCALAR_TYPE(element);
        int index = element.handle;
        Eigen::Vector3<ScalarT> x = element.variables(index);

        if (useEnergy0)
            return x.sum();
        else
            return x.prod();
    });

    std::cout << func.eval(x) << std::endl;
}

int main()
{
    Eigen::RowVector3d v0(0.0, 0.0, 0.0);
    Eigen::RowVector3d v1(1.0, 1.0, 1.0);
    Eigen::Matrix<double, 2, 3> vertexPositions;
    vertexPositions << v0, v1;

    my_algorithm(vertexPositions, true);
    my_algorithm(vertexPositions, false);
}

[Victorlouisdg]

I'll probably have to stick to the lambda approach for my use case, but thanks for the reply! :)