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/SorterBuilderInterface.hpp"
13 : #include "pairinteraction/ket/KetAtom.hpp"
14 : #include "pairinteraction/ket/KetPair.hpp"
15 :
16 : #include <nanobind/eigen/sparse.h>
17 : #include <nanobind/nanobind.h>
18 : #include <nanobind/stl/complex.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 : template <typename T>
27 8 : static void declare_basis(nb::module_ &m, std::string const &type_name) {
28 8 : std::string pyclass_name = "Basis" + type_name;
29 : using scalar_t = typename Basis<T>::scalar_t;
30 8 : nb::class_<Basis<T>, SorterBuilderInterface> pyclass(m, pyclass_name.c_str());
31 8 : pyclass.def("get_kets", &Basis<T>::get_kets)
32 8 : .def("get_ket", &Basis<T>::get_ket)
33 8 : .def("get_state", &Basis<T>::get_state)
34 8 : .def("get_number_of_states", &Basis<T>::get_number_of_states)
35 8 : .def("get_number_of_kets", &Basis<T>::get_number_of_kets)
36 8 : .def("has_quantum_number", &Basis<T>::has_quantum_number)
37 8 : .def("has_energy", &Basis<T>::has_energy)
38 8 : .def("is_sorted_by_energy", &Basis<T>::is_sorted_by_energy)
39 8 : .def("get_quantum_number", &Basis<T>::get_quantum_number)
40 8 : .def("get_energy", &Basis<T>::get_energy)
41 8 : .def("get_coefficients", &Basis<T>::get_coefficients)
42 8 : .def("copy_with_coefficients", &Basis<T>::copy_with_coefficients)
43 8 : .def("get_sorter", &Basis<T>::get_sorter)
44 16 : .def("get_indices_of_blocks", &Basis<T>::get_indices_of_blocks)
45 8 : .def("transformed",
46 8 : [](const Basis<T> &self,
47 : const Eigen::SparseMatrix<scalar_t, Eigen::RowMajor> &transformation) {
48 0 : return self.transformed(transformation);
49 : })
50 8 : .def("transformed",
51 8 : nb::overload_cast<const Eigen::PermutationMatrix<Eigen::Dynamic, Eigen::Dynamic> &>(
52 : &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 : // keep_alive because the creator stores a reference to the database
70 12 : pyclass.def(nb::init<Database &>(), nb::keep_alive<1, 2>())
71 4 : .def("set_species", &BasisAtomCreator<T>::set_species)
72 4 : .def("restrict_energy", &BasisAtomCreator<T>::restrict_energy)
73 4 : .def("restrict_quantum_number", &BasisAtomCreator<T>::restrict_quantum_number)
74 4 : .def("set_quantum_number_standard_deviation_factor",
75 4 : &BasisAtomCreator<T>::set_quantum_number_standard_deviation_factor)
76 4 : .def("add_ket", &BasisAtomCreator<T>::add_ket)
77 8 : .def("create", &BasisAtomCreator<T>::create, nb::call_guard<nb::gil_scoped_release>());
78 4 : }
79 :
80 : template <typename T>
81 4 : static void declare_basis_pair(nb::module_ &m, std::string const &type_name) {
82 4 : std::string pyclass_name = "BasisPair" + type_name;
83 4 : nb::class_<BasisPair<T>, Basis<BasisPair<T>>> pyclass(m, pyclass_name.c_str());
84 4 : pyclass.def("get_matrix_elements", &BasisPair<T>::get_matrix_elements)
85 4 : .def("get_basis1", &BasisPair<T>::get_basis1)
86 4 : .def("get_basis2", &BasisPair<T>::get_basis2);
87 4 : }
88 :
89 : template <typename T>
90 4 : static void declare_basis_pair_creator(nb::module_ &m, std::string const &type_name) {
91 4 : std::string pyclass_name = "BasisPairCreator" + type_name;
92 8 : nb::class_<BasisPairCreator<T>> pyclass(m, pyclass_name.c_str());
93 : pyclass
94 12 : .def(nb::init<std::shared_ptr<const BasisAtom<T>>, std::shared_ptr<const BasisAtom<T>>>())
95 4 : .def("restrict_energy", &BasisPairCreator<T>::restrict_energy)
96 4 : .def("restrict_quantum_number_m", &BasisPairCreator<T>::restrict_quantum_number_m)
97 4 : .def("restrict_parity_under_inversion",
98 4 : &BasisPairCreator<T>::restrict_parity_under_inversion)
99 4 : .def("restrict_parity_under_permutation",
100 4 : &BasisPairCreator<T>::restrict_parity_under_permutation)
101 8 : .def("create", &BasisPairCreator<T>::create, nb::call_guard<nb::gil_scoped_release>());
102 4 : }
103 :
104 2 : void bind_basis(nb::module_ &m) {
105 2 : declare_basis<BasisAtom<double>>(m, "BasisAtomReal");
106 2 : declare_basis<BasisAtom<std::complex<double>>>(m, "BasisAtomComplex");
107 2 : declare_basis_atom<double>(m, "Real");
108 2 : declare_basis_atom<std::complex<double>>(m, "Complex");
109 2 : declare_basis_atom_creator<double>(m, "Real");
110 2 : declare_basis_atom_creator<std::complex<double>>(m, "Complex");
111 :
112 2 : declare_basis<BasisPair<double>>(m, "BasisPairReal");
113 2 : declare_basis<BasisPair<std::complex<double>>>(m, "BasisPairComplex");
114 2 : declare_basis_pair<double>(m, "Real");
115 2 : declare_basis_pair<std::complex<double>>(m, "Complex");
116 2 : declare_basis_pair_creator<double>(m, "Real");
117 2 : declare_basis_pair_creator<std::complex<double>>(m, "Complex");
118 2 : }
|
