Using Skala in C++ with LibTorch

This page provides an example of how to use Skala in C++ with the LibTorch library. The example demonstrates how to load a Skala model, prepare input features, compute the exchange-correlation energy and its derivatives with respect to the input features, and access the results.

Note

This example evaluates the Skala model on precomputed density features and is meant as a starting point for calling the model from your own code. To add Skala to an existing electronic structure code, where the integration grid, the density features, and the exchange-correlation potential have to be computed from a density matrix, use GauXC instead. A Fortran version of this example is also available.

Setting up development environment

The repository’s locked cpp-integration Pixi environment provides everything needed for this example: LibTorch, nlohmann_json, CMake, Ninja, and the C and C++ toolchains, the hf CLI to download the model checkpoint, and the skala runtime with PySCF to generate the input features. From the repository root, activate the environment:

pixi shell --locked -e cpp-integration

All following commands are executed in this shell from the repository root. Configure and build the example with:

cmake -B build_example -S model/examples/cpp/cpp_integration -G Ninja
cmake --build build_example

Preparing the model and input features

To evaluate Skala, we download the model checkpoint from HuggingFace using the hf CLI from the huggingface_hub package:

hf download microsoft/skala-1.1 skala-1.1-rev1.fun --local-dir .

The example reads the input features from torch .pt files, one file per feature. They are generated for the H2 molecule with the def2-QZVP basis set by the model/examples/prepare_inputs.py script, which is shared by the C++ and Fortran examples:

python model/examples/prepare_inputs.py --output-dir features

Feature preparation script

model/examples/prepare_inputs.py

#!/usr/bin/env python3

import argparse
from pathlib import Path

import torch
from pyscf.dft import gen_grid
from skala.functional.traditional import LDA
from skala.pyscf.evaluation import FeatureSpec
from skala.pyscf.model_chunking import evaluate_model_features
from skala_model import SkalaFunctional

from pyscf import dft, gto


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output-dir",
        default=".",
        type=Path,
        help="Output directory for generated feature files.",
    )
    parser.add_argument(
        "--basis",
        default="def2-qzvp",
        type=str,
        help="Basis set.",
    )
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)

    molecule = gto.M(
        atom="H 0 0 0; H 0 0 1",
        basis=args.basis,
        verbose=0,
    )

    # Create a set of meta-GGA features for this molecule.
    dm = get_density_matrix(molecule)
    grid = gen_grid.Grids(molecule)
    grid.level = 3
    grid.build(sort_grids=False)
    features = evaluate_model_features(
        molecule, dm, grid, FeatureSpec(SkalaFunctional.features)
    )

    # Save all features as individual .pt files.
    for key, value in features.items():
        torch.save(value, str(args.output_dir / f"{key}.pt"))

    print(f"Saved features to {args.output_dir}")

    lda_exc = LDA().get_exc(features)
    print(f"For reference, LDAx Exc = {lda_exc.item()}")


def get_density_matrix(mol: gto.Mole) -> torch.Tensor:
    """Computes an example density matrix for a given molecule using PySCF."""
    ks = dft.RKS(mol, xc="b3lyp5")
    ks = ks.density_fit()
    ks.kernel()
    dm = torch.from_numpy(ks.make_rdm1())
    return dm


if __name__ == "__main__":
    main()

Tip

Here we load the features directly from torch .pt files, in your application you may want to create them directly from your own arrays, for example with torch::from_blob.

With the checkpoint and the features in place, run the example:

./build_example/skala_cpp_integration ./skala-1.1-rev1.fun ./features

Build system setup

For building the C++ application, we will use CMake. We will use the following directory structure and files:

model/examples/
├── prepare_inputs.py
└── cpp/
    └── cpp_integration/
        ├── CMakeLists.txt
        └── main.cpp

The CMakeLists.txt file sets up the project, finds LibTorch and nlohmann_json, and declares the executable:

CMakeLists.txt

cmake_minimum_required(VERSION 3.10 FATAL_ERROR)
project(
  "skala_cpp_integration"
  VERSION "0.0.1"
  LANGUAGES CXX
)

