Program Listing for File ObservablesTNCuda.hpp

↰ Return to documentation for file (pennylane_lightning/core/simulators/lightning_tensor/tncuda/observables/ObservablesTNCuda.hpp)

// Copyright 2024 Xanadu Quantum Technologies Inc. and contributors.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <tuple>
#include <unordered_set>
#include <vector>

#include "Constant.hpp"
#include "ConstantUtil.hpp" // lookup
#include "Error.hpp"
#include "Util.hpp"

#include "cuda_helpers.hpp"

namespace {
using namespace Pennylane::Util;
using namespace Pennylane::LightningGPU::Util;

template <class T> using vector1D = std::vector<T>;
template <class T> using vector2D = std::vector<vector1D<T>>;
template <class T> using vector3D = std::vector<vector2D<T>>;
} // namespace

namespace Pennylane::LightningTensor::TNCuda::Observables {
template <class TensorNetT> class ObservableTNCuda {
  public:
    using PrecisionT = typename TensorNetT::PrecisionT;
    using ComplexT = typename TensorNetT::ComplexT;
    using MetaDataT = std::tuple<std::string, std::vector<PrecisionT>,
                                 std::vector<ComplexT>>; // name, params, matrix

  protected:
    vector1D<PrecisionT> coeffs_;      // coefficients of each term
    vector1D<std::size_t> numTensors_; // number of tensors in each term
    vector2D<std::size_t>
        numStateModes_; // number of state modes of each tensor in each term
    vector3D<std::size_t>
        stateModes_;               // state modes of each tensor in each term
    vector2D<MetaDataT> metaData_; // meta data of each tensor in each term

  protected:
    ObservableTNCuda() = default;
    ObservableTNCuda(const ObservableTNCuda &) = default;
    ObservableTNCuda(ObservableTNCuda &&) noexcept = default;
    ObservableTNCuda &operator=(const ObservableTNCuda &) = default;
    ObservableTNCuda &operator=(ObservableTNCuda &&) noexcept = default;

  private:
    [[nodiscard]] virtual bool
    isEqual(const ObservableTNCuda<TensorNetT> &other) const = 0;

  public:
    virtual ~ObservableTNCuda() = default;

    [[nodiscard]] virtual auto getObsName() const -> std::string = 0;

    [[nodiscard]] virtual auto getWires() const -> std::vector<std::size_t> = 0;

    [[nodiscard]] auto getNumTensors() const -> const vector1D<std::size_t> & {
        return numTensors_;
    }

    [[nodiscard]] auto getNumStateModes() const
        -> const vector2D<std::size_t> & {
        return numStateModes_;
    }

    [[nodiscard]] auto getStateModes() const -> const vector3D<std::size_t> & {
        return stateModes_;
    }

    [[nodiscard]] auto getMetaData() const -> const vector2D<MetaDataT> & {
        return metaData_;
    }

    [[nodiscard]] auto getCoeffs() const -> const vector1D<PrecisionT> & {
        return coeffs_;
    };

    [[nodiscard]] auto
    operator==(const ObservableTNCuda<TensorNetT> &other) const -> bool {
        return typeid(*this) == typeid(other) && isEqual(other);
    }

    [[nodiscard]] auto
    operator!=(const ObservableTNCuda<TensorNetT> &other) const -> bool {
        return !(*this == other);
    }
};

template <class TensorNetT>
class NamedObsTNCuda : public ObservableTNCuda<TensorNetT> {
  public:
    using BaseType = ObservableTNCuda<TensorNetT>;
    using PrecisionT = typename TensorNetT::PrecisionT;
    using ComplexT = typename TensorNetT::ComplexT;

  private:
    std::string obs_name_;
    std::vector<std::size_t> wires_;
    std::vector<PrecisionT> params_;

    [[nodiscard]] auto isEqual(const ObservableTNCuda<TensorNetT> &other) const
        -> bool override {
        const auto &other_cast =
            static_cast<const NamedObsTNCuda<TensorNetT> &>(other);

        return (obs_name_ == other_cast.obs_name_) &&
               (wires_ == other_cast.wires_) && (params_ == other_cast.params_);
    }

