Line data    Source code

       1             : // SPDX-FileCopyrightText: 2024 PairInteraction Developers
       2             : // SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : #include "./Ket.py.hpp"
       5             : 
       6             : #include "pairinteraction/basis/BasisAtom.hpp"
       7             : #include "pairinteraction/database/Database.hpp"
       8             : #include "pairinteraction/ket/Ket.hpp"
       9             : #include "pairinteraction/ket/KetAtom.hpp"
      10             : #include "pairinteraction/ket/KetAtomCreator.hpp"
      11             : #include "pairinteraction/ket/KetPair.hpp"
      12             : #include "pairinteraction/utils/traits.hpp"
      13             : 
      14             : #include <complex>
      15             : #include <nanobind/nanobind.h>
      16             : #include <nanobind/operators.h>
      17             : #include <nanobind/stl/shared_ptr.h>
      18             : #include <nanobind/stl/string.h>
      19             : #include <nanobind/stl/vector.h>
      20             : #include <sstream>
      21             : 
      22             : namespace nb = nanobind;
      23             : using namespace pairinteraction;
      24             : 
      25           2 : static void declare_ket(nb::module_ &m) {
      26           2 :     std::string pyclass_name = "Ket";
      27           2 :     nb::class_<Ket> pyclass(m, pyclass_name.c_str());
      28           2 :     pyclass.def("get_energy", &Ket::get_energy)
      29           4 :         .def("get_label", &Ket::get_label)
      30           2 :         .def("__str__", [](Ket const &self) {
      31           0 :             std::stringstream ss;
      32           0 :             ss << self;
      33           0 :             return ss.str();
      34           0 :         });
      35           2 : }
      36             : 
      37           2 : static void declare_ket_atom(nb::module_ &m) {
      38           2 :     std::string pyclass_name = "KetAtom";
      39           2 :     nb::class_<KetAtom, Ket> pyclass(m, pyclass_name.c_str());
      40           2 :     pyclass.def("get_database", &KetAtom::get_database, nb::rv_policy::reference)
      41           2 :         .def("get_species", &KetAtom::get_species)
      42           2 :         .def("get_quantum_number", &KetAtom::get_quantum_number)
      43           2 :         .def("get_quantum_number_std", &KetAtom::get_quantum_number_std)
      44           2 :         .def(nb::self == nb::self) // NOLINT(misc-redundant-expression)
      45           0 :         .def("__hash__", [](const KetAtom &self) { return KetAtom::hash{}(self); });
      46           2 : }
      47             : 
      48           2 : static void declare_ket_atom_creator(nb::module_ &m) {
      49           2 :     std::string pyclass_name = "KetAtomCreator";
      50           4 :     nb::class_<KetAtomCreator> pyclass(m, pyclass_name.c_str());
      51           2 :     pyclass.def(nb::init<>())
      52           2 :         .def(nb::init<std::string, int, double, double, double>())
      53           2 :         .def("set_species", &KetAtomCreator::set_species)
      54           2 :         .def("set_energy", &KetAtomCreator::set_energy)
      55           2 :         .def("set_quantum_number", &KetAtomCreator::set_quantum_number)
      56           4 :         .def("create", &KetAtomCreator::create);
      57           2 : }
      58             : 
      59             : template <typename T>
      60           4 : static void declare_ket_pair(nb::module_ &m, std::string const &type_name) {
      61           4 :     std::string pyclass_name = "KetPair" + type_name;
      62           4 :     nb::class_<KetPair<T>, Ket> pyclass(m, pyclass_name.c_str());
      63           4 :     pyclass.def("get_atomic_states", &KetPair<T>::get_atomic_states)
      64           4 :         .def("get_quantum_number_m", &KetPair<T>::get_quantum_number_m)
      65           4 :         .def(nb::self == nb::self) // NOLINT(misc-redundant-expression)
      66           4 :         .def("__hash__", [](const KetPair<T> &self) { return typename KetPair<T>::hash{}(self); });
      67           4 : }
      68             : 
      69           2 : void bind_ket(nb::module_ &m) {
      70           2 :     declare_ket(m);
      71           2 :     declare_ket_atom(m);
      72           2 :     declare_ket_atom_creator(m);
      73           2 :     declare_ket_pair<double>(m, "Real");
      74           2 :     declare_ket_pair<std::complex<double>>(m, "Complex");
      75           2 : }