find_package(nlohmann_json REQUIRED)
find_package(Torch REQUIRED)

add_executable(
  ${PROJECT_NAME}
  ./main.cpp
)
target_link_libraries(
  ${PROJECT_NAME}
  "${TORCH_LIBRARIES}"
)
set_property(
  TARGET ${PROJECT_NAME}
  PROPERTY CXX_STANDARD 17
)

C++ application

We start the program with the necessary headers and some type aliases for the TorchScript calling convention:

main.cpp

#include <torch/script.h>
#include <torch/csrc/autograd/autograd.h>
#include <nlohmann/json.hpp>

#include <iostream>
#include <memory>

using json = nlohmann::json;
using IValueList = std::vector<c10::IValue>;
using IValueMap = std::unordered_map<std::string, c10::IValue>;
using FeatureDict = c10::Dict<std::string, at::Tensor>;

The Skala checkpoint is a TorchScript module. Next to the module, the checkpoint stores the list of required input features and the protocol version as extra files, which we read while loading the model. The get_exc_density method of the module evaluates the exchange-correlation energy density on the grid.

main.cpp

std::tuple<torch::jit::Method, std::vector<std::string>>
load_model(const std::string &filename, torch::DeviceType device)
{
  torch::jit::script::Module mod;
  torch::jit::ExtraFilesMap extra_files{{"features", ""}, {"protocol_version", ""}};
  std::vector<std::string> keys;

  try
  {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    mod = torch::jit::load(filename, device, extra_files);
  }
  catch (const c10::Error &e)
  {
    throw std::runtime_error("Error loading the model from " + filename + ": " + e.what());
  }

  auto version = json::parse(extra_files.at("protocol_version")).get<int>();
  if (version != 2)
  {
    throw std::runtime_error("Unsupported protocol version " + std::to_string(version));
  }

  auto features = json::parse(extra_files.at("features"));
  // check if features is array
  if (!features.is_array())
  {
    throw std::runtime_error("features is not an array");
  }
  for (const auto &feature : features)
  {
    if (!feature.is_string())
    {
      throw std::runtime_error("feature is not a string");
    }
    keys.push_back(feature.get<std::string>());
  }

  return std::make_tuple(mod.get_method("get_exc_density"), keys);
}

The input to the Skala model is passed via a dictionary of tensors. We load each required feature from its .pt file and insert it into the dictionary under its feature name:

main.cpp

at::Tensor
load_feature(const std::string &filename, torch::DeviceType device)
{
  std::ifstream input(filename, std::ios::binary);
  if (!input.is_open())
  {
    throw std::runtime_error("Failed to open feature file: " + filename);
  }
  std::vector<char> bytes(
      (std::istreambuf_iterator<char>(input)),
      (std::istreambuf_iterator<char>()));

  input.close();
  return torch::jit::pickle_load(bytes).toTensor().to(device);
}

FeatureDict
load_features(const std::string &prefix, const std::vector<std::string> &keys, torch::DeviceType device)
{
  FeatureDict featmap;
  for (const auto &key : keys)
  {
    featmap.insert(key, load_feature(prefix + "/" + key + ".pt", device));
  }
  return featmap;
}

With the model and the features available, we evaluate the exchange-correlation energy density by calling the model method with the feature dictionary as mol keyword argument:

main.cpp

at::Tensor
get_exc(const torch::jit::Method &exc_func, const FeatureDict &features)
{
  IValueList args;
  IValueMap kwargs;
  kwargs["mol"] = features;
  return exc_func(args, kwargs).toTensor();
}

To obtain the derivatives of the exchange-correlation energy with respect to the input features, which are needed to build the exchange-correlation potential, we enable gradients on the relevant feature tensors and use LibTorch’s autograd:

main.cpp

