From 969417fb1341d0728b7da58838a56aa1344b2d80 Mon Sep 17 00:00:00 2001 From: Alex McCaskey Date: Fri, 19 Jul 2024 19:40:29 +0000 Subject: [PATCH 1/7] initial port of gqe Signed-off-by: Alex McCaskey --- cudaqlib/gse.h | 4 +- cudaqlib/gse/CMakeLists.txt | 4 +- cudaqlib/gse/adapt/adapt.h | 12 +- cudaqlib/gse/adapt/operator_pool.h | 40 -- .../gse/adapt/pools/spin_complement_gsd.h | 152 -------- cudaqlib/gse/utils/CMakeLists.txt | 9 + cudaqlib/gse/utils/operator_pool.h | 70 ++++ cudaqlib/gse/utils/pools/CMakeLists.txt | 26 ++ .../gse/utils/pools/spin_complement_gsd.cpp | 155 ++++++++ .../gse/utils/pools/spin_complement_gsd.h | 26 ++ cudaqlib/gse/utils/pools/uccsd_pool.cpp | 71 ++++ cudaqlib/gse/utils/pools/uccsd_pool.h | 23 ++ cudaqlib/operators/chemistry/molecule.cpp | 8 +- cudaqlib/utils/extension_point.h | 2 +- python/CMakeLists.txt | 4 +- python/bindings/gse/py_gse.cpp | 20 +- python/cudaqlib/__init__.py | 1 + python/cudaqlib/algorithms/gqe.py | 358 ++++++++++++++++++ python/cudaqlib/algorithms/loss.py | 53 +++ python/cudaqlib/algorithms/transformer.py | 106 ++++++ .../pyscf/generators/gas_phase_generator.py | 1 - tests/CMakeLists.txt | 3 +- tests/gse/AdaptTester.cpp | 5 +- tests/kernels/UCCSDTester.cpp | 10 +- 24 files changed, 948 insertions(+), 215 deletions(-) delete mode 100644 cudaqlib/gse/adapt/operator_pool.h delete mode 100644 cudaqlib/gse/adapt/pools/spin_complement_gsd.h create mode 100644 cudaqlib/gse/utils/CMakeLists.txt create mode 100644 cudaqlib/gse/utils/operator_pool.h create mode 100644 cudaqlib/gse/utils/pools/CMakeLists.txt create mode 100644 cudaqlib/gse/utils/pools/spin_complement_gsd.cpp create mode 100644 cudaqlib/gse/utils/pools/spin_complement_gsd.h create mode 100644 cudaqlib/gse/utils/pools/uccsd_pool.cpp create mode 100644 cudaqlib/gse/utils/pools/uccsd_pool.h create mode 100644 python/cudaqlib/algorithms/gqe.py create mode 100644 python/cudaqlib/algorithms/loss.py create mode 100644 python/cudaqlib/algorithms/transformer.py diff --git a/cudaqlib/gse.h b/cudaqlib/gse.h index 73e9f5f..774b76b 100644 --- a/cudaqlib/gse.h +++ b/cudaqlib/gse.h @@ -10,5 +10,7 @@ #include "gse/adapt/adapt.h" #include "gse/vqe/vqe.h" +#include "gse/utils/operator_pool.h" +#include "gse/utils/pools/uccsd_pool.h" +#include "gse/utils/pools/spin_complement_gsd.h" -#include "gse/adapt/pools/spin_complement_gsd.h" diff --git a/cudaqlib/gse/CMakeLists.txt b/cudaqlib/gse/CMakeLists.txt index eaabdaa..53e86e7 100644 --- a/cudaqlib/gse/CMakeLists.txt +++ b/cudaqlib/gse/CMakeLists.txt @@ -8,13 +8,15 @@ set(LIBRARY_NAME cudaq-gse) +add_subdirectory(utils) + add_library(${LIBRARY_NAME} INTERFACE) target_include_directories(${LIBRARY_NAME} INTERFACE $ $ $) -target_link_libraries(${LIBRARY_NAME} INTERFACE cudaq::cudaq) +target_link_libraries(${LIBRARY_NAME} INTERFACE cudaq::cudaq cudaq-operator-pools cudaq-optim) install( TARGETS ${LIBRARY_NAME} diff --git a/cudaqlib/gse/adapt/adapt.h b/cudaqlib/gse/adapt/adapt.h index bc39c54..51d9f40 100644 --- a/cudaqlib/gse/adapt/adapt.h +++ b/cudaqlib/gse/adapt/adapt.h @@ -8,10 +8,9 @@ #pragma once -#include "cudaqlib/gse/vqe/vqe.h" #include "cudaq/qis/qubit_qis.h" - -#include "operator_pool.h" +#include "cudaqlib/gse/utils/operator_pool.h" +#include "cudaqlib/gse/vqe/vqe.h" namespace cudaq::gse { @@ -36,7 +35,7 @@ struct AdaptKernel { template auto adapt_vqe(const InitialState &initialState, const spin_op &H, - const operator_pool &operatorPool, optim::optimizer &optimizer, + const std::vector &poolList, optim::optimizer &optimizer, observe_gradient *gradient, const adapt_options options = adapt_options()) { @@ -44,7 +43,6 @@ auto adapt_vqe(const InitialState &initialState, const spin_op &H, std::vector trotterList; std::vector thetas; auto numQubits = H.num_qubits(); - auto poolList = operatorPool.generate(); double energy = 0.0, lastNorm = std::numeric_limits::max(); // Compute the [H,Oi] @@ -114,14 +112,14 @@ auto adapt_vqe(const InitialState &initialState, const spin_op &H, template auto adapt_vqe(const InitialState &initialState, const spin_op &H, - const operator_pool &operatorPool, optim::optimizer &optimizer, + const std::vector &operatorPool, optim::optimizer &optimizer, const adapt_options options = adapt_options()) { return adapt_vqe(initialState, H, operatorPool, optimizer, nullptr, options); } template auto adapt_vqe(const InitialState &initialState, const spin_op &H, - const operator_pool &operatorPool, + const std::vector &operatorPool, const adapt_options options = adapt_options()) { cudaq::optim::cobyla optimizer; return adapt_vqe(initialState, H, operatorPool, optimizer, options); diff --git a/cudaqlib/gse/adapt/operator_pool.h b/cudaqlib/gse/adapt/operator_pool.h deleted file mode 100644 index 375bb34..0000000 --- a/cudaqlib/gse/adapt/operator_pool.h +++ /dev/null @@ -1,40 +0,0 @@ -/****************************************************************-*- C++ -*-**** - * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates. * - * All rights reserved. * - * * - * This source code and the accompanying materials are made available under * - * the terms of the Apache License 2.0 which accompanies this distribution. * - ******************************************************************************/ - -#pragma once - -#include "cudaq/spin_op.h" - -namespace cudaq::gse { -class operator_pool { -private: - std::vector pool; - -public: - operator_pool() = default; - operator_pool(const std::vector &opPool) : pool(opPool) {} - virtual std::vector generate() const { return pool; } -}; - -namespace fermion { -inline spin_op adag(std::size_t numQubits, std::size_t j) { - spin_op zprod(numQubits); - for (std::size_t k = 0; k < j; k++) - zprod *= spin::z(k); - return 0.5 * zprod * (spin::x(j) - std::complex{0, 1} * spin::y(j)); -} - -inline spin_op a(std::size_t numQubits, std::size_t j) { - spin_op zprod(numQubits); - for (std::size_t k = 0; k < j; k++) - zprod *= spin::z(k); - return 0.5 * zprod * (spin::x(j) + std::complex{0, 1} * spin::y(j)); -} -} // namespace fermion - -} // namespace cudaq::gse diff --git a/cudaqlib/gse/adapt/pools/spin_complement_gsd.h b/cudaqlib/gse/adapt/pools/spin_complement_gsd.h deleted file mode 100644 index 286370f..0000000 --- a/cudaqlib/gse/adapt/pools/spin_complement_gsd.h +++ /dev/null @@ -1,152 +0,0 @@ -#pragma once - -#include "../operator_pool.h" - -namespace cudaq::gse { - -// adapted from -// https://github.com/mayhallgroup/adapt-vqe/blob/master/src/operator_pools.py - -class spin_complement_gsd : public operator_pool { -protected: - std::int64_t numOrbitals; - -public: - spin_complement_gsd(std::int64_t q) : numOrbitals(q) {} - std::vector generate() const override { - std::vector pool; - auto numQubits = 2 * numOrbitals; - std::vector alphaOrbs, betaOrbs; - for (auto i : range(numOrbitals)) { - alphaOrbs.push_back(2 * i); - betaOrbs.push_back(alphaOrbs.back() + 1); - } - - for (auto p : alphaOrbs) { - for (auto q : alphaOrbs) { - if (p >= q) - continue; - auto oneElectron = - fermion::adag(numQubits, q) * fermion::a(numQubits, p) - - fermion::adag(numQubits, p) * fermion::a(numQubits, q); - oneElectron += - fermion::adag(numQubits, q + 1) * fermion::a(numQubits, p + 1) - - fermion::adag(numQubits, p + 1) * fermion::a(numQubits, q + 1); - - std::unordered_map> terms; - oneElectron.for_each_term([&](spin_op &term) { - auto coeff = term.get_coefficient(); - if (std::fabs(coeff.real()) < 1e-12 && - std::fabs(coeff.imag()) < 1e-12) - return; - - if (std::fabs(coeff.real()) < 1e-12) - terms.insert({term.get_raw_data().first[0], - std::complex{0., coeff.imag()}}); - }); - - if (!terms.empty()) - pool.emplace_back(terms); - } - } - - int pq = 0; - for (auto p : alphaOrbs) { - for (auto q : alphaOrbs) { - if (p > q) - continue; - - int rs = 0; - for (auto r : alphaOrbs) { - for (auto s : alphaOrbs) { - if (r > s) - continue; - if (pq < rs) - continue; - - auto twoElectron = - fermion::adag(numQubits, r) * fermion::a(numQubits, p) * - fermion::adag(numQubits, s) * fermion::a(numQubits, q) - - fermion::adag(numQubits, q) * fermion::a(numQubits, s) * - fermion::adag(numQubits, p) * fermion::a(numQubits, r); - twoElectron += - fermion::adag(numQubits, r + 1) * fermion::a(numQubits, p + 1) * - fermion::adag(numQubits, s + 1) * - fermion::a(numQubits, q + 1) - - fermion::adag(numQubits, q + 1) * fermion::a(numQubits, s + 1) * - fermion::adag(numQubits, p + 1) * - fermion::a(numQubits, r + 1); - - std::unordered_map> - terms; - twoElectron.for_each_term([&](spin_op &term) { - auto coeff = term.get_coefficient(); - if (std::fabs(coeff.real()) < 1e-12 && - std::fabs(coeff.imag()) < 1e-12) - return; - - if (std::fabs(coeff.real()) < 1e-12) - terms.insert({term.get_raw_data().first[0], - std::complex{0., coeff.imag()}}); - }); - - if (!terms.empty()) - pool.push_back(terms); - rs++; - } - } - pq++; - } - } - - pq = 0; - for (auto p : alphaOrbs) { - for (auto q : betaOrbs) { - - int rs = 0; - for (auto r : alphaOrbs) { - for (auto s : betaOrbs) { - - if (pq < rs) - continue; - - auto twoElectron = - fermion::adag(numQubits, r) * fermion::a(numQubits, p) * - fermion::adag(numQubits, s) * fermion::a(numQubits, q) - - fermion::adag(numQubits, q) * fermion::a(numQubits, s) * - fermion::adag(numQubits, p) * fermion::a(numQubits, r); - if (p > q) - continue; - - twoElectron += - fermion::adag(numQubits, s - 1) * fermion::a(numQubits, q - 1) * - fermion::adag(numQubits, r + 1) * - fermion::a(numQubits, p + 1) - - fermion::adag(numQubits, p + 1) * fermion::a(numQubits, r + 1) * - fermion::adag(numQubits, q - 1) * - fermion::a(numQubits, s - 1); - std::unordered_map> - terms; - twoElectron.for_each_term([&](spin_op &term) { - auto coeff = term.get_coefficient(); - if (std::fabs(coeff.real()) < 1e-12 && - std::fabs(coeff.imag()) < 1e-12) - return; - - if (std::fabs(coeff.real()) < 1e-12) - terms.insert({term.get_raw_data().first[0], - std::complex{0., coeff.imag()}}); - }); - if (!terms.empty()) - pool.push_back(terms); - rs++; - } - } - pq++; - } - } - - return pool; - } -}; -} // namespace cudaq::gse \ No newline at end of file diff --git a/cudaqlib/gse/utils/CMakeLists.txt b/cudaqlib/gse/utils/CMakeLists.txt new file mode 100644 index 0000000..c4c5847 --- /dev/null +++ b/cudaqlib/gse/utils/CMakeLists.txt @@ -0,0 +1,9 @@ +# ============================================================================ # +# Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates. # +# All rights reserved. # +# # +# This source code and the accompanying materials are made available under # +# the terms of the Apache License 2.0 which accompanies this distribution. # +# ============================================================================ # + +add_subdirectory(pools) \ No newline at end of file diff --git a/cudaqlib/gse/utils/operator_pool.h b/cudaqlib/gse/utils/operator_pool.h new file mode 100644 index 0000000..212b068 --- /dev/null +++ b/cudaqlib/gse/utils/operator_pool.h @@ -0,0 +1,70 @@ +/****************************************************************-*- C++ -*-**** + * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#pragma once + +#include +#include +#include + +#include "cudaq/spin_op.h" +#include "cudaqlib/utils/extension_point.h" + +namespace cudaq { +inline std::optional::const_iterator> +findIter(const std::vector &possibleNames, + const std::unordered_map &m) { + for (auto &name : possibleNames) { + auto iter = m.find(name); + if (iter != m.end()) + return iter; + } + return std::nullopt; +} + +inline std::size_t getIntLike(const std::any &any) { + try { + return std::any_cast(any); + } catch (...) { + // If this throws then we'll just have an error + return std::any_cast(any); + } +} + +class operator_pool : public extension_point { +public: + operator_pool() = default; + virtual std::vector + generate(const std::unordered_map &config) const = 0; +}; + +CUDAQ_DEFINE_EXTENSION_IMPL(operator_pool) + +#define CUDAQ_REGISTER_OPERATOR_POOL(TYPE) \ + static inline const std::string class_identifier = #TYPE; \ + static std::unique_ptr create() { \ + return std::make_unique(); \ + } + +namespace fermion { +inline spin_op adag(std::size_t numQubits, std::size_t j) { + spin_op zprod(numQubits); + for (std::size_t k = 0; k < j; k++) + zprod *= spin::z(k); + return 0.5 * zprod * (spin::x(j) - std::complex{0, 1} * spin::y(j)); +} + +inline spin_op a(std::size_t numQubits, std::size_t j) { + spin_op zprod(numQubits); + for (std::size_t k = 0; k < j; k++) + zprod *= spin::z(k); + return 0.5 * zprod * (spin::x(j) + std::complex{0, 1} * spin::y(j)); +} +} // namespace fermion + +} // namespace cudaq diff --git a/cudaqlib/gse/utils/pools/CMakeLists.txt b/cudaqlib/gse/utils/pools/CMakeLists.txt new file mode 100644 index 0000000..f8ee5ec --- /dev/null +++ b/cudaqlib/gse/utils/pools/CMakeLists.txt @@ -0,0 +1,26 @@ +# ============================================================================ # +# Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. # +# All rights reserved. # +# # +# This source code and the accompanying materials are made available under # +# the terms of the Apache License 2.0 which accompanies this distribution. # +# ============================================================================ # + +add_compile_options(-Wno-attributes) +set(LIBRARY_NAME cudaq-operator-pools) +add_library(${LIBRARY_NAME} SHARED uccsd_pool.cpp spin_complement_gsd.cpp) +target_include_directories( + ${LIBRARY_NAME} + PUBLIC $) +target_link_libraries(${LIBRARY_NAME} PRIVATE cudaq::cudaq-spin) + +install( + TARGETS ${LIBRARY_NAME} + EXPORT ${LIBRARY_NAME}Targets + DESTINATION lib) + +install( + EXPORT ${LIBRARY_NAME}Targets + FILE ${LIBRARY_NAME}Targets.cmake + NAMESPACE cudaq:: + DESTINATION lib/cmake/gse) \ No newline at end of file diff --git a/cudaqlib/gse/utils/pools/spin_complement_gsd.cpp b/cudaqlib/gse/utils/pools/spin_complement_gsd.cpp new file mode 100644 index 0000000..144dc50 --- /dev/null +++ b/cudaqlib/gse/utils/pools/spin_complement_gsd.cpp @@ -0,0 +1,155 @@ +/****************************************************************-*- C++ -*-**** + * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#include "spin_complement_gsd.h" + +namespace cudaq { + +std::vector spin_complement_gsd::generate( + const std::unordered_map &config) const { + auto iter = config.find("num-orbitals"); + if (iter == config.end()) + throw std::runtime_error("spin_complement_gsd requires num-orbitals config " + "parameter of type std::size_t."); + + auto numOrbitals = std::any_cast(iter->second); + + + std::vector pool; + auto numQubits = 2 * numOrbitals; + std::vector alphaOrbs, betaOrbs; + for (auto i : range(numOrbitals)) { + alphaOrbs.push_back(2 * i); + betaOrbs.push_back(alphaOrbs.back() + 1); + } + + for (auto p : alphaOrbs) { + for (auto q : alphaOrbs) { + if (p >= q) + continue; + auto oneElectron = + fermion::adag(numQubits, q) * fermion::a(numQubits, p) - + fermion::adag(numQubits, p) * fermion::a(numQubits, q); + oneElectron += + fermion::adag(numQubits, q + 1) * fermion::a(numQubits, p + 1) - + fermion::adag(numQubits, p + 1) * fermion::a(numQubits, q + 1); + + std::unordered_map> terms; + oneElectron.for_each_term([&](spin_op &term) { + auto coeff = term.get_coefficient(); + if (std::fabs(coeff.real()) < 1e-12 && std::fabs(coeff.imag()) < 1e-12) + return; + + if (std::fabs(coeff.real()) < 1e-12) + terms.insert({term.get_raw_data().first[0], + std::complex{0., coeff.imag()}}); + }); + + if (!terms.empty()) + pool.emplace_back(terms); + } + } + + int pq = 0; + for (auto p : alphaOrbs) { + for (auto q : alphaOrbs) { + if (p > q) + continue; + + int rs = 0; + for (auto r : alphaOrbs) { + for (auto s : alphaOrbs) { + if (r > s) + continue; + if (pq < rs) + continue; + + auto twoElectron = + fermion::adag(numQubits, r) * fermion::a(numQubits, p) * + fermion::adag(numQubits, s) * fermion::a(numQubits, q) - + fermion::adag(numQubits, q) * fermion::a(numQubits, s) * + fermion::adag(numQubits, p) * fermion::a(numQubits, r); + twoElectron += + fermion::adag(numQubits, r + 1) * fermion::a(numQubits, p + 1) * + fermion::adag(numQubits, s + 1) * + fermion::a(numQubits, q + 1) - + fermion::adag(numQubits, q + 1) * fermion::a(numQubits, s + 1) * + fermion::adag(numQubits, p + 1) * + fermion::a(numQubits, r + 1); + + std::unordered_map> terms; + twoElectron.for_each_term([&](spin_op &term) { + auto coeff = term.get_coefficient(); + if (std::fabs(coeff.real()) < 1e-12 && + std::fabs(coeff.imag()) < 1e-12) + return; + + if (std::fabs(coeff.real()) < 1e-12) + terms.insert({term.get_raw_data().first[0], + std::complex{0., coeff.imag()}}); + }); + + if (!terms.empty()) + pool.push_back(terms); + rs++; + } + } + pq++; + } + } + + pq = 0; + for (auto p : alphaOrbs) { + for (auto q : betaOrbs) { + + int rs = 0; + for (auto r : alphaOrbs) { + for (auto s : betaOrbs) { + + if (pq < rs) + continue; + + auto twoElectron = + fermion::adag(numQubits, r) * fermion::a(numQubits, p) * + fermion::adag(numQubits, s) * fermion::a(numQubits, q) - + fermion::adag(numQubits, q) * fermion::a(numQubits, s) * + fermion::adag(numQubits, p) * fermion::a(numQubits, r); + if (p > q) + continue; + + twoElectron += + fermion::adag(numQubits, s - 1) * fermion::a(numQubits, q - 1) * + fermion::adag(numQubits, r + 1) * + fermion::a(numQubits, p + 1) - + fermion::adag(numQubits, p + 1) * fermion::a(numQubits, r + 1) * + fermion::adag(numQubits, q - 1) * + fermion::a(numQubits, s - 1); + std::unordered_map> terms; + twoElectron.for_each_term([&](spin_op &term) { + auto coeff = term.get_coefficient(); + if (std::fabs(coeff.real()) < 1e-12 && + std::fabs(coeff.imag()) < 1e-12) + return; + + if (std::fabs(coeff.real()) < 1e-12) + terms.insert({term.get_raw_data().first[0], + std::complex{0., coeff.imag()}}); + }); + if (!terms.empty()) + pool.push_back(terms); + rs++; + } + } + pq++; + } + } + + return pool; +} + +} // namespace cudaq \ No newline at end of file diff --git a/cudaqlib/gse/utils/pools/spin_complement_gsd.h b/cudaqlib/gse/utils/pools/spin_complement_gsd.h new file mode 100644 index 0000000..9176b85 --- /dev/null +++ b/cudaqlib/gse/utils/pools/spin_complement_gsd.h @@ -0,0 +1,26 @@ +/****************************************************************-*- C++ -*-**** + * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#pragma once + +#include "../operator_pool.h" + +namespace cudaq { + +// adapted from +// https://github.com/mayhallgroup/adapt-vqe/blob/master/src/operator_pools.py + +class spin_complement_gsd + : public details::operator_pool_impl { + +public: + std::vector generate( + const std::unordered_map &config) const override; + CUDAQ_REGISTER_OPERATOR_POOL(spin_complement_gsd) +}; +} // namespace cudaq \ No newline at end of file diff --git a/cudaqlib/gse/utils/pools/uccsd_pool.cpp b/cudaqlib/gse/utils/pools/uccsd_pool.cpp new file mode 100644 index 0000000..b15f57d --- /dev/null +++ b/cudaqlib/gse/utils/pools/uccsd_pool.cpp @@ -0,0 +1,71 @@ +/****************************************************************-*- C++ -*-**** + * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#include "uccsd_pool.h" +#include "cudaqlib/kernels/uccsd.h" +namespace cudaq { + +std::vector +uccsd::generate(const std::unordered_map &config) const { + auto iter = findIter({"num-qubits", "num_qubits"}, + config); // config.find("num-qubits"); + if (!iter.has_value()) + throw std::runtime_error("uccsd_pool requires num-qubits config " + "parameter of type std::size_t."); + auto numQubits = getIntLike(iter.value()->second); + + iter = findIter({"num-electrons", "num_electrons"}, config); + if (!iter.has_value()) + throw std::runtime_error("uccsd_pool requires num-electrons config " + "parameter of type std::size_t."); + auto numElectrons = getIntLike(iter.value()->second); + + auto [singlesAlpha, singlesBeta, doublesMixed, doublesAlpha, doublesBeta] = + cudaq::get_uccsd_excitations(numElectrons, numQubits); + + std::vector ops; + using namespace cudaq::spin; + + auto addSinglesExcitation = [numQubits](std::vector &ops, + std::size_t p, std::size_t q) { + spin_op o(numQubits); + ops.emplace_back(o * spin::y(p) * spin::x(q)); + ops.emplace_back(o * spin::x(p) * spin::y(q)); + }; + + auto addDoublesExcitation = [numQubits](std::vector &ops, + std::size_t p, std::size_t q, + std::size_t r, std::size_t s) { + spin_op o(numQubits); + ops.emplace_back(o * spin::x(p) * spin::x(q) * spin::x(r) * spin::y(s)); + ops.emplace_back(o * spin::x(p) * spin::x(q) * spin::y(r) * spin::x(s)); + ops.emplace_back(o * spin::x(p) * spin::y(q) * spin::y(r) * spin::y(s)); + ops.emplace_back(o * spin::x(p) * spin::y(q) * spin::x(r) * spin::x(s)); + + ops.emplace_back(o * spin::y(p) * spin::x(q) * spin::x(r) * spin::x(s)); + ops.emplace_back(o * spin::y(p) * spin::x(q) * spin::y(r) * spin::y(s)); + ops.emplace_back(o * spin::y(p) * spin::y(q) * spin::x(r) * spin::y(s)); + ops.emplace_back(o * spin::y(p) * spin::y(q) * spin::y(r) * spin::x(s)); + }; + + for (auto &sa : singlesAlpha) + addSinglesExcitation(ops, sa[0], sa[1]); + for (auto &sa : singlesBeta) + addSinglesExcitation(ops, sa[0], sa[1]); + + for (auto &d : doublesMixed) + addDoublesExcitation(ops, d[0], d[1], d[2], d[3]); + for (auto &d : doublesAlpha) + addDoublesExcitation(ops, d[0], d[1], d[2], d[3]); + for (auto &d : doublesBeta) + addDoublesExcitation(ops, d[0], d[1], d[2], d[3]); + + return ops; +} + +} // namespace cudaq \ No newline at end of file diff --git a/cudaqlib/gse/utils/pools/uccsd_pool.h b/cudaqlib/gse/utils/pools/uccsd_pool.h new file mode 100644 index 0000000..8c4523f --- /dev/null +++ b/cudaqlib/gse/utils/pools/uccsd_pool.h @@ -0,0 +1,23 @@ +/****************************************************************-*- C++ -*-**** + * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#pragma once + +#include "../operator_pool.h" + +namespace cudaq { + +class uccsd : public details::operator_pool_impl { + +public: + std::vector generate( + const std::unordered_map &config) const override; + CUDAQ_REGISTER_OPERATOR_POOL(uccsd) +}; + +} // namespace cudaq \ No newline at end of file diff --git a/cudaqlib/operators/chemistry/molecule.cpp b/cudaqlib/operators/chemistry/molecule.cpp index 17a6dc2..3e2a7e2 100644 --- a/cudaqlib/operators/chemistry/molecule.cpp +++ b/cudaqlib/operators/chemistry/molecule.cpp @@ -9,7 +9,7 @@ #include "MoleculePackageDriver.h" #include -#include +#include namespace cudaq::operators { @@ -84,8 +84,10 @@ void molecule_options::dump() { std::cout << "\tsymmetry: " << symmetry << "\n"; std::cout << "\tcycles: " << cycles << "\n"; std::cout << "\tinitguess: " << initguess << "\n"; - std::cout << "\tnele_cas: " << nele_cas.value() << "\n"; - std::cout << "\tnorb_cas: " << norb_cas.value() << "\n"; + std::cout << "\tnele_cas: " << (nele_cas.has_value() ? nele_cas.value() + : -1) << "\n"; + std::cout << "\tnorb_cas: " << (norb_cas.has_value() ? norb_cas.value() + : -1) << "\n"; std::cout << "\tUR: " << std::boolalpha << UR << "\n"; std::cout << "\tMP2: " << std::boolalpha << MP2 << "\n"; std::cout << "\tnatorb: " << std::boolalpha << natorb << "\n"; diff --git a/cudaqlib/utils/extension_point.h b/cudaqlib/utils/extension_point.h index dc65aa9..e4a9d41 100644 --- a/cudaqlib/utils/extension_point.h +++ b/cudaqlib/utils/extension_point.h @@ -27,7 +27,7 @@ class extension_point { auto ®istry = get_registry(); auto iter = registry.find(name); if (iter == registry.end()) - throw std::runtime_error("error"); + throw std::runtime_error("Cannot find extension with name = " + name); return iter->second(std::forward(args)...); } diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 445db99..efc4a44 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -43,7 +43,9 @@ target_link_libraries(_pycudaqlib PRIVATE cudaq::cudaq cudaq-operators - cudaq-optim) + cudaq-optim + cudaq-operator-pools + ) if (NOT SKBUILD) install(DIRECTORY cudaqlib DESTINATION python/cudaqlib) diff --git a/python/bindings/gse/py_gse.cpp b/python/bindings/gse/py_gse.cpp index 99ae94b..1d69e5e 100644 --- a/python/bindings/gse/py_gse.cpp +++ b/python/bindings/gse/py_gse.cpp @@ -10,10 +10,13 @@ #include #include -#include "cudaqlib/gse.h" +#include "cudaqlib/gse/utils/operator_pool.h" + #include "../utils/type_casters.h" #include "../utils/utils.h" +#include "cudaqlib/gse.h" + namespace py = pybind11; namespace cudaq::gse { @@ -21,6 +24,16 @@ namespace cudaq::gse { void bindGse(py::module &mod) { auto gse = mod.def_submodule("gse"); + gse.def("get_operator_pool", [](const std::string &name, py::kwargs config) { + std::unordered_map asCpp; + for (auto &[k, v] : config) { + std::string asStr = k.cast(); + + if (py::isinstance(v)) + asCpp.insert({k.cast(), v.cast()}); + } + return operator_pool::get(name)->generate(asCpp); + }); gse.def( "vqe", @@ -32,8 +45,7 @@ void bindGse(py::module &mod) { optOptions.max_iterations = cudaq::getValueOr( options, "max_iterations", -1); // Or is just a dummy - optOptions.verbose = - cudaq::getValueOr(options, "verbose", false); + optOptions.verbose = cudaq::getValueOr(options, "verbose", false); auto optimizerName = cudaq::getValueOr(options, "optimizer", "cobyla"); auto optimizer = cudaq::optim::optimizer::get(optimizerName); @@ -47,7 +59,7 @@ void bindGse(py::module &mod) { } auto gradientName = cudaq::getValueOr(options, "gradient", - "parameter_shift"); + "parameter_shift"); auto gradient = cudaq::optim::observe_gradient::get(gradientName); gradient->set_parameterized_kernel(kernelWrapper); gradient->set_spin_op(op); diff --git a/python/cudaqlib/__init__.py b/python/cudaqlib/__init__.py index e55ac27..c82b93e 100644 --- a/python/cudaqlib/__init__.py +++ b/python/cudaqlib/__init__.py @@ -10,3 +10,4 @@ import os, os.path from ._pycudaqlib import * from .kernels import * +from .algorithms.gqe import gqe \ No newline at end of file diff --git a/python/cudaqlib/algorithms/gqe.py b/python/cudaqlib/algorithms/gqe.py new file mode 100644 index 0000000..84bc161 --- /dev/null +++ b/python/cudaqlib/algorithms/gqe.py @@ -0,0 +1,358 @@ +from .transformer import Transformer +import torch +import lightning as L +from abc import ABC, abstractmethod +import math, os, json, sys +from torch.utils.data import Dataset +import matplotlib.pyplot as plt +import numpy as np +import time + +torch.set_float32_matmul_precision('high') + + +def key(distance): + v = str(distance).replace(".", "_") + return v + + +def plot_trajectory(trajectory_path): + data = [] + with open(trajectory_path, 'r') as file: + for line in file: + data.append(json.loads(line)) + + iterations = [] + losses = [] + avg_energies = [] + min_energies = [] + global_min = 0 + for entry in data: + iterations.append(entry['iter']) + losses.append(entry['loss']) + avg_energies.append(sum(entry['energies']) / len(entry['energies'])) + if min(entry['energies']) < global_min: + global_min = min(entry['energies']) + min_energies.append(global_min) + + plt.figure(figsize=(12, 8)) + + plt.subplot(3, 1, 1) + plt.plot(iterations, losses, marker='o', linestyle='-', color='b') + plt.title('Loss per Epoch') + plt.xlabel('Iteration') + plt.ylabel('Loss') + + plt.subplot(3, 1, 2) + plt.plot(iterations, avg_energies, marker='o', linestyle='-', color='g') + plt.title('Average Energy per Epoch') + plt.xlabel('Iteration') + plt.ylabel('Average Energy') + + plt.subplot(3, 1, 3) + plt.plot(iterations, min_energies, marker='o', linestyle='-', color='r') + plt.title('Minimum Energy per Epoch') + plt.xlabel('Iteration') + plt.ylabel('Minimum Energy') + + plt.tight_layout() + plt.savefig('trajectory_plot.png') + + +def pretrain_file(cfg): + return f"{cfg.save_dir}{cfg.name}_{cfg.seed}_checkpoint_pretrain.ckpt" + + +def train_file(cfg): + return f"{cfg.save_dir}{cfg.name}_{cfg.seed}.txt" + + +def trajectory_file(cfg, distance): + return f"{cfg.save_dir}{cfg.name}_trajectory_{distance}.ckpt" + + +def image_file(cfg, errors): + suffix = "" + if errors is not None: + suffix = "-detail" + return f"{cfg.save_dir}result-{cfg.name}{suffix}.pdf" + + +def ground_state_file(cfg): + return f"{cfg.save_dir}gs_{cfg.name}.txt" + + +def random_file(cfg, seed): + return f'{cfg.save_dir}{cfg.name}_random_{seed}.txt' + + +class TemperatureScheduler(ABC): + + @abstractmethod + def get(self, iter): + pass + + +class DefaultScheduler(TemperatureScheduler): + + def __init__(self, start, delta) -> None: + self.start = start + self.delta = delta + + def get(self, iter): + return self.start + iter * self.delta + + +class CosineScheduler(TemperatureScheduler): + + def __init__(self, minimum, maximum, frequency) -> None: + self.minimum = minimum + self.maximum = maximum + self.frequency = frequency + + def get(self, iter): + return (self.maximum + self.minimum) / 2 - ( + self.maximum - self.minimum) / 2 * math.cos( + 2 * math.pi * iter / self.frequency) + + +class TrajectoryData: + + def __init__(self, iter_num, loss, indices, energies): + self.iter_num = iter_num + self.loss = loss + self.indices = indices + self.energies = energies + + def to_json(self): + map = { + "iter": self.iter_num, + "loss": self.loss, + "indices": self.indices, + "energies": self.energies + } + return json.dumps(map) + + @classmethod + def from_json(self, string): + if string.startswith('"'): + string = string[1:len(string) - 1] + string = string.replace("\\", "") + map = json.loads(string) + return TrajectoryData(map["iter"], map["loss"], map["indices"], + map["energies"]) + + +class EnergyDataset(Dataset): + + def __init__(self, file_paths, threshold=sys.maxsize): + tensor_x = [] + tensor_y = [] + self.min_energy = sys.maxsize + self.min_indices = None + for file_path in file_paths: + with open(file_path) as f: + for l in f.readlines(): + data = TrajectoryData.from_json(l.rstrip()) + for indices, energy in zip(data.indices, data.energies): + if threshold < energy: + continue + if self.min_energy > energy: + self.min_energy = energy + self.min_indices = indices + result = [0] + result.extend(indices) + tensor_x.append(result) + tensor_y.append(energy) + self.tensors = (torch.tensor(tensor_x, dtype=torch.int64), + torch.tensor(tensor_y, dtype=torch.float)) + + def __getitem__(self, index): + result = self.tensors[0][index], self.tensors[1][index] + return result + + def __len__(self): + return self.tensors[0].size(0) + + +class FileMonitor: + + def __init__(self): + self.lines = [] + + def record(self, iter_num, loss, energies, indices): + energies = energies.cpu().numpy().tolist() + indices = indices.cpu().numpy().tolist() + data = TrajectoryData(iter_num, loss.item(), indices, energies) + self.lines.append(data.to_json()) + + def save(self, path): + with open(path, 'w') as f: + for l in self.lines: + f.write(f"{l}\n") + + +class GPTQETaskBase(): + + def __init__(self, + temperature_scheduler: TemperatureScheduler = None) -> None: + self.temperature_scheduler = temperature_scheduler + + # def train(self, cfg, model, optimizer, numQPUs=1): + # self.numQPUs = numQPUs + + # # Not yet... + # # from mpi4py import MPI + + # # Initialize MPI + # # self.mpiComm = MPI.COMM_WORLD + # # self.mpiRank = self.mpiComm.Get_rank() + # # self.mpiSize = self.mpiComm.Get_size() + # fabric = L.Fabric(accelerator="auto", devices=1) + # fabric.seed_everything(cfg.seed) + # fabric.launch() + + # min_indices_dict = {} + # distances = cfg.distances + # filename = train_file(cfg) + # m = {} + # if os.path.exists(filename): + # with open(filename) as f: + # for l in f.readlines(): + # items = l.rstrip().split('\t') + # if len(items) != 2: + # continue + # distance, energy = items + # distance = float(distance) + # energy = float(energy) + # m[distance] = energy + # with open(filename, 'w') as f: + # for distance in distances: + # print("distance:", distance) + # if cfg.cache and distance in m: + # min_energy = m[distance] + # print("already computed, skipped", distance) + # else: + # indices, min_energy = self.do_train(cfg, distance, fabric, model, optimizer) + # min_indices_dict[str(distance)] = indices + # f.write(f"{distance}\t{min_energy}\n") + # fabric.log('circuit', json.dumps(min_indices_dict)) + # # if (self.mpiComm): + # # MPI.Finalize() + # return min_indices_dict + + def train(self, cfg, model, optimizer): + fabric = L.Fabric(accelerator="auto") + fabric.seed_everything(cfg.seed) + fabric.launch() + + min_indices, energy = self.do_train(cfg, fabric, model, + optimizer) + fabric.log('circuit', json.dumps(min_indices)) + return min_indices, energy + + def do_train(self, cfg, fabric, model, optimizer): + print(cfg) + monitor = FileMonitor() + model, optimizer = fabric.setup(model, optimizer) + model.mark_forward_method(model.train_step) + + # if key(distance) in cfg.check_points: + # print("loaded from the checkpoint") + # cp = fabric.load(cfg.check_points[key(distance)]) + # model.load_state_dict(cp["model"]) + # optimizer.load_state_dict(cp["optimizer"]) + pytorch_total_params = sum( + p.numel() for p in model.parameters() if p.requires_grad) + print(f"total trainable params: {pytorch_total_params / 1e6:.2f}M") + min_energy = sys.maxsize + min_indices = None + for epoch in range(cfg.max_iters): + optimizer.zero_grad() + l = None + start = time.time() + loss, energies, indices, log_values = model.train_step() + # numQPUs=self.numQPUs, comm=self.mpiComm) + print('epoch', epoch, 'model.train_step time:', time.time() - start, torch.min(energies)) + if l is None: + l = loss + else: + l = l + loss + monitor.record(epoch, loss, energies, indices) + for e, indices in zip(energies, indices): + energy = e.item() + if energy < min_energy: + min_energy = e.item() + min_indices = indices + log_values[f"min_energy at"] = min_energy + log_values[f"temperature at"] = model.temperature + if cfg.verbose: + print(f"energies: {energies}") + print(f"temperature: {model.temperature}") + fabric.log_dict(log_values, step=epoch) + fabric.backward(loss) + fabric.clip_gradients(model, optimizer, max_norm=cfg.grad_norm_clip) + optimizer.step() + # scheduler.step() + if self.temperature_scheduler is not None: + model.temperature = self.temperature_scheduler.get(epoch) + else: + model.temperature += cfg.del_temperature + model.set_cost(None) + # state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} + # fabric.save(cfg.save_dir + f"checkpoint_{distance}.ckpt", state) + # if cfg.save_data: + # monitor.save(trajectory_file(cfg, distance)) + # if cfg.plot: + # plot_trajectory( + # str(cfg.save_dir) + str(cfg.name) + "_trajectory_" + + # str(distance) + ".ckpt") + indices = min_indices.cpu().numpy().tolist() + return min_energy, indices + + +def get_default_chemistry_configs(): + from ml_collections import ConfigDict + cfg = ConfigDict() + cfg.verbose = False + cfg.save_data = False + cfg.print_exact = True + cfg.name = "gptqe" + cfg.warmup_steps = 100 + cfg.num_samples = 5 # akin to batch size + cfg.max_iters = 100 + cfg.nqubit = 4 + cfg.ngates = 20 + cfg.nshot = 0 + cfg.seed = 3047 + cfg.distances = [0.5, 0.6, 0.7, 0.7414, 0.8, 0.9, 1.0, 1.5, + 2.0] # choices of the distance between two atoms + cfg.time_pool = [ + 0.003125, -0.003125, 0.00625, -0.00625, 0.0125, -0.0125, 0.025, -0.025, + 0.05, -0.05, 0.1, -0.1 + ] + cfg.time_factor = 1.0 + cfg.lr = 5e-7 + cfg.energy_offset = 0.0 + cfg.grad_norm_clip = 1.0 + cfg.temperature = 5.0 + cfg.del_temperature = 0.05 + cfg.resid_pdrop = 0.0 + cfg.embd_pdrop = 0.0 + cfg.attn_pdrop = 0.0 + cfg.check_points = {} + cfg.dry = False + cfg.small = False + cfg.cache = True + cfg.save_dir = "./output/" + return cfg + + +def gqe(cost, **kwargs): + cfg = get_default_chemistry_configs() + cfg.vocab_size = 12 + distance = .7474 + model = Transformer(cfg, distance, cost, loss='exp') + optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr) + gqe = GPTQETaskBase() + return gqe.train(cfg, model, optimizer)#, numQPUs=1) diff --git a/python/cudaqlib/algorithms/loss.py b/python/cudaqlib/algorithms/loss.py new file mode 100644 index 0000000..9520f46 --- /dev/null +++ b/python/cudaqlib/algorithms/loss.py @@ -0,0 +1,53 @@ +from abc import abstractmethod, ABC +import torch + + +class LogitMatchingLoss(ABC): + + @abstractmethod + def compute(self, energies, logits_tensor, temperature, log_values): + pass + + +class ExpLogitMatching(LogitMatchingLoss): + + def __init__(self, energy_offset, label) -> None: + self._label = label + self.energy_offset = energy_offset + self.loss_fn = torch.nn.MSELoss() + + def compute(self, energies, logits_tensor, temperature, log_values): + mean_logits = torch.mean(logits_tensor, 1) + log_values[f"mean_logits at {self._label}"] = torch.mean( + mean_logits - self.energy_offset) + log_values[f"mean energy at {self._label}"] = torch.mean(energies) + mean_logits = torch.mean(logits_tensor, 1) + device = mean_logits.device + return self.loss_fn( + torch.exp(-mean_logits), + torch.exp(-energies.to(device) - self.energy_offset)) + + +class GFlowLogitMatching(LogitMatchingLoss): + + def __init__(self, energy_offset, device, label, nn: torch.nn) -> None: + self._label = label + self.loss_fn = torch.nn.MSELoss() + self.energy_offset = energy_offset + self.normalization = 10**-5 + self.param = torch.nn.Parameter(torch.tensor([0.0]).to(device)) + nn.register_parameter(name="energy_offset", param=self.param) + + def compute(self, energies, logits_tensor, temperature, log_values): + mean_logits = torch.mean(logits_tensor, 1) + energy_offset = self.energy_offset + self.param / self.normalization + log_values[f"energy_offset at {self._label}"] = energy_offset + log_values[f"mean_logits at {self._label}"] = torch.mean(mean_logits - + energy_offset) + log_values[f"mean energy at {self._label}"] = torch.mean(energies) + mean_logits = torch.mean(logits_tensor, 1) + device = mean_logits.device + loss = self.loss_fn( + torch.exp(-mean_logits), + torch.exp(-(energies.to(device) + energy_offset.to(device)))) + return loss diff --git a/python/cudaqlib/algorithms/transformer.py b/python/cudaqlib/algorithms/transformer.py new file mode 100644 index 0000000..08b3c0b --- /dev/null +++ b/python/cudaqlib/algorithms/transformer.py @@ -0,0 +1,106 @@ +import torch +from torch.nn import functional as F +from transformers import GPT2LMHeadModel, GPT2Config +from lightning import LightningModule +from .loss import ExpLogitMatching, GFlowLogitMatching + +def get_device(): + if torch.cuda.is_available(): + return 'cuda' + elif torch.backends.mps.is_available(): + return 'mps' + return 'cpu' + + +class SmallConfig(GPT2Config): + + def __init__(self, **kwargs): + super().__init__(n_layer=6, n_head=6, **kwargs) + + +class Transformer(LightningModule): + + def __init__(self, cfg, label, cost, loss="exp"): + super().__init__() + self._label = label + self.cfg = cfg + gpt2cfg = GPT2Config( + **{k: cfg[k] for k in GPT2Config().to_dict().keys() & cfg.keys()}) + if cfg.small: + gpt2cfg = SmallConfig( + ** + {k: cfg[k] for k in GPT2Config().to_dict().keys() & cfg.keys()}) + self.transformer = GPT2LMHeadModel(gpt2cfg).to(get_device()) + self.ngates = cfg.ngates + self.num_samples = cfg.num_samples + self.temperature = cfg.temperature + self.save_hyperparameters() + self._starting_idx = torch.zeros(self.num_samples, + 1, + dtype=torch.int, + device=get_device()) + if loss == "exp": + self.loss = ExpLogitMatching(cfg.energy_offset, self._label) + print("ExpMatching") + else: + self.loss = GFlowLogitMatching(cfg.energy_offset, get_device(), + self._label, self) + print("GFlow") + self._cost = cost + + def generate_logits(self, idx): + logits = self.transformer(idx)[0] + return logits + + def set_cost(self, cost): + self._cost = cost + + def gather(self, idx, logits_base): + b_size = idx.shape[0] + return torch.gather(logits_base, 2, idx.reshape(b_size, -1, + 1)).reshape(b_size, -1) + + def computeCost(self, idx_output, **kwargs): + return torch.tensor([self._cost(row) for row in idx_output], + dtype=torch.float) + + def train_step(self, indices=None, energies=None, numQPUs=None, comm=None): + log_values = {} + if energies is not None: + assert indices is not None + idx_output = indices[:, 1:] + logits_base = self.generate_logits(idx_output) + else: + idx_output, logits_base = self.generate() + energies = self.computeCost(idx_output, numQPUs=numQPUs, comm=comm) + logits_tensor = self.gather(idx_output, logits_base) + allLogits = logits_tensor + + loss = self.loss.compute(energies, allLogits, self.temperature, + log_values) + log_values[f"loss at {self._label}"] = loss + return loss, energies, idx_output, log_values + + def generate(self, idx=None, ngates=None): + """ + Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete + the sequence max_new_tokens times, feeding the predictions back into the model each time. + Most likely you'll want to make sure to be in model.eval() mode of operation for this. + """ + if idx is None: + idx = self._starting_idx.clone() + condition_length = idx.size(dim=1) + if ngates is None: + ngates = self.ngates + for _ in range(ngates): + idx_cond = idx + logits_base = self.generate_logits(idx_cond) + logits = logits_base[:, -1, :] + probs = F.softmax(-self.temperature * logits, dim=-1) + idx_next = torch.multinomial(probs, num_samples=1) + idx = torch.cat((idx, idx_next), dim=1) + idx = idx[:, condition_length:] + return idx, logits_base + + def forward(self): + return self.train_step() diff --git a/python/cudaqlib/tools/chemistry/pyscf/generators/gas_phase_generator.py b/python/cudaqlib/tools/chemistry/pyscf/generators/gas_phase_generator.py index 209ee0b..5475f46 100644 --- a/python/cudaqlib/tools/chemistry/pyscf/generators/gas_phase_generator.py +++ b/python/cudaqlib/tools/chemistry/pyscf/generators/gas_phase_generator.py @@ -131,7 +131,6 @@ def get_spin_hamiltonian(self, xyz:str, spin:int, charge: int, basis:str, symmet "WARN: nele_cas is not None and norb_cas is None. nele_cas and norb_cas should be either both None\ or have values") - print("HOWDY: ", norb_cas, nele_cas) ######################################################################## # To add (coming soon) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a731767..0a017b2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -23,7 +23,7 @@ if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND NOT APPLE) endif() target_link_libraries(test_gse PRIVATE - cudaq-optim + cudaq-gse gtest_main) target_link_libraries(test_gse INTERFACE cudaq-gse) gtest_discover_tests(test_gse) @@ -46,6 +46,7 @@ target_link_libraries(test_kernels PRIVATE cudaq::cudaq-builder cudaq-operators + cudaq-operator-pools cudaq-optim gtest_main) gtest_discover_tests(test_kernels) diff --git a/tests/gse/AdaptTester.cpp b/tests/gse/AdaptTester.cpp index 5de8ab3..62d421f 100644 --- a/tests/gse/AdaptTester.cpp +++ b/tests/gse/AdaptTester.cpp @@ -25,14 +25,15 @@ TEST(GSETester, checkSimpleAdapt) { 1, 1, 3, 3, -0.0454063, -0, 15}; cudaq::spin_op h(h2_data, 4); - cudaq::gse::spin_complement_gsd pool(h.num_qubits() / 2); + auto pool = cudaq::operator_pool::get("spin_complement_gsd"); + auto poolList = pool->generate({{"num-orbitals", h.num_qubits() / 2}}); auto initialState = [&](cudaq::qvector<> &q) __qpu__ { for (std::size_t i = 0; i < 2; i++) x(q[i]); }; - auto energy = cudaq::gse::adapt_vqe(initialState, h, pool, + auto energy = cudaq::gse::adapt_vqe(initialState, h, poolList, {.grad_norm_tolerance = 1e-4}); EXPECT_NEAR(energy, -1.13, 1e-2); } \ No newline at end of file diff --git a/tests/kernels/UCCSDTester.cpp b/tests/kernels/UCCSDTester.cpp index d07171f..256d765 100644 --- a/tests/kernels/UCCSDTester.cpp +++ b/tests/kernels/UCCSDTester.cpp @@ -9,10 +9,10 @@ #include #include +#include "cudaq.h" #include "cudaqlib/gse.h" #include "cudaqlib/kernels/uccsd.h" #include "cudaqlib/operators.h" -#include "cudaq.h" std::vector> h2_hpq_data{ -1.24884680e+00, 0.00000000e+00, -9.24110683e-17, 0.00000000e+00, @@ -335,3 +335,11 @@ TEST(UCCSDTester, checkUCCSDAnsatz) { EXPECT_NEAR(result.energy, -1.137, 1e-3); } } + +TEST(UCCSDTester, checkOperatorPool) { + auto pool = cudaq::operator_pool::get("uccsd"); + auto ops = pool->generate({{"num-qubits", 4}, {"num-electrons", 2}}); + + for (auto o : ops) + o.dump(); +} From 77cf9dd451bfb07b20a18f237004ccbc11160689 Mon Sep 17 00:00:00 2001 From: Alex McCaskey Date: Fri, 19 Jul 2024 20:31:09 +0000 Subject: [PATCH 2/7] cleanup Signed-off-by: Alex McCaskey --- python/cudaqlib/algorithms/gqe.py | 144 ++-------------------- python/cudaqlib/algorithms/transformer.py | 12 +- 2 files changed, 19 insertions(+), 137 deletions(-) diff --git a/python/cudaqlib/algorithms/gqe.py b/python/cudaqlib/algorithms/gqe.py index 84bc161..aa41fe4 100644 --- a/python/cudaqlib/algorithms/gqe.py +++ b/python/cudaqlib/algorithms/gqe.py @@ -7,6 +7,7 @@ import matplotlib.pyplot as plt import numpy as np import time +from ml_collections import ConfigDict torch.set_float32_matmul_precision('high') @@ -15,50 +16,6 @@ def key(distance): v = str(distance).replace(".", "_") return v - -def plot_trajectory(trajectory_path): - data = [] - with open(trajectory_path, 'r') as file: - for line in file: - data.append(json.loads(line)) - - iterations = [] - losses = [] - avg_energies = [] - min_energies = [] - global_min = 0 - for entry in data: - iterations.append(entry['iter']) - losses.append(entry['loss']) - avg_energies.append(sum(entry['energies']) / len(entry['energies'])) - if min(entry['energies']) < global_min: - global_min = min(entry['energies']) - min_energies.append(global_min) - - plt.figure(figsize=(12, 8)) - - plt.subplot(3, 1, 1) - plt.plot(iterations, losses, marker='o', linestyle='-', color='b') - plt.title('Loss per Epoch') - plt.xlabel('Iteration') - plt.ylabel('Loss') - - plt.subplot(3, 1, 2) - plt.plot(iterations, avg_energies, marker='o', linestyle='-', color='g') - plt.title('Average Energy per Epoch') - plt.xlabel('Iteration') - plt.ylabel('Average Energy') - - plt.subplot(3, 1, 3) - plt.plot(iterations, min_energies, marker='o', linestyle='-', color='r') - plt.title('Minimum Energy per Epoch') - plt.xlabel('Iteration') - plt.ylabel('Minimum Energy') - - plt.tight_layout() - plt.savefig('trajectory_plot.png') - - def pretrain_file(cfg): return f"{cfg.save_dir}{cfg.name}_{cfg.seed}_checkpoint_pretrain.ckpt" @@ -198,70 +155,19 @@ def __init__(self, temperature_scheduler: TemperatureScheduler = None) -> None: self.temperature_scheduler = temperature_scheduler - # def train(self, cfg, model, optimizer, numQPUs=1): - # self.numQPUs = numQPUs - - # # Not yet... - # # from mpi4py import MPI - - # # Initialize MPI - # # self.mpiComm = MPI.COMM_WORLD - # # self.mpiRank = self.mpiComm.Get_rank() - # # self.mpiSize = self.mpiComm.Get_size() - # fabric = L.Fabric(accelerator="auto", devices=1) - # fabric.seed_everything(cfg.seed) - # fabric.launch() - - # min_indices_dict = {} - # distances = cfg.distances - # filename = train_file(cfg) - # m = {} - # if os.path.exists(filename): - # with open(filename) as f: - # for l in f.readlines(): - # items = l.rstrip().split('\t') - # if len(items) != 2: - # continue - # distance, energy = items - # distance = float(distance) - # energy = float(energy) - # m[distance] = energy - # with open(filename, 'w') as f: - # for distance in distances: - # print("distance:", distance) - # if cfg.cache and distance in m: - # min_energy = m[distance] - # print("already computed, skipped", distance) - # else: - # indices, min_energy = self.do_train(cfg, distance, fabric, model, optimizer) - # min_indices_dict[str(distance)] = indices - # f.write(f"{distance}\t{min_energy}\n") - # fabric.log('circuit', json.dumps(min_indices_dict)) - # # if (self.mpiComm): - # # MPI.Finalize() - # return min_indices_dict - - def train(self, cfg, model, optimizer): + def train(self, cfg, model, pool, optimizer): fabric = L.Fabric(accelerator="auto") fabric.seed_everything(cfg.seed) fabric.launch() - - min_indices, energy = self.do_train(cfg, fabric, model, + min_indices, energy = self.do_train(cfg, fabric, model,pool, optimizer) fabric.log('circuit', json.dumps(min_indices)) return min_indices, energy - def do_train(self, cfg, fabric, model, optimizer): - print(cfg) + def do_train(self, cfg, fabric, model, pool, optimizer): monitor = FileMonitor() model, optimizer = fabric.setup(model, optimizer) model.mark_forward_method(model.train_step) - - # if key(distance) in cfg.check_points: - # print("loaded from the checkpoint") - # cp = fabric.load(cfg.check_points[key(distance)]) - # model.load_state_dict(cp["model"]) - # optimizer.load_state_dict(cp["optimizer"]) pytorch_total_params = sum( p.numel() for p in model.parameters() if p.requires_grad) print(f"total trainable params: {pytorch_total_params / 1e6:.2f}M") @@ -271,8 +177,7 @@ def do_train(self, cfg, fabric, model, optimizer): optimizer.zero_grad() l = None start = time.time() - loss, energies, indices, log_values = model.train_step() - # numQPUs=self.numQPUs, comm=self.mpiComm) + loss, energies, indices, log_values = model.train_step(pool) print('epoch', epoch, 'model.train_step time:', time.time() - start, torch.min(energies)) if l is None: l = loss @@ -286,52 +191,30 @@ def do_train(self, cfg, fabric, model, optimizer): min_indices = indices log_values[f"min_energy at"] = min_energy log_values[f"temperature at"] = model.temperature - if cfg.verbose: - print(f"energies: {energies}") - print(f"temperature: {model.temperature}") fabric.log_dict(log_values, step=epoch) fabric.backward(loss) fabric.clip_gradients(model, optimizer, max_norm=cfg.grad_norm_clip) optimizer.step() - # scheduler.step() if self.temperature_scheduler is not None: model.temperature = self.temperature_scheduler.get(epoch) else: model.temperature += cfg.del_temperature model.set_cost(None) - # state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} + state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} # fabric.save(cfg.save_dir + f"checkpoint_{distance}.ckpt", state) # if cfg.save_data: - # monitor.save(trajectory_file(cfg, distance)) - # if cfg.plot: - # plot_trajectory( - # str(cfg.save_dir) + str(cfg.name) + "_trajectory_" + - # str(distance) + ".ckpt") + # monitor.save(trajectory_file(cfg, distance)) indices = min_indices.cpu().numpy().tolist() return min_energy, indices -def get_default_chemistry_configs(): - from ml_collections import ConfigDict +def get_default_config(): cfg = ConfigDict() cfg.verbose = False - cfg.save_data = False - cfg.print_exact = True - cfg.name = "gptqe" - cfg.warmup_steps = 100 cfg.num_samples = 5 # akin to batch size cfg.max_iters = 100 - cfg.nqubit = 4 cfg.ngates = 20 - cfg.nshot = 0 cfg.seed = 3047 - cfg.distances = [0.5, 0.6, 0.7, 0.7414, 0.8, 0.9, 1.0, 1.5, - 2.0] # choices of the distance between two atoms - cfg.time_pool = [ - 0.003125, -0.003125, 0.00625, -0.00625, 0.0125, -0.0125, 0.025, -0.025, - 0.05, -0.05, 0.1, -0.1 - ] - cfg.time_factor = 1.0 cfg.lr = 5e-7 cfg.energy_offset = 0.0 cfg.grad_norm_clip = 1.0 @@ -348,11 +231,10 @@ def get_default_chemistry_configs(): return cfg -def gqe(cost, **kwargs): - cfg = get_default_chemistry_configs() - cfg.vocab_size = 12 - distance = .7474 - model = Transformer(cfg, distance, cost, loss='exp') +def gqe(cost, pool, config=None, **kwargs): + cfg = get_default_config() if config == None else config + cfg.vocab_size = len(pool) + model = Transformer(cfg, cost, loss='exp') optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr) gqe = GPTQETaskBase() - return gqe.train(cfg, model, optimizer)#, numQPUs=1) + return gqe.train(cfg, model, pool, optimizer)#, numQPUs=1) diff --git a/python/cudaqlib/algorithms/transformer.py b/python/cudaqlib/algorithms/transformer.py index 08b3c0b..594cf0e 100644 --- a/python/cudaqlib/algorithms/transformer.py +++ b/python/cudaqlib/algorithms/transformer.py @@ -20,9 +20,9 @@ def __init__(self, **kwargs): class Transformer(LightningModule): - def __init__(self, cfg, label, cost, loss="exp"): + def __init__(self, cfg, cost, loss="exp"): super().__init__() - self._label = label + self._label = 'label_stand_in' self.cfg = cfg gpt2cfg = GPT2Config( **{k: cfg[k] for k in GPT2Config().to_dict().keys() & cfg.keys()}) @@ -60,11 +60,11 @@ def gather(self, idx, logits_base): return torch.gather(logits_base, 2, idx.reshape(b_size, -1, 1)).reshape(b_size, -1) - def computeCost(self, idx_output, **kwargs): - return torch.tensor([self._cost(row) for row in idx_output], + def computeCost(self, idx_output, pool, **kwargs): + return torch.tensor([self._cost([pool[i] for i in row]) for row in idx_output], dtype=torch.float) - def train_step(self, indices=None, energies=None, numQPUs=None, comm=None): + def train_step(self, pool, indices=None, energies=None, numQPUs=None, comm=None): log_values = {} if energies is not None: assert indices is not None @@ -72,7 +72,7 @@ def train_step(self, indices=None, energies=None, numQPUs=None, comm=None): logits_base = self.generate_logits(idx_output) else: idx_output, logits_base = self.generate() - energies = self.computeCost(idx_output, numQPUs=numQPUs, comm=comm) + energies = self.computeCost(idx_output, pool, numQPUs=numQPUs, comm=comm) logits_tensor = self.gather(idx_output, logits_base) allLogits = logits_tensor From 82d9aa148743d67edb53e33cc8f6c13ad44a5c5d Mon Sep 17 00:00:00 2001 From: Alex McCaskey Date: Mon, 22 Jul 2024 12:45:03 +0000 Subject: [PATCH 3/7] fix issue with multiple gpus Signed-off-by: Alex McCaskey --- python/cudaqlib/algorithms/gqe.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/cudaqlib/algorithms/gqe.py b/python/cudaqlib/algorithms/gqe.py index aa41fe4..58a4726 100644 --- a/python/cudaqlib/algorithms/gqe.py +++ b/python/cudaqlib/algorithms/gqe.py @@ -156,7 +156,7 @@ def __init__(self, self.temperature_scheduler = temperature_scheduler def train(self, cfg, model, pool, optimizer): - fabric = L.Fabric(accelerator="auto") + fabric = L.Fabric(accelerator="auto", devices=1) # 1 device for now fabric.seed_everything(cfg.seed) fabric.launch() min_indices, energy = self.do_train(cfg, fabric, model,pool, @@ -167,7 +167,7 @@ def train(self, cfg, model, pool, optimizer): def do_train(self, cfg, fabric, model, pool, optimizer): monitor = FileMonitor() model, optimizer = fabric.setup(model, optimizer) - model.mark_forward_method(model.train_step) + model.mark_forward_method('train_step')#model.train_step) pytorch_total_params = sum( p.numel() for p in model.parameters() if p.requires_grad) print(f"total trainable params: {pytorch_total_params / 1e6:.2f}M") @@ -200,7 +200,7 @@ def do_train(self, cfg, fabric, model, pool, optimizer): else: model.temperature += cfg.del_temperature model.set_cost(None) - state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} + # state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} # fabric.save(cfg.save_dir + f"checkpoint_{distance}.ckpt", state) # if cfg.save_data: # monitor.save(trajectory_file(cfg, distance)) @@ -237,4 +237,4 @@ def gqe(cost, pool, config=None, **kwargs): model = Transformer(cfg, cost, loss='exp') optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr) gqe = GPTQETaskBase() - return gqe.train(cfg, model, pool, optimizer)#, numQPUs=1) + return gqe.train(cfg, model, pool, optimizer) From 681f56504d6d785d3ad067a187115b37128945ec Mon Sep 17 00:00:00 2001 From: Alex McCaskey Date: Mon, 22 Jul 2024 12:47:49 +0000 Subject: [PATCH 4/7] add the example Signed-off-by: Alex McCaskey --- examples/python/gqe_h2.py | 56 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 examples/python/gqe_h2.py diff --git a/examples/python/gqe_h2.py b/examples/python/gqe_h2.py new file mode 100644 index 0000000..ae98763 --- /dev/null +++ b/examples/python/gqe_h2.py @@ -0,0 +1,56 @@ +import cudaq, cudaqlib, numpy as np + +# Define the molecule +geometry = [('H', (0., 0., 0.)), ('H', (0., 0., .7474))] +molecule = cudaqlib.operators.create_molecule(geometry, + 'sto-3g', + 0, + 0) + +# Get the number of qubits +hamiltonian = molecule.hamiltonian + +# Get the number of qubits +numQubits = molecule.hamiltonian.get_qubit_count() + +# Create the operator pool +# FIXME needs to better incorporate thetas. +pool = cudaqlib.gse.get_operator_pool('uccsd', num_qubits=4, num_electrons=2) +operatorCoeffs = [ + 0.003125, -0.003125, 0.00625, -0.00625, 0.0125, -0.0125, 0.025, -0.025, + 0.05, -0.05, 0.1, -0.1 + ] +assert len(pool) == len(operatorCoeffs + ) +# Need an initial state +@cudaq.kernel +def init(q: cudaq.qview): + x(q[0]) + x(q[1]) + + +# Define the GQE cost function +def cost(sampledPoolOperations : list): + """ + Cost should take operator pool indices and + return the associated cost. For the chemistry + example, we'll take uccsd pool indices and return + cudaq observe result + """ + asWords = [cudaq.pauli_word(op.to_string(False)) for op in sampledPoolOperations] + + @cudaq.kernel + def kernel(numQubits: int, coeffs: list[float], + words: list[cudaq.pauli_word]): + q = cudaq.qvector(numQubits) + init(q) + for i, word in enumerate(words): + exp_pauli(coeffs[i], q, word) + + return cudaq.observe(kernel, molecule.hamiltonian, numQubits, operatorCoeffs, asWords).expectation() + + +minE, optimPoolOps = cudaqlib.gqe(cost, pool) +print(f'Ground Energy = {minE}') +print('Ansatz Ops') +for idx in optimPoolOps: print(pool[idx].to_string(False)) \ No newline at end of file From 29a848d831e07ad979fe2b494b973bf357544cae Mon Sep 17 00:00:00 2001 From: Alex McCaskey Date: Mon, 22 Jul 2024 14:02:00 +0000 Subject: [PATCH 5/7] fix ci build error Signed-off-by: Alex McCaskey --- docker/base/Dockerfile | 2 +- python/cudaqlib/__init__.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docker/base/Dockerfile b/docker/base/Dockerfile index 97a743e..2d5c3ef 100644 --- a/docker/base/Dockerfile +++ b/docker/base/Dockerfile @@ -14,4 +14,4 @@ RUN cd .. && git clone -b cudaqlib https://github.com/amccaskey/cuda-quantum \ -DCUDAQ_ENABLE_PYTHON=TRUE -DCMAKE_INSTALL_PREFIX=$HOME/.cudaq \ && ninja install \ && apt-get update && apt-get install -y gfortran libblas-dev \ - && python3 -m pip install cppe \ No newline at end of file + && python3 -m pip install cppe torch lightning ml-collections \ No newline at end of file diff --git a/python/cudaqlib/__init__.py b/python/cudaqlib/__init__.py index c82b93e..ad56ba6 100644 --- a/python/cudaqlib/__init__.py +++ b/python/cudaqlib/__init__.py @@ -10,4 +10,8 @@ import os, os.path from ._pycudaqlib import * from .kernels import * -from .algorithms.gqe import gqe \ No newline at end of file +try: + from .algorithms.gqe import gqe +except: + print('[cudaqlib warning] Could not import GQE module. Install required modules (e.g. torch)') + pass \ No newline at end of file From 52f960f8db261f16c1ddda6fe42393e8e08d5bab Mon Sep 17 00:00:00 2001 From: Alex McCaskey Date: Mon, 22 Jul 2024 16:58:20 +0000 Subject: [PATCH 6/7] cleanup Signed-off-by: Alex McCaskey --- python/cudaqlib/algorithms/gqe.py | 114 +++++++++++++----------------- 1 file changed, 50 insertions(+), 64 deletions(-) diff --git a/python/cudaqlib/algorithms/gqe.py b/python/cudaqlib/algorithms/gqe.py index 58a4726..47a4c80 100644 --- a/python/cudaqlib/algorithms/gqe.py +++ b/python/cudaqlib/algorithms/gqe.py @@ -148,66 +148,6 @@ def save(self, path): for l in self.lines: f.write(f"{l}\n") - -class GPTQETaskBase(): - - def __init__(self, - temperature_scheduler: TemperatureScheduler = None) -> None: - self.temperature_scheduler = temperature_scheduler - - def train(self, cfg, model, pool, optimizer): - fabric = L.Fabric(accelerator="auto", devices=1) # 1 device for now - fabric.seed_everything(cfg.seed) - fabric.launch() - min_indices, energy = self.do_train(cfg, fabric, model,pool, - optimizer) - fabric.log('circuit', json.dumps(min_indices)) - return min_indices, energy - - def do_train(self, cfg, fabric, model, pool, optimizer): - monitor = FileMonitor() - model, optimizer = fabric.setup(model, optimizer) - model.mark_forward_method('train_step')#model.train_step) - pytorch_total_params = sum( - p.numel() for p in model.parameters() if p.requires_grad) - print(f"total trainable params: {pytorch_total_params / 1e6:.2f}M") - min_energy = sys.maxsize - min_indices = None - for epoch in range(cfg.max_iters): - optimizer.zero_grad() - l = None - start = time.time() - loss, energies, indices, log_values = model.train_step(pool) - print('epoch', epoch, 'model.train_step time:', time.time() - start, torch.min(energies)) - if l is None: - l = loss - else: - l = l + loss - monitor.record(epoch, loss, energies, indices) - for e, indices in zip(energies, indices): - energy = e.item() - if energy < min_energy: - min_energy = e.item() - min_indices = indices - log_values[f"min_energy at"] = min_energy - log_values[f"temperature at"] = model.temperature - fabric.log_dict(log_values, step=epoch) - fabric.backward(loss) - fabric.clip_gradients(model, optimizer, max_norm=cfg.grad_norm_clip) - optimizer.step() - if self.temperature_scheduler is not None: - model.temperature = self.temperature_scheduler.get(epoch) - else: - model.temperature += cfg.del_temperature - model.set_cost(None) - # state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} - # fabric.save(cfg.save_dir + f"checkpoint_{distance}.ckpt", state) - # if cfg.save_data: - # monitor.save(trajectory_file(cfg, distance)) - indices = min_indices.cpu().numpy().tolist() - return min_energy, indices - - def get_default_config(): cfg = ConfigDict() cfg.verbose = False @@ -230,11 +170,57 @@ def get_default_config(): cfg.save_dir = "./output/" return cfg +def __internal_run_gqe(temperature_scheduler: TemperatureScheduler, cfg : ConfigDict, model, pool, optimizer): + fabric = L.Fabric(accelerator="auto", devices=1) # 1 device for now + fabric.seed_everything(cfg.seed) + fabric.launch() + monitor = FileMonitor() + model, optimizer = fabric.setup(model, optimizer) + model.mark_forward_method('train_step')#model.train_step) + pytorch_total_params = sum( + p.numel() for p in model.parameters() if p.requires_grad) + print(f"total trainable params: {pytorch_total_params / 1e6:.2f}M") + min_energy = sys.maxsize + min_indices = None + for epoch in range(cfg.max_iters): + optimizer.zero_grad() + l = None + start = time.time() + loss, energies, indices, log_values = model.train_step(pool) + print('epoch', epoch, 'model.train_step time:', time.time() - start, torch.min(energies)) + if l is None: + l = loss + else: + l = l + loss + monitor.record(epoch, loss, energies, indices) + for e, indices in zip(energies, indices): + energy = e.item() + if energy < min_energy: + min_energy = e.item() + min_indices = indices + log_values[f"min_energy at"] = min_energy + log_values[f"temperature at"] = model.temperature + fabric.log_dict(log_values, step=epoch) + fabric.backward(loss) + fabric.clip_gradients(model, optimizer, max_norm=cfg.grad_norm_clip) + optimizer.step() + if temperature_scheduler is not None: + model.temperature = temperature_scheduler.get(epoch) + else: + model.temperature += cfg.del_temperature + model.set_cost(None) + # state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} + # fabric.save(cfg.save_dir + f"checkpoint_{distance}.ckpt", state) + # if cfg.save_data: + # monitor.save(trajectory_file(cfg, distance)) + min_indices = min_indices.cpu().numpy().tolist() + # return min_energy, indices + fabric.log('circuit', json.dumps(min_indices)) + return min_energy, min_indices def gqe(cost, pool, config=None, **kwargs): cfg = get_default_config() if config == None else config cfg.vocab_size = len(pool) - model = Transformer(cfg, cost, loss='exp') - optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr) - gqe = GPTQETaskBase() - return gqe.train(cfg, model, pool, optimizer) + model = Transformer(cfg, cost, loss='exp') if 'model' not in kwargs else kwargs['model'] + optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr) if 'optimizer' not in kwargs else kwargs['optimizer'] + return __internal_run_gqe(None, cfg, model, pool, optimizer) From 9cfba85e3fd1fb5d1b98ec8c2d42a9a4cc51d7a9 Mon Sep 17 00:00:00 2001 From: Alex McCaskey Date: Mon, 22 Jul 2024 18:15:08 +0000 Subject: [PATCH 7/7] format Signed-off-by: Alex McCaskey --- python/cudaqlib/algorithms/gqe.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/python/cudaqlib/algorithms/gqe.py b/python/cudaqlib/algorithms/gqe.py index 47a4c80..62d256b 100644 --- a/python/cudaqlib/algorithms/gqe.py +++ b/python/cudaqlib/algorithms/gqe.py @@ -16,6 +16,7 @@ def key(distance): v = str(distance).replace(".", "_") return v + def pretrain_file(cfg): return f"{cfg.save_dir}{cfg.name}_{cfg.seed}_checkpoint_pretrain.ckpt" @@ -148,6 +149,7 @@ def save(self, path): for l in self.lines: f.write(f"{l}\n") + def get_default_config(): cfg = ConfigDict() cfg.verbose = False @@ -170,13 +172,15 @@ def get_default_config(): cfg.save_dir = "./output/" return cfg -def __internal_run_gqe(temperature_scheduler: TemperatureScheduler, cfg : ConfigDict, model, pool, optimizer): - fabric = L.Fabric(accelerator="auto", devices=1) # 1 device for now + +def __internal_run_gqe(temperature_scheduler: TemperatureScheduler, + cfg: ConfigDict, model, pool, optimizer): + fabric = L.Fabric(accelerator="auto", devices=1) # 1 device for now fabric.seed_everything(cfg.seed) fabric.launch() monitor = FileMonitor() model, optimizer = fabric.setup(model, optimizer) - model.mark_forward_method('train_step')#model.train_step) + model.mark_forward_method('train_step') #model.train_step) pytorch_total_params = sum( p.numel() for p in model.parameters() if p.requires_grad) print(f"total trainable params: {pytorch_total_params / 1e6:.2f}M") @@ -187,7 +191,8 @@ def __internal_run_gqe(temperature_scheduler: TemperatureScheduler, cfg : Config l = None start = time.time() loss, energies, indices, log_values = model.train_step(pool) - print('epoch', epoch, 'model.train_step time:', time.time() - start, torch.min(energies)) + print('epoch', epoch, 'model.train_step time:', + time.time() - start, torch.min(energies)) if l is None: l = loss else: @@ -212,15 +217,19 @@ def __internal_run_gqe(temperature_scheduler: TemperatureScheduler, cfg : Config # state = {"model": model, "optimizer": optimizer, "hparams": model.hparams} # fabric.save(cfg.save_dir + f"checkpoint_{distance}.ckpt", state) # if cfg.save_data: - # monitor.save(trajectory_file(cfg, distance)) + # monitor.save(trajectory_file(cfg, distance)) min_indices = min_indices.cpu().numpy().tolist() # return min_energy, indices fabric.log('circuit', json.dumps(min_indices)) return min_energy, min_indices + def gqe(cost, pool, config=None, **kwargs): cfg = get_default_config() if config == None else config cfg.vocab_size = len(pool) - model = Transformer(cfg, cost, loss='exp') if 'model' not in kwargs else kwargs['model'] - optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr) if 'optimizer' not in kwargs else kwargs['optimizer'] + model = Transformer( + cfg, cost, loss='exp') if 'model' not in kwargs else kwargs['model'] + optimizer = torch.optim.AdamW( + model.parameters(), + lr=cfg.lr) if 'optimizer' not in kwargs else kwargs['optimizer'] return __internal_run_gqe(None, cfg, model, pool, optimizer)