  public:
    NamedObsTNCuda(std::string obs_name, std::vector<std::size_t> wires,
                   std::vector<PrecisionT> params = {})
        : obs_name_{obs_name}, wires_{wires}, params_{params} {
        BaseType::coeffs_.emplace_back(PrecisionT{1.0});
        BaseType::numTensors_.emplace_back(std::size_t{1});
        BaseType::numStateModes_.emplace_back(
            vector1D<std::size_t>{wires_.size()});
        BaseType::stateModes_.emplace_back(vector2D<std::size_t>{wires_});

        BaseType::metaData_.push_back(
            {std::make_tuple(obs_name, params_, std::vector<ComplexT>{})});
    }

    [[nodiscard]] auto getObsName() const -> std::string override {
        using Pennylane::Util::operator<<;
        std::ostringstream obs_stream;
        obs_stream << obs_name_ << wires_;
        return obs_stream.str();
    }

    [[nodiscard]] auto getWires() const -> std::vector<std::size_t> override {
        return wires_;
    }
};

template <class TensorNetT>
class HermitianObsTNCuda : public ObservableTNCuda<TensorNetT> {
  public:
    using BaseType = ObservableTNCuda<TensorNetT>;
    using PrecisionT = typename TensorNetT::PrecisionT;
    using ComplexT = typename TensorNetT::ComplexT;
    using MatrixT = std::vector<ComplexT>;

  private:
    inline static const MatrixHasher mh;
    MatrixT matrix_;
    std::vector<std::size_t> wires_;

    [[nodiscard]] auto isEqual(const ObservableTNCuda<TensorNetT> &other) const
        -> bool override {
        const auto &other_cast =
            static_cast<const HermitianObsTNCuda<TensorNetT> &>(other);

        return (matrix_ == other_cast.matrix_) && (wires_ == other_cast.wires_);
    }

  public:
    HermitianObsTNCuda(MatrixT matrix, std::vector<std::size_t> wires)
        : matrix_{std::move(matrix)}, wires_{std::move(wires)} {
        PL_ABORT_IF(wires_.size() != 1, "The number of Hermitian target wires "
                                        "must be 1 for Lightning-Tensor.");
        PL_ASSERT(matrix_.size() == Pennylane::Util::exp2(2 * wires_.size()));
        BaseType::coeffs_.emplace_back(PrecisionT{1.0});
        BaseType::numTensors_.emplace_back(std::size_t{1});
        BaseType::numStateModes_.emplace_back(
            vector1D<std::size_t>{wires_.size()});
        BaseType::stateModes_.emplace_back(vector2D<std::size_t>{wires_});

        BaseType::metaData_.push_back(
            {std::make_tuple("Hermitian", std::vector<PrecisionT>{}, matrix_)});
    }

    [[nodiscard]] auto getObsName() const -> std::string override {
        // To avoid collisions on cached GPU data, use matrix elements to
        // uniquely identify Hermitian
        // TODO: Replace with a performant hash function
        std::ostringstream obs_stream;
        obs_stream << "Hermitian" << mh(matrix_);
        return obs_stream.str();
    }

    [[nodiscard]] auto getWires() const -> std::vector<std::size_t> override {
        return wires_;
    }

    [[nodiscard]] auto getMatrix() const -> const MatrixT & { return matrix_; }
};

template <class TensorNetT>
class TensorProdObsTNCuda : public ObservableTNCuda<TensorNetT> {
  public:
    using BaseType = ObservableTNCuda<TensorNetT>;
    using PrecisionT = typename TensorNetT::PrecisionT;
    using MetaDataT = BaseType::MetaDataT;

  private:
    std::vector<std::shared_ptr<ObservableTNCuda<TensorNetT>>> obs_;
    std::vector<std::size_t> all_wires_;