std::tuple<at::Tensor, c10::Dict<std::string, at::Tensor>>
get_exc_and_grad(const torch::jit::Method &exc_func, const FeatureDict &features)
{
  // Create a mutable copy only for the tensors that need gradients
  FeatureDict features_with_grad;
  std::vector<at::Tensor> input_tensors;
  std::vector<std::string> tensor_keys;

  for (const auto& entry : features) {
    std::string key = entry.key();
    bool requires_grad = (key == "density" || key == "grad" || key == "kin" || key == "coarse_0_atomic_coords" || key == "grid_coords" || key == "grid_weights");

    auto tensor = entry.value().clone().requires_grad_(requires_grad);
    if (requires_grad) {
      input_tensors.push_back(tensor);
      tensor_keys.push_back(key);
    }
    features_with_grad.insert(key, tensor);
  }

  IValueList args;
  IValueMap kwargs;
  kwargs["mol"] = features_with_grad;

  auto exc_on_grid = exc_func(args, kwargs).toTensor();
  auto exc = (exc_on_grid * features_with_grad.at("grid_weights")).sum();

  auto gradients = torch::autograd::grad(
      {exc},                  // outputs
      input_tensors,          // inputs
      /*grad_outputs=*/{},    // grad_outputs (defaults to ones)
      /*retain_graph=*/false, // retain_graph, necessary for higher-order grads
      /*create_graph=*/false, // create_graph, necessary for higher-order grads
      /*allow_unused=*/true   // allow_unused
  );

  c10::Dict<std::string, at::Tensor> grad;
  for (size_t i = 0; i < tensor_keys.size(); ++i)
  {
    grad.insert(tensor_keys[i], gradients[i]);
  }

  return std::make_tuple(exc_on_grid, grad);
}

Finally, the main program reads the command line arguments, loads model and features, and computes the exchange-correlation energy by weighting the energy density with the grid weights and summing over the grid points:

main.cpp

int main(int argc, const char *argv[])
{
  if (argc != 3)
  {
    std::cerr << "usage: skala_cpp_integration <path-to-fun-file> <feature-file-directory>\n";
    return -1;
  }

  const torch::DeviceType device = torch::kCPU;

  const auto [exc_func, feature_keys] = load_model(std::string(argv[1]), device);
  const auto features = load_features(std::string(argv[2]), feature_keys, device);

  std::cout << "Compute Exc..." << std::endl;

  const auto exc_on_grid = get_exc(exc_func, features);
  const auto exc = (exc_on_grid * features.at("grid_weights")).sum();

  std::cout << "Exc = " << exc.item() << std::endl;

  std::cout << "Compute Exc and dExc/dfeat..." << std::endl;

  const auto [exc_on_grid2, grad] = get_exc_and_grad(exc_func, features);
  const auto exc2 = (exc_on_grid2 * features.at("grid_weights")).sum();

  std::cout << "Exc = " << exc2.item() << std::endl;
  for (const auto &kv : grad)
  {
    std::cout << "|dExc/d(" << kv.key() << ")| = " << kv.value().norm().item() << std::endl;
  }

  return 0;
}

Running the example

Rebuild the application after changing the sources:

cmake --build build_example

And run it, passing the path to the Skala model and the feature directory as command line arguments:

./build_example/skala_cpp_integration ./skala-1.1-rev1.fun ./features

The output for the H2 molecule with the def2-QZVP basis set should look like this:

Compute Exc...
Exc = -0.624491
Compute Exc and dExc/dfeat...
Exc = -0.624491
|dExc/d(density)| = 1.39944
|dExc/d(kin)| = 0.846093
|dExc/d(grad)| = 0.0237195
|dExc/d(grid_coords)| = 0.000438993
|dExc/d(grid_weights)| = 9.72895
|dExc/d(coarse_0_atomic_coords)| = 0.00108825

Note

You are expected to add the D3(BJ) dispersion correction (using B3LYP damping parameters) to the final energy of Skala. You can use e.g. simple-dftd3 for that.

Tip

The performance tuning guide from Intel provides useful tips on how to tune the performance of PyTorch models on CPU.

Summary

In this example, we have demonstrated how to use Skala in C++ with the LibTorch library. We have shown how to set up the development environment, build the application using CMake, and run the application to compute exchange-correlation energies and their derivatives using a Skala model. This example serves as a starting point for integrating Skala into your C++ applications using LibTorch, and can be extended to include more complex features and functionalities as needed. For integrating Skala into an existing electronic structure code, see also the GauXC integration.