Examples

In this section, we show basic examples on how to define a quantum operator, quantum state, and then compute the action of the quantum operator on a quantum state, and, optionally, backward-differentiate the operator action (compute gradients) with respect to user-provided real parameters parameterizing the operator. We also show examples on how to propagate an MPS state in time using the TDVP method (both the single-site update and the two-site update that adapts the MPS bond extents), how to apply a matrix product operator (MPO) to an MPS via variational ALS fitting, how to compute the extreme eigenspectrum of a given operator, and how to compute the ground state of an MPO acting on an MPS using the split-scope DMRG method (both the 1-site and 2-site variants). For clarity, the quantum operator for each example is defined inside a separate C++ header, specifically transverse_ising_full_fused.h, transverse_ising_full_fused_noisy.h and transverse_ising_full_fused_noisy_grad.h, where it is wrapped in a helper C++ class UserDefinedLiouvillian. We also provide a utility header helpers.h containing convenient GPU array creation/destruction, initialization, copying, and printing helper functions.

Building code

Assuming cuQuantum has been extracted in CUQUANTUM_ROOT and cuTENSOR is in CUTENSOR_ROOT, we update the library path as follows:

export LD_LIBRARY_PATH=${CUQUANTUM_ROOT}/lib:${CUTENSOR_ROOT}/lib:${LD_LIBRARY_PATH}

A serial sample code discussed below (operator_action_example.cpp) can be built via the following command:

nvcc operator_action_example.cpp -I${CUQUANTUM_ROOT}/include -L${CUQUANTUM_ROOT}/lib -L${CUTENSOR_ROOT}/lib -lcudensitymat -lcutensornet -lcutensor -lcusolver -lcublas -lcurand -o operator_action_example

For static linking against the cuDensityMat library, use the following command:

nvcc operator_action_example.cpp -I${CUQUANTUM_ROOT}/include ${CUQUANTUM_ROOT}/lib/libcudensitymat_static.a ${CUQUANTUM_ROOT}/lib/libcutensornet_static.a -L${CUTENSOR_ROOT}/lib -lcutensor -lcusolver -lcublas -lcurand -o operator_action_example

In order to build a parallel (MPI) version of the example operator_action_mpi_example.cpp, one will need to have a CUDA-aware MPI library installed (e.g., recent OpenMPI, MPICH or MVAPICH) and then set the environment variable $CUDENSITYMAT_COMM_LIB to the path to the MPI interface wrapper shared library libcudensitymat_distributed_interface_mpi.so. The MPI interface wrapper shared library libcudensitymat_distributed_interface_mpi.so can be built inside the ${CUQUANTUM_ROOT}/distributed_interfaces folder by calling the build script provided there. In order to link the executable to a CUDA-aware MPI library, one will need to add -I${MPI_PATH}/include and -L${MPI_PATH}/lib -lmpi to the build command:

nvcc operator_action_mpi_example.cpp -DMPI_ENABLED -I${CUQUANTUM_ROOT}/include -I${MPI_PATH}/include -L${CUQUANTUM_ROOT}/lib -L${CUTENSOR_ROOT}/lib -L${MPI_PATH}/lib -lcudensitymat -lcutensornet -lcutensor -lcusolver -lcublas -lcurand -lmpi -o operator_action_mpi_example

In order to build a parallel (NCCL) version of the example operator_action_nccl_example.cpp, one will additionally need the NCCL library installed and the environment variable $CUDENSITYMAT_COMM_LIB set to the path to the NCCL interface wrapper shared library libcudensitymat_distributed_interface_nccl.so. The NCCL interface wrapper shared library libcudensitymat_distributed_interface_nccl.so can be built inside the ${CUQUANTUM_ROOT}/distributed_interfaces folder by calling the build script provided there. Note that MPI is still required for bootstrapping (e.g., broadcasting the ncclUniqueId). In order to link the executable, one will need to add -I${NCCL_PATH}/include and -L${NCCL_PATH}/lib -lnccl in addition to the MPI flags:

nvcc operator_action_nccl_example.cpp -DNCCL_ENABLED -DMPI_ENABLED -I${CUQUANTUM_ROOT}/include -I${MPI_PATH}/include -I${NCCL_PATH}/include -L${CUQUANTUM_ROOT}/lib -L${CUTENSOR_ROOT}/lib -L${MPI_PATH}/lib -L${NCCL_PATH}/lib -lcudensitymat -lcutensornet -lcutensor -lcusolver -lcublas -lcurand -lmpi -lnccl -o operator_action_nccl_example

Warning

When running operator_action_mpi_example.cpp with a non-CUDA-aware MPI library, the program will crash.

Note

Depending on the installation of the cuQuantum SDK package, you may need to replace lib above by lib64, depending which folder name is used inside your cuQuantum SDK package.

Code example (serial execution on a single GPU)

The following code example illustrates the common steps necessary to use the cuDensityMat library to compute the action of a quantum many-body operator on a quantum state. The full sample code can be found in the NVIDIA/cuQuantum repository (main serial code and operator definition as well as the utility code).

First let’s introduce a helper class to construct a specific quantum many-body operator, for example, the transverse field Ising Hamiltonian with fused ZZ terms and an additional noise term. Here we choose to make the f(t) coefficient depend on time and a single user-provided real parameter Omega. We use a CPU-side user-defined scalar callback function to define the dependence of the f(t) coefficient on time and the user-provided real parameter Omega. Note that inside the callback function definition, we explicitly expect the data type to be CUDA_C_64F (double-precision complex numbers), which applies to the scalar coefficient f(t) set by the callback function in-place.

/* Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#pragma once

#include <cudensitymat.h> // cuDensityMat library header
#include "helpers.h"      // GPU helper functions

#include <cmath>
#include <complex>
#include <vector>
#include <iostream>
#include <cassert>


/* DESCRIPTION:
   Time-dependent transverse-field Ising Hamiltonian operator
   with ordered and fused ZZ terms, plus fused unitary dissipation terms:
    H = sum_{i} {h_i * X_i}                // transverse field sum of X_i operators with static h_i coefficients 
      + f(t) * sum_{i < j} {g_ij * ZZ_ij}  // modulated sum of the fused ordered {Z_i * Z_j} terms with static g_ij coefficients
      + d * sum_{i} {Y_i * {..} * Y_i}     // scaled sum of the dissipation terms {Y_i * {..} * Y_i} fused into the YY_ii super-operators
   where {..} is the placeholder for the density matrix to show that the Y_i operators act from different sides.
*/

/** Define the numerical type and data type for the GPU computations (same) */
using NumericalType = std::complex<double>;      // do not change
constexpr cudaDataType_t dataType = CUDA_C_64F;  // do not change


/** Example of a user-provided scalar CPU callback C function
 *  defining a time-dependent coefficient inside the Hamiltonian:
 *  f(t) = exp(i * Omega * t) = cos(Omega * t) + i * sin(Omega * t)
 */