    [[nodiscard]] auto isEqual(const ObservableTNCuda<TensorNetT> &other) const
        -> bool override {
        const auto &other_cast =
            static_cast<const TensorProdObsTNCuda<TensorNetT> &>(other);

        if (obs_.size() != other_cast.obs_.size()) {
            return false;
        }

        for (std::size_t i = 0; i < obs_.size(); i++) {
            if (*obs_[i] != *other_cast.obs_[i]) {
                return false;
            }
        }
        return true;
    }

  public:
    template <typename... Ts>
    explicit TensorProdObsTNCuda(Ts &&...arg) : obs_{std::forward<Ts>(arg)...} {
        if (obs_.size() == 1 &&
            obs_[0]->getObsName().find('@') != std::string::npos) {
            // This would prevent the misuse of this constructor for creating
            // TensorProdObs(TensorProdObs).
            PL_ABORT("A new TensorProdObs observable cannot be created "
                     "from a single TensorProdObs.");
        }

        for (const auto &ob : obs_) {
            PL_ABORT_IF(ob->getObsName().find("Hamiltonian") !=
                            std::string::npos,
                        "A TensorProdObs observable cannot be created from a "
                        "Hamiltonian.");
        }

        BaseType::coeffs_.push_back(PrecisionT{1.0});
        BaseType::numTensors_.push_back(obs_.size());

        vector1D<std::size_t> numStateModesLocal;
        vector2D<std::size_t> stateModesLocal;
        vector1D<MetaDataT> dataLocal;

        for (const auto &ob : obs_) {
            numStateModesLocal.insert(numStateModesLocal.end(),
                                      ob->getNumStateModes().front().begin(),
                                      ob->getNumStateModes().front().end());

            stateModesLocal.insert(stateModesLocal.end(),
                                   ob->getStateModes().front().begin(),
                                   ob->getStateModes().front().end());

            dataLocal.insert(dataLocal.end(), ob->getMetaData().front().begin(),
                             ob->getMetaData().front().end());
        }

        BaseType::numStateModes_.emplace_back(numStateModesLocal);
        BaseType::stateModes_.emplace_back(stateModesLocal);
        BaseType::metaData_.emplace_back(dataLocal);

        std::unordered_set<std::size_t> wires;
        for (const auto &ob : obs_) {
            const auto ob_wires = ob->getWires();
            for (const auto wire : ob_wires) {
                PL_ABORT_IF(wires.contains(wire),
                            "All wires in observables must be disjoint.");
                wires.insert(wire);
            }
        }
        all_wires_ = std::vector<std::size_t>(wires.begin(), wires.end());
        std::sort(all_wires_.begin(), all_wires_.end());
    }

    static auto
    create(std::initializer_list<std::shared_ptr<ObservableTNCuda<TensorNetT>>>
               obs) -> std::shared_ptr<TensorProdObsTNCuda<TensorNetT>> {
        return std::shared_ptr<TensorProdObsTNCuda<TensorNetT>>{
            new TensorProdObsTNCuda(std::move(obs))};
    }

    static auto
    create(std::vector<std::shared_ptr<ObservableTNCuda<TensorNetT>>> obs)
        -> std::shared_ptr<TensorProdObsTNCuda<TensorNetT>> {
        return std::shared_ptr<TensorProdObsTNCuda<TensorNetT>>{
            new TensorProdObsTNCuda(std::move(obs))};
    }

    [[nodiscard]] auto getSize() const -> std::size_t { return obs_.size(); }

    [[nodiscard]] auto getWires() const -> std::vector<std::size_t> override {
        return all_wires_;
    }

    [[nodiscard]] auto getObsName() const -> std::string override {
        using Pennylane::Util::operator<<;
        std::ostringstream obs_stream;
        const auto obs_size = obs_.size();
        for (std::size_t idx = 0; idx < obs_size; idx++) {
            obs_stream << obs_[idx]->getObsName();
            if (idx != obs_size - 1) {
                obs_stream << " @ ";
            }
        }
        return obs_stream.str();
    }

