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/stl/shared_ptr.h>
      17             : #include <nanobind/stl/string.h>
      18             : #include <nanobind/stl/vector.h>
      19             : #include <sstream>
      20             : 
      21             : namespace nb = nanobind;
      22             : using namespace pairinteraction;
      23             : 
      24           1 : static void declare_ket(nb::module_ &m) {
      25           1 :     std::string pyclass_name = "Ket";
      26           1 :     nb::class_<Ket> pyclass(m, pyclass_name.c_str());
      27           1 :     pyclass.def("get_energy", &Ket::get_energy)
      28           1 :         .def("get_quantum_number_f", &Ket::get_quantum_number_f)
      29           1 :         .def("get_quantum_number_m", &Ket::get_quantum_number_m)
      30           1 :         .def("get_parity", &Ket::get_parity)
      31           2 :         .def("get_label", &Ket::get_label)
      32           1 :         .def("__str__", [](Ket const &self) {
      33           0 :             std::stringstream ss;
      34           0 :             ss << self;
      35           0 :             return ss.str();
      36           0 :         });
      37           1 : }
      38             : 
      39           1 : static void declare_ket_atom(nb::module_ &m) {
      40           1 :     std::string pyclass_name = "KetAtom";
      41           1 :     nb::class_<KetAtom, Ket> pyclass(m, pyclass_name.c_str());
      42           1 :     pyclass.def("get_species", &KetAtom::get_species)
      43           1 :         .def("get_quantum_number_n", &KetAtom::get_quantum_number_n)
      44           1 :         .def("get_quantum_number_nu", &KetAtom::get_quantum_number_nu)
      45           1 :         .def("get_quantum_number_nui", &KetAtom::get_quantum_number_nui)
      46           1 :         .def("get_quantum_number_l", &KetAtom::get_quantum_number_l)
      47           1 :         .def("get_quantum_number_s", &KetAtom::get_quantum_number_s)
      48           1 :         .def("get_quantum_number_j", &KetAtom::get_quantum_number_j)
      49           1 :         .def("get_quantum_number_l_ryd", &KetAtom::get_quantum_number_l_ryd)
      50           1 :         .def("get_quantum_number_j_ryd", &KetAtom::get_quantum_number_j_ryd)
      51           1 :         .def("get_quantum_number_nui_std", &KetAtom::get_quantum_number_nui_std)
      52           1 :         .def("get_quantum_number_l_std", &KetAtom::get_quantum_number_l_std)
      53           1 :         .def("get_quantum_number_s_std", &KetAtom::get_quantum_number_s_std)
      54           1 :         .def("get_quantum_number_j_std", &KetAtom::get_quantum_number_j_std)
      55           1 :         .def("get_quantum_number_l_ryd_std", &KetAtom::get_quantum_number_l_ryd_std)
      56           1 :         .def("get_quantum_number_j_ryd_std", &KetAtom::get_quantum_number_j_ryd_std)
      57           1 :         .def("is_j_total_momentum", &KetAtom::is_j_total_momentum)
      58           1 :         .def("is_calculated_with_mqdt", &KetAtom::is_calculated_with_mqdt)
      59             :         .def("get_underspecified_channel_contribution",
      60           1 :              &KetAtom::get_underspecified_channel_contribution);
      61           1 : }
      62             : 
      63           1 : static void declare_ket_atom_creator(nb::module_ &m) {
      64           1 :     std::string pyclass_name = "KetAtomCreator";
      65           2 :     nb::class_<KetAtomCreator> pyclass(m, pyclass_name.c_str());
      66           1 :     pyclass.def(nb::init<>())
      67           2 :         .def(nb::init<std::string, int, double, double, double>())
      68           1 :         .def("set_species", &KetAtomCreator::set_species)
      69           1 :         .def("set_energy", &KetAtomCreator::set_energy)
      70           1 :         .def("set_quantum_number_f", &KetAtomCreator::set_quantum_number_f)
      71           1 :         .def("set_quantum_number_m", &KetAtomCreator::set_quantum_number_m)
      72           1 :         .def("set_parity", &KetAtomCreator::set_parity)
      73           1 :         .def("set_quantum_number_n", &KetAtomCreator::set_quantum_number_n)
      74           1 :         .def("set_quantum_number_nu", &KetAtomCreator::set_quantum_number_nu)
      75           1 :         .def("set_quantum_number_nui", &KetAtomCreator::set_quantum_number_nui)
      76           1 :         .def("set_quantum_number_l", &KetAtomCreator::set_quantum_number_l)
      77           1 :         .def("set_quantum_number_s", &KetAtomCreator::set_quantum_number_s)
      78           1 :         .def("set_quantum_number_j", &KetAtomCreator::set_quantum_number_j)
      79           1 :         .def("set_quantum_number_l_ryd", &KetAtomCreator::set_quantum_number_l_ryd)
      80           1 :         .def("set_quantum_number_j_ryd", &KetAtomCreator::set_quantum_number_j_ryd)
      81           1 :         .def("create", &KetAtomCreator::create);
      82           1 : }
      83             : 
      84             : template <typename T>
      85           2 : static void declare_ket_pair(nb::module_ &m, std::string const &type_name) {
      86           2 :     std::string pyclass_name = "KetPair" + type_name;
      87           2 :     nb::class_<KetPair<T>, Ket> pyclass(m, pyclass_name.c_str());
      88           2 :     pyclass.def("get_atomic_states", &KetPair<T>::get_atomic_states);
      89           2 : }
      90             : 
      91           1 : void bind_ket(nb::module_ &m) {
      92           1 :     declare_ket(m);
      93           1 :     declare_ket_atom(m);
      94           1 :     declare_ket_atom_creator(m);
      95           1 :     declare_ket_pair<double>(m, "Real");
      96           1 :     declare_ket_pair<std::complex<double>>(m, "Complex");
      97           1 : }