extern "C"
int32_t fCoefComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarStorage,    //inout: CPU-accessible storage for the returned coefficient value(s) of shape [0:batchSize-1]
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    auto * tdCoef = static_cast<cuDoubleComplex*>(scalarStorage); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t i = 0; i < batchSize; ++i) {
      const auto omega = params[i * numParams + 0]; // params[0][i]: 0-th parameter for i-th instance of the batch
      tdCoef[i] = make_cuDoubleComplex(std::cos(omega * time), std::sin(omega * time)); // value of the i-th instance of the coefficients batch
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** Convenience class which encapsulates a user-defined Liouvillian operator (system Hamiltonian + dissipation terms):
 *  - Constructor constructs the desired Liouvillian operator (`cudensitymatOperator_t`)
 *  - Method `get()` returns a reference to the constructed Liouvillian operator
 *  - Destructor releases all resources used by the Liouvillian operator
 */
class UserDefinedLiouvillian final
{
private:
  // Data members
  cudensitymatHandle_t handle;             // library context handle
  int64_t stateBatchSize;                  // quantum state batch size
  const std::vector<int64_t> spaceShape;   // Hilbert space shape (extents of the modes of the composite Hilbert space)
  void * spinXelems {nullptr};             // elements of the X spin operator in GPU RAM (F-order storage)
  void * spinYYelems {nullptr};            // elements of the fused YY two-spin operator in GPU RAM (F-order storage)
  void * spinZZelems {nullptr};            // elements of the fused ZZ two-spin operator in GPU RAM (F-order storage)
  cudensitymatElementaryOperator_t spinX;  // X spin operator (elementary tensor operator)
  cudensitymatElementaryOperator_t spinYY; // fused YY two-spin operator (elementary tensor operator)
  cudensitymatElementaryOperator_t spinZZ; // fused ZZ two-spin operator (elementary tensor operator)
  cudensitymatOperatorTerm_t oneBodyTerm;  // operator term: H1 = sum_{i} {h_i * X_i} (one-body term)
  cudensitymatOperatorTerm_t twoBodyTerm;  // operator term: H2 = f(t) * sum_{i < j} {g_ij * ZZ_ij} (two-body term)
  cudensitymatOperatorTerm_t noiseTerm;    // operator term: D1 = d * sum_{i} {YY_ii}  // Y_i operators act from different sides on the density matrix (two-body mixed term)
  cudensitymatOperator_t liouvillian;      // full operator: (-i * (H1 + H2) * {..}) + (i * {..} * (H1 + H2)) + D1{..} (super-operator)

public:

  // Constructor constructs a user-defined Liouvillian operator
  UserDefinedLiouvillian(cudensitymatHandle_t contextHandle,             // library context handle
                         const std::vector<int64_t> & hilbertSpaceShape, // Hilbert space shape
                         int64_t batchSize):                             // batch size for the quantum state
    handle(contextHandle), stateBatchSize(batchSize), spaceShape(hilbertSpaceShape)
  {
    // Define the necessary operator tensors in GPU memory (F-order storage!)
    spinXelems = createInitializeArrayGPU<NumericalType>(  // X[i0; j0]
                  {{0.0, 0.0}, {1.0, 0.0},   // 1st column of matrix X
                   {1.0, 0.0}, {0.0, 0.0}}); // 2nd column of matrix X

    spinYYelems = createInitializeArrayGPU<NumericalType>(  // YY[i0, i1; j0, j1] := Y[i0; j0] * Y[i1; j1]
                    {{0.0, 0.0},  {0.0, 0.0}, {0.0, 0.0}, {-1.0, 0.0},  // 1st column of matrix YY
                     {0.0, 0.0},  {0.0, 0.0}, {1.0, 0.0}, {0.0, 0.0},   // 2nd column of matrix YY
                     {0.0, 0.0},  {1.0, 0.0}, {0.0, 0.0}, {0.0, 0.0},   // 3rd column of matrix YY
                     {-1.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}}); // 4th column of matrix YY

    spinZZelems = createInitializeArrayGPU<NumericalType>(  // ZZ[i0, i1; j0, j1] := Z[i0; j0] * Z[i1; j1]
                    {{1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {0.0, 0.0},   // 1st column of matrix ZZ
                     {0.0, 0.0}, {-1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},   // 2nd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {-1.0, 0.0}, {0.0, 0.0},   // 3rd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {1.0, 0.0}}); // 4th column of matrix ZZ

    // Construct the necessary Elementary Tensor Operators
    //  X_i operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        1,                                   // one-body operator
                        std::vector<int64_t>({2}).data(),    // acts in tensor space of shape {2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinXelems,                          // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinX));                            // the created elementary tensor operator
    //  ZZ_ij = Z_i * Z_j fused operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        2,                                   // two-body operator
                        std::vector<int64_t>({2,2}).data(),  // acts in tensor space of shape {2,2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinZZelems,                         // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinZZ));                           // the created elementary tensor operator
    //  YY_ii = Y_i * {..} * Y_i fused operator (note action from different sides)
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        2,                                   // two-body operator
                        std::vector<int64_t>({2,2}).data(),  // acts in tensor space of shape {2,2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinYYelems,                         // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinYY));                           // the created elementary tensor operator

    // Construct the necessary Operator Terms from tensor products of Elementary Tensor Operators
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &oneBodyTerm));                      // the created empty operator term
    //  Define the operator term: H1 = sum_{i} {h_i * X_i}
    for (int32_t i = 0; i < spaceShape.size(); ++i) {
      const double h_i = 1.0 / static_cast<double>(i+1); // assign some value to the time-independent h_i coefficient
      HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                          oneBodyTerm,
                          1,                                                             // number of elementary tensor operators in the product
                          std::vector<cudensitymatElementaryOperator_t>({spinX}).data(), // elementary tensor operators forming the product
                          std::vector<int32_t>({i}).data(),                              // space modes acted on by the operator product
                          std::vector<int32_t>({0}).data(),                              // space mode action duality (0: from the left; 1: from the right)
                          make_cuDoubleComplex(h_i, 0.0),                                // h_i constant coefficient: Always 64-bit-precision complex number
                          cudensitymatScalarCallbackNone,                                // no time-dependent coefficient associated with this operator product
                          cudensitymatScalarGradientCallbackNone));                      // no coefficient gradient associated with this operator product
    }
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &twoBodyTerm));                      // the created empty operator term
    //  Define the operator term: H2 = f(t) * sum_{i < j} {g_ij * ZZ_ij}
    for (int32_t i = 0; i < spaceShape.size() - 1; ++i) {
      for (int32_t j = (i + 1); j < spaceShape.size(); ++j) {
        const double g_ij = -1.0 / static_cast<double>(i + j + 1); // assign some value to the time-independent g_ij coefficient
        HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                            twoBodyTerm,
                            1,                                                              // number of elementary tensor operators in the product
                            std::vector<cudensitymatElementaryOperator_t>({spinZZ}).data(), // elementary tensor operators forming the product
                            std::vector<int32_t>({i, j}).data(),                            // space modes acted on by the operator product
                            std::vector<int32_t>({0, 0}).data(),                            // space mode action duality (0: from the left; 1: from the right)
                            make_cuDoubleComplex(g_ij, 0.0),                                // g_ij constant coefficient: Always 64-bit-precision complex number
                            cudensitymatScalarCallbackNone,                                 // no time-dependent coefficient associated with this operator product
                            cudensitymatScalarGradientCallbackNone));                       // no coefficient gradient associated with this operator product
      }
    }
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &noiseTerm));                        // the created empty operator term
    //  Define the operator term: D1 = d * sum_{i} {YY_ii}
    for (int32_t i = 0; i < spaceShape.size(); ++i) {
      HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                          noiseTerm,
                          1,                                                              // number of elementary tensor operators in the product
                          std::vector<cudensitymatElementaryOperator_t>({spinYY}).data(), // elementary tensor operators forming the product
                          std::vector<int32_t>({i, i}).data(),                            // space modes acted on by the operator product (from different sides)
                          std::vector<int32_t>({0, 1}).data(),                            // space mode action duality (0: from the left; 1: from the right)
                          make_cuDoubleComplex(1.0, 0.0),                                 // default coefficient: Always 64-bit-precision complex number
                          cudensitymatScalarCallbackNone,                                 // no time-dependent coefficient associated with this operator product
                          cudensitymatScalarGradientCallbackNone));                       // no coefficient gradient associated with this operator product
    }

    // Construct the full Liouvillian operator as a sum of the operator terms
    //  Create an empty operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                        spaceShape.size(),                // Hilbert space rank (number of modes)
                        spaceShape.data(),                // Hilbert space shape (modes extents)
                        &liouvillian));                   // the created empty operator (super-operator)
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        oneBodyTerm,                      // appended operator term
                        0,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, -1.0),  // -i constant
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        twoBodyTerm,                     // appended operator term
                        0,                               // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, -1.0), // -i constant
                        {fCoefComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr}, // CPU scalar callback function defining the time-dependent coefficient associated with this operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with this operator term as a whole
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        oneBodyTerm,                      // appended operator term
                        1,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, 1.0),   // i constant
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        twoBodyTerm,                     // appended operator term
                        1,                               // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, 1.0),  // i constant
                        {fCoefComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr}, // CPU scalar callback function defining the time-dependent coefficient associated with this operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with this operator term as a whole
    //  Append an operator term to the operator (super-operator)
    const double d = 0.42; // assign some value to the time-independent coefficient
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        noiseTerm,                        // appended operator term
                        0,                                // operator term action duality as a whole (no duality reversing in this case)
                        make_cuDoubleComplex(d, 0.0),     // constant coefficient associated with the operator term as a whole
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
  }

  // Destructor destructs the user-defined Liouvillian operator
  ~UserDefinedLiouvillian()
  {
    // Destroy the Liouvillian operator
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(liouvillian));

    // Destroy operator terms
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(noiseTerm));
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(twoBodyTerm));
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(oneBodyTerm));

    // Destroy elementary tensor operators
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinYY));
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinZZ));
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinX));

    // Destroy operator tensors
    destroyArrayGPU(spinYYelems);
    destroyArrayGPU(spinZZelems);
    destroyArrayGPU(spinXelems);
  }

  // Disable copy constructor/assignment (GPU resources are private, no deep copy)
  UserDefinedLiouvillian(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian & operator=(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian(UserDefinedLiouvillian &&) = delete;
  UserDefinedLiouvillian & operator=(UserDefinedLiouvillian &&) = delete;

  /** Returns the number of externally provided Hamiltonian parameters. */
  int32_t getNumParameters() const
  {
    return 1; // one parameter Omega
  }

  /** Get access to the constructed Liouvillian operator. */
  cudensitymatOperator_t & get()
  {
    return liouvillian;
  }

};

Now we can use this parameterized quantum many-body operator in our main code to compute the action of the operator on a mixed quantum state (density matrix).

/* Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <cudensitymat.h>  // cuDensityMat library header
#include "helpers.h"       // helper functions


// Transverse Ising Hamiltonian with double summation ordering
// and spin-operator fusion, plus fused dissipation terms
#include "transverse_ising_full_fused_noisy.h"  // user-defined Liouvillian operator example

#include <cmath>
#include <complex>
#include <vector>
#include <chrono>
#include <iostream>
#include <cassert>


// Number of times to perform operator action on a quantum state
constexpr int NUM_REPEATS = 2;

// Logging verbosity
bool verbose = true;


// Example workflow
void exampleWorkflow(cudensitymatHandle_t handle)
{
  // Define the composite Hilbert space shape and
  // quantum state batch size (number of individual quantum states in a batched simulation)
  const std::vector<int64_t> spaceShape({2,2,2,2,2,2,2,2}); // dimensions of quantum degrees of freedom
  const int64_t batchSize = 1;                              // number of quantum states per batch (default is 1)

  if (verbose) {
    std::cout << "Hilbert space rank = " << spaceShape.size() << "; Shape = (";
    for (const auto & dimsn: spaceShape)
      std::cout << dimsn << ",";
    std::cout << ")" << std::endl;
    std::cout << "Quantum state batch size = " << batchSize << std::endl;
  }

  // Construct a user-defined Liouvillian operator using a convenience C++ class
  UserDefinedLiouvillian liouvillian(handle, spaceShape, batchSize);
  if (verbose)
    std::cout << "Constructed the Liouvillian operator\n";

  // Set and place external user-provided Hamiltonian parameters in GPU memory
  const int32_t numParams = liouvillian.getNumParameters(); // number of external user-provided Hamiltonian parameters
  std::vector<double> cpuHamParams(numParams * batchSize);
  for (int64_t j = 0; j < batchSize; ++j) {
    for (int32_t i = 0; i < numParams; ++i) {
      cpuHamParams[j * numParams + i] = double(i+1) / double(j+1); // just setting some parameter values for each instance of the batch
    }
  }
  auto * hamiltonianParams = static_cast<double *>(createInitializeArrayGPU(cpuHamParams));

  // Declare the input quantum state
  cudensitymatState_t inputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &inputState));

  // Query the size of the quantum state storage
  std::size_t storageSize {0}; // only one storage component (tensor) is needed (no tensor factorization)
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      inputState,
                      1,               // only one storage component (tensor)
                      &storageSize));  // storage size in bytes
  const std::size_t stateVolume = storageSize / sizeof(NumericalType);  // quantum state tensor volume (number of elements)
  if (verbose)
    std::cout << "Quantum state storage size (bytes) = " << storageSize << std::endl;

  // Prepare some initial value for the input quantum state batch
  std::vector<NumericalType> inputStateValue(stateVolume);
  if constexpr (std::is_same_v<NumericalType, float>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0f / float(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, double>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0 / double(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<float>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0f / float(i+1), -1.0f / float(i+2)}; // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<double>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0 / double(i+1), -1.0 / double(i+2)}; // just some value
    }
  } else {
    std::cerr << "Error: Unsupported data type!\n";
    std::exit(1);
  }
  // Allocate initialized GPU storage for the input quantum state with prepared values
  auto * inputStateElems = createInitializeArrayGPU(inputStateValue);
  if (verbose)
    std::cout << "Allocated input quantum state storage and initialized it to some value\n";

  // Attach initialized GPU storage to the input quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      inputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({inputStateElems}).data(),      // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed input quantum state\n";

  // Declare the output quantum state of the same shape
  cudensitymatState_t outputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &outputState));

  // Allocate initialized GPU storage for the output quantum state
  auto * outputStateElems = createArrayGPU<NumericalType>(stateVolume);
  if (verbose)
    std::cout << "Allocated output quantum state storage\n";

  // Attach GPU storage to the output quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      outputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({outputStateElems}).data(),     // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed output quantum state\n";

  // Declare a workspace descriptor
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  // Query free GPU memory
  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95); // take 95% of the free memory for the workspace buffer
  if (verbose)
    std::cout << "Max workspace buffer size (bytes) = " << freeMem << std::endl;

  // Prepare the Liouvillian operator action on a quantum state (needs to be done only once)
  const auto startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorPrepareAction(handle,
                      liouvillian.get(),
                      inputState,
                      outputState,
                      CUDENSITYMAT_COMPUTE_64F,  // GPU compute type
                      freeMem,                   // max available GPU free memory for the workspace
                      workspaceDescr,            // workspace descriptor
                      0x0));                     // default CUDA stream
  const auto finishTime = std::chrono::high_resolution_clock::now();
  const std::chrono::duration<double> timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator action preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  std::size_t requiredBufferSize {0};
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  // Allocate GPU storage for the workspace buffer
  const std::size_t bufferVolume = requiredBufferSize / sizeof(NumericalType);
  auto * workspaceBuffer = createArrayGPU<NumericalType>(bufferVolume);
  if (verbose)
    std::cout << "Allocated workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      requiredBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Apply the Liouvillian operator to the input quatum state
  // and accumulate its action into the output quantum state (note accumulative += semantics)
  for (int32_t repeat = 0; repeat < NUM_REPEATS; ++repeat) { // repeat multiple times for accurate timing
    // Zero out the output quantum state
    HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle,
                        outputState,
                        0x0));
    if (verbose)
      std::cout << "Initialized the output quantum state to zero\n";
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    const auto startTime = std::chrono::high_resolution_clock::now();
    HANDLE_CUDM_ERROR(cudensitymatOperatorComputeAction(handle,
                        liouvillian.get(),
                        0.3,                                   // time point (some value)
                        batchSize,                             // user-defined batch size
                        numParams,                             // number of external user-defined Hamiltonian parameters
                        hamiltonianParams,                     // external Hamiltonian parameters in GPU memory
                        inputState,                            // input quantum state
                        outputState,                           // output quantum state
                        workspaceDescr,                        // workspace descriptor
                        0x0));                                 // default CUDA stream
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    const auto finishTime = std::chrono::high_resolution_clock::now();
    const std::chrono::duration<double> timeSec = finishTime - startTime;
    if (verbose)
      std::cout << "Operator action computation time (sec) = " << timeSec.count() << std::endl;
  }

  // Compute the squared norm of the output quantum state
  void * norm2 = createInitializeArrayGPU(std::vector<double>(batchSize, 0.0));
  HANDLE_CUDM_ERROR(cudensitymatStateComputeNorm(handle,
                      outputState,
                      norm2,
                      0x0));
  if (verbose) {
    std::cout << "Computed the output quantum state norm:\n";
    printArrayGPU<double>(norm2, batchSize);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Destroy the norm2 array
  destroyArrayGPU(norm2);

  // Destroy workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));

  // Destroy workspace buffer storage
  destroyArrayGPU(workspaceBuffer);

  // Destroy quantum states
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(outputState));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(inputState));

  // Destroy quantum state storage
  destroyArrayGPU(outputStateElems);
  destroyArrayGPU(inputStateElems);

  // Destroy external Hamiltonian parameters
  destroyArrayGPU(static_cast<void *>(hamiltonianParams));

  if (verbose)
    std::cout << "Destroyed resources\n" << std::flush;
}


int main(int argc, char ** argv)
{
  // Assign a GPU to the process
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose)
    std::cout << "Set active device\n";

  // Create a library handle
  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created a library handle\n";

  // Run the example
  exampleWorkflow(handle);

  // Destroy the library handle
  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed the library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());

  // Done
  return 0;
}

Code example (parallel execution on multiple GPUs)

It is straightforward to adapt the main serial code and enable parallel execution across multiple/many GPU devices (across multiple/many nodes). Two distributed communication backends are supported: MPI and NCCL (experimental).

MPI backend

We will illustrate parallel execution with an example using the Message Passing Interface (MPI) as the communication layer. Below we show the minor additions that need to be made in order to enable distributed parallel execution without making any changes to the original serial source code.

The full sample code can be found in the NVIDIA/cuQuantum repository (main MPI code and operator definition as well as the utility code).

Here is the updated main code for multi-GPU runs using MPI.

/* Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <cudensitymat.h>  // cuDensityMat library header
#include "helpers.h"       // helper functions


// Transverse Ising Hamiltonian with double summation ordering
// and spin-operator fusion, plus fused dissipation terms
#include "transverse_ising_full_fused_noisy.h"  // user-defined Liouvillian operator example


// MPI library (optional)
#ifdef MPI_ENABLED
#include <mpi.h>
#endif

#include <cmath>
#include <complex>
#include <vector>
#include <chrono>
#include <iostream>
#include <cassert>


// Number of times to perform operator action on a quantum state
constexpr int NUM_REPEATS = 2;

// Logging verbosity
bool verbose = true;


// Example workflow
void exampleWorkflow(cudensitymatHandle_t handle)
{
  // Define the composite Hilbert space shape and
  // quantum state batch size (number of individual quantum states in a batched simulation)
  const std::vector<int64_t> spaceShape({2,2,2,2,2,2,2,2}); // dimensions of quantum degrees of freedom
  const int64_t batchSize = 1;                              // number of quantum states per batch (default is 1)

  if (verbose) {
    std::cout << "Hilbert space rank = " << spaceShape.size() << "; Shape = (";
    for (const auto & dimsn: spaceShape)
      std::cout << dimsn << ",";
    std::cout << ")" << std::endl;
    std::cout << "Quantum state batch size = " << batchSize << std::endl;
  }

  // Construct a user-defined Liouvillian operator using a convenience C++ class
  UserDefinedLiouvillian liouvillian(handle, spaceShape, batchSize);
  if (verbose)
    std::cout << "Constructed the Liouvillian operator\n";

  // Set and place external user-provided Hamiltonian parameters in GPU memory
  const int32_t numParams = liouvillian.getNumParameters(); // number of external user-provided Hamiltonian parameters
  std::vector<double> cpuHamParams(numParams * batchSize);
  for (int64_t j = 0; j < batchSize; ++j) {
    for (int32_t i = 0; i < numParams; ++i) {
      cpuHamParams[j * numParams + i] = double(i+1) / double(j+1); // just setting some parameter values for each instance of the batch
    }
  }
  auto * hamiltonianParams = static_cast<double *>(createInitializeArrayGPU(cpuHamParams));

  // Declare the input quantum state
  cudensitymatState_t inputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &inputState));

  // Query the size of the quantum state storage
  std::size_t storageSize {0}; // only one storage component (tensor) is needed (no tensor factorization)
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      inputState,
                      1,               // only one storage component (tensor)
                      &storageSize));  // storage size in bytes
  const std::size_t stateVolume = storageSize / sizeof(NumericalType);  // quantum state tensor volume (number of elements)
  if (verbose)
    std::cout << "Quantum state storage size (bytes) = " << storageSize << std::endl;

  // Prepare some initial value for the input quantum state batch
  std::vector<NumericalType> inputStateValue(stateVolume);
  if constexpr (std::is_same_v<NumericalType, float>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0f / float(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, double>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0 / double(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<float>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0f / float(i+1), -1.0f / float(i+2)}; // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<double>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0 / double(i+1), -1.0 / double(i+2)}; // just some value
    }
  } else {
    std::cerr << "Error: Unsupported data type!\n";
    std::exit(1);
  }
  // Allocate initialized GPU storage for the input quantum state with prepared values
  auto * inputStateElems = createInitializeArrayGPU(inputStateValue);
  if (verbose)
    std::cout << "Allocated input quantum state storage and initialized it to some value\n";

  // Attach initialized GPU storage to the input quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      inputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({inputStateElems}).data(),      // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed input quantum state\n";

  // Declare the output quantum state of the same shape
  cudensitymatState_t outputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &outputState));

  // Allocate initialized GPU storage for the output quantum state
  auto * outputStateElems = createArrayGPU<NumericalType>(stateVolume);
  if (verbose)
    std::cout << "Allocated output quantum state storage\n";

  // Attach GPU storage to the output quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      outputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({outputStateElems}).data(),     // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed output quantum state\n";

  // Declare a workspace descriptor
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  // Query free GPU memory
  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95); // take 95% of the free memory for the workspace buffer
  if (verbose)
    std::cout << "Max workspace buffer size (bytes) = " << freeMem << std::endl;

  // Prepare the Liouvillian operator action on a quantum state (needs to be done only once)
  const auto startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorPrepareAction(handle,
                      liouvillian.get(),
                      inputState,
                      outputState,
                      CUDENSITYMAT_COMPUTE_64F,  // GPU compute type
                      freeMem,                   // max available GPU free memory for the workspace
                      workspaceDescr,            // workspace descriptor
                      0x0));                     // default CUDA stream
  const auto finishTime = std::chrono::high_resolution_clock::now();
  const std::chrono::duration<double> timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator action preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  std::size_t requiredBufferSize {0};
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  // Allocate GPU storage for the workspace buffer
  const std::size_t bufferVolume = requiredBufferSize / sizeof(NumericalType);
  auto * workspaceBuffer = createArrayGPU<NumericalType>(bufferVolume);
  if (verbose)
    std::cout << "Allocated workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      requiredBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Apply the Liouvillian operator to the input quatum state
  // and accumulate its action into the output quantum state (note accumulative += semantics)
  for (int32_t repeat = 0; repeat < NUM_REPEATS; ++repeat) { // repeat multiple times for accurate timing
    // Zero out the output quantum state
    HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle,
                        outputState,
                        0x0));
    if (verbose)
      std::cout << "Initialized the output quantum state to zero\n";
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    const auto startTime = std::chrono::high_resolution_clock::now();
    HANDLE_CUDM_ERROR(cudensitymatOperatorComputeAction(handle,
                        liouvillian.get(),
                        0.3,                                   // time point (some value)
                        batchSize,                             // user-defined batch size
                        numParams,                             // number of external user-defined Hamiltonian parameters
                        hamiltonianParams,                     // external Hamiltonian parameters in GPU memory
                        inputState,                            // input quantum state
                        outputState,                           // output quantum state
                        workspaceDescr,                        // workspace descriptor
                        0x0));                                 // default CUDA stream
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    const auto finishTime = std::chrono::high_resolution_clock::now();
    const std::chrono::duration<double> timeSec = finishTime - startTime;
    if (verbose)
      std::cout << "Operator action computation time (sec) = " << timeSec.count() << std::endl;
  }

  // Compute the squared norm of the output quantum state
  void * norm2 = createInitializeArrayGPU(std::vector<double>(batchSize, 0.0));
  HANDLE_CUDM_ERROR(cudensitymatStateComputeNorm(handle,
                      outputState,
                      norm2,
                      0x0));
  if (verbose) {
    std::cout << "Computed the output quantum state norm:\n";
    printArrayGPU<double>(norm2, batchSize);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Destroy the norm2 array
  destroyArrayGPU(norm2);

  // Destroy workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));

  // Destroy workspace buffer storage
  destroyArrayGPU(workspaceBuffer);

  // Destroy quantum states
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(outputState));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(inputState));

  // Destroy quantum state storage
  destroyArrayGPU(outputStateElems);
  destroyArrayGPU(inputStateElems);

  // Destroy external Hamiltonian parameters
  destroyArrayGPU(static_cast<void *>(hamiltonianParams));

  if (verbose)
    std::cout << "Destroyed resources\n" << std::flush;
}


int main(int argc, char ** argv)
{
  // Initialize MPI library (if needed)
#ifdef MPI_ENABLED
  HANDLE_MPI_ERROR(MPI_Init(&argc, &argv));
  int procRank {-1};
  HANDLE_MPI_ERROR(MPI_Comm_rank(MPI_COMM_WORLD, &procRank));
  int numProcs {0};
  HANDLE_MPI_ERROR(MPI_Comm_size(MPI_COMM_WORLD, &numProcs));
  if (procRank != 0) verbose = false;
  if (verbose)
    std::cout << "Initialized MPI library\n";
#else
  const int procRank {0};
  const int numProcs {1};
#endif

  // Assign a GPU to the process
  int numDevices {0};
  HANDLE_CUDA_ERROR(cudaGetDeviceCount(&numDevices));
  const int deviceId = procRank % numDevices;
  HANDLE_CUDA_ERROR(cudaSetDevice(deviceId));
  if (verbose)
    std::cout << "Set active device\n";

  // Create a library handle
  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created a library handle\n";

  // Reset distributed configuration (once)
#ifdef MPI_ENABLED
  MPI_Comm comm;
  HANDLE_MPI_ERROR(MPI_Comm_dup(MPI_COMM_WORLD, &comm));
  HANDLE_CUDM_ERROR(cudensitymatResetDistributedConfiguration(handle,
                      CUDENSITYMAT_DISTRIBUTED_PROVIDER_MPI,
                      &comm, sizeof(comm)));
#endif

  // Run the example
  exampleWorkflow(handle);

  // Synchronize MPI processes
#ifdef MPI_ENABLED
  HANDLE_MPI_ERROR(MPI_Barrier(MPI_COMM_WORLD));
#endif

  // Destroy the library handle
  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed the library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());

  // Finalize the MPI library
#ifdef MPI_ENABLED
  HANDLE_MPI_ERROR(MPI_Finalize());
  if (verbose)
    std::cout << "Finalized MPI library\n";
#endif

  // Done
  return 0;
}

NCCL backend (experimental)

NCCL (NVIDIA Collective Communications Library) can provide better performance for GPU-to-GPU communication, especially within a single node with NVLink connectivity. Below we show an example using NCCL as the communication layer. Note that the NCCL backend is currently experimental. MPI is used for process spawning and bootstrapping (e.g., broadcasting the ncclUniqueId), but all GPU-to-GPU communication uses NCCL.

The full sample code can be found in the NVIDIA/cuQuantum repository (main NCCL code and operator definition as well as the utility code).

Here is the main code for multi-GPU runs using NCCL.

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <cudensitymat.h>  // cuDensityMat library header
#include "helpers.h"       // helper functions


// Transverse Ising Hamiltonian with double summation ordering
// and spin-operator fusion, plus fused dissipation terms
#include "transverse_ising_full_fused_noisy.h"  // user-defined Liouvillian operator example


// NCCL library (required for this example)
#ifdef NCCL_ENABLED
#include <nccl.h>
#endif

// MPI library (used for bootstrapping NCCL communicator)
#ifdef MPI_ENABLED
#include <mpi.h>
#endif

#include <cmath>
#include <complex>
#include <vector>
#include <chrono>
#include <iostream>
#include <cassert>


// Number of times to perform operator action on a quantum state
constexpr int NUM_REPEATS = 2;

// Logging verbosity
bool verbose = true;


#ifdef NCCL_ENABLED
// Error handling macro for NCCL
#define HANDLE_NCCL_ERROR(x)                                 \
{                                                            \
  const ncclResult_t err = x;                                \
  if (err != ncclSuccess)                                    \
  {                                                          \
    printf("NCCL Error: %s in line %d\n",                    \
           ncclGetErrorString(err), __LINE__);               \
    fflush(stdout);                                          \
    std::abort();                                            \
  }                                                          \
};
#endif


// Example workflow
void exampleWorkflow(cudensitymatHandle_t handle)
{
  // Define the composite Hilbert space shape and
  // quantum state batch size (number of individual quantum states in a batched simulation)
  const std::vector<int64_t> spaceShape({2,2,2,2,2,2,2,2}); // dimensions of quantum degrees of freedom
  const int64_t batchSize = 1;                              // number of quantum states per batch (default is 1)

  if (verbose) {
    std::cout << "Hilbert space rank = " << spaceShape.size() << "; Shape = (";
    for (const auto & dimsn: spaceShape)
      std::cout << dimsn << ",";
    std::cout << ")" << std::endl;
    std::cout << "Quantum state batch size = " << batchSize << std::endl;
  }

  // Construct a user-defined Liouvillian operator using a convenience C++ class
  UserDefinedLiouvillian liouvillian(handle, spaceShape, batchSize);
  if (verbose)
    std::cout << "Constructed the Liouvillian operator\n";

  // Set and place external user-provided Hamiltonian parameters in GPU memory
  const int32_t numParams = liouvillian.getNumParameters(); // number of external user-provided Hamiltonian parameters
  std::vector<double> cpuHamParams(numParams * batchSize);
  for (int64_t j = 0; j < batchSize; ++j) {
    for (int32_t i = 0; i < numParams; ++i) {
      cpuHamParams[j * numParams + i] = double(i+1) / double(j+1); // just setting some parameter values for each instance of the batch
    }
  }
  auto * hamiltonianParams = static_cast<double *>(createInitializeArrayGPU(cpuHamParams));

  // Declare the input quantum state
  cudensitymatState_t inputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &inputState));

  // Query the size of the quantum state storage
  std::size_t storageSize {0}; // only one storage component (tensor) is needed (no tensor factorization)
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      inputState,
                      1,               // only one storage component (tensor)
                      &storageSize));  // storage size in bytes
  const std::size_t stateVolume = storageSize / sizeof(NumericalType);  // quantum state tensor volume (number of elements)
  if (verbose)
    std::cout << "Quantum state storage size (bytes) = " << storageSize << std::endl;

  // Prepare some initial value for the input quantum state batch
  std::vector<NumericalType> inputStateValue(stateVolume);
  if constexpr (std::is_same_v<NumericalType, float>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0f / float(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, double>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0 / double(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<float>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0f / float(i+1), -1.0f / float(i+2)}; // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<double>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0 / double(i+1), -1.0 / double(i+2)}; // just some value
    }
  } else {
    std::cerr << "Error: Unsupported data type!\n";
    std::exit(1);
  }
  // Allocate initialized GPU storage for the input quantum state with prepared values
  auto * inputStateElems = createInitializeArrayGPU(inputStateValue);
  if (verbose)
    std::cout << "Allocated input quantum state storage and initialized it to some value\n";

  // Attach initialized GPU storage to the input quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      inputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({inputStateElems}).data(),      // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed input quantum state\n";

  // Declare the output quantum state of the same shape
  cudensitymatState_t outputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &outputState));

  // Allocate initialized GPU storage for the output quantum state
  auto * outputStateElems = createArrayGPU<NumericalType>(stateVolume);
  if (verbose)
    std::cout << "Allocated output quantum state storage\n";

  // Attach GPU storage to the output quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      outputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({outputStateElems}).data(),     // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed output quantum state\n";

  // Declare a workspace descriptor
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  // Query free GPU memory
  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95); // take 95% of the free memory for the workspace buffer
  if (verbose)
    std::cout << "Max workspace buffer size (bytes) = " << freeMem << std::endl;

  // Prepare the Liouvillian operator action on a quantum state (needs to be done only once)
  const auto startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorPrepareAction(handle,
                      liouvillian.get(),
                      inputState,
                      outputState,
                      CUDENSITYMAT_COMPUTE_64F,  // GPU compute type
                      freeMem,                   // max available GPU free memory for the workspace
                      workspaceDescr,            // workspace descriptor
                      0x0));                     // default CUDA stream
  const auto finishTime = std::chrono::high_resolution_clock::now();
  const std::chrono::duration<double> timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator action preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  std::size_t requiredBufferSize {0};
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  // Allocate GPU storage for the workspace buffer
  const std::size_t bufferVolume = requiredBufferSize / sizeof(NumericalType);
  auto * workspaceBuffer = createArrayGPU<NumericalType>(bufferVolume);
  if (verbose)
    std::cout << "Allocated workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      requiredBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Apply the Liouvillian operator to the input quatum state
  // and accumulate its action into the output quantum state (note accumulative += semantics)
  for (int32_t repeat = 0; repeat < NUM_REPEATS; ++repeat) { // repeat multiple times for accurate timing
    // Zero out the output quantum state
    HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle,
                        outputState,
                        0x0));
    if (verbose)
      std::cout << "Initialized the output quantum state to zero\n";
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    const auto startTime = std::chrono::high_resolution_clock::now();
    HANDLE_CUDM_ERROR(cudensitymatOperatorComputeAction(handle,
                        liouvillian.get(),
                        0.3,                                   // time point (some value)
                        batchSize,                             // user-defined batch size
                        numParams,                             // number of external user-defined Hamiltonian parameters
                        hamiltonianParams,                     // external Hamiltonian parameters in GPU memory
                        inputState,                            // input quantum state
                        outputState,                           // output quantum state
                        workspaceDescr,                        // workspace descriptor
                        0x0));                                 // default CUDA stream
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    const auto finishTime = std::chrono::high_resolution_clock::now();
    const std::chrono::duration<double> timeSec = finishTime - startTime;
    if (verbose)
      std::cout << "Operator action computation time (sec) = " << timeSec.count() << std::endl;
  }

  // Compute the squared norm of the output quantum state
  void * norm2 = createInitializeArrayGPU(std::vector<double>(batchSize, 0.0));
  HANDLE_CUDM_ERROR(cudensitymatStateComputeNorm(handle,
                      outputState,
                      norm2,
                      0x0));
  if (verbose) {
    std::cout << "Computed the output quantum state norm:\n";
    printArrayGPU<double>(norm2, batchSize);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Destroy the norm2 array
  destroyArrayGPU(norm2);

  // Destroy workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));

  // Destroy workspace buffer storage
  destroyArrayGPU(workspaceBuffer);

  // Destroy quantum states
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(outputState));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(inputState));

  // Destroy quantum state storage
  destroyArrayGPU(outputStateElems);
  destroyArrayGPU(inputStateElems);

  // Destroy external Hamiltonian parameters
  destroyArrayGPU(static_cast<void *>(hamiltonianParams));

  if (verbose)
    std::cout << "Destroyed resources\n" << std::flush;
}


int main(int argc, char ** argv)
{
#if defined(NCCL_ENABLED) && defined(MPI_ENABLED)
  // Initialize MPI library (used to bootstrap NCCL)
  HANDLE_MPI_ERROR(MPI_Init(&argc, &argv));
  int procRank {-1};
  HANDLE_MPI_ERROR(MPI_Comm_rank(MPI_COMM_WORLD, &procRank));
  int numProcs {0};
  HANDLE_MPI_ERROR(MPI_Comm_size(MPI_COMM_WORLD, &numProcs));
  if (procRank != 0) verbose = false;
  if (verbose)
    std::cout << "Initialized MPI library (for NCCL bootstrap)\n";

  // Assign a GPU to the process
  int numDevices {0};
  HANDLE_CUDA_ERROR(cudaGetDeviceCount(&numDevices));
  const int deviceId = procRank % numDevices;
  HANDLE_CUDA_ERROR(cudaSetDevice(deviceId));
  if (verbose)
    std::cout << "Set active device to GPU " << deviceId << "\n";

  // Initialize NCCL communicator
  // Step 1: Generate unique ID on rank 0 and broadcast to all ranks
  ncclUniqueId ncclId;
  if (procRank == 0) {
    HANDLE_NCCL_ERROR(ncclGetUniqueId(&ncclId));
  }
  HANDLE_MPI_ERROR(MPI_Bcast(&ncclId, sizeof(ncclId), MPI_BYTE, 0, MPI_COMM_WORLD));

  // Step 2: Initialize NCCL communicator with the shared unique ID
  ncclComm_t ncclComm;
  HANDLE_NCCL_ERROR(ncclCommInitRank(&ncclComm, numProcs, ncclId, procRank));
  if (verbose)
    std::cout << "Initialized NCCL communicator with " << numProcs << " ranks\n";

  // Create a library handle
  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created a library handle\n";

  // Reset distributed configuration with NCCL communicator
  // The barrier buffer is now managed internally by cuDensityMat
  HANDLE_CUDM_ERROR(cudensitymatResetDistributedConfiguration(handle,
                      CUDENSITYMAT_DISTRIBUTED_PROVIDER_NCCL,
                      &ncclComm, sizeof(ncclComm)));
  if (verbose)
    std::cout << "Configured distributed execution with NCCL\n";

  // Run the example
  exampleWorkflow(handle);

  // Synchronize processes via NCCL barrier (uses allreduce internally)
  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  HANDLE_MPI_ERROR(MPI_Barrier(MPI_COMM_WORLD));

  // Destroy the library handle
  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed the library handle\n";

  // Finalize NCCL communicator
  HANDLE_NCCL_ERROR(ncclCommFinalize(ncclComm));
  HANDLE_NCCL_ERROR(ncclCommDestroy(ncclComm));
  if (verbose)
    std::cout << "Finalized NCCL communicator\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());

  // Finalize the MPI library
  HANDLE_MPI_ERROR(MPI_Finalize());
  if (verbose)
    std::cout << "Finalized MPI library\n";

#else
  // Fallback for when NCCL or MPI is not enabled
  (void)argc;
  (void)argv;
  std::cerr << "This example requires both NCCL_ENABLED and MPI_ENABLED to be defined.\n";
  std::cerr << "NCCL uses MPI for bootstrapping (sharing ncclUniqueId across processes).\n";
  std::cerr << "Build with: -DENABLE_NCCL=TRUE -DENABLE_MPI=TRUE\n";
  return 1;
#endif

  // Done
  return 0;
}

Code example (serial execution with backward differentiation)

The following code example illustrates how to use the cuDensityMat library to not only compute the action of a quantum many-body operator on a quantum state, but also backward-differentiate it (compute gradients) with respect to user-provided real parameters parameterizing the operator (one real parameter Omega in this example). The full sample code can be found in the NVIDIA/cuQuantum repository (main serial gradient code and operator definition for gradient as well as the utility code).

First let’s construct a specific quantum many-body operator which, in this case, is a slightly modified version of the quantum many-body operator used in main serial code. Here we make both the h(t) and f(t) scalar coefficients depend on time and a single user-provided real parameter Omega via different (made-up) functional forms. In order to backward-differentiate the operator action with respect to this single user-provided real parameter Omega, we need to manually define a gradient callback function for each regular callback function we have (for h(t) and f(t) in this example). In our example, we define two CPU-side scalar gradient callback functions which compute the vector-jacobian product (VJP) of the scalar adjoint of h(t) and f(t) with respect to the user-provided real parameter Omega, respectively. A gradient callback function is expected to accumulate the VJP result into the paramsGrad output array, the final value of which will be the gradient(s) of the user-defined cost function with respect to the user-provided real parameters parameterizing the operator. As before, all regular and gradient callback functions used in our example explicitly expect the data type to be CUDA_C_64F (double-precision complex numbers).

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#pragma once

#include <cudensitymat.h> // cuDensityMat library header
#include "helpers.h"      // GPU helper functions

#include <cmath>
#include <complex>
#include <vector>
#include <iostream>
#include <cassert>


/* DESCRIPTION:
   Time-dependent transverse-field Ising Hamiltonian operator
   with ordered and fused ZZ terms, plus fused unitary dissipation terms:
    H = sum_{i} {h_i(t) * X_i}             // transverse field sum of X_i operators with time-dependent h_i(t) coefficients 
      + f(t) * sum_{i < j} {g_ij * ZZ_ij}  // modulated sum of the fused ordered {Z_i * Z_j} terms with static g_ij coefficients
      + d * sum_{i} {Y_i * {..} * Y_i}     // scaled sum of the dissipation terms {Y_i * {..} * Y_i} fused into the YY_ii super-operators
   where {..} is the placeholder for the density matrix to show that the Y_i operators act from different sides.
*/

/** Define the numerical type and data type for the GPU computations (same) */
using NumericalType = std::complex<double>;      // do not change
constexpr cudaDataType_t dataType = CUDA_C_64F;  // do not change


/** Example of a user-provided scalar CPU callback C function
 *  defining a time-dependent coefficient h_i(t) inside the Hamiltonian:
 *  h_i(t) = exp(-Omega * t)
 */
extern "C"
int32_t hCoefComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarStorage,    //inout: CPU-accessible storage for the returned coefficient value(s) of shape [0:batchSize-1]
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    auto * tdCoef = static_cast<cuDoubleComplex *>(scalarStorage); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t i = 0; i < batchSize; ++i) {
      const auto omega = params[i * numParams + 0]; // params[0][i]: 0-th parameter for i-th instance of the batch
      tdCoef[i] = make_cuDoubleComplex(std::exp((-omega) * time), 0.0); // value of the i-th instance of the coefficients batch
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** User-provided gradient callback function for the user-provided
 *  scalar callback function with respect to its single parameter Omega.
 *  It accumulates a partial derivative 2*Re(dCost/dOmega) = 2*Re(dCost/dCoef * dCoef/dOmega),
 *  where:
 *  - Cost is some user-defined real scalar cost function,
 *  - dCost/dCoef is the adjoint of the cost function with respect to the coefficient (or their batch) associated with the callback function,
 *  - dCoef/dOmega is the gradient of the coefficient (or their batch) with respect to the parameter Omega:
 *    dCoef/dOmega = -time * exp(-Omega * time)
 */
extern "C"
int32_t hCoefGradComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarGrad,       //inout: CPU-accessible storage for the adjoint(s) of the coefficient(s) of shape [0:batchSize-1]
  double * paramsGrad,     //inout: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of the returned gradient(s) of the parameter(s)
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    const auto * tdCoefAdjoint = static_cast<const cuDoubleComplex *>(scalarGrad); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t i = 0; i < batchSize; ++i) {
      const auto omega = params[i * numParams + 0]; // params[0][i]: 0-th parameter for i-th instance of the batch
      paramsGrad[i * numParams + 0] += // IMPORTANT: Accumulate the partial derivative for the i-th instance of the batch, not overwrite it!
        2.0 * cuCreal(cuCmul(tdCoefAdjoint[i], make_cuDoubleComplex(std::exp((-omega) * time) * (-time), 0.0)));
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** Example of a user-provided scalar CPU callback C function
 *  defining a time-dependent coefficient f(t) inside the Hamiltonian:
 *  f(t) = exp(i * Omega * t) = cos(Omega * t) + i * sin(Omega * t)
 */
extern "C"
int32_t fCoefComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarStorage,    //inout: CPU-accessible storage for the returned coefficient value(s) of shape [0:batchSize-1]
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    auto * tdCoef = static_cast<cuDoubleComplex *>(scalarStorage); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t i = 0; i < batchSize; ++i) {
      const auto omega = params[i * numParams + 0]; // params[0][i]: 0-th parameter for i-th instance of the batch
      tdCoef[i] = make_cuDoubleComplex(std::cos(omega * time), std::sin(omega * time)); // value of the i-th instance of the coefficients batch
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** User-provided gradient callback function for the user-provided
 *  scalar callback function with respect to its single parameter Omega.
 *  It accumulates a partial derivative 2*Re(dCost/dOmega) = 2*Re(dCost/dCoef * dCoef/dOmega),
 *  where:
 *  - Cost is some user-defined real scalar cost function,
 *  - dCost/dCoef is the adjoint of the cost function with respect to the coefficient associated with the callback function,
 *  - dCoef/dOmega is the gradient of the coefficient with respect to the parameter Omega:
 *    dCoef/dOmega = -i * time * exp(i * Omega * time)
 *                 = -i * time * (cos(Omega * time) + i * sin(Omega * time)) =
 *                 = -time * sin(Omega * time) + i * time * cos(Omega * time)
 */
extern "C"
int32_t fCoefGradComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarGrad,       //inout: CPU-accessible storage for the adjoint(s) of the coefficient(s) of shape [0:batchSize-1]
  double * paramsGrad,     //inout: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of the returned gradient(s) of the parameter(s)
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    const auto * tdCoefAdjoint = static_cast<const cuDoubleComplex *>(scalarGrad); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t i = 0; i < batchSize; ++i) {
      const auto omega = params[i * numParams + 0]; // params[0][i]: 0-th parameter for i-th instance of the batch
      paramsGrad[i * numParams + 0] += // IMPORTANT: Accumulate the partial derivative for the i-th instance of the batch, not overwrite it!
        2.0 * cuCreal(cuCmul(tdCoefAdjoint[i], make_cuDoubleComplex(-std::sin(omega * time) * time, std::cos(omega * time) * time)));
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** Convenience class which encapsulates a user-defined Liouvillian operator (system Hamiltonian + dissipation terms):
 *  - Constructor constructs the desired Liouvillian operator (`cudensitymatOperator_t`)
 *  - Method `get()` returns a reference to the constructed Liouvillian operator
 *  - Destructor releases all resources used by the Liouvillian operator
 */
class UserDefinedLiouvillian final
{
private:
  // Data members
  cudensitymatHandle_t handle;             // library context handle
  int64_t stateBatchSize;                  // quantum state batch size
  const std::vector<int64_t> spaceShape;   // Hilbert space shape (extents of the modes of the composite Hilbert space)
  void * spinXelems {nullptr};             // elements of the X spin operator in GPU RAM (F-order storage)
  void * spinYYelems {nullptr};            // elements of the fused YY two-spin operator in GPU RAM (F-order storage)
  void * spinZZelems {nullptr};            // elements of the fused ZZ two-spin operator in GPU RAM (F-order storage)
  cudensitymatElementaryOperator_t spinX;  // X spin operator (elementary tensor operator)
  cudensitymatElementaryOperator_t spinYY; // fused YY two-spin operator (elementary tensor operator)
  cudensitymatElementaryOperator_t spinZZ; // fused ZZ two-spin operator (elementary tensor operator)
  cudensitymatOperatorTerm_t oneBodyTerm;  // operator term: H1 = sum_{i} {h_i(t) * X_i} (one-body term)
  cudensitymatOperatorTerm_t twoBodyTerm;  // operator term: H2 = f(t) * sum_{i < j} {g_ij * ZZ_ij} (two-body term)
  cudensitymatOperatorTerm_t noiseTerm;    // operator term: D1 = d * sum_{i} {YY_ii}  // Y_i operators act from different sides on the density matrix (two-body mixed term)
  cudensitymatOperator_t liouvillian;      // full operator: (-i * (H1 + H2) * {..}) + (i * {..} * (H1 + H2)) + D1{..} (super-operator)

public:

  // Constructor constructs a user-defined Liouvillian operator
  UserDefinedLiouvillian(cudensitymatHandle_t contextHandle,             // library context handle
                         const std::vector<int64_t> & hilbertSpaceShape, // Hilbert space shape
                         int64_t batchSize):                             // batch size for the quantum state
    handle(contextHandle), stateBatchSize(batchSize), spaceShape(hilbertSpaceShape)
  {
    // Define the necessary operator tensors in GPU memory (F-order storage!)
    spinXelems = createInitializeArrayGPU<NumericalType>(  // X[i0; j0]
                  {{0.0, 0.0}, {1.0, 0.0},   // 1st column of matrix X
                   {1.0, 0.0}, {0.0, 0.0}}); // 2nd column of matrix X

    spinYYelems = createInitializeArrayGPU<NumericalType>(  // YY[i0, i1; j0, j1] := Y[i0; j0] * Y[i1; j1]
                    {{0.0, 0.0},  {0.0, 0.0}, {0.0, 0.0}, {-1.0, 0.0},  // 1st column of matrix YY
                     {0.0, 0.0},  {0.0, 0.0}, {1.0, 0.0}, {0.0, 0.0},   // 2nd column of matrix YY
                     {0.0, 0.0},  {1.0, 0.0}, {0.0, 0.0}, {0.0, 0.0},   // 3rd column of matrix YY
                     {-1.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}}); // 4th column of matrix YY

    spinZZelems = createInitializeArrayGPU<NumericalType>(  // ZZ[i0, i1; j0, j1] := Z[i0; j0] * Z[i1; j1]
                    {{1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {0.0, 0.0},   // 1st column of matrix ZZ
                     {0.0, 0.0}, {-1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},   // 2nd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {-1.0, 0.0}, {0.0, 0.0},   // 3rd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {1.0, 0.0}}); // 4th column of matrix ZZ

    // Construct the necessary Elementary Tensor Operators
    //  X_i operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        1,                                   // one-body operator
                        std::vector<int64_t>({2}).data(),    // acts in tensor space of shape {2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinXelems,                          // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinX));                            // the created elementary tensor operator
    //  ZZ_ij = Z_i * Z_j fused operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        2,                                   // two-body operator
                        std::vector<int64_t>({2,2}).data(),  // acts in tensor space of shape {2,2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinZZelems,                         // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinZZ));                           // the created elementary tensor operator
    //  YY_ii = Y_i * {..} * Y_i fused operator (note action from different sides)
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        2,                                   // two-body operator
                        std::vector<int64_t>({2,2}).data(),  // acts in tensor space of shape {2,2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinYYelems,                         // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinYY));                           // the created elementary tensor operator

    // Construct the necessary Operator Terms from tensor products of Elementary Tensor Operators
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &oneBodyTerm));                      // the created empty operator term
    //  Define the operator term: H1 = sum_{i} {h_i(t) * X_i}
    for (int32_t i = 0; i < spaceShape.size(); ++i) {
      HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                          oneBodyTerm,
                          1,                                                             // number of elementary tensor operators in the product
                          std::vector<cudensitymatElementaryOperator_t>({spinX}).data(), // elementary tensor operators forming the product
                          std::vector<int32_t>({i}).data(),                              // space modes acted on by the operator product
                          std::vector<int32_t>({0}).data(),                              // space mode action duality (0: from the left; 1: from the right)
                          make_cuDoubleComplex(1.0, 0.0),                                // static coefficient part: Always 64-bit-precision complex number
                          {hCoefComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr},   // CPU scalar callback function defining the time-dependent coefficient associated with this operator product
                          {hCoefGradComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr,
                           CUDENSITYMAT_DIFFERENTIATION_DIR_BACKWARD})); // CPU scalar gradient callback function defining the gradient of the coefficient with respect to the parameter Omega
    }
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &twoBodyTerm));                      // the created empty operator term
    //  Define the operator term: H2 = f(t) * sum_{i < j} {g_ij * ZZ_ij}
    for (int32_t i = 0; i < spaceShape.size() - 1; ++i) {
      for (int32_t j = (i + 1); j < spaceShape.size(); ++j) {
        const double g_ij = -1.0 / static_cast<double>(i + j + 1); // assign some value to the time-independent g_ij coefficient
        HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                            twoBodyTerm,
                            1,                                                              // number of elementary tensor operators in the product
                            std::vector<cudensitymatElementaryOperator_t>({spinZZ}).data(), // elementary tensor operators forming the product
                            std::vector<int32_t>({i, j}).data(),                            // space modes acted on by the operator product
                            std::vector<int32_t>({0, 0}).data(),                            // space mode action duality (0: from the left; 1: from the right)
                            make_cuDoubleComplex(g_ij, 0.0),                                // g_ij static coefficient: Always 64-bit-precision complex number
                            cudensitymatScalarCallbackNone,                                 // no time-dependent coefficient associated with this operator product
                            cudensitymatScalarGradientCallbackNone));                       // no coefficient gradient associated with this operator product
      }
    }
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &noiseTerm));                        // the created empty operator term
    //  Define the operator term: D1 = d * sum_{i} {YY_ii}
    for (int32_t i = 0; i < spaceShape.size(); ++i) {
      HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                          noiseTerm,
                          1,                                                              // number of elementary tensor operators in the product
                          std::vector<cudensitymatElementaryOperator_t>({spinYY}).data(), // elementary tensor operators forming the product
                          std::vector<int32_t>({i, i}).data(),                            // space modes acted on by the operator product (from different sides)
                          std::vector<int32_t>({0, 1}).data(),                            // space mode action duality (0: from the left; 1: from the right)
                          make_cuDoubleComplex(1.0, 0.0),                                 // default coefficient: Always 64-bit-precision complex number
                          cudensitymatScalarCallbackNone,                                 // no time-dependent coefficient associated with this operator product
                          cudensitymatScalarGradientCallbackNone));                       // no coefficient gradient associated with this operator product
    }

    // Construct the full Liouvillian operator as a sum of the created operator terms
    //  Create an empty operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                        spaceShape.size(),                // Hilbert space rank (number of modes)
                        spaceShape.data(),                // Hilbert space shape (modes extents)
                        &liouvillian));                   // the created empty operator (super-operator)
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        oneBodyTerm,                      // appended operator term
                        0,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, -1.0),  // -i constant
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        twoBodyTerm,                      // appended operator term
                        0,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, -1.0),  // -i constant
                        {fCoefComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr}, // CPU scalar callback function defining the time-dependent coefficient associated with this operator term as a whole
                        {fCoefGradComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr,
                         CUDENSITYMAT_DIFFERENTIATION_DIR_BACKWARD})); // CPU scalar gradient callback function defining the gradient of the coefficient with respect to the parameter Omega
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        oneBodyTerm,                      // appended operator term
                        1,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, +1.0),  // +i constant
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        twoBodyTerm,                      // appended operator term
                        1,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, 1.0),   // +i constant
                        {fCoefComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr}, // CPU scalar callback function defining the time-dependent coefficient associated with this operator term as a whole
                        {fCoefGradComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr,
                         CUDENSITYMAT_DIFFERENTIATION_DIR_BACKWARD})); // CPU scalar gradient callback function defining the gradient of the coefficient with respect to the parameter Omega
    //  Append an operator term to the operator (super-operator)
    const double d = 1.0; // assign some value to the time-independent coefficient
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        noiseTerm,                        // appended operator term
                        0,                                // operator term action duality as a whole (no duality reversing in this case)
                        make_cuDoubleComplex(d, 0.0),     // static coefficient associated with the operator term as a whole
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
  }

  // Destructor destructs the user-defined Liouvillian operator
  ~UserDefinedLiouvillian()
  {
    // Destroy the Liouvillian operator
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(liouvillian));

    // Destroy operator terms
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(noiseTerm));
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(twoBodyTerm));
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(oneBodyTerm));

    // Destroy elementary tensor operators
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinYY));
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinZZ));
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinX));

    // Destroy operator tensors
    destroyArrayGPU(spinYYelems);
    destroyArrayGPU(spinZZelems);
    destroyArrayGPU(spinXelems);
  }

  // Disable copy constructor/assignment (GPU resources are private, no deep copy)
  UserDefinedLiouvillian(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian & operator=(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian(UserDefinedLiouvillian &&) = delete;
  UserDefinedLiouvillian & operator=(UserDefinedLiouvillian &&) = delete;

  /** Returns the number of externally provided Hamiltonian parameters. */
  int32_t getNumParameters() const
  {
    return 1; // one parameter Omega
  }

  /** Get access to the constructed Liouvillian operator. */
  cudensitymatOperator_t & get()
  {
    return liouvillian;
  }

};

Now we can use the defined quantum many-body operator in our main code to compute its action on a mixed quantum state and then backward-differentiate it (compute gradients) with respect to the user-provided real parameter Omega. For the sake of simplicity, we pass a made-up adjoint of the output quantum state to the cudensitymatOperatorComputeActionBackwardDiff() call, which is just the output quantum state itself (in real scenarios, the adjoint of the output quantum state will depend on the user-chosen cost function and will be provided by the user). Upon completion of the cudensitymatOperatorComputeActionBackwardDiff() call, the paramsGrad output argument will contain the gradient of the user-defined cost function with respect to the user-provided real parameter Omega. Additionally, the backward-differentiation API call will also return the adjoint of the input quantum state for cases where the input quantum state implicitly depends on the user-provided real parameters (for example, cases where the input quantum state comes from a previous operator action step, which is typical for time-integration of quantum dynamics master equations). Note that both output arguments, namely paramsGrad and stateInAdj, are accumulative, i.e., they will be accumulated into (it is user’s responsibility to zero them out before the first call!).

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <cudensitymat.h>  // cuDensityMat library header
#include "helpers.h"       // helper functions


// Transverse Ising Hamiltonian with double summation ordering
// and spin-operator fusion, plus fused dissipation terms
#include "transverse_ising_full_fused_noisy_grad.h" // user-defined Liouvillian operator example

#include <cmath>
#include <complex>
#include <vector>
#include <chrono>
#include <iostream>
#include <cassert>


// Number of times to perform operator action on a quantum state
constexpr int NUM_REPEATS = 2;

// Logging verbosity
bool verbose = true;


// Example workflow
void exampleWorkflow(cudensitymatHandle_t handle)
{
  // Define the composite Hilbert space shape and
  // quantum state batch size (number of individual quantum states in a batched simulation)
  const std::vector<int64_t> spaceShape({2,2,2,2}); // dimensions of quantum degrees of freedom
  const int64_t batchSize = 1;                      // number of quantum states per batch

  if (verbose) {
    std::cout << "Hilbert space rank = " << spaceShape.size() << "; Shape = (";
    for (const auto & dimsn: spaceShape)
      std::cout << dimsn << ",";
    std::cout << ")" << std::endl;
    std::cout << "Quantum state batch size = " << batchSize << std::endl;
  }

  // Construct a user-defined Liouvillian operator using a convenience C++ class
  UserDefinedLiouvillian liouvillian(handle, spaceShape, batchSize);
  if (verbose)
    std::cout << "Constructed the Liouvillian operator\n";

  // Set and place external user-provided Hamiltonian parameters in GPU memory
  const int32_t numParams = liouvillian.getNumParameters(); // number of external user-provided Hamiltonian parameters
  if (verbose)
    std::cout << "Number of external user-provided Hamiltonian parameters = " << numParams << std::endl;
  std::vector<double> cpuHamParams(numParams * batchSize);
  for (int64_t j = 0; j < batchSize; ++j) {
    for (int32_t i = 0; i < numParams; ++i) {
      cpuHamParams[j * numParams + i] = double(i+1) / double(j+1); // just setting some parameter values for each instance of the batch
    }
  }
  auto * hamiltonianParams = static_cast<double *>(createInitializeArrayGPU(cpuHamParams));
  if (verbose)
    std::cout << "Created an array of external user-provided Hamiltonian parameters in GPU memory\n";

  // Create an array of gradients for the user-provided Hamiltonian parameters in GPU memory
  std::vector<double> cpuHamParamsGrad(numParams * batchSize, 0.0);
  auto * hamiltonianParamsGrad = static_cast<double *>(createInitializeArrayGPU(cpuHamParamsGrad));
  if (verbose)
    std::cout << "Created an array of gradients for the external user-provided Hamiltonian parameters in GPU memory\n";

  // Declare the input quantum state
  cudensitymatState_t inputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &inputState));

  // Query the size of the quantum state storage
  std::size_t storageSize {0}; // only one storage component (tensor) is needed (no tensor factorization)
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      inputState,
                      1,               // only one storage component (tensor)
                      &storageSize));  // storage size in bytes
  const std::size_t stateVolume = storageSize / sizeof(NumericalType);  // quantum state tensor volume (number of elements)
  if (verbose)
    std::cout << "Quantum state storage size (bytes) = " << storageSize << std::endl;

  // Prepare some initial value for the input quantum state batch
  std::vector<NumericalType> inputStateValue(stateVolume);
  if constexpr (std::is_same_v<NumericalType, float>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0f / float(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, double>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0 / double(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<float>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0f / float(i+1), -1.0f / float(i+2)}; // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<double>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0 / double(i+1), -1.0 / double(i+2)}; // just some value
    }
  } else {
    std::cerr << "Error: Unsupported data type!\n";
    std::exit(1);
  }
  // Allocate initialized GPU storage for the input quantum state with prepared values
  auto * inputStateElems = createInitializeArrayGPU(inputStateValue);
  if (verbose)
    std::cout << "Allocated input quantum state storage and initialized it to some value\n";

  // Attach initialized GPU storage to the input quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      inputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({inputStateElems}).data(),      // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed input quantum state\n";

  // Declare the output quantum state of the same shape
  cudensitymatState_t outputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &outputState));

  // Allocate GPU storage for the output quantum state
  auto * outputStateElems = createArrayGPU<NumericalType>(stateVolume);
  if (verbose)
    std::cout << "Allocated output quantum state storage\n";

  // Attach GPU storage to the output quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      outputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({outputStateElems}).data(),     // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed output quantum state\n";

  // Declare the adjoint input quantum state of the same shape
  cudensitymatState_t inputStateAdj;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,  // data type must match
                      &inputStateAdj));

  // Allocate GPU storage for the adjoint input quantum state
  auto * inputStateAdjElems = createArrayGPU<NumericalType>(stateVolume);
  if (verbose)
    std::cout << "Allocated adjoint input quantum state storage\n";

  // Attach GPU storage to the adjoint input quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      inputStateAdj,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({inputStateAdjElems}).data(),   // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed adjoint input quantum state\n";

  // Declare a workspace descriptor
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  // Query free GPU memory
  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95); // take 95% of the free memory as the workspace budget
  if (verbose)
    std::cout << "Max workspace buffer size (bytes) = " << freeMem << std::endl;

  // Prepare the Liouvillian operator action on a quantum state (needs to be done only once)
  auto startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorPrepareAction(handle,
                      liouvillian.get(),
                      inputState,
                      outputState,
                      CUDENSITYMAT_COMPUTE_64F,  // GPU compute type
                      freeMem,                   // max available GPU free memory for the workspace
                      workspaceDescr,            // workspace descriptor
                      0x0));                     // default CUDA stream
  auto finishTime = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator action preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  std::size_t requiredBufferSize {0};
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  if (requiredBufferSize > freeMem) {
    std::cerr << "Error: Required workspace buffer size is greater than the available GPU free memory!\n";
    std::exit(1);
  }

  // Allocate GPU storage for the workspace buffer
  std::size_t workspaceBufferSize = requiredBufferSize;
  void * workspaceBuffer = createArrayGPU<NumericalType>(workspaceBufferSize / sizeof(NumericalType));
  if (verbose)
    std::cout << "Allocated workspace buffer of size (bytes) = " << workspaceBufferSize << std::endl;

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      workspaceBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << workspaceBufferSize << std::endl;

  // Apply the Liouvillian operator to the input quatum state
  // and accumulate its action into the output quantum state (note the accumulative += semantics)
  for (int32_t repeat = 0; repeat < NUM_REPEATS; ++repeat) { // repeat multiple times for accurate timing
    // Zero out the output quantum state
    HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle,
                        outputState,
                        0x0));
    if (verbose)
      std::cout << "Initialized the output quantum state to zero\n";
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    startTime = std::chrono::high_resolution_clock::now();
    HANDLE_CUDM_ERROR(cudensitymatOperatorComputeAction(handle,
                        liouvillian.get(),
                        0.3,                // time point (some value)
                        batchSize,          // user-defined batch size
                        numParams,          // number of external user-defined Hamiltonian parameters
                        hamiltonianParams,  // external Hamiltonian parameters in GPU memory
                        inputState,         // input quantum state
                        outputState,        // output quantum state
                        workspaceDescr,     // workspace descriptor
                        0x0));              // default CUDA stream
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    finishTime = std::chrono::high_resolution_clock::now();
    timeSec = finishTime - startTime;
    if (verbose)
      std::cout << "Operator action computation time (sec) = " << timeSec.count() << std::endl;
  }

  // Compute the squared norm of the output quantum state
  void * norm2 = createInitializeArrayGPU(std::vector<double>(batchSize, 0.0));
  HANDLE_CUDM_ERROR(cudensitymatStateComputeNorm(handle,
                      outputState,
                      norm2,
                      0x0));
  if (verbose) {
    std::cout << "Computed the output quantum state norm:\n";
    printArrayGPU<double>(norm2, batchSize);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Prepare the Liouvillian operator action backward differentiation (needs to be done only once)
  startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorPrepareActionBackwardDiff(handle,
                      liouvillian.get(),
                      inputState,
                      outputState,               // adjoint output quantum state is always congruent to the output quantum state
                      CUDENSITYMAT_COMPUTE_64F,  // GPU compute type
                      freeMem,                   // max available GPU free memory for the workspace buffer
                      workspaceDescr,            // workspace descriptor
                      0x0));                     // default CUDA stream
  finishTime = std::chrono::high_resolution_clock::now();
  timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator action backward differentiation preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  requiredBufferSize = 0;
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  if (requiredBufferSize > freeMem) {
    std::cerr << "Error: Required workspace buffer size is greater than the available GPU free memory!\n";
    std::exit(1);
  }

  // Reallocate the workspace buffer if the backward pass requires more memory
  if (requiredBufferSize > workspaceBufferSize) {
    destroyArrayGPU(workspaceBuffer);
    workspaceBufferSize = requiredBufferSize;
    workspaceBuffer = createArrayGPU<NumericalType>(workspaceBufferSize / sizeof(NumericalType));
    if (verbose)
      std::cout << "Re-allocated workspace buffer of size (bytes) = " << workspaceBufferSize << std::endl;
  }

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      requiredBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Liouvillian operator action backward differentiation:
  // The adjoint output quantum state, which is always congruent to the output quantum state,
  // depends on the user-defined cost function, so here we simply pass the previously computed output quantum state.
  // In real-life applications, the user will pass their adjoint output quantum state, computed for their cost function.
  for (int32_t repeat = 0; repeat < NUM_REPEATS; ++repeat) { // repeat multiple times for accurate timing
    // Zero out the adjoint input quantum state and gradients
    HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle,
                        inputStateAdj,
                        0x0));
    initializeArrayGPU(std::vector<double>(numParams * batchSize, 0.0), hamiltonianParamsGrad);
    if (verbose)
      std::cout << "Initialized the adjoint input quantum state and gradients to zero\n";
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    startTime = std::chrono::high_resolution_clock::now();
    HANDLE_CUDM_ERROR(cudensitymatOperatorComputeActionBackwardDiff(handle,
                        liouvillian.get(),
                        0.3,                    // time point (some value)
                        batchSize,              // user-defined batch size
                        numParams,              // number of external user-defined Hamiltonian parameters
                        hamiltonianParams,      // external Hamiltonian parameters in GPU memory
                        inputState,             // input quantum state
                        outputState,            // adjoint output quantum state (here we just pass the previously computed output quantum state for simplicity)
                        inputStateAdj,          // adjoint input quantum state
                        hamiltonianParamsGrad,  // partial gradients with respect to the user-defined real parameters
                        workspaceDescr,         // workspace descriptor
                        0x0));                  // default CUDA stream
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    finishTime = std::chrono::high_resolution_clock::now();
    timeSec = finishTime - startTime;
    if (verbose)
      std::cout << "Operator action backward differentiation computation time (sec) = " << timeSec.count() << std::endl;
  }

  // Compute the squared norm of the adjoint input quantum state
  initializeArrayGPU(std::vector<double>(batchSize, 0.0), norm2);
  HANDLE_CUDM_ERROR(cudensitymatStateComputeNorm(handle,
                      inputStateAdj,
                      norm2,
                      0x0));
  if (verbose) {
    std::cout << "Computed the adjoint input quantum state norm:\n";
    printArrayGPU<double>(norm2, batchSize);
    std::cout << "Hamiltonian parameters gradients:\n";
    printArrayGPU<double>(hamiltonianParamsGrad, numParams * batchSize);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Destroy the norm2 array
  destroyArrayGPU(norm2);

  // Destroy workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));

  // Destroy workspace buffer storage
  destroyArrayGPU(workspaceBuffer);

  // Destroy quantum states
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(inputStateAdj));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(outputState));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(inputState));

  // Destroy quantum state storage
  destroyArrayGPU(inputStateAdjElems);
  destroyArrayGPU(outputStateElems);
  destroyArrayGPU(inputStateElems);

  // Destroy external Hamiltonian parameters
  destroyArrayGPU(static_cast<void *>(hamiltonianParamsGrad));
  destroyArrayGPU(static_cast<void *>(hamiltonianParams));

  if (verbose)
    std::cout << "Destroyed resources\n" << std::flush;
}