    [[nodiscard]] auto getObs() const
        -> std::vector<std::shared_ptr<ObservableTNCuda<TensorNetT>>> {
        return obs_;
    };
};

template <class TensorNetT>
class HamiltonianTNCuda : public ObservableTNCuda<TensorNetT> {
  public:
    using BaseType = ObservableTNCuda<TensorNetT>;
    using PrecisionT = typename TensorNetT::PrecisionT;

  private:
    std::vector<PrecisionT> coeffs_ham_;
    std::vector<std::shared_ptr<ObservableTNCuda<TensorNetT>>> obs_;

    [[nodiscard]] bool
    isEqual(const ObservableTNCuda<TensorNetT> &other) const override {
        const auto &other_cast =
            static_cast<const HamiltonianTNCuda<TensorNetT> &>(other);

        if (coeffs_ham_ != other_cast.coeffs_ham_) {
            return false;
        }

        for (std::size_t i = 0; i < obs_.size(); i++) {
            if (*obs_[i] != *other_cast.obs_[i]) {
                return false;
            }
        }
        return true;
    }

  public:
    template <typename T1, typename T2>
    HamiltonianTNCuda(T1 &&coeffs, T2 &&obs)
        : coeffs_ham_{std::forward<T1>(coeffs)}, obs_{std::forward<T2>(obs)} {
        BaseType::coeffs_ = coeffs_ham_;
        PL_ASSERT(BaseType::coeffs_.size() == obs_.size());

        for (std::size_t term_idx = 0; term_idx < BaseType::coeffs_.size();
             term_idx++) {
            auto ob = obs_[term_idx];
            // This is aligned with statevector backends
            PL_ABORT_IF(ob->getObsName().find("Hamiltonian") !=
                            std::string::npos,
                        "A Hamiltonian observable cannot be created from "
                        "another Hamiltonian.");
            BaseType::numTensors_.emplace_back(ob->getNumTensors().front());
            BaseType::numStateModes_.emplace_back(
                ob->getNumStateModes().front());
            BaseType::stateModes_.emplace_back(ob->getStateModes().front());
            BaseType::metaData_.emplace_back(ob->getMetaData().front());
        }
    }

    static auto
    create(std::initializer_list<PrecisionT> coeffs,
           std::initializer_list<std::shared_ptr<ObservableTNCuda<TensorNetT>>>
               obs) -> std::shared_ptr<HamiltonianTNCuda<TensorNetT>> {
        return std::shared_ptr<HamiltonianTNCuda<TensorNetT>>(
            new HamiltonianTNCuda<TensorNetT>{std::move(coeffs),
                                              std::move(obs)});
    }

    [[nodiscard]] auto getWires() const -> std::vector<std::size_t> override {
        std::unordered_set<std::size_t> wires;

        for (const auto &ob : obs_) {
            const auto ob_wires = ob->getWires();
            wires.insert(ob_wires.begin(), ob_wires.end());
        }
        auto all_wires = std::vector<std::size_t>(wires.begin(), wires.end());
        std::sort(all_wires.begin(), all_wires.end());
        return all_wires;
    }

    [[nodiscard]] auto getObsName() const -> std::string override {
        using Pennylane::Util::operator<<;
        std::ostringstream ss;
        ss << "Hamiltonian: { 'coeffs' : " << BaseType::coeffs_
           << ", 'observables' : [";
        const auto term_size = BaseType::coeffs_.size();
        for (std::size_t t = 0; t < term_size; t++) {
            ss << obs_[t]->getObsName();
            if (t != term_size - 1) {
                ss << ", ";
            }
        }
        ss << "]}";
        return ss.str();
    }
    [[nodiscard]] auto getObs() const
        -> std::vector<std::shared_ptr<ObservableTNCuda<TensorNetT>>> {
        return obs_;
    };

    [[nodiscard]] auto getCoeffs() const -> std::vector<PrecisionT> {
        return BaseType::getCoeffs();
    };
};
} // namespace Pennylane::LightningTensor::TNCuda::Observables

api/program_listing_file_pennylane_lightning_core_simulators_lightning_tensor_tncuda_observables_ObservablesTNCuda.hpp

 

Download Python script

 

Download Notebook

 

View on GitHub