Line data    Source code

       1             : // SPDX-FileCopyrightText: 2024 PairInteraction Developers
       2             : // SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : #include "./Basis.py.hpp"
       5             : 
       6             : #include "pairinteraction/basis/Basis.hpp"
       7             : #include "pairinteraction/basis/BasisAtom.hpp"
       8             : #include "pairinteraction/basis/BasisAtomCreator.hpp"
       9             : #include "pairinteraction/basis/BasisPair.hpp"
      10             : #include "pairinteraction/basis/BasisPairCreator.hpp"
      11             : #include "pairinteraction/database/Database.hpp"
      12             : #include "pairinteraction/interfaces/TransformationBuilderInterface.hpp"
      13             : #include "pairinteraction/ket/KetAtom.hpp"
      14             : #include "pairinteraction/ket/KetPair.hpp"
      15             : #include "pairinteraction/system/SystemAtom.hpp"
      16             : 
      17             : #include <nanobind/eigen/sparse.h>
      18             : #include <nanobind/nanobind.h>
      19             : #include <nanobind/stl/complex.h>
      20             : #include <nanobind/stl/set.h>
      21             : #include <nanobind/stl/shared_ptr.h>
      22             : #include <nanobind/stl/string.h>
      23             : #include <nanobind/stl/vector.h>
      24             : 
      25             : namespace nb = nanobind;
      26             : using namespace pairinteraction;
      27             : 
      28             : template <typename T>
      29           8 : static void declare_basis(nb::module_ &m, std::string const &type_name) {
      30           8 :     std::string pyclass_name = "Basis" + type_name;
      31             :     using scalar_t = typename Basis<T>::scalar_t;
      32           8 :     nb::class_<Basis<T>, TransformationBuilderInterface<scalar_t>> pyclass(m, pyclass_name.c_str());
      33           8 :     pyclass.def("get_kets", &Basis<T>::get_kets)
      34           8 :         .def("get_ket", &Basis<T>::get_ket)
      35           8 :         .def("get_state", &Basis<T>::get_state)
      36           8 :         .def("get_number_of_states", &Basis<T>::get_number_of_states)
      37           8 :         .def("get_number_of_kets", &Basis<T>::get_number_of_kets)
      38           8 :         .def("get_quantum_number_f", &Basis<T>::get_quantum_number_f)
      39           8 :         .def("get_quantum_number_m", &Basis<T>::get_quantum_number_m)
      40           8 :         .def("get_parity", &Basis<T>::get_parity)
      41           8 :         .def("get_coefficients", &Basis<T>::get_coefficients)
      42           8 :         .def("copy_with_coefficients", &Basis<T>::copy_with_coefficients)
      43           8 :         .def("get_transformation", &Basis<T>::get_transformation)
      44           8 :         .def("get_sorter", &Basis<T>::get_sorter)
      45           8 :         .def("get_indices_of_blocks", &Basis<T>::get_indices_of_blocks)
      46           8 :         .def("get_sorter_without_checks", &Basis<T>::get_sorter_without_checks)
      47           8 :         .def("get_indices_of_blocks_without_checks",
      48           8 :              &Basis<T>::get_indices_of_blocks_without_checks)
      49           8 :         .def(
      50             :             "transformed",
      51           8 :             nb::overload_cast<const Transformation<scalar_t> &>(&Basis<T>::transformed, nb::const_))
      52           8 :         .def("transformed", nb::overload_cast<const Sorting &>(&Basis<T>::transformed, nb::const_))
      53           8 :         .def("canonicalized", &Basis<T>::canonicalized)
      54           8 :         .def("is_canonical", &Basis<T>::is_canonical)
      55           8 :         .def("merge", &Basis<T>::merge);
      56           8 : }
      57             : 
      58             : template <typename T>
      59           4 : static void declare_basis_atom(nb::module_ &m, std::string const &type_name) {
      60           4 :     std::string pyclass_name = "BasisAtom" + type_name;
      61           4 :     nb::class_<BasisAtom<T>, Basis<BasisAtom<T>>> pyclass(m, pyclass_name.c_str());
      62           4 :     pyclass.def("get_matrix_elements", &BasisAtom<T>::get_matrix_elements);
      63           4 : }
      64             : 
      65             : template <typename T>
      66           4 : static void declare_basis_atom_creator(nb::module_ &m, std::string const &type_name) {
      67           4 :     std::string pyclass_name = "BasisAtomCreator" + type_name;
      68           8 :     nb::class_<BasisAtomCreator<T>> pyclass(m, pyclass_name.c_str());
      69          12 :     pyclass.def(nb::init<>())
      70           4 :         .def("set_species", &BasisAtomCreator<T>::set_species)
      71           4 :         .def("restrict_energy", &BasisAtomCreator<T>::restrict_energy)
      72           4 :         .def("restrict_quantum_number", &BasisAtomCreator<T>::restrict_quantum_number)
      73           4 :         .def("set_quantum_number_standard_deviation_factor",
      74           4 :              &BasisAtomCreator<T>::set_quantum_number_standard_deviation_factor)
      75           4 :         .def("add_ket", &BasisAtomCreator<T>::add_ket)
      76           8 :         .def("create", &BasisAtomCreator<T>::create, nb::call_guard<nb::gil_scoped_release>());
      77           4 : }
      78             : 
      79             : template <typename T>
      80           4 : static void declare_basis_pair(nb::module_ &m, std::string const &type_name) {
      81           4 :     std::string pyclass_name = "BasisPair" + type_name;
      82           4 :     nb::class_<BasisPair<T>, Basis<BasisPair<T>>> pyclass(m, pyclass_name.c_str());
      83           4 :     pyclass.def("get_matrix_elements", &BasisPair<T>::get_matrix_elements)
      84           4 :         .def("get_basis1", &BasisPair<T>::get_basis1)
      85           4 :         .def("get_basis2", &BasisPair<T>::get_basis2);
      86           4 : }
      87             : 
      88             : template <typename T>
      89           4 : static void declare_basis_pair_creator(nb::module_ &m, std::string const &type_name) {
      90           4 :     std::string pyclass_name = "BasisPairCreator" + type_name;
      91           8 :     nb::class_<BasisPairCreator<T>> pyclass(m, pyclass_name.c_str());
      92             :     pyclass
      93          12 :         .def(nb::init<>())
      94             :         // keep_alive because add() stores only a reference to the system, which must outlive the
      95             :         // creator (the system is dereferenced in create())
      96           4 :         .def("add", &BasisPairCreator<T>::add, nb::keep_alive<1, 2>())
      97           4 :         .def("restrict_energy", &BasisPairCreator<T>::restrict_energy)
      98           4 :         .def("restrict_quantum_number_m", &BasisPairCreator<T>::restrict_quantum_number_m)
      99           4 :         .def("restrict_parity_under_inversion",
     100           4 :              &BasisPairCreator<T>::restrict_parity_under_inversion)
     101           4 :         .def("restrict_parity_under_permutation",
     102           4 :              &BasisPairCreator<T>::restrict_parity_under_permutation)
     103           8 :         .def("create", &BasisPairCreator<T>::create, nb::call_guard<nb::gil_scoped_release>());
     104           4 : }
     105             : 
     106           2 : void bind_basis(nb::module_ &m) {
     107           2 :     declare_basis<BasisAtom<double>>(m, "BasisAtomReal");
     108           2 :     declare_basis<BasisAtom<std::complex<double>>>(m, "BasisAtomComplex");
     109           2 :     declare_basis_atom<double>(m, "Real");
     110           2 :     declare_basis_atom<std::complex<double>>(m, "Complex");
     111           2 :     declare_basis_atom_creator<double>(m, "Real");
     112           2 :     declare_basis_atom_creator<std::complex<double>>(m, "Complex");
     113             : 
     114           2 :     declare_basis<BasisPair<double>>(m, "BasisPairReal");
     115           2 :     declare_basis<BasisPair<std::complex<double>>>(m, "BasisPairComplex");
     116           2 :     declare_basis_pair<double>(m, "Real");
     117           2 :     declare_basis_pair<std::complex<double>>(m, "Complex");
     118           2 :     declare_basis_pair_creator<double>(m, "Real");
     119           2 :     declare_basis_pair_creator<std::complex<double>>(m, "Complex");
     120           2 : }