int main(int argc, char ** argv)
{
  // Assign a GPU to the process
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose)
    std::cout << "Set active device\n";

  // Create a library handle
  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created a library handle\n";

  // Run the example
  exampleWorkflow(handle);

  // Destroy the library handle
  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed the library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());

  // Done
  return 0;
}

Code example (serial batched execution with backward differentiation)

The following code example extends backward differentiation to batched operators and quantum states. Here the Hamiltonian operator contains batched coefficients such that each instance of a batched quantum state is acted on by a different instance of the batched operator (with different coefficient values). The full sample code can be found in the NVIDIA/cuQuantum repository (main serial batched gradient code and operator definition for batched gradient as well as the utility code).

The Hamiltonian definition makes both the h(t) and f(t) scalar coefficients batched, requiring user-supplied vector storage for their static and dynamic (total) values. The corresponding scalar and scalar gradient callback functions also operate on a batch instead of a single instance. Furthermore, both params and paramsGrad arrays become truly two-dimensional arrays, with the first dimension corresponding to the number of user-provided real parameters and the second dimension corresponding to the batch size.

/* Copyright (c) 2026-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#pragma once

#include <cudensitymat.h> // cuDensityMat library header
#include "helpers.h"      // GPU helper functions

#include <cmath>
#include <complex>
#include <vector>
#include <iostream>
#include <cassert>


/* DESCRIPTION:
   Batched time-dependent transverse-field Ising Hamiltonian operator
   with ordered and fused ZZ terms, plus fused unitary dissipation terms:
    H[k] = sum_{i} {h_i(t)[k] * X_i}          // transverse-field sum of X_i operators with batched time-dependent h_i(t)[k] coefficients 
      + f(t)[k] * sum_{i < j} {g_ij * ZZ_ij}  // batched modulated sum of the fused ordered {Z_i * Z_j} terms with static g_ij coefficients
      + d * sum_{i} {Y_i * {..} * Y_i}        // scaled sum of the dissipation terms {Y_i * {..} * Y_i} fused into the YY_ii super-operators
   where {..} is the placeholder for the density matrix to show that the Y_i operators act from different sides.
*/

/** Define the numerical type and data type for the GPU computations (same) */
using NumericalType = std::complex<double>;      // do not change
constexpr cudaDataType_t dataType = CUDA_C_64F;  // do not change


/** User-provided batched scalar CPU callback C function
 *  defining a batched time-dependent coefficient h_i(t) for all instances
 *  of the batch inside the Hamiltonian:
 *  h_i(t)[k] = exp(-Omega[k] * t) for k = 0, ..., batchSize-1
 */
