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/KetNotUniqueError.hpp"
      12             : #include "pairinteraction/ket/KetPair.hpp"
      13             : #include "pairinteraction/ket/QuantumNumberNotAvailableError.hpp"
      14             : #include "pairinteraction/utils/traits.hpp"
      15             : 
      16             : #include <complex>
      17             : #include <nanobind/nanobind.h>
      18             : #include <nanobind/operators.h>
      19             : #include <nanobind/stl/shared_ptr.h>
      20             : #include <nanobind/stl/string.h>
      21             : #include <nanobind/stl/vector.h>
      22             : 
      23             : namespace nb = nanobind;
      24             : using namespace pairinteraction;
      25             : 
      26           2 : static void declare_ket(nb::module_ &m) {
      27           2 :     std::string pyclass_name = "Ket";
      28           2 :     nb::class_<Ket> pyclass(m, pyclass_name.c_str());
      29           2 :     pyclass.def("get_energy", &Ket::get_energy)
      30           2 :         .def("has_quantum_number", &Ket::has_quantum_number)
      31           2 :         .def("get_quantum_number", &Ket::get_quantum_number);
      32           2 : }
      33             : 
      34           2 : static void declare_ket_atom(nb::module_ &m) {
      35           2 :     std::string pyclass_name = "KetAtom";
      36           2 :     nb::class_<KetAtom, Ket> pyclass(m, pyclass_name.c_str());
      37           2 :     pyclass.def("get_database", &KetAtom::get_database, nb::rv_policy::reference)
      38           2 :         .def("get_species", &KetAtom::get_species)
      39           2 :         .def("get_quantum_number_std", &KetAtom::get_quantum_number_std)
      40           2 :         .def(nb::self == nb::self) // NOLINT(misc-redundant-expression)
      41        4420 :         .def("__hash__", [](const KetAtom &self) { return KetAtom::hash{}(self); });
      42           2 : }
      43             : 
      44           2 : static void declare_ket_atom_creator(nb::module_ &m) {
      45           2 :     std::string pyclass_name = "KetAtomCreator";
      46           4 :     nb::class_<KetAtomCreator> pyclass(m, pyclass_name.c_str());
      47             :     // keep_alive because the creator stores a reference to the database
      48           6 :     pyclass.def(nb::init<Database &>(), nb::keep_alive<1, 2>())
      49           2 :         .def("set_species", &KetAtomCreator::set_species)
      50           2 :         .def("set_energy", &KetAtomCreator::set_energy)
      51           2 :         .def("set_quantum_number", &KetAtomCreator::set_quantum_number)
      52           4 :         .def("create", &KetAtomCreator::create);
      53           2 : }
      54             : 
      55             : template <typename T>
      56           4 : static void declare_ket_pair(nb::module_ &m, std::string const &type_name) {
      57           4 :     std::string pyclass_name = "KetPair" + type_name;
      58           4 :     nb::class_<KetPair<T>, Ket> pyclass(m, pyclass_name.c_str());
      59           4 :     pyclass.def("get_atomic_states", &KetPair<T>::get_atomic_states)
      60           4 :         .def(nb::self == nb::self) // NOLINT(misc-redundant-expression)
      61         228 :         .def("__hash__", [](const KetPair<T> &self) { return typename KetPair<T>::hash{}(self); });
      62           4 : }
      63             : 
      64           2 : static void declare_ket_not_unique_error(nb::module_ &m) {
      65           2 :     static nb::object exc_type = nb::exception<KetNotUniqueError>(m, "KetNotUniqueError");
      66           2 :     nb::register_exception_translator(
      67           8 :         [](const std::exception_ptr &p, void *payload) {
      68             :             try {
      69          12 :                 std::rethrow_exception(p);
      70           6 :             } catch (const KetNotUniqueError &e) {
      71           0 :                 auto *type = static_cast<PyObject *>(payload);
      72           0 :                 nb::object exc = nb::borrow(type)(e.what());
      73           0 :                 exc.attr("kets") = nb::cast(e.get_kets());
      74           0 :                 PyErr_SetObject(type, exc.ptr());
      75           0 :             }
      76           0 :         },
      77           2 :         exc_type.ptr());
      78           2 : }
      79             : 
      80           2 : static void declare_quantum_number_not_available_error(nb::module_ &m) {
      81           4 :     static nb::object exc_type = nb::exception<QuantumNumberNotAvailableError>(
      82           4 :         m, "QuantumNumberNotAvailableError", PyExc_ValueError);
      83           2 :     nb::register_exception_translator(
      84           8 :         [](const std::exception_ptr &p, void *payload) {
      85             :             try {
      86          12 :                 std::rethrow_exception(p);
      87           6 :             } catch (const QuantumNumberNotAvailableError &e) {
      88           0 :                 auto *type = static_cast<PyObject *>(payload);
      89           0 :                 nb::object exc = nb::borrow(type)(e.what());
      90           0 :                 exc.attr("name") = nb::cast(e.get_name());
      91           0 :                 PyErr_SetObject(type, exc.ptr());
      92           0 :             }
      93           0 :         },
      94           2 :         exc_type.ptr());
      95           2 : }
      96             : 
      97           2 : void bind_ket(nb::module_ &m) {
      98           2 :     declare_ket(m);
      99           2 :     declare_ket_atom(m);
     100           2 :     declare_ket_atom_creator(m);
     101           2 :     declare_ket_pair<double>(m, "Real");
     102           2 :     declare_ket_pair<std::complex<double>>(m, "Complex");
     103           2 :     declare_ket_not_unique_error(m);
     104           2 :     declare_quantum_number_not_available_error(m);
     105           2 : }