extern "C"
int32_t hCoefBatchComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarStorage,    //inout: CPU-accessible storage for the returned batched coefficient values of shape [0:batchSize-1]
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    auto * tdCoef = static_cast<cuDoubleComplex *>(scalarStorage); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t k = 0; k < batchSize; ++k) { // for each instance of the batch
      const auto omega = params[k * numParams + 0]; // params[0][k]: 0-th parameter for k-th instance of the batch
      tdCoef[k] = make_cuDoubleComplex(std::exp((-omega) * time), 0.0); // value of the k-th instance of the batched coefficients
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** User-provided batched scalar gradient callback function (CPU-side) for the user-provided
 *  batched scalar callback function hCoefBatchComplex64, defining the gradients with respect
 *  to its single (batched) parameter Omega. It accumulates a partial derivative:
 *    2*Re(dCost/dOmega[k]) = 2*Re(dCost/dCoef[k] * dCoef[k]/dOmega[k]),
 *  where:
 *  - Cost is some user-defined real scalar cost function,
 *  - dCost/dCoef[k] is the adjoint of the cost function with respect to the k-th instance of the batched coefficient associated with the callback function,
 *  - dCoef[k]/dOmega[k] is the gradient of the k-th instance of the batched coefficient with respect to the parameter Omega[k]:
 *    dCoef[k]/dOmega[k] = -t * exp(-Omega[k] * t) for k = 0, ..., batchSize-1
 */
extern "C"
int32_t hCoefBatchGradComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarGrad,       //in: CPU-accessible storage for the batched adjoint of the batched coefficient of shape [0:batchSize-1]
  double * paramsGrad,     //inout: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of the returned gradients of the parameter(s) for all instances of the batch
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    const auto * tdCoefAdjoint = static_cast<const cuDoubleComplex *>(scalarGrad); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t k = 0; k < batchSize; ++k) { // for each instance of the batch
      const auto omega = params[k * numParams + 0]; // params[0][k]: 0-th parameter for k-th instance of the batch
      paramsGrad[k * numParams + 0] += // IMPORTANT: Accumulate the partial derivative for the k-th instance of the batch, not overwrite it!
        2.0 * cuCreal(cuCmul(tdCoefAdjoint[k], make_cuDoubleComplex(std::exp((-omega) * time) * (-time), 0.0)));
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** User-provided batched scalar CPU callback C function
 *  defining a batched time-dependent coefficient f(t) for all instances
 *  of the batch inside the Hamiltonian:
 *  f(t)[k] = exp(i * Omega[k] * t)
 *          = cos(Omega[k] * t) + i * sin(Omega[k] * t) for k = 0, ..., batchSize-1
 */
extern "C"
int32_t fCoefBatchComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarStorage,    //inout: CPU-accessible storage for the returned batched coefficient values of shape [0:batchSize-1]
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    auto * tdCoef = static_cast<cuDoubleComplex *>(scalarStorage); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t k = 0; k < batchSize; ++k) { // for each instance of the batch
      const auto omega = params[k * numParams + 0]; // params[0][k]: 0-th parameter for k-th instance of the batch
      tdCoef[k] = make_cuDoubleComplex(std::cos(omega * time), std::sin(omega * time)); // value of the k-th instance of the batched coefficients
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** User-provided batched scalar gradient callback function (CPU-side) for the user-provided
 *  batched scalar callback function fCoefBatchComplex64, defining the gradients with respect
 *  to its single (batched) parameter Omega. It accumulates a partial derivative:
 *    2*Re(dCost/dOmega[k]) = 2*Re(dCost/dCoef[k] * dCoef[k]/dOmega[k]),
 *  where:
 *  - Cost is some user-defined real scalar cost function,
 *  - dCost/dCoef[k] is the adjoint of the cost function with respect to the k-th instance of the batched coefficient associated with the callback function,
 *  - dCoef[k]/dOmega[k] is the gradient of the k-th instance of the batched coefficient with respect to the parameter Omega[k]:
 *    dCoef[k]/dOmega[k] = i * t * exp(i * Omega[k] * t)
 *                       = i * t * cos(Omega[k] * t) - t * sin(Omega[k] * t) for k = 0, ..., batchSize-1
 */
extern "C"
int32_t fCoefBatchGradComplex64(
  double time,             //in: time point
  int64_t batchSize,       //in: user-defined batch size (number of coefficients in the batch)
  int32_t numParams,       //in: number of external user-provided Hamiltonian parameters (this function expects one parameter, Omega)
  const double * params,   //in: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of user-provided Hamiltonian parameters for all instances of the batch
  cudaDataType_t dataType, //in: data type (expecting CUDA_C_64F in this specific callback function)
  void * scalarGrad,       //in: CPU-accessible storage for the batched adjoint of the batched coefficient of shape [0:batchSize-1]
  double * paramsGrad,     //inout: params[0:numParams-1][0:batchSize-1]: GPU-accessible F-ordered array of the returned gradients of the parameter(s) for all instances of the batch
  cudaStream_t stream)     //in: CUDA stream (default is 0x0)
{
  if (dataType == CUDA_C_64F) {
    const auto * tdCoefAdjoint = static_cast<const cuDoubleComplex *>(scalarGrad); // casting to cuDoubleComplex because this callback function expects CUDA_C_64F data type
    for (int64_t k = 0; k < batchSize; ++k) {
      const auto omega = params[k * numParams + 0]; // params[0][k]: 0-th parameter for k-th instance of the batch
      paramsGrad[k * numParams + 0] += // IMPORTANT: Accumulate the partial derivative for the k-th instance of the batch, not overwrite it!
        2.0 * cuCreal(cuCmul(tdCoefAdjoint[k], make_cuDoubleComplex(-std::sin(omega * time) * time, std::cos(omega * time) * time)));
    }
  } else {
    return 1; // error code (1: Error)
  }
  return 0; // error code (0: Success)
}


/** Convenience class which encapsulates a user-defined Liouvillian operator (system Hamiltonian + dissipation terms):
 *  - Constructor constructs the desired Liouvillian operator (`cudensitymatOperator_t`)
 *  - Method `get()` returns a reference to the constructed Liouvillian operator
 *  - Destructor releases all resources used by the Liouvillian operator
 */
class UserDefinedLiouvillian final
{
private:
  // Data members
  cudensitymatHandle_t handle;             // library context handle
  int64_t operBatchSize;                   // batch size for the super-operator
  const std::vector<int64_t> spaceShape;   // Hilbert space shape (extents of the modes of the composite Hilbert space)
  void * spinXelems {nullptr};             // elements of the X spin operator in GPU RAM (F-order storage)
  void * spinYYelems {nullptr};            // elements of the fused YY two-spin operator in GPU RAM (F-order storage)
  void * spinZZelems {nullptr};            // elements of the fused ZZ two-spin operator in GPU RAM (F-order storage)
  cudensitymatElementaryOperator_t spinX;  // X spin operator (elementary tensor operator)
  cudensitymatElementaryOperator_t spinYY; // fused YY two-spin operator (elementary tensor operator)
  cudensitymatElementaryOperator_t spinZZ; // fused ZZ two-spin operator (elementary tensor operator)
  cudensitymatOperatorTerm_t oneBodyTerm;  // operator term: H1 = sum_{i} {h_i(t) * X_i} (one-body term)
  cudensitymatOperatorTerm_t twoBodyTerm;  // operator term: H2 = f(t) * sum_{i < j} {g_ij * ZZ_ij} (two-body term)
  cudensitymatOperatorTerm_t noiseTerm;    // operator term: D1 = d * sum_{i} {YY_ii}  // Y_i operators act from different sides on the density matrix (two-body mixed term)
  // Batched coefficients
  cuDoubleComplex * hCoefsStatic {nullptr}; // static part of the h(t) batched coefficients in the one-body term (of length operBatchSize)
  cuDoubleComplex * hCoefsTotal {nullptr};  // total h(t) batched coefficients in the one-body term (of length operBatchSize)
  cuDoubleComplex * fCoefsStaticMinus {nullptr}; // static part of the f(t) batched coefficients in the two-body term (of length operBatchSize)
  cuDoubleComplex * fCoefsTotalMinus {nullptr};  // total f(t) batched coefficients in the two-body term (of length operBatchSize)
  cuDoubleComplex * fCoefsStaticPlus {nullptr};  // static part of the f(t) batched coefficients in the dual two-body term (of length operBatchSize)
  cuDoubleComplex * fCoefsTotalPlus {nullptr};   // total f(t) batched coefficients in the dual two-body term (of length operBatchSize)
  // Final Liouvillian operator
  cudensitymatOperator_t liouvillian; // full operator: (-i * (H1 + H2) * {..}) + (i * {..} * (H1 + H2)) + D1{..} (super-operator)

public:

  // Constructor constructs a user-defined Liouvillian operator
  UserDefinedLiouvillian(cudensitymatHandle_t contextHandle,             // library context handle
                         const std::vector<int64_t> & hilbertSpaceShape, // Hilbert space shape
                         int64_t batchSize):                             // batch size for the super-operator
    handle(contextHandle), operBatchSize(batchSize), spaceShape(hilbertSpaceShape)
  {
    // Define the necessary operator tensors in GPU memory (F-order storage!)
    spinXelems = createInitializeArrayGPU<NumericalType>(  // X[i0; j0]
                  {{0.0, 0.0}, {1.0, 0.0},   // 1st column of matrix X
                   {1.0, 0.0}, {0.0, 0.0}}); // 2nd column of matrix X

    spinYYelems = createInitializeArrayGPU<NumericalType>(  // YY[i0, i1; j0, j1] := Y[i0; j0] * Y[i1; j1]
                    {{0.0, 0.0},  {0.0, 0.0}, {0.0, 0.0}, {-1.0, 0.0},  // 1st column of matrix YY
                     {0.0, 0.0},  {0.0, 0.0}, {1.0, 0.0}, {0.0, 0.0},   // 2nd column of matrix YY
                     {0.0, 0.0},  {1.0, 0.0}, {0.0, 0.0}, {0.0, 0.0},   // 3rd column of matrix YY
                     {-1.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}}); // 4th column of matrix YY

    spinZZelems = createInitializeArrayGPU<NumericalType>(  // ZZ[i0, i1; j0, j1] := Z[i0; j0] * Z[i1; j1]
                    {{1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {0.0, 0.0},   // 1st column of matrix ZZ
                     {0.0, 0.0}, {-1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},   // 2nd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {-1.0, 0.0}, {0.0, 0.0},   // 3rd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {1.0, 0.0}}); // 4th column of matrix ZZ

    // Construct the necessary Elementary Tensor Operators
    //  X_i operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        1,                                   // one-body operator
                        std::vector<int64_t>({2}).data(),    // acts in tensor space of shape {2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinXelems,                          // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinX));                            // the created elementary tensor operator
    //  ZZ_ij = Z_i * Z_j fused operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        2,                                   // two-body operator
                        std::vector<int64_t>({2,2}).data(),  // acts in tensor space of shape {2,2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinZZelems,                         // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinZZ));                           // the created elementary tensor operator
    //  YY_ii = Y_i * {..} * Y_i fused operator (note action from different sides)
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        2,                                   // two-body operator
                        std::vector<int64_t>({2,2}).data(),  // acts in tensor space of shape {2,2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinYYelems,                         // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinYY));                           // the created elementary tensor operator

    // Construct the necessary Operator Terms from tensor products of Elementary Tensor Operators
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &oneBodyTerm));                      // the created empty operator term
    //  Define the batched operator term: H1[k] = sum_{i} {h_i(t)[k] * X_i}
    hCoefsStatic = static_cast<cuDoubleComplex *>(createInitializeArrayGPU(
      std::vector<cuDoubleComplex>(operBatchSize, make_cuDoubleComplex(1.0, 0.0)))); // 1.0 constant for all coefficient instances in the batch
    hCoefsTotal = static_cast<cuDoubleComplex *>(createInitializeArrayGPU(
      std::vector<cuDoubleComplex>(operBatchSize, make_cuDoubleComplex(0.0, 0.0)))); // storage for the total coefficients for all instances of the batch
    for (int32_t i = 0; i < spaceShape.size(); ++i) {
      HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProductBatch(handle,
                          oneBodyTerm,
                          1,                                                             // number of elementary tensor operators in the product
                          std::vector<cudensitymatElementaryOperator_t>({spinX}).data(), // elementary tensor operators forming the product
                          std::vector<int32_t>({i}).data(),                              // space modes acted on by the operator product
                          std::vector<int32_t>({0}).data(),                              // space mode action duality (0: from the left; 1: from the right)
                          operBatchSize,                                                 // batch size
                          hCoefsStatic,                                                  // static part of the h(t) batched coefficients in the one-body term
                          hCoefsTotal,                                                   // total h(t) batched coefficients in the one-body term
                          {hCoefBatchComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr}, // CPU batched scalar callback function defining the time-dependent coefficient associated with this operator product
                          {hCoefBatchGradComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr})); // CPU batched scalar gradient callback function defining the gradient of the coefficient with respect to the parameter Omega
    }
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &twoBodyTerm));                      // the created empty operator term
    //  Define the operator term: H2 = f(t) * sum_{i < j} {g_ij * ZZ_ij}
    for (int32_t i = 0; i < spaceShape.size() - 1; ++i) {
      for (int32_t j = (i + 1); j < spaceShape.size(); ++j) {
        const double g_ij = -1.0 / static_cast<double>(i + j + 1); // assign some value to the time-independent g_ij coefficient
        HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                            twoBodyTerm,
                            1,                                                              // number of elementary tensor operators in the product
                            std::vector<cudensitymatElementaryOperator_t>({spinZZ}).data(), // elementary tensor operators forming the product
                            std::vector<int32_t>({i, j}).data(),                            // space modes acted on by the operator product
                            std::vector<int32_t>({0, 0}).data(),                            // space mode action duality (0: from the left; 1: from the right)
                            make_cuDoubleComplex(g_ij, 0.0),                                // g_ij static coefficient: Always 64-bit-precision complex number
                            cudensitymatScalarCallbackNone,                                 // no time-dependent coefficient associated with this operator product
                            cudensitymatScalarGradientCallbackNone));                       // no coefficient gradient associated with this operator product
      }
    }
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &noiseTerm));                        // the created empty operator term
    //  Define the operator term: D1 = d * sum_{i} {YY_ii}
    for (int32_t i = 0; i < spaceShape.size(); ++i) {
      HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                          noiseTerm,
                          1,                                                              // number of elementary tensor operators in the product
                          std::vector<cudensitymatElementaryOperator_t>({spinYY}).data(), // elementary tensor operators forming the product
                          std::vector<int32_t>({i, i}).data(),                            // space modes acted on by the operator product (from different sides)
                          std::vector<int32_t>({0, 1}).data(),                            // space mode action duality (0: from the left; 1: from the right)
                          make_cuDoubleComplex(1.0, 0.0),                                 // default coefficient: Always 64-bit-precision complex number
                          cudensitymatScalarCallbackNone,                                 // no time-dependent coefficient associated with this operator product
                          cudensitymatScalarGradientCallbackNone));                       // no coefficient gradient associated with this operator product
    }

    // Construct the full Liouvillian operator as a sum of the created operator terms
    //  Create an empty operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                        spaceShape.size(),                // Hilbert space rank (number of modes)
                        spaceShape.data(),                // Hilbert space shape (modes extents)
                        &liouvillian));                   // the created empty operator (super-operator)
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        oneBodyTerm,                      // appended operator term
                        0,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, -1.0),  // -i constant
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
    //  Append an operator term to the operator (super-operator)
    fCoefsStaticMinus = static_cast<cuDoubleComplex *>(createInitializeArrayGPU(
      std::vector<cuDoubleComplex>(operBatchSize, make_cuDoubleComplex(0.0, -1.0)))); // -i constant for all coefficient instances in the batch
    fCoefsTotalMinus = static_cast<cuDoubleComplex *>(createInitializeArrayGPU(
      std::vector<cuDoubleComplex>(operBatchSize, make_cuDoubleComplex(0.0, 0.0)))); // storage for the total coefficients for all instances of the batch
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTermBatch(handle,
                        liouvillian,
                        twoBodyTerm,                      // appended operator term
                        0,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        operBatchSize,                    // number of instances of the operator term in the batch (they differ by the coefficient value)
                        fCoefsStaticMinus,                // static part of the f(t) batched coefficients in the two-body term
                        fCoefsTotalMinus,                 // total f(t) batched coefficients in the two-body term
                        {fCoefBatchComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr}, // CPU batched scalar callback function defining the time-dependent coefficient associated with this operator term as a whole
                        {fCoefBatchGradComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr})); // CPU batched scalar gradient callback function defining the gradient of the coefficient with respect to parameter Omega
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        oneBodyTerm,                      // appended operator term
                        1,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(0.0, +1.0),  // +i constant
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
    //  Append an operator term to the operator (super-operator)
    fCoefsStaticPlus = static_cast<cuDoubleComplex *>(createInitializeArrayGPU(
      std::vector<cuDoubleComplex>(operBatchSize, make_cuDoubleComplex(0.0, +1.0)))); // +i constant for all coefficient instances in the batch
    fCoefsTotalPlus = static_cast<cuDoubleComplex *>(createInitializeArrayGPU(
      std::vector<cuDoubleComplex>(operBatchSize, make_cuDoubleComplex(0.0, 0.0)))); // storage for the total coefficients for all instances of the batch
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTermBatch(handle,
                        liouvillian,
                        twoBodyTerm,                      // appended operator term
                        1,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        operBatchSize,                    // number of instances of the operator term in the batch (they differ by the coefficient value)
                        fCoefsStaticPlus,                 // static part of the f(t) batched coefficients in the dual two-body term
                        fCoefsTotalPlus,                  // total f(t) batched coefficients in the dual two-body term
                        {fCoefBatchComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr}, // CPU batched scalar callback function defining the time-dependent coefficient associated with this operator term as a whole
                        {fCoefBatchGradComplex64, CUDENSITYMAT_CALLBACK_DEVICE_CPU, nullptr})); // CPU batched scalar gradient callback function defining the gradient of the coefficient with respect to parameter Omega
    //  Append an operator term to the operator (super-operator)
    const double d = 1.0; // assign some value to the time-independent coefficient
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        noiseTerm,                        // appended operator term
                        0,                                // operator term action duality as a whole (no duality reversing in this case)
                        make_cuDoubleComplex(d, 0.0),     // static coefficient associated with the operator term as a whole
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
  }

  // Destructor destructs the user-defined Liouvillian operator
  ~UserDefinedLiouvillian()
  {
    // Destroy the Liouvillian operator
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(liouvillian));

    // Destroy operator terms
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(noiseTerm));
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(twoBodyTerm));
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(oneBodyTerm));

    // Destroy elementary tensor operators
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinYY));
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinZZ));
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinX));

    // Destroy the batched coefficients
    destroyArrayGPU(fCoefsTotalPlus);
    destroyArrayGPU(fCoefsStaticPlus);
    destroyArrayGPU(fCoefsTotalMinus);
    destroyArrayGPU(fCoefsStaticMinus);
    destroyArrayGPU(hCoefsTotal);
    destroyArrayGPU(hCoefsStatic);

    // Destroy operator tensors
    destroyArrayGPU(spinYYelems);
    destroyArrayGPU(spinZZelems);
    destroyArrayGPU(spinXelems);
  }

  // Disable copy constructor/assignment (GPU resources are private, no deep copy)
  UserDefinedLiouvillian(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian & operator=(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian(UserDefinedLiouvillian &&) = delete;
  UserDefinedLiouvillian & operator=(UserDefinedLiouvillian &&) = delete;

  /** Returns the number of externally provided Hamiltonian parameters. */
  int32_t getNumParameters() const
  {
    return 1; // one parameter Omega
  }

  /** Get access to the constructed Liouvillian operator. */
  cudensitymatOperator_t & get()
  {
    return liouvillian;
  }

};

Once the batched Liouvillian operator has been defined, the rest of the code logic is largely identical to the non-batched case.

/* Copyright (c) 2026-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <cudensitymat.h>  // cuDensityMat library header
#include "helpers.h"       // helper functions


// Batched time-dependent transverse-field Ising Hamiltonian operator
// with ordered and fused ZZ terms, plus fused unitary dissipation terms
#include "transverse_ising_full_fused_noisy_batch_grad.h" // user-defined batched Liouvillian operator example

#include <cmath>
#include <complex>
#include <vector>
#include <chrono>
#include <iostream>
#include <cassert>


// Number of times to perform operator action on a quantum state
constexpr int NUM_REPEATS = 2;

// Logging verbosity
bool verbose = true;


// Example workflow
void exampleWorkflow(cudensitymatHandle_t handle)
{
  // Define the composite Hilbert space shape and
  // quantum state batch size (number of individual quantum states in a batched simulation)
  const std::vector<int64_t> spaceShape({2,2,2,2}); // dimensions of quantum degrees of freedom
  const int64_t batchSize = 3;                      // number of quantum states per batch

  if (verbose) {
    std::cout << "Hilbert space rank = " << spaceShape.size() << "; Shape = (";
    for (const auto & dimsn: spaceShape)
      std::cout << dimsn << ",";
    std::cout << ")" << std::endl;
    std::cout << "Quantum state batch size = " << batchSize << std::endl;
  }

  // Construct a user-defined batched Liouvillian operator using a convenience C++ class
  // Note that the constructed Liouvillian operator has some batched coefficients
  UserDefinedLiouvillian liouvillian(handle, spaceShape, batchSize);
  if (verbose)
    std::cout << "Constructed the Liouvillian operator\n";

  // Set and place external user-provided Hamiltonian parameters in GPU memory
  const int32_t numParams = liouvillian.getNumParameters(); // number of external user-provided Hamiltonian parameters
  if (verbose)
    std::cout << "Number of external user-provided Hamiltonian parameters = " << numParams << std::endl;
  std::vector<double> cpuHamParams(numParams * batchSize);
  for (int64_t j = 0; j < batchSize; ++j) {
    for (int32_t i = 0; i < numParams; ++i) {
      cpuHamParams[j * numParams + i] = double(i+1) / double(j+1); // just setting some parameter values for each instance of the batch
    }
  }
  auto * hamiltonianParams = static_cast<double *>(createInitializeArrayGPU(cpuHamParams));
  if (verbose)
    std::cout << "Created an array of external user-provided Hamiltonian parameters in GPU memory\n";

  // Create an array of gradients for the user-provided Hamiltonian parameters in GPU memory
  std::vector<double> cpuHamParamsGrad(numParams * batchSize, 0.0);
  auto * hamiltonianParamsGrad = static_cast<double *>(createInitializeArrayGPU(cpuHamParamsGrad));
  if (verbose)
    std::cout << "Created an array of gradients for the external user-provided Hamiltonian parameters in GPU memory\n";

  // Declare the input quantum state
  cudensitymatState_t inputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &inputState));

  // Query the size of the quantum state storage
  std::size_t storageSize {0}; // only one storage component (tensor) is needed (no tensor factorization)
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      inputState,
                      1,               // only one storage component (tensor)
                      &storageSize));  // storage size in bytes
  const std::size_t stateVolume = storageSize / sizeof(NumericalType);  // quantum state tensor volume (number of elements)
  if (verbose)
    std::cout << "Quantum state storage size (bytes) = " << storageSize << std::endl;

  // Prepare some initial value for the input quantum state batch
  std::vector<NumericalType> inputStateValue(stateVolume);
  if constexpr (std::is_same_v<NumericalType, float>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0f / float(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, double>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = 1.0 / double(i+1); // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<float>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0f / float(i+1), -1.0f / float(i+2)}; // just some value
    }
  } else if constexpr (std::is_same_v<NumericalType, std::complex<double>>) {
    for (std::size_t i = 0; i < stateVolume; ++i) {
      inputStateValue[i] = NumericalType{1.0 / double(i+1), -1.0 / double(i+2)}; // just some value
    }
  } else {
    std::cerr << "Error: Unsupported data type!\n";
    std::exit(1);
  }
  // Allocate initialized GPU storage for the input quantum state with prepared values
  auto * inputStateElems = createInitializeArrayGPU(inputStateValue);
  if (verbose)
    std::cout << "Allocated input quantum state storage and initialized it to some value\n";

  // Attach initialized GPU storage to the input quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      inputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({inputStateElems}).data(),      // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed input quantum state\n";

  // Declare the output quantum state of the same shape
  cudensitymatState_t outputState;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,
                      &outputState));

  // Allocate GPU storage for the output quantum state
  auto * outputStateElems = createArrayGPU<NumericalType>(stateVolume);
  if (verbose)
    std::cout << "Allocated output quantum state storage\n";

  // Attach GPU storage to the output quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      outputState,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({outputStateElems}).data(),     // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed output quantum state\n";

  // Declare the adjoint input quantum state of the same shape
  cudensitymatState_t inputStateAdj;
  HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                      CUDENSITYMAT_STATE_PURITY_MIXED,  // pure (state vector) or mixed (density matrix) state
                      spaceShape.size(),
                      spaceShape.data(),
                      batchSize,
                      dataType,  // data type must match
                      &inputStateAdj));

  // Allocate GPU storage for the adjoint input quantum state
  auto * inputStateAdjElems = createArrayGPU<NumericalType>(stateVolume);
  if (verbose)
    std::cout << "Allocated adjoint input quantum state storage\n";

  // Attach GPU storage to the adjoint input quantum state
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      inputStateAdj,
                      1,                                                 // only one storage component (tensor)
                      std::vector<void*>({inputStateAdjElems}).data(),   // pointer to the GPU storage for the quantum state
                      std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
  if (verbose)
    std::cout << "Constructed adjoint input quantum state\n";

  // Declare a workspace descriptor
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  // Query free GPU memory
  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.45); // take 45% of the free memory for the workspace buffer
  if (verbose)
    std::cout << "Max workspace buffer size (bytes) = " << freeMem << std::endl;

  // Allocate GPU storage for the workspace buffer
  const std::size_t bufferVolume = freeMem / sizeof(NumericalType);
  auto * workspaceBuffer = createArrayGPU<NumericalType>(bufferVolume);
  if (verbose)
    std::cout << "Allocated workspace buffer of size (bytes) = " << freeMem << std::endl;

  // Prepare the Liouvillian operator action on a quantum state (needs to be done only once)
  auto startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorPrepareAction(handle,
                      liouvillian.get(),
                      inputState,
                      outputState,
                      CUDENSITYMAT_COMPUTE_64F,  // GPU compute type
                      freeMem,                   // max available GPU free memory for the workspace
                      workspaceDescr,            // workspace descriptor
                      0x0));                     // default CUDA stream
  auto finishTime = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator action preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  std::size_t requiredBufferSize {0};
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  if (requiredBufferSize > freeMem) {
    std::cerr << "Error: Required workspace buffer size is greater than the available GPU free memory!\n";
    std::exit(1);
  }

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      requiredBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Apply the Liouvillian operator to the input quatum state
  // and accumulate its action into the output quantum state (note the accumulative += semantics)
  for (int32_t repeat = 0; repeat < NUM_REPEATS; ++repeat) { // repeat multiple times for accurate timing
    // Zero out the output quantum state
    HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle,
                        outputState,
                        0x0));
    if (verbose)
      std::cout << "Initialized the output quantum state to zero\n";
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    startTime = std::chrono::high_resolution_clock::now();
    HANDLE_CUDM_ERROR(cudensitymatOperatorComputeAction(handle,
                        liouvillian.get(),
                        0.3,                // time point (some value)
                        batchSize,          // user-defined batch size
                        numParams,          // number of external user-defined Hamiltonian parameters
                        hamiltonianParams,  // external Hamiltonian parameters in GPU memory
                        inputState,         // input quantum state
                        outputState,        // output quantum state
                        workspaceDescr,     // workspace descriptor
                        0x0));              // default CUDA stream
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    finishTime = std::chrono::high_resolution_clock::now();
    timeSec = finishTime - startTime;
    if (verbose)
      std::cout << "Operator action computation time (sec) = " << timeSec.count() << std::endl;
  }

  // Compute the squared norm of the output quantum state
  void * norm2 = createInitializeArrayGPU(std::vector<double>(batchSize, 0.0));
  HANDLE_CUDM_ERROR(cudensitymatStateComputeNorm(handle,
                      outputState,
                      norm2,
                      0x0));
  if (verbose) {
    std::cout << "Computed the output quantum state norm:\n";
    printArrayGPU<double>(norm2, batchSize);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Prepare the Liouvillian operator action backward differentiation (needs to be done only once)
  startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorPrepareActionBackwardDiff(handle,
                      liouvillian.get(),
                      inputState,
                      outputState,               // adjoint output quantum state is always congruent to the output quantum state
                      CUDENSITYMAT_COMPUTE_64F,  // GPU compute type
                      freeMem,                   // max available GPU free memory for the workspace buffer
                      workspaceDescr,            // workspace descriptor
                      0x0));                     // default CUDA stream
  finishTime = std::chrono::high_resolution_clock::now();
  timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator action backward differentiation preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  requiredBufferSize = 0;
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  if (requiredBufferSize > freeMem) {
    std::cerr << "Error: Required workspace buffer size is greater than the available GPU free memory!\n";
    std::exit(1);
  }

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      requiredBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Liouvillian operator action backward differentiation:
  // The adjoint output quantum state, which is always congruent to the output quantum state,
  // depends on the user-defined cost function, so here we simply pass the previously computed output quantum state.
  // In real-life applications, the user will pass their adjoint output quantum state, computed for their cost function.
  for (int32_t repeat = 0; repeat < NUM_REPEATS; ++repeat) { // repeat multiple times for accurate timing
    // Zero out the adjoint input quantum state and gradients
    HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle,
                        inputStateAdj,
                        0x0));
    initializeArrayGPU(std::vector<double>(numParams * batchSize, 0.0), hamiltonianParamsGrad);
    if (verbose)
      std::cout << "Initialized the adjoint input quantum state and gradients to zero\n";
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    startTime = std::chrono::high_resolution_clock::now();
    HANDLE_CUDM_ERROR(cudensitymatOperatorComputeActionBackwardDiff(handle,
                        liouvillian.get(),
                        0.3,                    // time point (some value)
                        batchSize,              // user-defined batch size
                        numParams,              // number of external user-defined Hamiltonian parameters
                        hamiltonianParams,      // external Hamiltonian parameters in GPU memory
                        inputState,             // input quantum state
                        outputState,            // adjoint output quantum state (here we just pass the previously computed output quantum state for simplicity)
                        inputStateAdj,          // adjoint input quantum state
                        hamiltonianParamsGrad,  // partial gradients with respect to the user-defined real parameters
                        workspaceDescr,         // workspace descriptor
                        0x0));                  // default CUDA stream
    HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
    finishTime = std::chrono::high_resolution_clock::now();
    timeSec = finishTime - startTime;
    if (verbose)
      std::cout << "Operator action backward differentiation computation time (sec) = " << timeSec.count() << std::endl;
  }

  // Compute the squared norm of the adjoint input quantum state
  initializeArrayGPU(std::vector<double>(batchSize, 0.0), norm2);
  HANDLE_CUDM_ERROR(cudensitymatStateComputeNorm(handle,
                      inputStateAdj,
                      norm2,
                      0x0));
  if (verbose) {
    std::cout << "Computed the adjoint input quantum state norm:\n";
    printArrayGPU<double>(norm2, batchSize);
    std::cout << "Hamiltonian parameters gradients:\n";
    printArrayGPU<double>(hamiltonianParamsGrad, numParams * batchSize);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Destroy the norm2 array
  destroyArrayGPU(norm2);

  // Destroy workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));

  // Destroy workspace buffer storage
  destroyArrayGPU(workspaceBuffer);

  // Destroy quantum states
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(inputStateAdj));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(outputState));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(inputState));

  // Destroy quantum state storage
  destroyArrayGPU(inputStateAdjElems);
  destroyArrayGPU(outputStateElems);
  destroyArrayGPU(inputStateElems);

  // Destroy external Hamiltonian parameters
  destroyArrayGPU(static_cast<void *>(hamiltonianParamsGrad));
  destroyArrayGPU(static_cast<void *>(hamiltonianParams));

  if (verbose)
    std::cout << "Destroyed resources\n" << std::flush;
}


int main(int argc, char ** argv)
{
  // Assign a GPU to the process
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose)
    std::cout << "Set active device\n";

  // Create a library handle
  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created a library handle\n";

  // Run the example
  exampleWorkflow(handle);

  // Destroy the library handle
  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed the library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());

  // Done
  return 0;
}

Code example (serial execution of MPS-TDVP time propagation)

The following code example illustrates how to use the cuDensityMat library to propagate a pure MPS state in time under a transverse-field Ising Hamiltonian encoded as an MPO, using the split-scope TDVP method with Krylov subspace exponentiation. The example builds the Hamiltonian MPO, creates input and output MPS states, initializes the input MPS to a Neel state, configures the TDVP and Krylov solver parameters, and then runs a time-stepping loop. The full sample code can be found in the NVIDIA/cuQuantum repository (main serial MPS-TDVP code as well as the utility code).

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

// MPS-TDVP time propagation example.
//
// Demonstrates propagating a pure MPS state under a transverse-field Ising
// Hamiltonian encoded as an MPO, using the split-scope TDVP method with
// Krylov subspace exponentiation.
//
// Workflow:
//  1. Build Hamiltonian as MPO (nearest-neighbor ZZ + transverse X field)
//  2. Create input and output MPS states (cudensitymatCreateStateMPS)
//  3. Initialize input MPS to a Neel state |0101...>
//  4. Create TDVP time propagation object
//  5. Configure TDVP / Krylov parameters
//  6. Prepare propagation and allocate workspace
//  7. Time-stepping loop
//  8. Clean up all resources

#include <cudensitymat.h>
#include "helpers.h"

#include <algorithm>
#include <cmath>
#include <complex>
#include <vector>
#include <numeric>
#include <chrono>
#include <iostream>
#include <cassert>


using Complex = std::complex<double>;
constexpr cudaDataType_t DATA_TYPE = CUDA_C_64F;

constexpr bool verbose = true;

// --- Simulation parameters ---
constexpr int32_t  NUM_SITES    = 40;
constexpr int64_t  PHYS_DIM     = 2;
constexpr int64_t  MAX_BOND_DIM = 128;
constexpr int64_t  MPO_BOND_DIM = 3;
constexpr int32_t  NUM_STEPS    = 3;
constexpr double   DT           = 0.01;
constexpr double   J_COUPLING   = 1.0;
constexpr double   H_FIELD      = 0.5;


// ============================================================================
// Transverse-field Ising MPO builder
// ============================================================================
//
// H = -J * sum_{i} Z_i Z_{i+1}  +  h * sum_{i} X_i
//
// Standard MPO representation with bond dimension 3.
// Site tensor mode ordering follows cuDensityMat convention (column-major):
//   Left boundary (site 0):     [phys_ket(d), right_bond(bR), phys_bra(d)]
//   Interior sites:             [left_bond(bL), phys_ket(d), right_bond(bR), phys_bra(d)]
//   Right boundary (site N-1):  [left_bond(bL), phys_ket(d), phys_bra(d)]
//
// The bulk MPO matrix (indexed by bond dimensions aL, aR) is:
//
//        |  I      0     0  |
//   W =  |  Z      0     0  |
//        | h*X   -J*Z    I  |
//
// Left boundary row-vector  (d x 3 x d):  W_L = [ h*X  -J*Z  I ]
// Right boundary col-vector (3 x d x d):  W_R = [ I; Z; h*X ]^T
//
// Pauli matrices: I = diag(1,1),  X = [[0,1],[1,0]],  Z = diag(1,-1)

struct IsingMPO {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build() {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0};
    const Complex mone{-1.0, 0.0};
    const Complex jc{-J_COUPLING, 0.0};
    const Complex hc{H_FIELD, 0.0};

    // Pauli matrices stored column-major: mat[col*2 + row]
    auto I_mat = [&](int r, int c) -> Complex { return (r == c) ? one : zero; };
    auto X_mat = [&](int r, int c) -> Complex { return (r != c) ? one : zero; };
    auto Z_mat = [&](int r, int c) -> Complex { return (r == c) ? ((r == 0) ? one : mone) : zero; };

    hostTensors.resize(NUM_SITES);
    gpuPtrs.resize(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL = (site == 0) ? 1 : MPO_BOND_DIM;
      const int64_t bR = (site == NUM_SITES - 1) ? 1 : MPO_BOND_DIM;
      const int64_t vol = bL * PHYS_DIM * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      // Library mode ordering (column-major):
      //   [left_bond, phys_ket, right_bond, phys_bra]
      // Boundary sites omit the missing bond; bL=1 or bR=1 makes
      // the degenerate dimension transparent in the flat index.
      auto idx = [&](int64_t aL, int64_t ket, int64_t aR, int64_t bra) -> int64_t {
        return aL + bL * (ket + PHYS_DIM * (aR + bR * bra));
      };

      if (NUM_SITES == 1) {
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket)
            hostTensors[site][idx(0, ket, 0, bra)] = hc * X_mat(bra, ket);
      } else if (site == 0) {
        // Left boundary: row-vector [h*X, -J*Z, I]
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = hc * X_mat(bra, ket);
            hostTensors[site][idx(0, ket, 1, bra)] = jc * Z_mat(bra, ket);
            hostTensors[site][idx(0, ket, 2, bra)] = I_mat(bra, ket);
          }
      } else if (site == NUM_SITES - 1) {
        // Right boundary: col-vector [I; Z; h*X]^T
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = I_mat(bra, ket);
            hostTensors[site][idx(1, ket, 0, bra)] = Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 0, bra)] = hc * X_mat(bra, ket);
          }
      } else {
        // Bulk:
        //   row 0: [I, 0, 0]
        //   row 1: [Z, 0, 0]
        //   row 2: [h*X, -J*Z, I]
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = I_mat(bra, ket);
            hostTensors[site][idx(1, ket, 0, bra)] = Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 0, bra)] = hc * X_mat(bra, ket);
            hostTensors[site][idx(2, ket, 1, bra)] = jc * Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 2, bra)] = I_mat(bra, ket);
          }
      }

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Build a Neel-state MPS  |0,1,0,1,...>  with given bond dimensions
// ============================================================================
// Each MPS tensor A[site] has shape (bondL, phys, bondR) in column-major.
// For a product state, only the (0,sigma,0) slice is nonzero:
//   A[0,sigma,0] = delta(sigma, site%2)
// Remaining bond indices are zero-padded.

struct NeelMPS {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build(const std::vector<int64_t>& bondDims) {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0};

    hostTensors.resize(NUM_SITES);
    gpuPtrs.resize(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL = (site == 0) ? 1 : bondDims[site - 1];
      const int64_t bR = (site == NUM_SITES - 1) ? 1 : bondDims[site];
      const int64_t vol = bL * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      // T[aL, sigma, aR]  column-major
      auto idx = [&](int64_t aL, int64_t sigma, int64_t aR) -> int64_t {
        return aL + bL * (sigma + PHYS_DIM * aR);
      };

      const int64_t neelSpin = site % 2;  // 0 for even sites, 1 for odd
      hostTensors[site][idx(0, neelSpin, 0)] = one;

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Example workflow
// ============================================================================

void exampleWorkflow(cudensitymatHandle_t handle)
{
  // --- 1. Build the transverse-field Ising MPO ---
  IsingMPO mpo;
  mpo.build();
  if (verbose)
    std::cout << "Built transverse-field Ising MPO (bond dim " << MPO_BOND_DIM << ")\n";

  const std::vector<int64_t> spaceShape(NUM_SITES, PHYS_DIM);
  std::vector<int64_t> mpoBondDims(NUM_SITES - 1, MPO_BOND_DIM);

  cudensitymatMatrixProductOperator_t mpoHandle;
  std::vector<cudensitymatWrappedTensorCallback_t> mpoCallbacks(NUM_SITES, cudensitymatTensorCallbackNone);
  std::vector<cudensitymatWrappedTensorGradientCallback_t> mpoGradCallbacks(NUM_SITES, cudensitymatTensorGradientCallbackNone);
  HANDLE_CUDM_ERROR(cudensitymatCreateMatrixProductOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpoBondDims.data(),
                      DATA_TYPE,
                      mpo.gpuPtrs.data(),
                      mpoCallbacks.data(),
                      mpoGradCallbacks.data(),
                      &mpoHandle));
  if (verbose)
    std::cout << "Created MPO handle\n";

  // --- 2. Build the Operator from the MPO ---
  cudensitymatOperatorTerm_t operatorTerm;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      &operatorTerm));

  std::vector<int32_t> modesActedOn(NUM_SITES);
  std::iota(modesActedOn.begin(), modesActedOn.end(), 0);
  std::vector<int32_t> modeDuality(NUM_SITES, 0);
  std::vector<int32_t> mpoConjugation = {0};

  HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendMPOProduct(handle,
                      operatorTerm,
                      1,
                      &mpoHandle,
                      mpoConjugation.data(),
                      modesActedOn.data(),
                      modeDuality.data(),
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));

  cudensitymatOperator_t hamiltonian;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      &hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                      hamiltonian,
                      operatorTerm,
                      0,
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));
  if (verbose)
    std::cout << "Constructed Hamiltonian operator from MPO\n";

  // --- 3. Create input and output MPS states ---
  const int64_t batchSize = 1;

  // For OBC, cap bond dimensions by the exact Hilbert space dimension on each side.
  std::vector<int64_t> mpsBondDims(NUM_SITES - 1);
  for (int32_t i = 0; i < NUM_SITES - 1; ++i) {
    int64_t leftDim = 1;
    for (int32_t j = 0; j <= i; ++j) leftDim *= spaceShape[j];
    int64_t rightDim = 1;
    for (int32_t j = i + 1; j < NUM_SITES; ++j) rightDim *= spaceShape[j];
    mpsBondDims[i] = std::min({MAX_BOND_DIM, leftDim, rightDim});
  }

  cudensitymatState_t stateIn, stateOut;
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpsBondDims.data(),
                      DATA_TYPE,
                      batchSize,
                      &stateIn));
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpsBondDims.data(),
                      DATA_TYPE,
                      batchSize,
                      &stateOut));

  // Query number of MPS components (= NUM_SITES tensors)
  int32_t numComponents = 0;
  HANDLE_CUDM_ERROR(cudensitymatStateGetNumComponents(handle, stateIn, &numComponents));
  assert(numComponents == NUM_SITES);

  // Query storage sizes for each MPS tensor
  std::vector<std::size_t> componentSizes(numComponents);
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      stateIn, numComponents, componentSizes.data()));

  if (verbose) {
    std::cout << "MPS state has " << numComponents << " components, sizes (bytes):";
    for (auto s : componentSizes) std::cout << " " << s;
    std::cout << "\n";
  }

  // --- 4. Allocate GPU storage and initialise Neel MPS for stateIn ---
  NeelMPS neelMps;
  neelMps.build(mpsBondDims);

  // Allocate empty storage for stateOut
  std::vector<void*> stateOutBufs(numComponents, nullptr);
  for (int32_t c = 0; c < numComponents; ++c)
    stateOutBufs[c] = createArrayGPU<Complex>(componentSizes[c] / sizeof(Complex));

  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      stateIn, numComponents,
                      neelMps.gpuPtrs.data(), componentSizes.data()));
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      stateOut, numComponents,
                      stateOutBufs.data(), componentSizes.data()));

  if (verbose)
    std::cout << "Initialized MPS states (Neel state |0101...>)\n";

  // --- 5. Create TDVP time propagation object ---
  cudensitymatTimePropagation_t timeProp;
  HANDLE_CUDM_ERROR(cudensitymatCreateTimePropagation(handle,
                      hamiltonian,
                      1,  // Hermitian
                      CUDENSITYMAT_PROPAGATION_SCOPE_SPLIT,
                      CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV,
                      &timeProp));
  if (verbose)
    std::cout << "Created TDVP time propagation object\n";

  // --- 6. Configure TDVP and Krylov parameters ---
  // TDVP config: 2nd-order, 1-site
  cudensitymatTimePropagationScopeSplitTDVPConfig_t tdvpConfig;
  HANDLE_CUDM_ERROR(cudensitymatCreateTimePropagationScopeSplitTDVPConfig(handle, &tdvpConfig));
  {
    const int32_t order = 2;
    HANDLE_CUDM_ERROR(cudensitymatTimePropagationScopeSplitTDVPConfigSetAttribute(handle,
                        tdvpConfig,
                        CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_ORDER,
                        &order, sizeof(order)));
  }
  HANDLE_CUDM_ERROR(cudensitymatTimePropagationConfigure(handle,
                      timeProp,
                      CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_CONFIG,
                      &tdvpConfig, sizeof(tdvpConfig)));

  // Krylov config
  cudensitymatTimePropagationApproachKrylovConfig_t krylovConfig;
  HANDLE_CUDM_ERROR(cudensitymatCreateTimePropagationApproachKrylovConfig(handle, &krylovConfig));
  {
    const int32_t maxDim = 10;
    HANDLE_CUDM_ERROR(cudensitymatTimePropagationApproachKrylovConfigSetAttribute(handle,
                        krylovConfig,
                        CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV_MAX_DIM,
                        &maxDim, sizeof(maxDim)));
    const double tol = 1e-8;
    HANDLE_CUDM_ERROR(cudensitymatTimePropagationApproachKrylovConfigSetAttribute(handle,
                        krylovConfig,
                        CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV_TOLERANCE,
                        &tol, sizeof(tol)));
  }
  HANDLE_CUDM_ERROR(cudensitymatTimePropagationConfigure(handle,
                      timeProp,
                      CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV_CONFIG,
                      &krylovConfig, sizeof(krylovConfig)));
  if (verbose)
    std::cout << "Configured TDVP (order=2, 1-site) with Krylov (max_dim=10, tol=1e-8)\n";

  // --- 7. Prepare propagation and allocate workspace ---
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95);
  if (verbose)
    std::cout << "Available workspace memory (bytes) = " << freeMem << "\n";

  HANDLE_CUDM_ERROR(cudensitymatTimePropagationPrepare(handle,
                      timeProp,
                      stateIn,
                      stateOut,
                      CUDENSITYMAT_COMPUTE_64F,
                      freeMem,
                      workspaceDescr,
                      0x0));
  if (verbose)
    std::cout << "Prepared time propagation\n";

  // Query and allocate scratch workspace
  std::size_t scratchSize = 0;
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &scratchSize));
  void * scratchBuf = nullptr;
  if (scratchSize > 0) {
    HANDLE_CUDA_ERROR(cudaMalloc(&scratchBuf, scratchSize));
    HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                        workspaceDescr,
                        CUDENSITYMAT_MEMSPACE_DEVICE,
                        CUDENSITYMAT_WORKSPACE_SCRATCH,
                        scratchBuf, scratchSize));
  }
  if (verbose)
    std::cout << "Scratch workspace (bytes) = " << scratchSize << "\n";

  // --- 8. Time-stepping loop ---
  if (verbose) {
    std::cout << "\nStarting TDVP propagation: " << NUM_STEPS << " steps, dt = " << DT << "\n";
    std::cout << "Hamiltonian: H = -" << J_COUPLING << " sum Z_i Z_{i+1} + "
              << H_FIELD << " sum X_i\n\n";
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  const auto wallStart = std::chrono::high_resolution_clock::now();

  for (int32_t step = 0; step < NUM_STEPS; ++step) {
    const double currentTime = step * DT;

    HANDLE_CUDM_ERROR(cudensitymatTimePropagationCompute(handle,
                        timeProp,
                        DT,      // timeStepReal
                        0.0,     // timeStepImag
                        currentTime,
                        batchSize,
                        0,       // numParams
                        nullptr, // params
                        stateIn,
                        stateOut,
                        workspaceDescr,
                        0x0));

    // Swap input/output for next step (double-buffering)
    std::swap(stateIn, stateOut);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  const auto wallEnd = std::chrono::high_resolution_clock::now();
  const std::chrono::duration<double> elapsed = wallEnd - wallStart;
  if (verbose)
    std::cout << "\nTotal propagation wall time (sec) = " << elapsed.count() << "\n";

  // --- 9. Clean up ---
  if (scratchBuf)
    HANDLE_CUDA_ERROR(cudaFree(scratchBuf));
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));
  HANDLE_CUDM_ERROR(cudensitymatDestroyTimePropagation(timeProp));
  HANDLE_CUDM_ERROR(cudensitymatDestroyTimePropagationApproachKrylovConfig(krylovConfig));
  HANDLE_CUDM_ERROR(cudensitymatDestroyTimePropagationScopeSplitTDVPConfig(tdvpConfig));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(operatorTerm));
  HANDLE_CUDM_ERROR(cudensitymatDestroyMatrixProductOperator(mpoHandle));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(stateOut));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(stateIn));

  for (auto buf : stateOutBufs)
    destroyArrayGPU(buf);
  neelMps.destroy();
  mpo.destroy();

  if (verbose)
    std::cout << "Destroyed all resources\n";
}


int main(int argc, char ** argv)
{
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose)
    std::cout << "Set active device\n";

  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created library handle\n";

  exampleWorkflow(handle);

  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());
  return 0;
}

Code example (serial execution of two-site MPS-TDVP time propagation)

The following code example illustrates the two-site variant of the split-scope TDVP method, selected via the CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_NUM_SITES attribute set to 2. Unlike the single-site update, the two-site update can grow the MPS bond extents dynamically: each evolved two-site block is re-split via a truncated SVD whose policy is supplied through an attached cudensitymatSVDConfig_t (here with a global cap via CUDENSITYMAT_SVD_CONFIG_MAX_EXTENT). The example starts from a polarized product state with small initial current bond extents recorded via cudensitymatStateMPSSetCurrentBondExtents(), and after each step queries cudensitymatStateMPSGetCurrentBondExtents() and compares the current extents against the maximum (buffer) extents to monitor bond growth and truncation. The full sample code can be found in the NVIDIA/cuQuantum repository (main serial two-site MPS-TDVP code as well as the utility code).

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

// Two-site MPS-TDVP time propagation example.
//
// Demonstrates propagating a pure MPS state under a transverse-field Ising
// Hamiltonian encoded as an MPO, using the split-scope TDVP method with the
// TWO-site local update (CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_NUM_SITES = 2)
// and Krylov subspace exponentiation.
//
// Unlike the single-site companion (mps_tdvp_example.cpp), the two-site sweep
// can dynamically grow the MPS bond dimensions: each two-site block is evolved
// and then re-split via a truncated SVD. This example therefore additionally:
//   * starts from small current bond extents (> 1, below the buffer maximum)
//     via cudensitymatStateMPSSetCurrentBondExtents;
//   * attaches a TDVP SVD configuration with a maximum-extent cap
//     (CUDENSITYMAT_SVD_CONFIG_MAX_EXTENT); and
//   * after each step queries cudensitymatStateMPSGetCurrentBondExtents and
//     compares against the maximum extents to monitor bond growth.
//
// Workflow:
//  1. Build Hamiltonian as MPO (nearest-neighbor ZZ + transverse X field)
//  2. Create input and output MPS states (cudensitymatCreateStateMPS)
//  3. Initialize input MPS to a polarized product state |000...0> and set small
//     current bond extents
//  4. Create TDVP time propagation object
//  5. Configure TDVP (num_sites = 2 + SVD config) / Krylov parameters
//  6. Prepare propagation and allocate workspace
//  7. Time-stepping loop; after each step monitor current vs. maximum extents
//  8. Clean up all resources

#include <cudensitymat.h>
#include "helpers.h"

#include <algorithm>
#include <cmath>
#include <complex>
#include <vector>
#include <numeric>
#include <chrono>
#include <iostream>
#include <cassert>


using Complex = std::complex<double>;
constexpr cudaDataType_t DATA_TYPE = CUDA_C_64F;

constexpr bool verbose = true;

// --- Simulation parameters ---
constexpr int32_t  NUM_SITES        = 8;
constexpr int64_t  PHYS_DIM         = 2;
constexpr int64_t  MAX_BOND_DIM     = 8;    // maximum (buffer) bond dimension cap
constexpr int64_t  MPO_BOND_DIM     = 3;
constexpr int64_t  INITIAL_BOND_DIM = 2;    // small initial current bond dimension (> 1)
constexpr int32_t  NUM_STEPS        = 5;
constexpr double   DT               = 0.05;
constexpr double   J_COUPLING       = 1.0;
constexpr double   H_FIELD          = 0.5;


// ============================================================================
// Transverse-field Ising MPO builder
// ============================================================================
//
// H = -J * sum_{i} Z_i Z_{i+1}  +  h * sum_{i} X_i
//
// Standard MPO representation with bond dimension 3.
// Site tensor mode ordering follows cuDensityMat convention (column-major):
//   Left boundary (site 0):     [phys_ket(d), right_bond(bR), phys_bra(d)]
//   Interior sites:             [left_bond(bL), phys_ket(d), right_bond(bR), phys_bra(d)]
//   Right boundary (site N-1):  [left_bond(bL), phys_ket(d), phys_bra(d)]
//
// The bulk MPO matrix (indexed by bond dimensions aL, aR) is:
//
//        |  I      0     0  |
//   W =  |  Z      0     0  |
//        | h*X   -J*Z    I  |
//
// Pauli matrices: I = diag(1,1),  X = [[0,1],[1,0]],  Z = diag(1,-1)

struct IsingMPO {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build() {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0};
    const Complex mone{-1.0, 0.0};
    const Complex jc{-J_COUPLING, 0.0};
    const Complex hc{H_FIELD, 0.0};

    auto I_mat = [&](int r, int c) -> Complex { return (r == c) ? one : zero; };
    auto X_mat = [&](int r, int c) -> Complex { return (r != c) ? one : zero; };
    auto Z_mat = [&](int r, int c) -> Complex { return (r == c) ? ((r == 0) ? one : mone) : zero; };

    hostTensors.resize(NUM_SITES);
    gpuPtrs.resize(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL = (site == 0) ? 1 : MPO_BOND_DIM;
      const int64_t bR = (site == NUM_SITES - 1) ? 1 : MPO_BOND_DIM;
      const int64_t vol = bL * PHYS_DIM * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      // Library mode ordering (column-major): [left_bond, phys_ket, right_bond, phys_bra]
      auto idx = [&](int64_t aL, int64_t ket, int64_t aR, int64_t bra) -> int64_t {
        return aL + bL * (ket + PHYS_DIM * (aR + bR * bra));
      };

      if (site == 0) {
        // Left boundary: row-vector [h*X, -J*Z, I]
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = hc * X_mat(bra, ket);
            hostTensors[site][idx(0, ket, 1, bra)] = jc * Z_mat(bra, ket);
            hostTensors[site][idx(0, ket, 2, bra)] = I_mat(bra, ket);
          }
      } else if (site == NUM_SITES - 1) {
        // Right boundary: col-vector [I; Z; h*X]^T
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = I_mat(bra, ket);
            hostTensors[site][idx(1, ket, 0, bra)] = Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 0, bra)] = hc * X_mat(bra, ket);
          }
      } else {
        // Bulk:
        //   row 0: [I, 0, 0]
        //   row 1: [Z, 0, 0]
        //   row 2: [h*X, -J*Z, I]
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = I_mat(bra, ket);
            hostTensors[site][idx(1, ket, 0, bra)] = Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 0, bra)] = hc * X_mat(bra, ket);
            hostTensors[site][idx(2, ket, 1, bra)] = jc * Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 2, bra)] = I_mat(bra, ket);
          }
      }

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Build a polarized product-state MPS  |000...0>  with given (maximum/buffer)
// bond extents
// ============================================================================
// Each MPS tensor A[site] has shape (bondL, phys, bondR) in column-major over
// the MAXIMUM (buffer) extents. For a product state only the (0,sigma,0) slice
// is nonzero; the smaller current bond extents (set separately) leave the rest
// of the buffer zero-padded.

struct ProductMPS {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build(const std::vector<int64_t>& maxBondDims) {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0};

    hostTensors.resize(NUM_SITES);
    gpuPtrs.resize(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL = (site == 0) ? 1 : maxBondDims[site - 1];
      const int64_t bR = (site == NUM_SITES - 1) ? 1 : maxBondDims[site];
      const int64_t vol = bL * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      // T[aL, sigma, aR]  column-major over the maximum extents.
      auto idx = [&](int64_t aL, int64_t sigma, int64_t aR) -> int64_t {
        return aL + bL * (sigma + PHYS_DIM * aR);
      };

      hostTensors[site][idx(0, 0, 0)] = one;  // |0> on each site (product state)

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Example workflow
// ============================================================================

void exampleWorkflow(cudensitymatHandle_t handle)
{
  // --- 1. Build the transverse-field Ising MPO ---
  IsingMPO mpo;
  mpo.build();
  if (verbose)
    std::cout << "Built transverse-field Ising MPO (bond dim " << MPO_BOND_DIM << ")\n";

  const std::vector<int64_t> spaceShape(NUM_SITES, PHYS_DIM);
  std::vector<int64_t> mpoBondDims(NUM_SITES - 1, MPO_BOND_DIM);

  cudensitymatMatrixProductOperator_t mpoHandle;
  std::vector<cudensitymatWrappedTensorCallback_t> mpoCallbacks(NUM_SITES, cudensitymatTensorCallbackNone);
  std::vector<cudensitymatWrappedTensorGradientCallback_t> mpoGradCallbacks(NUM_SITES, cudensitymatTensorGradientCallbackNone);
  HANDLE_CUDM_ERROR(cudensitymatCreateMatrixProductOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpoBondDims.data(),
                      DATA_TYPE,
                      mpo.gpuPtrs.data(),
                      mpoCallbacks.data(),
                      mpoGradCallbacks.data(),
                      &mpoHandle));
  if (verbose)
    std::cout << "Created MPO handle\n";

  // --- 2. Build the Operator from the MPO ---
  cudensitymatOperatorTerm_t operatorTerm;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      &operatorTerm));

  std::vector<int32_t> modesActedOn(NUM_SITES);
  std::iota(modesActedOn.begin(), modesActedOn.end(), 0);
  std::vector<int32_t> modeDuality(NUM_SITES, 0);
  std::vector<int32_t> mpoConjugation = {0};

  HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendMPOProduct(handle,
                      operatorTerm,
                      1,
                      &mpoHandle,
                      mpoConjugation.data(),
                      modesActedOn.data(),
                      modeDuality.data(),
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));

  cudensitymatOperator_t hamiltonian;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      &hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                      hamiltonian,
                      operatorTerm,
                      0,
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));
  if (verbose)
    std::cout << "Constructed Hamiltonian operator from MPO\n";

  // --- 3. Create input and output MPS states ---
  const int64_t batchSize = 1;

  // Maximum (buffer) bond extents, fixed at state creation: capped by the exact
  // Hilbert space dimension on each side and by MAX_BOND_DIM.
  std::vector<int64_t> maxBondDims(NUM_SITES - 1);
  for (int32_t i = 0; i < NUM_SITES - 1; ++i) {
    int64_t leftDim = 1;
    for (int32_t j = 0; j <= i; ++j) leftDim *= spaceShape[j];
    int64_t rightDim = 1;
    for (int32_t j = i + 1; j < NUM_SITES; ++j) rightDim *= spaceShape[j];
    maxBondDims[i] = std::min({MAX_BOND_DIM, leftDim, rightDim});
  }

  // Initial current bond extents: > 1, clamped to the maximum, so the two-site
  // sweep starts below the buffer maximum and can grow the bonds.
  std::vector<int64_t> initialBondDims(NUM_SITES - 1);
  for (int32_t i = 0; i < NUM_SITES - 1; ++i)
    initialBondDims[i] = std::min<int64_t>(INITIAL_BOND_DIM, maxBondDims[i]);

  cudensitymatState_t stateIn, stateOut;
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      maxBondDims.data(),
                      DATA_TYPE,
                      batchSize,
                      &stateIn));
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      maxBondDims.data(),
                      DATA_TYPE,
                      batchSize,
                      &stateOut));

  // Query number of MPS components (= NUM_SITES tensors)
  int32_t numComponents = 0;
  HANDLE_CUDM_ERROR(cudensitymatStateGetNumComponents(handle, stateIn, &numComponents));
  assert(numComponents == NUM_SITES);

  // Query storage sizes for each MPS tensor (sized for the MAXIMUM extents)
  std::vector<std::size_t> componentSizes(numComponents);
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      stateIn, numComponents, componentSizes.data()));

  if (verbose) {
    std::cout << "MPS state has " << numComponents << " components, sizes (bytes):";
    for (auto s : componentSizes) std::cout << " " << s;
    std::cout << "\n";
  }

  // --- 4. Allocate GPU storage and initialise the product-state MPS for stateIn ---
  ProductMPS productMps;
  productMps.build(maxBondDims);

  // Allocate empty storage for stateOut
  std::vector<void*> stateOutBufs(numComponents, nullptr);
  for (int32_t c = 0; c < numComponents; ++c)
    stateOutBufs[c] = createArrayGPU<Complex>(componentSizes[c] / sizeof(Complex));

  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      stateIn, numComponents,
                      productMps.gpuPtrs.data(), componentSizes.data()));
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      stateOut, numComponents,
                      stateOutBufs.data(), componentSizes.data()));

  // Record the small initial CURRENT (valid) bond extents (metadata only).
  HANDLE_CUDM_ERROR(cudensitymatStateMPSSetCurrentBondExtents(handle,
                      stateIn, initialBondDims.data()));
  HANDLE_CUDM_ERROR(cudensitymatStateMPSSetCurrentBondExtents(handle,
                      stateOut, initialBondDims.data()));

  if (verbose) {
    std::cout << "Initialized MPS states (polarized product state |000...0>)\n";
    std::cout << "Maximum (buffer) bond extents:";
    for (auto b : maxBondDims) std::cout << " " << b;
    std::cout << "\nInitial current bond extents:";
    for (auto b : initialBondDims) std::cout << " " << b;
    std::cout << "\n";
  }

  // --- 5. Create TDVP time propagation object ---
  cudensitymatTimePropagation_t timeProp;
  HANDLE_CUDM_ERROR(cudensitymatCreateTimePropagation(handle,
                      hamiltonian,
                      1,  // Hermitian
                      CUDENSITYMAT_PROPAGATION_SCOPE_SPLIT,
                      CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV,
                      &timeProp));
  if (verbose)
    std::cout << "Created TDVP time propagation object\n";

  // --- 6. Configure TDVP and Krylov parameters ---
  // TDVP config: 2nd-order, TWO-site, with an attached SVD truncation policy.
  cudensitymatTimePropagationScopeSplitTDVPConfig_t tdvpConfig;
  HANDLE_CUDM_ERROR(cudensitymatCreateTimePropagationScopeSplitTDVPConfig(handle, &tdvpConfig));
  {
    const int32_t order = 2;
    HANDLE_CUDM_ERROR(cudensitymatTimePropagationScopeSplitTDVPConfigSetAttribute(handle,
                        tdvpConfig,
                        CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_ORDER,
                        &order, sizeof(order)));
    const int32_t numSites = 2;
    HANDLE_CUDM_ERROR(cudensitymatTimePropagationScopeSplitTDVPConfigSetAttribute(handle,
                        tdvpConfig,
                        CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_NUM_SITES,
                        &numSites, sizeof(numSites)));
  }

  // SVD truncation policy for the two-site splits (captured by value at Configure
  // time, so it can be destroyed right after). Here just a global bond cap.
  cudensitymatSVDConfig_t svdConfig;
  HANDLE_CUDM_ERROR(cudensitymatCreateSVDConfig(handle, &svdConfig));
  {
    const int64_t svdMaxExtent = MAX_BOND_DIM;
    HANDLE_CUDM_ERROR(cudensitymatSVDConfigSetAttribute(handle, svdConfig,
                        CUDENSITYMAT_SVD_CONFIG_MAX_EXTENT,
                        &svdMaxExtent, sizeof(svdMaxExtent)));
  }
  HANDLE_CUDM_ERROR(cudensitymatTimePropagationScopeSplitTDVPConfigSetAttribute(handle,
                      tdvpConfig,
                      CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_SVD_CONFIG,
                      &svdConfig, sizeof(svdConfig)));

  HANDLE_CUDM_ERROR(cudensitymatTimePropagationConfigure(handle,
                      timeProp,
                      CUDENSITYMAT_PROPAGATION_SPLIT_SCOPE_TDVP_CONFIG,
                      &tdvpConfig, sizeof(tdvpConfig)));

  HANDLE_CUDM_ERROR(cudensitymatDestroySVDConfig(svdConfig));

  // Krylov config
  cudensitymatTimePropagationApproachKrylovConfig_t krylovConfig;
  HANDLE_CUDM_ERROR(cudensitymatCreateTimePropagationApproachKrylovConfig(handle, &krylovConfig));
  {
    const int32_t maxDim = 10;
    HANDLE_CUDM_ERROR(cudensitymatTimePropagationApproachKrylovConfigSetAttribute(handle,
                        krylovConfig,
                        CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV_MAX_DIM,
                        &maxDim, sizeof(maxDim)));
    const double tol = 1e-8;
    HANDLE_CUDM_ERROR(cudensitymatTimePropagationApproachKrylovConfigSetAttribute(handle,
                        krylovConfig,
                        CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV_TOLERANCE,
                        &tol, sizeof(tol)));
  }
  HANDLE_CUDM_ERROR(cudensitymatTimePropagationConfigure(handle,
                      timeProp,
                      CUDENSITYMAT_PROPAGATION_APPROACH_KRYLOV_CONFIG,
                      &krylovConfig, sizeof(krylovConfig)));
  if (verbose)
    std::cout << "Configured TDVP (order=2, num_sites=2, SVD max_extent=" << MAX_BOND_DIM
              << ") with Krylov (max_dim=10, tol=1e-8)\n";

  // --- 7. Prepare propagation and allocate workspace ---
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95);
  if (verbose)
    std::cout << "Available workspace memory (bytes) = " << freeMem << "\n";

  HANDLE_CUDM_ERROR(cudensitymatTimePropagationPrepare(handle,
                      timeProp,
                      stateIn,
                      stateOut,
                      CUDENSITYMAT_COMPUTE_64F,
                      freeMem,
                      workspaceDescr,
                      0x0));
  if (verbose)
    std::cout << "Prepared time propagation\n";

  // Query and allocate scratch workspace
  std::size_t scratchSize = 0;
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &scratchSize));
  void * scratchBuf = nullptr;
  if (scratchSize > 0) {
    HANDLE_CUDA_ERROR(cudaMalloc(&scratchBuf, scratchSize));
    HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                        workspaceDescr,
                        CUDENSITYMAT_MEMSPACE_DEVICE,
                        CUDENSITYMAT_WORKSPACE_SCRATCH,
                        scratchBuf, scratchSize));
  }
  if (verbose)
    std::cout << "Scratch workspace (bytes) = " << scratchSize << "\n";

  // --- 8. Time-stepping loop ---
  if (verbose) {
    std::cout << "\nStarting 2-site TDVP propagation: " << NUM_STEPS << " steps, dt = " << DT << "\n";
    std::cout << "Hamiltonian: H = -" << J_COUPLING << " sum Z_i Z_{i+1} + "
              << H_FIELD << " sum X_i\n\n";
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  const auto wallStart = std::chrono::high_resolution_clock::now();

  const int32_t numBonds = NUM_SITES - 1;
  std::vector<int64_t> currentBonds(numBonds, 0);

  for (int32_t step = 0; step < NUM_STEPS; ++step) {
    const double currentTime = step * DT;

    HANDLE_CUDM_ERROR(cudensitymatTimePropagationCompute(handle,
                        timeProp,
                        DT,      // timeStepReal
                        0.0,     // timeStepImag
                        currentTime,
                        batchSize,
                        0,       // numParams
                        nullptr, // params
                        stateIn,
                        stateOut,
                        workspaceDescr,
                        0x0));

    // Compare the evolved state's current bond extents against the maximum
    // (buffer) extents to see how the bonds grow / get truncated.
    HANDLE_CUDM_ERROR(cudensitymatStateMPSGetCurrentBondExtents(handle,
                        stateOut, currentBonds.data()));
    if (verbose) {
      std::cout << "Step " << (step + 1) << "/" << NUM_STEPS
                << "  current/maximum bond extents:";
      bool anyTruncated = false;
      for (int32_t b = 0; b < numBonds; ++b) {
        std::cout << " " << currentBonds[b] << "/" << maxBondDims[b];
        if (currentBonds[b] < maxBondDims[b]) anyTruncated = true;
      }
      std::cout << (anyTruncated ? "  (some bonds below maximum)" : "  (all bonds at maximum)") << "\n";
    }

    // Swap input/output for next step (double-buffering)
    std::swap(stateIn, stateOut);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  const auto wallEnd = std::chrono::high_resolution_clock::now();
  const std::chrono::duration<double> elapsed = wallEnd - wallStart;
  if (verbose)
    std::cout << "\nTotal propagation wall time (sec) = " << elapsed.count() << "\n";

  // --- 9. Clean up ---
  if (scratchBuf)
    HANDLE_CUDA_ERROR(cudaFree(scratchBuf));
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));
  HANDLE_CUDM_ERROR(cudensitymatDestroyTimePropagation(timeProp));
  HANDLE_CUDM_ERROR(cudensitymatDestroyTimePropagationApproachKrylovConfig(krylovConfig));
  HANDLE_CUDM_ERROR(cudensitymatDestroyTimePropagationScopeSplitTDVPConfig(tdvpConfig));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(operatorTerm));
  HANDLE_CUDM_ERROR(cudensitymatDestroyMatrixProductOperator(mpoHandle));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(stateOut));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(stateIn));

  for (auto buf : stateOutBufs)
    destroyArrayGPU(buf);
  productMps.destroy();
  mpo.destroy();

  if (verbose)
    std::cout << "Destroyed all resources\n";
}


int main(int argc, char ** argv)
{
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose)
    std::cout << "Set active device\n";

  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created library handle\n";

  exampleWorkflow(handle);

  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());
  return 0;
}

Code example (serial execution of MPS-MPO operator action via ALS)

The following code example illustrates how to use the cuDensityMat library to compute the action of a transverse-field Ising Hamiltonian encoded as an MPO on a pure MPS state, fitted via the split-scope variational ALS path (CUDENSITYMAT_FITTING_SCOPE_SPLIT + CUDENSITYMAT_FITTING_APPROACH_LINSOLVE). The example builds the Hamiltonian MPO, creates input and output MPS states, initializes the input MPS to a Néel state, zero-initializes the output MPS via cudensitymatStateInitializeZero (which puts the ALS fitter into the assignment-mode starting-point path), configures the ALS sub-config (max_sweeps, tolerance), prepares the action and computes one step, then reports the output MPS bond extents. The full sample code can be found in the NVIDIA/cuQuantum repository (main serial MPS-MPO action code as well as the utility code).

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

// MPS-MPO operator-action example.
//
// Demonstrates applying a 1-D transverse-field Ising MPO to a pure MPS state
// via the split-scope variational ALS fitting path:
//
//   stateOut += alpha * H * stateIn
//
// All quantum-state inputs are matrix product states (MPS); the Hamiltonian
// is encoded as a matrix product operator (MPO). The output MPS is fitted
// via 1-site alternating least squares (ALS), so the user controls the
// truncation budget by choosing the output MPS bond dimensions.
//
// Workflow:
//  1. Build Hamiltonian as MPO (nearest-neighbor ZZ + transverse X field)
//  2. Build single-term Operator wrapping the MPO
//  3. Create input and output MPS states (cudensitymatCreateStateMPS)
//  4. Initialize input MPS to a Neel state |0101...> and zero-initialise
//     stateOut via cudensitymatStateInitializeZero
//  5. Create OperatorAction (FITTING_SCOPE_SPLIT + FITTING_APPROACH_LINSOLVE)
//  6. (Optional) Configure ALS sub-config (num_sites, max_sweeps, tolerance)
//  7. Prepare action and allocate workspace
//  8. Compute one operator action at t = 0 (stateOut <- alpha * H * stateIn)
//  9. Report the output MPS bond extents
// 10. Clean up all resources
//
// Assignment semantics. The split-scope operator action fits stateOut to
// alpha * H * stateIn accumulated on top of the prior contents of stateOut.
// Zero-initialising stateOut with cudensitymatStateInitializeZero before the
// Compute call therefore makes the result a plain assignment
// stateOut = alpha * H * stateIn (up to FP64 fitting accuracy). The example
// reports the allocated output MPS component shapes; a schedule-agnostic norm
// check should use the public MPS norm API once it is available.

#include <cudensitymat.h>
#include "helpers.h"

#include <algorithm>
#include <cmath>
#include <complex>
#include <vector>
#include <numeric>
#include <iostream>
#include <cassert>


using Complex = std::complex<double>;
constexpr cudaDataType_t DATA_TYPE = CUDA_C_64F;

constexpr bool verbose = true;

// --- Simulation parameters ---
// 8 qubits, TFIM (J = 1.0, h = 0.5). Output MPS capped at bond dim 16 which
// saturates the natural max bond ([2,4,8,16,8,4,2]) for an 8-site qubit chain
// and is comfortably above chi_MPO * chi_MPS_in = 3 * 1 = 3.
constexpr int32_t  NUM_SITES        = 8;
constexpr int64_t  PHYS_DIM         = 2;
constexpr int64_t  MAX_OUT_BOND_DIM = 16;
constexpr int64_t  MPO_BOND_DIM     = 3;
constexpr double   J_COUPLING       = 1.0;
constexpr double   H_FIELD          = 0.5;


// ============================================================================
// Transverse-field Ising MPO builder (bond dim 3)
// ============================================================================
//
// H = -J * sum_{i} Z_i Z_{i+1}  +  h * sum_{i} X_i
//
// The bulk MPO matrix (indexed by bond dimensions aL, aR) is:
//
//        |  I      0     0  |
//   W =  |  Z      0     0  |
//        | h*X   -J*Z    I  |
//
// Library mode ordering (column-major):
//   Left boundary (site 0):     [phys_ket(d), right_bond(bR), phys_bra(d)]   (bL omitted, treated as 1)
//   Interior sites:             [left_bond(bL), phys_ket(d), right_bond(bR), phys_bra(d)]
//   Right boundary (site N-1):  [left_bond(bL), phys_ket(d), phys_bra(d)]    (bR omitted, treated as 1)

struct IsingMPO {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build() {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0}, mone{-1.0, 0.0};
    const Complex jc{-J_COUPLING, 0.0};
    const Complex hc{H_FIELD, 0.0};

    auto I_mat = [&](int r, int c) -> Complex { return (r == c) ? one : zero; };
    auto X_mat = [&](int r, int c) -> Complex { return (r != c) ? one : zero; };
    auto Z_mat = [&](int r, int c) -> Complex { return (r == c) ? ((r == 0) ? one : mone) : zero; };

    hostTensors.resize(NUM_SITES);
    gpuPtrs.resize(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL = (site == 0) ? 1 : MPO_BOND_DIM;
      const int64_t bR = (site == NUM_SITES - 1) ? 1 : MPO_BOND_DIM;
      const int64_t vol = bL * PHYS_DIM * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      auto idx = [&](int64_t aL, int64_t ket, int64_t aR, int64_t bra) -> int64_t {
        return aL + bL * (ket + PHYS_DIM * (aR + bR * bra));
      };

      if (NUM_SITES == 1) {
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket)
            hostTensors[site][idx(0, ket, 0, bra)] = hc * X_mat(bra, ket);
      } else if (site == 0) {
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = hc * X_mat(bra, ket);
            hostTensors[site][idx(0, ket, 1, bra)] = jc * Z_mat(bra, ket);
            hostTensors[site][idx(0, ket, 2, bra)] = I_mat(bra, ket);
          }
      } else if (site == NUM_SITES - 1) {
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = I_mat(bra, ket);
            hostTensors[site][idx(1, ket, 0, bra)] = Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 0, bra)] = hc * X_mat(bra, ket);
          }
      } else {
        for (int bra = 0; bra < PHYS_DIM; ++bra)
          for (int ket = 0; ket < PHYS_DIM; ++ket) {
            hostTensors[site][idx(0, ket, 0, bra)] = I_mat(bra, ket);
            hostTensors[site][idx(1, ket, 0, bra)] = Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 0, bra)] = hc * X_mat(bra, ket);
            hostTensors[site][idx(2, ket, 1, bra)] = jc * Z_mat(bra, ket);
            hostTensors[site][idx(2, ket, 2, bra)] = I_mat(bra, ket);
          }
      }

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Build a Neel-state MPS  |0,1,0,1,...>  with given bond dimensions
// ============================================================================
// Each MPS tensor A[site] has shape (bondL, phys, bondR) in column-major.
// For a product state, only the (0, site%2, 0) slice is nonzero; remaining
// bond indices are zero-padded.

struct NeelMPS {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build(const std::vector<int64_t>& bondDims) {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0};

    hostTensors.resize(NUM_SITES);
    gpuPtrs.resize(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL = (site == 0) ? 1 : bondDims[site - 1];
      const int64_t bR = (site == NUM_SITES - 1) ? 1 : bondDims[site];
      const int64_t vol = bL * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      auto idx = [&](int64_t aL, int64_t sigma, int64_t aR) -> int64_t {
        return aL + bL * (sigma + PHYS_DIM * aR);
      };

      const int64_t neelSpin = site % 2;
      hostTensors[site][idx(0, neelSpin, 0)] = one;

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Helper: cap MPS bond dimensions by the exact Hilbert-space dimension on
// either side. For an OBC qubit chain, the max admissible bond at site i is
// min(2^(i+1), 2^(N-i-1)).
// ============================================================================
static std::vector<int64_t> makeMaxBondDims(int64_t maxBondCap)
{
  std::vector<int64_t> dims(NUM_SITES - 1);
  for (int32_t i = 0; i < NUM_SITES - 1; ++i) {
    int64_t leftDim = 1;
    for (int32_t j = 0; j <= i; ++j) leftDim *= PHYS_DIM;
    int64_t rightDim = 1;
    for (int32_t j = i + 1; j < NUM_SITES; ++j) rightDim *= PHYS_DIM;
    dims[i] = std::min({maxBondCap, leftDim, rightDim});
  }
  return dims;
}


// ============================================================================
// Example workflow
// ============================================================================

void exampleWorkflow(cudensitymatHandle_t handle)
{
  // --- 1. Build the transverse-field Ising MPO ---
  IsingMPO mpo;
  mpo.build();
  if (verbose)
    std::cout << "Built transverse-field Ising MPO (bond dim " << MPO_BOND_DIM << ")\n";

  const std::vector<int64_t> spaceShape(NUM_SITES, PHYS_DIM);
  std::vector<int64_t> mpoBondDims(NUM_SITES - 1, MPO_BOND_DIM);

  cudensitymatMatrixProductOperator_t mpoHandle;
  std::vector<cudensitymatWrappedTensorCallback_t> mpoCallbacks(NUM_SITES, cudensitymatTensorCallbackNone);
  std::vector<cudensitymatWrappedTensorGradientCallback_t> mpoGradCallbacks(NUM_SITES, cudensitymatTensorGradientCallbackNone);
  HANDLE_CUDM_ERROR(cudensitymatCreateMatrixProductOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpoBondDims.data(),
                      DATA_TYPE,
                      mpo.gpuPtrs.data(),
                      mpoCallbacks.data(),
                      mpoGradCallbacks.data(),
                      &mpoHandle));
  if (verbose)
    std::cout << "Created MPO handle\n";

  // --- 2. Wrap the MPO in a single-term Operator ---
  // Split-scope operator actions require the operator to wrap exactly one
  // term containing a single MPO product.
  cudensitymatOperatorTerm_t operatorTerm;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      &operatorTerm));

  std::vector<int32_t> modesActedOn(NUM_SITES);
  std::iota(modesActedOn.begin(), modesActedOn.end(), 0);
  std::vector<int32_t> modeDuality(NUM_SITES, 0);
  std::vector<int32_t> mpoConjugation = {0};

  HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendMPOProduct(handle,
                      operatorTerm,
                      1,
                      &mpoHandle,
                      mpoConjugation.data(),
                      modesActedOn.data(),
                      modeDuality.data(),
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));

  cudensitymatOperator_t hamiltonian;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      &hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                      hamiltonian,
                      operatorTerm,
                      0,
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));
  if (verbose)
    std::cout << "Constructed Hamiltonian operator from MPO\n";

  // --- 3. Create input and output MPS states ---
  const int64_t batchSize = 1;

  // Both stateIn and stateOut are sized to the maximum admissible bond
  // dimension for the given physical dimensions. The Neel input is
  // mathematically a bond-1 state; we zero-pad it into the maximum-bond
  // tensors. The output is sized to MAX_OUT_BOND_DIM = 16, which is well
  // above the bond dimension needed to represent H * stateIn exactly
  // (chi_MPO * chi_MPS_in = 3).
  std::vector<int64_t> mpsBondDims = makeMaxBondDims(MAX_OUT_BOND_DIM);

  cudensitymatState_t stateIn, stateOut;
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpsBondDims.data(),
                      DATA_TYPE,
                      batchSize,
                      &stateIn));
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpsBondDims.data(),
                      DATA_TYPE,
                      batchSize,
                      &stateOut));

  int32_t numComponents = 0;
  HANDLE_CUDM_ERROR(cudensitymatStateGetNumComponents(handle, stateIn, &numComponents));
  assert(numComponents == NUM_SITES);

  std::vector<std::size_t> componentSizes(numComponents);
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                      stateIn, numComponents, componentSizes.data()));

  if (verbose) {
    std::cout << "MPS state has " << numComponents << " components, sizes (bytes):";
    for (auto s : componentSizes) std::cout << " " << s;
    std::cout << "\n";
    std::cout << "MPS bond dims:";
    for (auto b : mpsBondDims) std::cout << " " << b;
    std::cout << "\n";
  }

  // --- 4. Allocate GPU storage; initialise stateIn to Neel, stateOut to 0 ---
  //
  // The per-site GPU buffers for stateIn must be allocated AND populated with
  // their initial data before attaching them to the state, so we upload the
  // Neel data on the host first via NeelMPS::build. For stateOut a plain
  // cudaMalloc is enough; cudensitymatStateInitializeZero (called below)
  // writes zeros into the attached buffers so the operator action behaves as
  // an assignment.
  NeelMPS neelMps;
  neelMps.build(mpsBondDims);

  std::vector<void *> stateOutBuffers(numComponents, nullptr);
  for (int32_t c = 0; c < numComponents; ++c) {
    HANDLE_CUDA_ERROR(cudaMalloc(&stateOutBuffers[c], componentSizes[c]));
  }

  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      stateIn, numComponents,
                      neelMps.gpuPtrs.data(), componentSizes.data()));
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                      stateOut, numComponents,
                      stateOutBuffers.data(), componentSizes.data()));

  // Zero-initialise stateOut so that the subsequent operator action assigns
  // alpha * H * stateIn into it (rather than accumulating onto prior data).
  HANDLE_CUDM_ERROR(cudensitymatStateInitializeZero(handle, stateOut, /*stream=*/0));

  if (verbose)
    std::cout << "Initialized MPS states (input = Neel |0101...>, "
                 "output = zero MPS via cudensitymatStateInitializeZero)\n";

  // --- 5. Create the OperatorAction (split scope + LinSolve approach) ---
  cudensitymatOperator_t operators[] = {hamiltonian};
  cudensitymatOperatorAction_t operatorAction;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorAction(handle,
                      /*numOperators=*/1,
                      operators,
                      CUDENSITYMAT_FITTING_SCOPE_SPLIT,
                      CUDENSITYMAT_FITTING_APPROACH_LINSOLVE,
                      &operatorAction));
  if (verbose)
    std::cout << "Created OperatorAction (scope = SPLIT, approach = LINSOLVE)\n";

  // --- 6. (Optional) Configure ALS sub-config ---
  // The defaults (num_sites = 1, max_sweeps = 20, tolerance = 1e-10) already
  // suffice for a TFIM demo. We pass an explicit configuration here only to
  // show the wiring; tightening max_sweeps below 20 may degrade convergence.
  cudensitymatStateFittingScopeSplitALSConfig_t alsConfig{nullptr};
  HANDLE_CUDM_ERROR(cudensitymatCreateStateFittingScopeSplitALSConfig(handle, &alsConfig));
  {
    const int32_t numSites = 1;
    HANDLE_CUDM_ERROR(cudensitymatStateFittingScopeSplitALSConfigSetAttribute(handle,
                        alsConfig,
                        CUDENSITYMAT_FITTING_SPLIT_SCOPE_ALS_NUM_SITES,
                        &numSites, sizeof(numSites)));
    const int32_t maxSweeps = 20;
    HANDLE_CUDM_ERROR(cudensitymatStateFittingScopeSplitALSConfigSetAttribute(handle,
                        alsConfig,
                        CUDENSITYMAT_FITTING_SPLIT_SCOPE_ALS_MAX_SWEEPS,
                        &maxSweeps, sizeof(maxSweeps)));
    const double tolerance = 1e-10;
    HANDLE_CUDM_ERROR(cudensitymatStateFittingScopeSplitALSConfigSetAttribute(handle,
                        alsConfig,
                        CUDENSITYMAT_FITTING_SPLIT_SCOPE_ALS_TOLERANCE,
                        &tolerance, sizeof(tolerance)));
  }
  HANDLE_CUDM_ERROR(cudensitymatOperatorActionConfigure(handle,
                      operatorAction,
                      CUDENSITYMAT_FITTING_SPLIT_SCOPE_ALS_CONFIG,
                      &alsConfig, sizeof(alsConfig)));
  // The sub-config is deep-copied at Configure; it is safe to destroy now.
  HANDLE_CUDM_ERROR(cudensitymatDestroyStateFittingScopeSplitALSConfig(alsConfig));
  if (verbose)
    std::cout << "Configured ALS (num_sites = 1, max_sweeps = 20, tolerance = 1e-10)\n";

  // --- 7. Prepare action and allocate workspace ---
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95);
  if (verbose)
    std::cout << "Available workspace memory (bytes) = " << freeMem << "\n";

  HANDLE_CUDM_ERROR(cudensitymatOperatorActionPrepare(handle,
                      operatorAction,
                      &stateIn,
                      stateOut,
                      CUDENSITYMAT_COMPUTE_64F,
                      freeMem,
                      workspaceDescr,
                      /*stream=*/0));
  if (verbose)
    std::cout << "Prepared OperatorAction\n";

  std::size_t scratchSize = 0;
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &scratchSize));
  void * scratchBuf = nullptr;
  if (scratchSize > 0) {
    HANDLE_CUDA_ERROR(cudaMalloc(&scratchBuf, scratchSize));
    HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                        workspaceDescr,
                        CUDENSITYMAT_MEMSPACE_DEVICE,
                        CUDENSITYMAT_WORKSPACE_SCRATCH,
                        scratchBuf, scratchSize));
  }
  if (verbose)
    std::cout << "Scratch workspace (bytes) = " << scratchSize << "\n";

  // --- 8. Compute one operator action (stateOut <- alpha * H * stateIn) ---
  // Because stateOut was zero-initialised above, the split-scope fit assigns
  // alpha * H * stateIn into stateOut rather than accumulating onto prior
  // contents, so the result equals H |Neel> for this fixture (alpha = 1).
  if (verbose)
    std::cout << "\nApplying H to Neel state via OperatorActionCompute (t = 0)\n";

  HANDLE_CUDM_ERROR(cudensitymatOperatorActionCompute(handle,
                      operatorAction,
                      /*time=*/0.0,
                      batchSize,
                      /*numParams=*/0,
                      /*params=*/nullptr,
                      &stateIn,
                      stateOut,
                      workspaceDescr,
                      /*stream=*/0));
  HANDLE_CUDA_ERROR(cudaStreamSynchronize(0));
  if (verbose)
    std::cout << "OperatorActionCompute completed\n";

  // --- 9. Report the output MPS bond extents ---
  // Bond extents reported here are the as-allocated ones; the ALS fit cannot
  // exceed them but may use a subset internally.
  if (verbose) {
    std::cout << "\nOutput MPS component shapes (column-major modes):\n";
    for (int32_t c = 0; c < numComponents; ++c) {
      int32_t globalId = 0;
      int32_t numModes = 0;
      int32_t batchModeLocation = 0;
      HANDLE_CUDM_ERROR(cudensitymatStateGetComponentNumModes(handle,
                          stateOut, c, &globalId, &numModes, &batchModeLocation));
      std::vector<int64_t> extents(numModes), offsets(numModes);
      HANDLE_CUDM_ERROR(cudensitymatStateGetComponentInfo(handle,
                          stateOut, c, &globalId,
                          &numModes,
                          extents.data(),
                          offsets.data()));
      std::cout << "  site " << c << " : [";
      for (int32_t m = 0; m < numModes; ++m) {
        std::cout << extents[m] << (m + 1 < numModes ? ", " : "");
      }
      std::cout << "]\n";
    }
  }

  // --- 10. Clean up ---
  if (scratchBuf)
    HANDLE_CUDA_ERROR(cudaFree(scratchBuf));
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorAction(operatorAction));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(operatorTerm));
  HANDLE_CUDM_ERROR(cudensitymatDestroyMatrixProductOperator(mpoHandle));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(stateOut));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(stateIn));

  for (auto * buf : stateOutBuffers) {
    if (buf) HANDLE_CUDA_ERROR(cudaFree(buf));
  }
  neelMps.destroy();
  mpo.destroy();

  if (verbose)
    std::cout << "\nDestroyed all resources\n";
}


int main(int argc, char ** argv)
{
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose)
    std::cout << "Set active device\n";

  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created library handle\n";

  exampleWorkflow(handle);

  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());
  return 0;
}

Code example (serial execution of operator eigenspectrum computation)

The following code example illustrates how to use the cuDensityMat library for computing the extreme eigenspectrum of a given operator. The full sample code can be found in the NVIDIA/cuQuantum repository (main serial eigenspectrum code and operator definition as well as the utility code).

First, similarly to the above examples, we define a transverse-field Ising Hamiltonian with fused ZZ terms the eigenspectrum of which we want to compute. Specifically, in this case, we want to compute a number of the smallest real eigenvalues and their corresponding eigenvectors (pure quantum states). For simplicity, we made our operator time-independent (static).

/* Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#pragma once

#include <cudensitymat.h> // cuDensityMat library header
#include "helpers.h"      // GPU helper functions

#include <cmath>
#include <complex>
#include <vector>
#include <iostream>
#include <cassert>


/* DESCRIPTION:
   Transverse-field Ising Hamiltonian operator with ordered and fused ZZ terms:
    H = sum_{i} {h_i * X_i}         // transverse field sum of X_i operators with static h_i coefficients 
      + sum_{i < j} {g_ij * ZZ_ij}  // sum of the fused ordered {Z_i * Z_j} terms with static g_ij coefficients
*/

/** Define the numerical type and data type for the GPU computations (same) */
using NumericalType = std::complex<double>;      // do not change
constexpr cudaDataType_t dataType = CUDA_C_64F;  // do not change


/** Convenience class which encapsulates a user-defined Liouvillian operator (system Hamiltonian + dissipation terms):
 *  - Constructor constructs the desired Liouvillian operator (`cudensitymatOperator_t`)
 *  - Method `get()` returns a reference to the constructed Liouvillian operator
 *  - Destructor releases all resources used by the Liouvillian operator
 */
class UserDefinedLiouvillian final
{
private:
  // Data members
  cudensitymatHandle_t handle;             // library context handle
  int64_t stateBatchSize;                  // quantum state batch size
  const std::vector<int64_t> spaceShape;   // Hilbert space shape (extents of the modes of the composite Hilbert space)
  void * spinXelems {nullptr};             // elements of the X spin operator in GPU RAM (F-order storage)
  void * spinZZelems {nullptr};            // elements of the fused ZZ two-spin operator in GPU RAM (F-order storage)
  cudensitymatElementaryOperator_t spinX;  // X spin operator (elementary tensor operator)
  cudensitymatElementaryOperator_t spinZZ; // fused ZZ two-spin operator (elementary tensor operator)
  cudensitymatOperatorTerm_t oneBodyTerm;  // operator term: H1 = sum_{i} {h_i * X_i} (one-body term)
  cudensitymatOperatorTerm_t twoBodyTerm;  // operator term: H2 = sum_{i < j} {g_ij * ZZ_ij} (two-body term)
  cudensitymatOperator_t liouvillian;      // full operator: H = H1 + H2

public:

  // Constructor constructs a user-defined Liouvillian operator
  UserDefinedLiouvillian(cudensitymatHandle_t contextHandle,             // library context handle
                         const std::vector<int64_t> & hilbertSpaceShape, // Hilbert space shape
                         int64_t batchSize):                             // batch size for the quantum state
    handle(contextHandle), stateBatchSize(batchSize), spaceShape(hilbertSpaceShape)
  {
    // Define the necessary operator tensors in GPU memory (F-order storage!)
    spinXelems = createInitializeArrayGPU<NumericalType>(  // X[i0; j0]
                  {{0.0, 0.0}, {1.0, 0.0},   // 1st column of matrix X
                   {1.0, 0.0}, {0.0, 0.0}}); // 2nd column of matrix X

    spinZZelems = createInitializeArrayGPU<NumericalType>(  // ZZ[i0, i1; j0, j1] := Z[i0; j0] * Z[i1; j1]
                    {{1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {0.0, 0.0},   // 1st column of matrix ZZ
                     {0.0, 0.0}, {-1.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},   // 2nd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {-1.0, 0.0}, {0.0, 0.0},   // 3rd column of matrix ZZ
                     {0.0, 0.0}, {0.0, 0.0},  {0.0, 0.0},  {1.0, 0.0}}); // 4th column of matrix ZZ

    // Construct the necessary Elementary Tensor Operators
    //  X_i operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        1,                                   // one-body operator
                        std::vector<int64_t>({2}).data(),    // acts in tensor space of shape {2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinXelems,                          // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinX));                            // the created elementary tensor operator
    //  ZZ_ij = Z_i * Z_j fused operator
    HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(handle,
                        2,                                   // two-body operator
                        std::vector<int64_t>({2,2}).data(),  // acts in tensor space of shape {2,2}
                        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, // dense tensor storage
                        0,                                   // 0 for dense tensors
                        nullptr,                             // nullptr for dense tensors
                        dataType,                            // data type
                        spinZZelems,                         // tensor elements in GPU memory
                        cudensitymatTensorCallbackNone,      // no tensor callback function (tensor is not time-dependent)
                        cudensitymatTensorGradientCallbackNone, // no tensor gradient callback function
                        &spinZZ));                           // the created elementary tensor operator

    // Construct the necessary Operator Terms from tensor products of Elementary Tensor Operators
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &oneBodyTerm));                      // the created empty operator term
    //  Define the operator term: H1 = sum_{i} {h_i * X_i}
    for (int32_t i = 0; i < spaceShape.size(); ++i) {
      const double h_i = 1.0 / static_cast<double>(i+1); // assign some value to the time-independent h_i coefficient
      HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                          oneBodyTerm,
                          1,                                                             // number of elementary tensor operators in the product
                          std::vector<cudensitymatElementaryOperator_t>({spinX}).data(), // elementary tensor operators forming the product
                          std::vector<int32_t>({i}).data(),                              // space modes acted on by the operator product
                          std::vector<int32_t>({0}).data(),                              // space mode action duality (0: from the left; 1: from the right)
                          make_cuDoubleComplex(h_i, 0.0),                                // h_i constant coefficient: Always 64-bit-precision complex number
                          cudensitymatScalarCallbackNone,                                // no time-dependent coefficient associated with this operator product
                          cudensitymatScalarGradientCallbackNone));                      // no coefficient gradient associated with this operator product
    }
    //  Create an empty operator term
    HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                        spaceShape.size(),                   // Hilbert space rank (number of modes)
                        spaceShape.data(),                   // Hilbert space shape (mode extents)
                        &twoBodyTerm));                      // the created empty operator term
    //  Define the operator term: H2 = sum_{i < j} {g_ij * ZZ_ij}
    for (int32_t i = 0; i < spaceShape.size() - 1; ++i) {
      for (int32_t j = (i + 1); j < spaceShape.size(); ++j) {
        const double g_ij = -1.0 / static_cast<double>(i + j + 1); // assign some value to the time-independent g_ij coefficient
        HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(handle,
                            twoBodyTerm,
                            1,                                                              // number of elementary tensor operators in the product
                            std::vector<cudensitymatElementaryOperator_t>({spinZZ}).data(), // elementary tensor operators forming the product
                            std::vector<int32_t>({i, j}).data(),                            // space modes acted on by the operator product
                            std::vector<int32_t>({0, 0}).data(),                            // space mode action duality (0: from the left; 1: from the right)
                            make_cuDoubleComplex(g_ij, 0.0),                                // g_ij constant coefficient: Always 64-bit-precision complex number
                            cudensitymatScalarCallbackNone,                                 // no time-dependent coefficient associated with this operator product
                            cudensitymatScalarGradientCallbackNone));                       // no coefficient gradient associated with this operator product
      }
    }

    // Construct the full Liouvillian operator as a sum of the operator terms
    //  Create an empty operator
    HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                        spaceShape.size(),                // Hilbert space rank (number of modes)
                        spaceShape.data(),                // Hilbert space shape (modes extents)
                        &liouvillian));                   // the created empty operator
    //  Append an operator term to the operator
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        oneBodyTerm,                      // appended operator term
                        0,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(1.0, 0.0),   // constant coefficient associated with the operator term as a whole
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with the operator term as a whole
    //  Append an operator term to the operator (super-operator)
    HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                        liouvillian,
                        twoBodyTerm,                      // appended operator term
                        0,                                // operator term action duality as a whole (0: acting from the left; 1: acting from the right)
                        make_cuDoubleComplex(1.0, 0.0),   // constant coefficient associated with the operator term as a whole
                        cudensitymatScalarCallbackNone,   // no time-dependent coefficient associated with the operator term as a whole
                        cudensitymatScalarGradientCallbackNone)); // no coefficient gradient associated with this operator term as a whole
  }

  // Destructor destructs the user-defined Liouvillian operator
  ~UserDefinedLiouvillian()
  {
    // Destroy the Liouvillian operator
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(liouvillian));

    // Destroy operator terms
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(twoBodyTerm));
    HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(oneBodyTerm));

    // Destroy elementary tensor operators
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinZZ));
    HANDLE_CUDM_ERROR(cudensitymatDestroyElementaryOperator(spinX));

    // Destroy operator tensors
    destroyArrayGPU(spinZZelems);
    destroyArrayGPU(spinXelems);
  }

  // Disable copy constructor/assignment (GPU resources are private, no deep copy)
  UserDefinedLiouvillian(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian & operator=(const UserDefinedLiouvillian &) = delete;
  UserDefinedLiouvillian(UserDefinedLiouvillian &&) = delete;
  UserDefinedLiouvillian & operator=(UserDefinedLiouvillian &&) = delete;

  /** Returns the number of externally provided Hamiltonian parameters. */
  int32_t getNumParameters() const
  {
    return 0; // no free parameters
  }

  /** Get access to the constructed Liouvillian operator. */
  cudensitymatOperator_t & get()
  {
    return liouvillian;
  }

};

Once the operator has been defined, we can follow standard steps to create necessary quantum states which will store the eigenstates of the defined operator. Then we can prepare the operator eigenspectrum computation, and, finally, compute the eigenspectrum. Note that the leading subset of the quantum states passed to the eigenspectrum compute call to store the computed eigenstates will also be used as the initial guesses for the first Krylov subspace block (if the block size is smaller than the number of requested eigenstates, only the leading quantum states will be used).

/* Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <cudensitymat.h>  // cuDensityMat library header
#include "helpers.h"       // helper functions


// Transverse Ising Hamiltonian with double summation ordering and spin-operator fusion
#include "transverse_ising_full_fused.h"  // user-defined Liouvillian operator example

#include <cmath>
#include <complex>
#include <vector>
#include <chrono>
#include <iostream>
#include <cassert>


// Logging verbosity
bool verbose = true;


// Example workflow
void exampleWorkflow(cudensitymatHandle_t handle)
{
  // Define the composite Hilbert space shape and
  // quantum state batch size (number of individual quantum states in a batched simulation)
  const std::vector<int64_t> spaceShape({2,2,2,2,2,2,2,2,2,2}); // dimensions of quantum degrees of freedom
  const int64_t batchSize = 1;        // number of quantum states per batch (currently only 1 state per batch)
  const int32_t numEigenStates = 4;   // number of eigenstates to compute

  if (verbose) {
    std::cout << "Hilbert space rank = " << spaceShape.size() << "; Shape = (";
    for (const auto & dimsn: spaceShape)
      std::cout << dimsn << ",";
    std::cout << ")" << std::endl;
    std::cout << "Quantum state batch size = " << batchSize << std::endl;
  }

  // Construct a user-defined Liouvillian operator using a convenience C++ class
  UserDefinedLiouvillian liouvillian(handle, spaceShape, batchSize);
  if (verbose)
    std::cout << "Constructed the Liouvillian operator\n";

  // Create quantum states to store the eigenstates
  std::size_t stateVolume {0};
  std::vector<cudensitymatState_t> eigenStates(numEigenStates);
  std::vector<void *> eigenStatesElems(numEigenStates);
  for (int32_t id = 0; id < numEigenStates; ++id) {

    // Declare the quantum state
    HANDLE_CUDM_ERROR(cudensitymatCreateState(handle,
                        CUDENSITYMAT_STATE_PURITY_PURE,  // pure (state vector)
                        spaceShape.size(),
                        spaceShape.data(),
                        batchSize,
                        dataType,
                        &eigenStates[id]));

    // Query the size of the quantum state storage
    std::size_t storageSize {0}; // only one storage component (tensor) is needed (no tensor factorization)
    HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle,
                        eigenStates[id],
                        1,               // only one storage component (tensor)
                        &storageSize));  // storage size in bytes
    stateVolume = storageSize / sizeof(NumericalType);  // quantum state tensor volume (number of elements)
    if (verbose)
      std::cout << "Quantum state storage size (bytes) = " << storageSize << std::endl;

    // Prepare some initial value for the quantum state
    std::vector<NumericalType> stateValue(stateVolume);
    if constexpr (std::is_same_v<NumericalType, double>) {
      for (std::size_t i = 0; i < stateVolume; ++i) {
        stateValue[i] = 1.0 / double(id*5 + i+1); // just some value
      }
    } else if constexpr (std::is_same_v<NumericalType, std::complex<double>>) {
      for (std::size_t i = 0; i < stateVolume; ++i) {
        stateValue[i] = NumericalType{1.0 / double(id*5 + i+1), -1.0 / double(id*3 + i+2)}; // just some value
      }
    } else {
      std::cerr << "Error: Unsupported data type!\n";
      std::exit(1);
    }
    // Allocate initialized GPU storage for the quantum state with prepared values
    eigenStatesElems[id] = createInitializeArrayGPU(stateValue);
    if (verbose)
      std::cout << "Allocated quantum state storage and initialized it to some value\n";

    // Attach initialized GPU storage to the quantum state
    HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle,
                        eigenStates[id],
                        1,                                                 // only one storage component (tensor)
                        std::vector<void*>({eigenStatesElems[id]}).data(), // pointer to the GPU storage for the quantum state
                        std::vector<std::size_t>({storageSize}).data()));  // size of the GPU storage for the quantum state
    if (verbose)
      std::cout << "Constructed quantum state\n";
  }

  // Allocate storage for the eigenvalues and convergence tolerances
  void * eigenvalues = createArrayGPU<NumericalType>(numEigenStates * batchSize);
  std::vector<double> tolerances(numEigenStates * batchSize, 1e-6);

  // Declare a workspace descriptor
  cudensitymatWorkspaceDescriptor_t workspaceDescr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  // Query free GPU memory
  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.95); // take 95% of the free memory for the workspace buffer
  if (verbose)
    std::cout << "Max workspace buffer size (bytes) = " << freeMem << std::endl;

  // Create the operator eigenspectrum computation object
  cudensitymatOperatorSpectrum_t spectrum;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorSpectrum(handle,
                      liouvillian.get(),
                      1,  // Hermitian operator
                      CUDENSITYMAT_OPERATOR_SPECTRUM_SMALLEST_REAL,
                      &spectrum));

  // Prepare the operator eigenspectrum computation (needs to be done only once)
  auto startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorSpectrumPrepare(handle,
                      spectrum,
                      numEigenStates,
                      eigenStates[0],
                      CUDENSITYMAT_COMPUTE_64F,
                      freeMem,
                      workspaceDescr,
                      0x0));
  auto finishTime = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator eigenspectrum preparation time (sec) = " << timeSec.count() << std::endl;

  // Query the required workspace buffer size (bytes)
  std::size_t requiredBufferSize {0};
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &requiredBufferSize));
  if (verbose)
    std::cout << "Required workspace buffer size (bytes) = " << requiredBufferSize << std::endl;

  // Allocate GPU storage for the workspace buffer
  const std::size_t bufferVolume = requiredBufferSize / sizeof(NumericalType);
  auto * workspaceBuffer = createArrayGPU<NumericalType>(bufferVolume);
  if (verbose)
    std::cout << "Allocated workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Attach the workspace buffer to the workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle,
                      workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      workspaceBuffer,
                      requiredBufferSize));
  if (verbose)
    std::cout << "Attached workspace buffer of size (bytes) = " << requiredBufferSize << std::endl;

  // Compute the operator eigenspectrum
  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  startTime = std::chrono::high_resolution_clock::now();
  HANDLE_CUDM_ERROR(cudensitymatOperatorSpectrumCompute(handle,
                      spectrum,
                      0.0,
                      batchSize,
                      0,
                      nullptr,
                      numEigenStates,
                      eigenStates.data(),
                      eigenvalues,
                      tolerances.data(),
                      workspaceDescr,
                      0x0));
  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  finishTime = std::chrono::high_resolution_clock::now();
  timeSec = finishTime - startTime;
  if (verbose)
    std::cout << "Operator eigenspectrum computation time (sec) = " << timeSec.count() << std::endl;

  // Print the eigenvalues
  if (verbose) {
    std::cout << "Eigenvalues:\n";
    printArrayGPU<NumericalType>(eigenvalues, numEigenStates);
  }

  // Print the residual norms
  if (verbose) {
    std::cout << "Residual norms:\n";
    printArrayCPU<double>(tolerances.data(), numEigenStates);
  }

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  // Destroy workspace descriptor
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));

  // Destroy workspace buffer storage
  destroyArrayGPU(workspaceBuffer);

  // Destroy operator eigenspectrum computation object
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorSpectrum(spectrum));

  // Destroy eigenvalues storage
  destroyArrayGPU(eigenvalues);

  // Destroy quantum states
  for (int32_t id = 0; id < numEigenStates; ++id)
    HANDLE_CUDM_ERROR(cudensitymatDestroyState(eigenStates[id]));

  // Destroy quantum state storage
  for (int32_t id = 0; id < numEigenStates; ++id)
    destroyArrayGPU(eigenStatesElems[id]);

  if (verbose)
    std::cout << "Destroyed resources\n" << std::flush;
}


int main(int argc, char ** argv)
{
  // Assign a GPU to the process
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose)
    std::cout << "Set active device\n";

  // Create a library handle
  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose)
    std::cout << "Created a library handle\n";

  // Run the example
  exampleWorkflow(handle);

  // Destroy the library handle
  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose)
    std::cout << "Destroyed the library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());

  // Done
  return 0;
}

Code example (serial execution of 1-site MPS-DMRG ground-state eigensolver)

The following code example illustrates how to use the cuDensityMat library to compute the ground state (smallest-real eigenpair) of a Heisenberg Hamiltonian encoded as an MPO, acting on a pure MPS state, using the split-scope 1-site DMRG method (CUDENSITYMAT_EIGEN_SCOPE_SPLIT + CUDENSITYMAT_EIGEN_APPROACH_KRYLOV + CUDENSITYMAT_EIGEN_SPECTRUM_SMALLEST_REAL) at the default 1-site update granularity (CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_NUM_SITES of 1). The example builds the Hamiltonian MPO, creates an initial MPS state, configures the DMRG and Krylov sub-configs, prepares and computes the ground state, then reports the energy and residual. The full sample code can be found in the NVIDIA/cuQuantum repository (main serial 1-site MPS-DMRG code as well as the utility code).

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

// DMRG eigensolver example.
//
// Computes the ground state of the 1-D spin-1/2 antiferromagnetic Heisenberg
// chain  H = J * sum_i ( Sx_i Sx_{i+1} + Sy_i Sy_{i+1} + Sz_i Sz_{i+1} )  with
// the 1-site DMRG solver (SCOPE_SPLIT + APPROACH_KRYLOV + SMALLEST_REAL) on an
// MPS state.  The exact N=8 open-boundary energy (J=1) is E0 = -3.374932598688;
// the example checks the DMRG result against it.
//
// Workflow:
//  1. Build the Heisenberg MPO (bond dim 5)
//  2. Create an initial MPS state (noisy Neel product)
//  3. Create the EigenDecomposition (SMALLEST_REAL + SCOPE_SPLIT + KRYLOV)
//  4. Configure the DMRG and Krylov sub-configs
//  5. Prepare + attach workspace + Compute
//  6. Report the ground-state energy and residual vs the exact value
//  7. Clean up

#include <cudensitymat.h>
#include "helpers.h"

#include <algorithm>
#include <cmath>
#include <complex>
#include <vector>
#include <numeric>
#include <random>
#include <iostream>
#include <iomanip>
#include <cassert>


using Complex = std::complex<double>;
constexpr cudaDataType_t kDataType = CUDA_C_64F;

constexpr bool verbose = true;

// --- Simulation parameters ---
constexpr int32_t  NUM_SITES  = 8;
constexpr int64_t  PHYS_DIM   = 2;     // spin-1/2
constexpr int64_t  MPO_BOND   = 5;     // Heisenberg MPO bond dim
constexpr int64_t  MAX_BOND   = 16;    // MPS bond-dim cap (= full Schmidt rank at N=8)
constexpr int32_t  MAX_SWEEPS = 50;
constexpr int32_t  KRYLOV_MAX_DIM = 3; // local Krylov subspace cap (boundary-dim constraint)
constexpr double   ENERGY_TOL = 1e-10;
constexpr double   RESIDUAL   = 1e-6;
constexpr double   NOISE_EPS  = 1e-2;
constexpr uint64_t RNG_SEED   = 1729;

constexpr double   J_COUPLING = 1.0;   // antiferromagnetic

// Exact ground-state energy for the N=8 open-boundary Heisenberg chain (J=1),
// from dense diagonalization of the 2^8 x 2^8 Hamiltonian.  Printed reference.
constexpr double   E0_EXACT   = -3.374932598688;


// ============================================================================
// Heisenberg MPO builder  (spin-1/2, bond dimension 5)
// ============================================================================
//
// Standard W-matrix: column 0 holds the finish operators (I, Sx, Sy, Sz), the
// bottom row the start operators (J*Sx, J*Sy, J*Sz, I); a term J*S^a_i S^a_{i+1}
// starts at site i and finishes at i+1, with no on-site term.  Left boundary is
// the bottom row, right boundary is column 0.
// Site tensor layout (b_L, d, b_R, d): index = aL + bL*(ket + d*(aR + bR*bra)).

struct HeisenbergMPO {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  // d x d spin-1/2 operators, entry (bra, ket) = <bra | S | ket>.
  static Complex I_op(int bra, int ket)  { return (bra == ket) ? Complex(1, 0) : Complex(0, 0); }
  static Complex Sx_op(int bra, int ket) { return (bra != ket) ? Complex(0.5, 0) : Complex(0, 0); }
  static Complex Sy_op(int bra, int ket) {
    if (bra == 0 && ket == 1) return Complex(0, -0.5);   // -i/2
    if (bra == 1 && ket == 0) return Complex(0,  0.5);   // +i/2
    return Complex(0, 0);
  }
  static Complex Sz_op(int bra, int ket) {
    if (bra == ket) return Complex(bra == 0 ? 0.5 : -0.5, 0);
    return Complex(0, 0);
  }

  void build() {
    const Complex zero{0.0, 0.0};
    hostTensors.assign(NUM_SITES, {});
    gpuPtrs.assign(NUM_SITES, nullptr);

    auto idx_full = [&](int64_t bL, int64_t bR, int64_t aL, int64_t ket,
                        int64_t aR, int64_t bra) {
      return aL + bL * (ket + PHYS_DIM * (aR + bR * bra));
    };
    // Place a d x d block at MPO corner (aL, aR); M is a callable (bra, ket).
    auto place = [&](std::vector<Complex>& T, int64_t bL, int64_t bR,
                     int64_t aL, int64_t aR, auto M) {
      for (int bra = 0; bra < PHYS_DIM; ++bra)
        for (int ket = 0; ket < PHYS_DIM; ++ket)
          T[idx_full(bL, bR, aL, ket, aR, bra)] += M(bra, ket);
    };

    const double J = J_COUPLING;
    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL  = (site == 0) ? 1 : MPO_BOND;
      const int64_t bR  = (site == NUM_SITES - 1) ? 1 : MPO_BOND;
      const int64_t vol = bL * PHYS_DIM * bR * PHYS_DIM;
      hostTensors[site].assign(vol, zero);
      auto& T = hostTensors[site];

      if (site == 0) {
        // Bottom row of W:  [0, J*Sx, J*Sy, J*Sz, I],  aL = 0.
        place(T, bL, bR, 0, 1, [&](int b, int k){ return Complex(J, 0) * Sx_op(b, k); });
        place(T, bL, bR, 0, 2, [&](int b, int k){ return Complex(J, 0) * Sy_op(b, k); });
        place(T, bL, bR, 0, 3, [&](int b, int k){ return Complex(J, 0) * Sz_op(b, k); });
        place(T, bL, bR, 0, 4, [&](int b, int k){ return I_op(b, k); });
      } else if (site == NUM_SITES - 1) {
        // First column of W:  [I, Sx, Sy, Sz, 0]^T,  aR = 0.
        place(T, bL, bR, 0, 0, [&](int b, int k){ return I_op(b, k); });
        place(T, bL, bR, 1, 0, [&](int b, int k){ return Sx_op(b, k); });
        place(T, bL, bR, 2, 0, [&](int b, int k){ return Sy_op(b, k); });
        place(T, bL, bR, 3, 0, [&](int b, int k){ return Sz_op(b, k); });
      } else {
        // Full bulk W: first column (finish) + bottom row (start).
        place(T, bL, bR, 0, 0, [&](int b, int k){ return I_op(b, k); });
        place(T, bL, bR, 1, 0, [&](int b, int k){ return Sx_op(b, k); });
        place(T, bL, bR, 2, 0, [&](int b, int k){ return Sy_op(b, k); });
        place(T, bL, bR, 3, 0, [&](int b, int k){ return Sz_op(b, k); });
        place(T, bL, bR, 4, 1, [&](int b, int k){ return Complex(J, 0) * Sx_op(b, k); });
        place(T, bL, bR, 4, 2, [&](int b, int k){ return Complex(J, 0) * Sy_op(b, k); });
        place(T, bL, bR, 4, 3, [&](int b, int k){ return Complex(J, 0) * Sz_op(b, k); });
        place(T, bL, bR, 4, 4, [&](int b, int k){ return I_op(b, k); });
      }

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Noisy Neel initial MPS  ( |up,down,up,...> + epsilon * Gaussian noise )
// ============================================================================
//
// Neel lies in the total-Sz = 0 sector of the singlet ground state; the noise
// gives DMRG a generic (non-exact) starting point.

struct NoisyNeelMPS {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build(const std::vector<int64_t>& bondDims, double epsilon, uint64_t seed) {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0};
    std::mt19937_64 rng(seed);
    std::normal_distribution<double> noise(0.0, epsilon);

    hostTensors.assign(NUM_SITES, {});
    gpuPtrs.assign(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL  = (site == 0) ? 1 : bondDims[site - 1];
      const int64_t bR  = (site == NUM_SITES - 1) ? 1 : bondDims[site];
      const int64_t vol = bL * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      auto idx = [&](int64_t aL, int64_t sigma, int64_t aR) {
        return aL + bL * (sigma + PHYS_DIM * aR);
      };

      // Neel seed: site i occupies |i mod 2>  (alternating up / down).
      const int64_t spin = site % 2;
      hostTensors[site][idx(0, spin, 0)] = one;
      // Per-element Gaussian noise (Re, Im) iid N(0, epsilon^2).
      for (auto & v : hostTensors[site]) {
        v += Complex(noise(rng), noise(rng));
      }

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Example workflow
// ============================================================================

void exampleWorkflow(cudensitymatHandle_t handle)
{
  if (verbose) {
    std::cout << "DMRG eigensolver example\n"
              << "  Model:           spin-1/2 antiferromagnetic Heisenberg chain\n"
              << "  H = J * sum_i ( Sx_i Sx_{i+1} + Sy_i Sy_{i+1} + Sz_i Sz_{i+1} )\n"
              << "  N (sites)       = " << NUM_SITES  << "\n"
              << "  d (phys dim)    = " << PHYS_DIM   << "\n"
              << "  J (coupling)    = " << J_COUPLING << "\n"
              << "  chi (bond cap)  = " << MAX_BOND   << "\n"
              << "  max sweeps      = " << MAX_SWEEPS << "\n"
              << "  energy tol      = " << ENERGY_TOL << "\n\n";
  }

  // --- 1. Build the Heisenberg MPO ---
  HeisenbergMPO mpo;
  mpo.build();
  if (verbose) std::cout << "Built Heisenberg MPO (bond dim " << MPO_BOND << ")\n";

  const std::vector<int64_t> spaceShape(NUM_SITES, PHYS_DIM);
  std::vector<int64_t> mpoBondDims(NUM_SITES - 1, MPO_BOND);

  cudensitymatMatrixProductOperator_t mpoHandle = nullptr;
  std::vector<cudensitymatWrappedTensorCallback_t> mpoCB(
      NUM_SITES, cudensitymatTensorCallbackNone);
  std::vector<cudensitymatWrappedTensorGradientCallback_t> mpoGCB(
      NUM_SITES, cudensitymatTensorGradientCallbackNone);
  HANDLE_CUDM_ERROR(cudensitymatCreateMatrixProductOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpoBondDims.data(),
                      kDataType,
                      mpo.gpuPtrs.data(),
                      mpoCB.data(),
                      mpoGCB.data(),
                      &mpoHandle));
  if (verbose) std::cout << "Created MPO handle\n";

  // --- 2. Wrap the MPO into an Operator (one term, one MPO product) ---
  cudensitymatOperatorTerm_t operatorTerm = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                      NUM_SITES, spaceShape.data(), &operatorTerm));

  std::vector<int32_t> modesActedOn(NUM_SITES);
  std::iota(modesActedOn.begin(), modesActedOn.end(), 0);
  std::vector<int32_t> modeDuality(NUM_SITES, 0);
  std::vector<int32_t> mpoConj = {0};
  HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendMPOProduct(handle,
                      operatorTerm,
                      1, &mpoHandle, mpoConj.data(),
                      modesActedOn.data(), modeDuality.data(),
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));

  cudensitymatOperator_t hamiltonian = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                      NUM_SITES, spaceShape.data(), &hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                      hamiltonian, operatorTerm, 0,
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));
  if (verbose) std::cout << "Created Hamiltonian operator from MPO\n";

  // --- 3. Build the initial MPS state ---
  //
  // Bond dims are capped at the natural Schmidt-rank bound min(d^i, d^(N-i)).
  std::vector<int64_t> mpsBondDims(NUM_SITES - 1);
  for (int32_t i = 0; i < NUM_SITES - 1; ++i) {
    int64_t leftDim = 1;
    for (int32_t j = 0; j <= i; ++j) leftDim *= spaceShape[j];
    int64_t rightDim = 1;
    for (int32_t j = i + 1; j < NUM_SITES; ++j) rightDim *= spaceShape[j];
    mpsBondDims[i] = std::min<int64_t>({MAX_BOND, leftDim, rightDim});
  }

  cudensitymatState_t mpsState = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpsBondDims.data(),
                      kDataType,
                      /*batchSize=*/1,
                      &mpsState));

  NoisyNeelMPS mps;
  mps.build(mpsBondDims, NOISE_EPS, RNG_SEED);

  int32_t numComponents = 0;
  HANDLE_CUDM_ERROR(cudensitymatStateGetNumComponents(handle, mpsState, &numComponents));
  std::vector<std::size_t> componentSizes(numComponents);
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle, mpsState,
                      numComponents, componentSizes.data()));
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle, mpsState,
                      numComponents, mps.gpuPtrs.data(), componentSizes.data()));
  if (verbose) std::cout << "Built and attached noisy-Neel MPS (alternating up/down + eps*Gaussian, seed="
                         << RNG_SEED << ")\n";

  // --- 4. Create EigenDecomposition (SCOPE_SPLIT + APPROACH_KRYLOV + SMALLEST_REAL) ---
  cudensitymatEigenDecomposition_t eigen = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateEigenDecomposition(handle,
                      hamiltonian,
                      /*isHermitian=*/1,
                      CUDENSITYMAT_EIGEN_SPECTRUM_SMALLEST_REAL,
                      CUDENSITYMAT_EIGEN_SCOPE_SPLIT,
                      CUDENSITYMAT_EIGEN_APPROACH_KRYLOV,
                      &eigen));
  if (verbose) std::cout << "Created EigenDecomposition object\n";

  // --- 5. Configure DMRG attributes (MAX_SWEEPS, ENERGY_TOLERANCE).
  // Sub-config values are deep-copied at Configure, so the handle can be freed now.
  cudensitymatEigenDecompositionScopeSplitDMRGConfig_t dmrgCfg = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateEigenDecompositionScopeSplitDMRGConfig(handle, &dmrgCfg));
  {
    const int32_t maxSweeps = MAX_SWEEPS;
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionScopeSplitDMRGConfigSetAttribute(
        handle, dmrgCfg,
        CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_MAX_SWEEPS,
        &maxSweeps, sizeof(maxSweeps)));
    const double energyTol = ENERGY_TOL;
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionScopeSplitDMRGConfigSetAttribute(
        handle, dmrgCfg,
        CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_ENERGY_TOLERANCE,
        &energyTol, sizeof(energyTol)));
  }
  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionConfigure(handle, eigen,
                      CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_CONFIG,
                      &dmrgCfg, sizeof(dmrgCfg)));
  HANDLE_CUDM_ERROR(cudensitymatDestroyEigenDecompositionScopeSplitDMRGConfig(dmrgCfg));
  dmrgCfg = nullptr;
  if (verbose) std::cout << "Configured DMRG sub-config (max_sweeps="
                         << MAX_SWEEPS << ", energy_tol=" << ENERGY_TOL << ")\n";

  // --- 5b. Configure the Krylov sub-config.
  // The 1-site local problem is only bL*d*bR = 4 at the boundary; the engine
  // needs dim >= (Krylov max_dim + min block size), so cap max_dim at 3.
  cudensitymatEigenDecompositionApproachKrylovConfig_t krylovCfg = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateEigenDecompositionApproachKrylovConfig(handle, &krylovCfg));
  {
    const int32_t krylovMaxDim = KRYLOV_MAX_DIM;
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionApproachKrylovConfigSetAttribute(
        handle, krylovCfg,
        CUDENSITYMAT_EIGEN_APPROACH_KRYLOV_MAX_DIM,
        &krylovMaxDim, sizeof(krylovMaxDim)));
  }
  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionConfigure(handle, eigen,
                      CUDENSITYMAT_EIGEN_APPROACH_KRYLOV_CONFIG,
                      &krylovCfg, sizeof(krylovCfg)));
  HANDLE_CUDM_ERROR(cudensitymatDestroyEigenDecompositionApproachKrylovConfig(krylovCfg));
  krylovCfg = nullptr;
  if (verbose) std::cout << "Configured Krylov sub-config (max_dim=" << KRYLOV_MAX_DIM << ")\n";

  // --- 6. Prepare + attach workspace ---
  cudensitymatWorkspaceDescriptor_t workspaceDescr = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.85);
  if (verbose) std::cout << "Available workspace memory (bytes) = " << freeMem << "\n";

  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionPrepare(handle, eigen,
                      /*maxEigenStates=*/1, mpsState, CUDENSITYMAT_COMPUTE_64F,
                      freeMem, workspaceDescr, /*stream=*/0x0));

  std::size_t scratchSize = 0;
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle, workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &scratchSize));
  void * scratchBuf = nullptr;
  if (scratchSize > 0) {
    HANDLE_CUDA_ERROR(cudaMalloc(&scratchBuf, scratchSize));
    HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle, workspaceDescr,
                        CUDENSITYMAT_MEMSPACE_DEVICE,
                        CUDENSITYMAT_WORKSPACE_SCRATCH,
                        scratchBuf, scratchSize));
  }
  if (verbose) std::cout << "Prepared DMRG plan; scratch workspace = " << scratchSize << " bytes\n";

  // --- 7. Compute the ground-state pair (eigenvalue, MPS) ---
  void * eigenvalueGpu = createArrayGPU<Complex>(1);
  double residual = RESIDUAL;
  cudensitymatState_t eigenstatesArr[1] = { mpsState };

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionCompute(handle, eigen,
                      /*time=*/0.0,
                      /*batchSize=*/1,
                      /*numParams=*/0, /*params=*/nullptr,
                      /*numEigenStates=*/1, eigenstatesArr,
                      eigenvalueGpu, &residual,
                      workspaceDescr, /*stream=*/0x0));
  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  cuDoubleComplex eigenvalueHost{0.0, 0.0};
  HANDLE_CUDA_ERROR(cudaMemcpy(&eigenvalueHost, eigenvalueGpu,
                               sizeof(cuDoubleComplex), cudaMemcpyDeviceToHost));

  const double eDmrg     = eigenvalueHost.x;
  const double energyErr = std::abs(eDmrg - E0_EXACT);
  const bool   converged = (energyErr < 1e-6) && (residual < RESIDUAL);

  std::cout << "\n=================================================================\n"
            << "DMRG ground-state energy : " << std::scientific << std::setprecision(12)
            << eDmrg << "\n"
            << "Exact (dense ED, N=" << NUM_SITES << ")  : " << E0_EXACT << "\n"
            << "Energy error             : " << std::scientific << std::setprecision(3)
            << energyErr << "\n"
            << "Final residual           : " << residual << "\n"
            << "Result                   : " << (converged ? "PASS" : "FAIL") << "\n"
            << "=================================================================\n\n";

  // --- 8. Cleanup (reverse Create order) ---
  destroyArrayGPU(eigenvalueGpu);
  if (scratchBuf) HANDLE_CUDA_ERROR(cudaFree(scratchBuf));
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));
  HANDLE_CUDM_ERROR(cudensitymatDestroyEigenDecomposition(eigen));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(mpsState));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(operatorTerm));
  HANDLE_CUDM_ERROR(cudensitymatDestroyMatrixProductOperator(mpoHandle));
  mps.destroy();
  mpo.destroy();

  if (verbose) std::cout << "Destroyed all resources\n";
}


int main(int /*argc*/, char ** /*argv*/)
{
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose) std::cout << "Set active CUDA device 0\n";

  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose) std::cout << "Created cuDensityMat library handle\n";

  exampleWorkflow(handle);

  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose) std::cout << "Destroyed cuDensityMat library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());
  return 0;
}

Code example (serial execution of 2-site MPS-DMRG ground-state eigensolver)

The following code example illustrates how to use the cuDensityMat library to compute the ground state of a transverse-field Ising Hamiltonian encoded as an MPO, acting on a pure MPS state, using the split-scope 2-site DMRG method with CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_NUM_SITES set to 2. The 2-site variant solves two adjacent sites jointly and re-splits the result with an SVD truncation, so the intervening bond dimension adapts during the computation. The example attaches an SVD configuration via CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_SVD_CONFIG (setting CUDENSITYMAT_SVD_CONFIG_MAX_EXTENT), seeds the initial current bond extents strictly between 1 and the buffer maximum via cudensitymatStateMPSSetCurrentBondExtents, computes the ground state, then queries cudensitymatStateMPSGetCurrentBondExtents to report the adapted bond extents. The full sample code can be found in the NVIDIA/cuQuantum repository (main serial 2-site MPS-DMRG code as well as the utility code).

/* Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

// Two-site DMRG eigensolver example.
//
// Computes the ground state of the 1-D transverse-field Ising model
//   H = -J * sum_i Z_i Z_{i+1}  -  h * sum_i X_i   (open boundary)
// with the 2-site DMRG solver (SCOPE_SPLIT + APPROACH_KRYLOV + SMALLEST_REAL)
// on an MPS state.  Two adjacent MPS sites are merged, locally eigensolved,
// SVD-split with adaptive bond truncation, and synced back per step.
//
// The chain is blocked: each MPS site carries a pair of qubits (phys dim 4,
// basis {|00>, |01>, |10>, |11>}), so 3 MPS sites model a 6-qubit chain.  A
// small SVD MAX_EXTENT cap and a partial-fill initial bond layout (current
// bond extents strictly between 1 and the buffer maximum) exercise the 2-site
// memory-layout and adaptive-growth contracts.
//
// 3 MPS sites is intentionally small: both bonds are boundary-adjacent, so the
// partial-fill current<max layout works here pending a cuTensorNet update that
// enables N>=4 interior-bond growth.
//
// Workflow:
//  1. Build the blocked-TFIM MPO (bond dim 3)
//  2. Wrap the MPO into a Hamiltonian operator
//  3. Create an initial MPS state (noisy product) + partial-fill bond layout
//  4. Create the EigenDecomposition (SMALLEST_REAL + SCOPE_SPLIT + KRYLOV)
//  5. Configure DMRG (NUM_SITES=2 + SVD config), Krylov, Prepare, attach
//  6. Compute the ground-state pair
//  7. Report the energy and the adapted (current vs max) bond extents
//  8. Clean up

#include <cudensitymat.h>
#include "helpers.h"

#include <algorithm>
#include <cmath>
#include <complex>
#include <vector>
#include <numeric>
#include <random>
#include <iostream>
#include <iomanip>
#include <cassert>


using Complex = std::complex<double>;
constexpr cudaDataType_t kDataType = CUDA_C_64F;

constexpr bool verbose = true;

// --- Simulation parameters ---
constexpr int32_t  NUM_SITES   = 3;     // blocked MPS sites (== 6 qubits); small by design
constexpr int64_t  PHYS_DIM    = 4;     // 2 qubits per blocked site
constexpr int64_t  MPO_BOND    = 3;     // standard TFIM MPO bond dim
constexpr int64_t  MAX_BOND    = 8;     // MPS bond-dim buffer cap
constexpr int32_t  NUM_DMRG_SITES = 2;  // 2-site DMRG
constexpr int64_t  SVD_MAX_EXTENT = 3;  // modest SVD truncation cap (< buffer extent)
constexpr int64_t  INIT_CURRENT_BOND = 2; // partial fill: strictly between 1 and the maximum
constexpr int32_t  MAX_SWEEPS  = 30;
constexpr int32_t  KRYLOV_MAX_DIM = 5;  // local Krylov subspace cap
constexpr double   ENERGY_TOL  = 1e-10;
constexpr double   RESIDUAL    = 1e-6;
constexpr double   NOISE_EPS   = 1e-2;
constexpr uint64_t RNG_SEED    = 1729;

constexpr double   J_COUPLING  = 1.0;   // ZZ coupling
constexpr double   H_FIELD     = 0.5;   // transverse field

// Exact ground-state energy for the 6-qubit OBC TFIM (J=1, h=0.5), from dense
// diagonalization of the 2^6 x 2^6 Hamiltonian.  Printed reference.
constexpr double   E0_EXACT    = -5.522029570800221;


// ============================================================================
// Blocked-TFIM MPO  (phys dim 4 per site, bond dimension 3)
// ============================================================================
//
// Each MPS site carries two qubits (q_left, q_right), basis index i = 2*q_left
// + q_right.  The W-matrix carries the standard 3-state TFIM finite automaton:
// column 0 holds the finish operators, the bottom row the start operators, plus
// the on-site block term B = -J*ZZ - h*(Xleft + Xright).
// Site tensor layout (aL, ket, aR, bra): idx = aL + bL*(ket + d*(aR + bR*bra)).

struct BlockedTFIMMPO {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  // 4x4 operators on a 2-qubit block; entry (bra, ket).
  static Complex I4_(int r, int c) { return (r == c) ? Complex(1, 0) : Complex(0, 0); }
  static Complex Zleft_(int r, int c) {
    if (r != c) return Complex(0, 0);
    return (r < 2) ? Complex(1, 0) : Complex(-1, 0);     // Z on first qubit: diag(+,+,-,-)
  }
  static Complex Zright_(int r, int c) {
    if (r != c) return Complex(0, 0);
    return ((r & 1) == 0) ? Complex(1, 0) : Complex(-1, 0);   // Z on second qubit: diag(+,-,+,-)
  }
  static Complex ZZ_(int r, int c) {
    if (r != c) return Complex(0, 0);
    const int rL = (r >> 1) & 1, rR = r & 1;
    return (rL == rR) ? Complex(1, 0) : Complex(-1, 0);  // Z⊗Z: diag(+,-,-,+)
  }
  static Complex Xleft_(int r, int c)  { return (r == (c ^ 0b10)) ? Complex(1, 0) : Complex(0, 0); }
  static Complex Xright_(int r, int c) { return (r == (c ^ 0b01)) ? Complex(1, 0) : Complex(0, 0); }
  // Within-block on-site Hamiltonian B = -J*ZZ - h*(Xleft + Xright).
  static Complex B_(int r, int c) {
    return -Complex(J_COUPLING, 0) * ZZ_(r, c)
           - Complex(H_FIELD, 0) * (Xleft_(r, c) + Xright_(r, c));
  }

  void build() {
    const Complex zero{0.0, 0.0};
    hostTensors.assign(NUM_SITES, {});
    gpuPtrs.assign(NUM_SITES, nullptr);

    // Left boundary site 0 — shape (bL=1, phys=4, bR=3, phys=4).
    {
      const int64_t bL = 1, bR = MPO_BOND;
      hostTensors[0].assign(bL * PHYS_DIM * bR * PHYS_DIM, zero);
      auto idx = [&](int64_t aL, int64_t ket, int64_t aR, int64_t bra) {
        return aL + bL * (ket + PHYS_DIM * (aR + bR * bra));
      };
      for (int bra = 0; bra < PHYS_DIM; ++bra) {
        for (int ket = 0; ket < PHYS_DIM; ++ket) {
          hostTensors[0][idx(0, ket, 0, bra)] = B_(bra, ket);
          hostTensors[0][idx(0, ket, 1, bra)] = Zright_(bra, ket);
          hostTensors[0][idx(0, ket, 2, bra)] = I4_(bra, ket);
        }
      }
      gpuPtrs[0] = createInitializeArrayGPU(hostTensors[0]);
    }

    // Bulk sites — shape (bL=3, phys=4, bR=3, phys=4).
    for (int32_t site = 1; site < NUM_SITES - 1; ++site) {
      const int64_t bL = MPO_BOND, bR = MPO_BOND;
      hostTensors[site].assign(bL * PHYS_DIM * bR * PHYS_DIM, zero);
      auto idx = [&](int64_t aL, int64_t ket, int64_t aR, int64_t bra) {
        return aL + bL * (ket + PHYS_DIM * (aR + bR * bra));
      };
      for (int bra = 0; bra < PHYS_DIM; ++bra) {
        for (int ket = 0; ket < PHYS_DIM; ++ket) {
          hostTensors[site][idx(0, ket, 0, bra)] = I4_(bra, ket);
          hostTensors[site][idx(1, ket, 0, bra)] = -Complex(J_COUPLING, 0) * Zleft_(bra, ket);
          hostTensors[site][idx(2, ket, 0, bra)] = B_(bra, ket);
          hostTensors[site][idx(2, ket, 1, bra)] = Zright_(bra, ket);
          hostTensors[site][idx(2, ket, 2, bra)] = I4_(bra, ket);
        }
      }
      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }

    // Right boundary site N-1 — shape (bL=3, phys=4, bR=1, phys=4).
    {
      const int32_t last = NUM_SITES - 1;
      const int64_t bL = MPO_BOND, bR = 1;
      hostTensors[last].assign(bL * PHYS_DIM * bR * PHYS_DIM, zero);
      auto idx = [&](int64_t aL, int64_t ket, int64_t aR, int64_t bra) {
        return aL + bL * (ket + PHYS_DIM * (aR + bR * bra));
      };
      for (int bra = 0; bra < PHYS_DIM; ++bra) {
        for (int ket = 0; ket < PHYS_DIM; ++ket) {
          hostTensors[last][idx(0, ket, 0, bra)] = I4_(bra, ket);
          hostTensors[last][idx(1, ket, 0, bra)] = -Complex(J_COUPLING, 0) * Zleft_(bra, ket);
          hostTensors[last][idx(2, ket, 0, bra)] = B_(bra, ket);
        }
      }
      gpuPtrs[last] = createInitializeArrayGPU(hostTensors[last]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Noisy product initial MPS  ( |00,00,...> + epsilon * Gaussian noise )
// ============================================================================
//
// A low-energy ferromagnetic product seed gives DMRG a generic, non-exact
// starting point.

struct NoisyProductMPS {

  std::vector<std::vector<Complex>> hostTensors;
  std::vector<void*> gpuPtrs;

  void build(const std::vector<int64_t>& bondDims, double epsilon, uint64_t seed) {
    const Complex zero{0.0, 0.0}, one{1.0, 0.0};
    std::mt19937_64 rng(seed);
    std::normal_distribution<double> noise(0.0, epsilon);

    hostTensors.assign(NUM_SITES, {});
    gpuPtrs.assign(NUM_SITES, nullptr);

    for (int32_t site = 0; site < NUM_SITES; ++site) {
      const int64_t bL  = (site == 0) ? 1 : bondDims[site - 1];
      const int64_t bR  = (site == NUM_SITES - 1) ? 1 : bondDims[site];
      const int64_t vol = bL * PHYS_DIM * bR;
      hostTensors[site].assign(vol, zero);

      auto idx = [&](int64_t aL, int64_t sigma, int64_t aR) {
        return aL + bL * (sigma + PHYS_DIM * aR);
      };

      // Product seed: every block in qubit-pair state |00> (index 0).
      hostTensors[site][idx(0, 0, 0)] = one;
      // Per-element Gaussian noise (Re, Im) iid N(0, epsilon^2).
      for (auto & v : hostTensors[site]) {
        v += Complex(noise(rng), noise(rng));
      }

      gpuPtrs[site] = createInitializeArrayGPU(hostTensors[site]);
    }
  }

  void destroy() {
    for (auto & ptr : gpuPtrs) {
      if (ptr) { destroyArrayGPU(ptr); ptr = nullptr; }
    }
  }
};


// ============================================================================
// Example workflow
// ============================================================================

void exampleWorkflow(cudensitymatHandle_t handle)
{
  if (verbose) {
    std::cout << "Two-site DMRG eigensolver example\n"
              << "  Model:           blocked transverse-field Ising chain (phys dim 4)\n"
              << "  H = -J * sum_i Z_i Z_{i+1}  -  h * sum_i X_i  (open boundary)\n"
              << "  N (MPS sites)   = " << NUM_SITES      << "   (== " << 2 * NUM_SITES << " qubits)\n"
              << "  d (phys dim)    = " << PHYS_DIM       << "\n"
              << "  J (coupling)    = " << J_COUPLING     << "\n"
              << "  h (field)       = " << H_FIELD        << "\n"
              << "  chi (bond cap)  = " << MAX_BOND       << "\n"
              << "  DMRG num_sites  = " << NUM_DMRG_SITES << "\n"
              << "  SVD max_extent  = " << SVD_MAX_EXTENT << "\n"
              << "  max sweeps      = " << MAX_SWEEPS     << "\n"
              << "  energy tol      = " << ENERGY_TOL     << "\n\n";
  }

  // --- 1. Build the blocked-TFIM MPO ---
  BlockedTFIMMPO mpo;
  mpo.build();
  if (verbose) std::cout << "Built blocked-TFIM MPO (bond dim " << MPO_BOND << ")\n";

  const std::vector<int64_t> spaceShape(NUM_SITES, PHYS_DIM);
  std::vector<int64_t> mpoBondDims(NUM_SITES - 1, MPO_BOND);

  cudensitymatMatrixProductOperator_t mpoHandle = nullptr;
  std::vector<cudensitymatWrappedTensorCallback_t> mpoCB(
      NUM_SITES, cudensitymatTensorCallbackNone);
  std::vector<cudensitymatWrappedTensorGradientCallback_t> mpoGCB(
      NUM_SITES, cudensitymatTensorGradientCallbackNone);
  HANDLE_CUDM_ERROR(cudensitymatCreateMatrixProductOperator(handle,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpoBondDims.data(),
                      kDataType,
                      mpo.gpuPtrs.data(),
                      mpoCB.data(),
                      mpoGCB.data(),
                      &mpoHandle));
  if (verbose) std::cout << "Created MPO handle\n";

  // --- 2. Wrap the MPO into an Operator (one term, one MPO product) ---
  cudensitymatOperatorTerm_t operatorTerm = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle,
                      NUM_SITES, spaceShape.data(), &operatorTerm));

  std::vector<int32_t> modesActedOn(NUM_SITES);
  std::iota(modesActedOn.begin(), modesActedOn.end(), 0);
  std::vector<int32_t> modeDuality(NUM_SITES, 0);
  std::vector<int32_t> mpoConj = {0};
  HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendMPOProduct(handle,
                      operatorTerm,
                      1, &mpoHandle, mpoConj.data(),
                      modesActedOn.data(), modeDuality.data(),
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));

  cudensitymatOperator_t hamiltonian = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateOperator(handle,
                      NUM_SITES, spaceShape.data(), &hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(handle,
                      hamiltonian, operatorTerm, 0,
                      make_cuDoubleComplex(1.0, 0.0),
                      cudensitymatScalarCallbackNone,
                      cudensitymatScalarGradientCallbackNone));
  if (verbose) std::cout << "Created Hamiltonian operator from MPO\n";

  // --- 3. Build the initial MPS state ---
  //
  // Buffer (maximum) bond dims are capped at the natural Schmidt-rank bound.
  std::vector<int64_t> mpsBondDims(NUM_SITES - 1);
  for (int32_t i = 0; i < NUM_SITES - 1; ++i) {
    int64_t leftDim = 1;
    for (int32_t j = 0; j <= i; ++j) leftDim *= spaceShape[j];
    int64_t rightDim = 1;
    for (int32_t j = i + 1; j < NUM_SITES; ++j) rightDim *= spaceShape[j];
    mpsBondDims[i] = std::min<int64_t>({MAX_BOND, leftDim, rightDim});
  }

  cudensitymatState_t mpsState = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateStateMPS(handle,
                      CUDENSITYMAT_STATE_PURITY_PURE,
                      NUM_SITES,
                      spaceShape.data(),
                      CUDENSITYMAT_BOUNDARY_CONDITION_OPEN,
                      mpsBondDims.data(),
                      kDataType,
                      /*batchSize=*/1,
                      &mpsState));

  NoisyProductMPS mps;
  mps.build(mpsBondDims, NOISE_EPS, RNG_SEED);

  int32_t numComponents = 0;
  HANDLE_CUDM_ERROR(cudensitymatStateGetNumComponents(handle, mpsState, &numComponents));
  std::vector<std::size_t> componentSizes(numComponents);
  HANDLE_CUDM_ERROR(cudensitymatStateGetComponentStorageSize(handle, mpsState,
                      numComponents, componentSizes.data()));
  HANDLE_CUDM_ERROR(cudensitymatStateAttachComponentStorage(handle, mpsState,
                      numComponents, mps.gpuPtrs.data(), componentSizes.data()));
  if (verbose) std::cout << "Built and attached noisy-product MPS (seed=" << RNG_SEED << ")\n";

  // --- 3b. Partial-fill bond layout: set initial CURRENT bond extents strictly
  // between 1 and the buffer maximum, exercising the 2-site partial-fill layout.
  const int32_t numBonds = NUM_SITES - 1;
  std::vector<int64_t> initialCurrentBonds(numBonds, INIT_CURRENT_BOND);
  HANDLE_CUDM_ERROR(cudensitymatStateMPSSetCurrentBondExtents(handle, mpsState,
                      initialCurrentBonds.data()));
  if (verbose) {
    std::cout << "Set initial current bond extents (partial fill) = {";
    for (int32_t b = 0; b < numBonds; ++b)
      std::cout << initialCurrentBonds[b] << (b + 1 < numBonds ? "," : "");
    std::cout << "}  (buffer max = {";
    for (int32_t b = 0; b < numBonds; ++b)
      std::cout << mpsBondDims[b] << (b + 1 < numBonds ? "," : "");
    std::cout << "})\n";
  }

  // --- 4. Create EigenDecomposition (SCOPE_SPLIT + APPROACH_KRYLOV + SMALLEST_REAL) ---
  cudensitymatEigenDecomposition_t eigen = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateEigenDecomposition(handle,
                      hamiltonian,
                      /*isHermitian=*/1,
                      CUDENSITYMAT_EIGEN_SPECTRUM_SMALLEST_REAL,
                      CUDENSITYMAT_EIGEN_SCOPE_SPLIT,
                      CUDENSITYMAT_EIGEN_APPROACH_KRYLOV,
                      &eigen));
  if (verbose) std::cout << "Created EigenDecomposition object\n";

  // --- 5. Configure DMRG attributes (NUM_SITES=2, SVD config, MAX_SWEEPS, ENERGY_TOLERANCE).
  // Sub-config values are deep-copied at Configure, so the handles can be freed afterward.
  cudensitymatEigenDecompositionScopeSplitDMRGConfig_t dmrgCfg = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateEigenDecompositionScopeSplitDMRGConfig(handle, &dmrgCfg));
  {
    const int32_t dmrgNumSites = NUM_DMRG_SITES;
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionScopeSplitDMRGConfigSetAttribute(
        handle, dmrgCfg,
        CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_NUM_SITES,
        &dmrgNumSites, sizeof(dmrgNumSites)));

    // SVD config: a modest MAX_EXTENT cap plus value cutoffs. The cap is below
    // the buffer maximum, so the truncated current bond extents adapt below it.
    cudensitymatSVDConfig_t svd = nullptr;
    HANDLE_CUDM_ERROR(cudensitymatCreateSVDConfig(handle, &svd));
    const int64_t maxExtent = SVD_MAX_EXTENT;
    const double absCutoff = 1e-12, relCutoff = 1e-12;
    HANDLE_CUDM_ERROR(cudensitymatSVDConfigSetAttribute(handle, svd,
        CUDENSITYMAT_SVD_CONFIG_MAX_EXTENT, &maxExtent, sizeof(maxExtent)));
    HANDLE_CUDM_ERROR(cudensitymatSVDConfigSetAttribute(handle, svd,
        CUDENSITYMAT_SVD_CONFIG_ABS_CUTOFF, &absCutoff, sizeof(absCutoff)));
    HANDLE_CUDM_ERROR(cudensitymatSVDConfigSetAttribute(handle, svd,
        CUDENSITYMAT_SVD_CONFIG_REL_CUTOFF, &relCutoff, sizeof(relCutoff)));
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionScopeSplitDMRGConfigSetAttribute(
        handle, dmrgCfg,
        CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_SVD_CONFIG, &svd, sizeof(svd)));
    HANDLE_CUDM_ERROR(cudensitymatDestroySVDConfig(svd));   // safe once attached (deep-cloned)

    const int32_t maxSweeps = MAX_SWEEPS;
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionScopeSplitDMRGConfigSetAttribute(
        handle, dmrgCfg,
        CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_MAX_SWEEPS,
        &maxSweeps, sizeof(maxSweeps)));
    const double energyTol = ENERGY_TOL;
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionScopeSplitDMRGConfigSetAttribute(
        handle, dmrgCfg,
        CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_ENERGY_TOLERANCE,
        &energyTol, sizeof(energyTol)));
  }
  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionConfigure(handle, eigen,
                      CUDENSITYMAT_EIGEN_SPLIT_SCOPE_DMRG_CONFIG,
                      &dmrgCfg, sizeof(dmrgCfg)));
  HANDLE_CUDM_ERROR(cudensitymatDestroyEigenDecompositionScopeSplitDMRGConfig(dmrgCfg));
  dmrgCfg = nullptr;
  if (verbose) std::cout << "Configured DMRG sub-config (num_sites=" << NUM_DMRG_SITES
                         << ", svd_max_extent=" << SVD_MAX_EXTENT
                         << ", max_sweeps=" << MAX_SWEEPS << ", energy_tol=" << ENERGY_TOL << ")\n";

  // --- 5b. Configure the Krylov sub-config.
  cudensitymatEigenDecompositionApproachKrylovConfig_t krylovCfg = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateEigenDecompositionApproachKrylovConfig(handle, &krylovCfg));
  {
    const int32_t krylovMaxDim = KRYLOV_MAX_DIM;
    HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionApproachKrylovConfigSetAttribute(
        handle, krylovCfg,
        CUDENSITYMAT_EIGEN_APPROACH_KRYLOV_MAX_DIM,
        &krylovMaxDim, sizeof(krylovMaxDim)));
  }
  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionConfigure(handle, eigen,
                      CUDENSITYMAT_EIGEN_APPROACH_KRYLOV_CONFIG,
                      &krylovCfg, sizeof(krylovCfg)));
  HANDLE_CUDM_ERROR(cudensitymatDestroyEigenDecompositionApproachKrylovConfig(krylovCfg));
  krylovCfg = nullptr;
  if (verbose) std::cout << "Configured Krylov sub-config (max_dim=" << KRYLOV_MAX_DIM << ")\n";

  // --- 6. Prepare + attach workspace ---
  cudensitymatWorkspaceDescriptor_t workspaceDescr = nullptr;
  HANDLE_CUDM_ERROR(cudensitymatCreateWorkspace(handle, &workspaceDescr));

  std::size_t freeMem = 0, totalMem = 0;
  HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
  freeMem = static_cast<std::size_t>(static_cast<double>(freeMem) * 0.85);
  if (verbose) std::cout << "Available workspace memory (bytes) = " << freeMem << "\n";

  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionPrepare(handle, eigen,
                      /*maxEigenStates=*/1, mpsState, CUDENSITYMAT_COMPUTE_64F,
                      freeMem, workspaceDescr, /*stream=*/0x0));

  std::size_t scratchSize = 0;
  HANDLE_CUDM_ERROR(cudensitymatWorkspaceGetMemorySize(handle, workspaceDescr,
                      CUDENSITYMAT_MEMSPACE_DEVICE,
                      CUDENSITYMAT_WORKSPACE_SCRATCH,
                      &scratchSize));
  void * scratchBuf = nullptr;
  if (scratchSize > 0) {
    HANDLE_CUDA_ERROR(cudaMalloc(&scratchBuf, scratchSize));
    HANDLE_CUDM_ERROR(cudensitymatWorkspaceSetMemory(handle, workspaceDescr,
                        CUDENSITYMAT_MEMSPACE_DEVICE,
                        CUDENSITYMAT_WORKSPACE_SCRATCH,
                        scratchBuf, scratchSize));
  }
  if (verbose) std::cout << "Prepared DMRG plan; scratch workspace = " << scratchSize << " bytes\n";

  // --- 7. Compute the ground-state pair (eigenvalue, MPS) ---
  void * eigenvalueGpu = createArrayGPU<Complex>(1);
  double residual = RESIDUAL;
  cudensitymatState_t eigenstatesArr[1] = { mpsState };

  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());
  HANDLE_CUDM_ERROR(cudensitymatEigenDecompositionCompute(handle, eigen,
                      /*time=*/0.0,
                      /*batchSize=*/1,
                      /*numParams=*/0, /*params=*/nullptr,
                      /*numEigenStates=*/1, eigenstatesArr,
                      eigenvalueGpu, &residual,
                      workspaceDescr, /*stream=*/0x0));
  HANDLE_CUDA_ERROR(cudaDeviceSynchronize());

  cuDoubleComplex eigenvalueHost{0.0, 0.0};
  HANDLE_CUDA_ERROR(cudaMemcpy(&eigenvalueHost, eigenvalueGpu,
                               sizeof(cuDoubleComplex), cudaMemcpyDeviceToHost));

  // Query the adapted (current) bond extents after Compute.
  std::vector<int64_t> finalCurrentBonds(numBonds, 0);
  HANDLE_CUDM_ERROR(cudensitymatStateMPSGetCurrentBondExtents(handle, mpsState,
                      finalCurrentBonds.data()));

  const double eDmrg     = eigenvalueHost.x;
  const double energyErr = std::abs(eDmrg - E0_EXACT);
  const bool   converged = (residual < RESIDUAL);

  std::cout << "\n=================================================================\n"
            << "DMRG ground-state energy : " << std::scientific << std::setprecision(12)
            << eDmrg << "\n"
            << "Exact (dense ED, N=" << 2 * NUM_SITES << " qubits) : " << E0_EXACT << "\n"
            << "Energy error             : " << std::scientific << std::setprecision(3)
            << energyErr << "\n"
            << "Final residual           : " << residual << "\n";
  std::cout << "Bond extents (current/max):";
  for (int32_t b = 0; b < numBonds; ++b)
    std::cout << " " << finalCurrentBonds[b] << "/" << mpsBondDims[b];
  std::cout << "\n"
            << "Result                   : " << (converged ? "PASS" : "FAIL") << "\n"
            << "=================================================================\n\n";

  // --- 8. Cleanup (reverse Create order) ---
  destroyArrayGPU(eigenvalueGpu);
  if (scratchBuf) HANDLE_CUDA_ERROR(cudaFree(scratchBuf));
  HANDLE_CUDM_ERROR(cudensitymatDestroyWorkspace(workspaceDescr));
  HANDLE_CUDM_ERROR(cudensitymatDestroyEigenDecomposition(eigen));
  HANDLE_CUDM_ERROR(cudensitymatDestroyState(mpsState));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperator(hamiltonian));
  HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(operatorTerm));
  HANDLE_CUDM_ERROR(cudensitymatDestroyMatrixProductOperator(mpoHandle));
  mps.destroy();
  mpo.destroy();

  if (verbose) std::cout << "Destroyed all resources\n";
}


int main(int /*argc*/, char ** /*argv*/)
{
  HANDLE_CUDA_ERROR(cudaSetDevice(0));
  if (verbose) std::cout << "Set active CUDA device 0\n";

  cudensitymatHandle_t handle;
  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
  if (verbose) std::cout << "Created cuDensityMat library handle\n";

  exampleWorkflow(handle);

  HANDLE_CUDM_ERROR(cudensitymatDestroy(handle));
  if (verbose) std::cout << "Destroyed cuDensityMat library handle\n";

  HANDLE_CUDA_ERROR(cudaDeviceReset());
  return 0;
}

Useful tips

• For debugging, one can set the environment variable CUDENSITYMAT_LOG_LEVEL=n. The level n = 0, 1, …, 5 corresponds to the logger level as described in the table below. The environment variable CUDENSITYMAT_LOG_FILE=<filepath> can be used to redirect the log output to a custom file at <filepath> instead of stdout.

Level	Summary	Long Description
0	Off	Logging is disabled (default)
1	Errors	Only errors will be logged
2	Performance Trace	API calls that launch CUDA kernels will log their parameters and important information
3	Performance Hints	Hints that can potentially improve the application’s performance
4	Heuristics Trace	Provides general information about the library execution, may contain details about heuristic status
5	API Trace	API calls will log their parameter and important information