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", nb::overload_cast<>(&Basis<T>::get_coefficients, nb::const_))
42 8 : .def("set_coefficients", &Basis<T>::set_coefficients)
43 8 : .def("get_transformation", &Basis<T>::get_transformation)
44 8 : .def("get_rotator", &Basis<T>::get_rotator)
45 8 : .def("get_sorter", &Basis<T>::get_sorter)
46 8 : .def("get_indices_of_blocks", &Basis<T>::get_indices_of_blocks)
47 8 : .def("get_sorter_without_checks", &Basis<T>::get_sorter_without_checks)
48 8 : .def("get_indices_of_blocks_without_checks",
49 8 : &Basis<T>::get_indices_of_blocks_without_checks)
50 8 : .def(
51 : "transformed",
52 8 : nb::overload_cast<const Transformation<scalar_t> &>(&Basis<T>::transformed, nb::const_))
53 8 : .def("transformed", nb::overload_cast<const Sorting &>(&Basis<T>::transformed, nb::const_))
54 8 : .def("get_amplitudes",
55 8 : nb::overload_cast<std::shared_ptr<const typename Basis<T>::ket_t>>(
56 : &Basis<T>::get_amplitudes, nb::const_))
57 8 : .def("get_amplitudes",
58 8 : nb::overload_cast<std::shared_ptr<const T>>(&Basis<T>::get_amplitudes, nb::const_))
59 8 : .def("get_overlaps",
60 8 : nb::overload_cast<std::shared_ptr<const typename Basis<T>::ket_t>>(
61 : &Basis<T>::get_overlaps, nb::const_))
62 8 : .def("get_overlaps",
63 8 : nb::overload_cast<std::shared_ptr<const T>>(&Basis<T>::get_overlaps, nb::const_))
64 8 : .def("get_matrix_elements",
65 8 : nb::overload_cast<std::shared_ptr<const typename Basis<T>::ket_t>, OperatorType, int>(
66 : &Basis<T>::get_matrix_elements, nb::const_))
67 8 : .def("get_matrix_elements",
68 8 : nb::overload_cast<std::shared_ptr<const T>, OperatorType, int>(
69 : &Basis<T>::get_matrix_elements, nb::const_))
70 8 : .def("get_corresponding_state",
71 8 : nb::overload_cast<size_t>(&Basis<T>::get_corresponding_state, nb::const_))
72 8 : .def("get_corresponding_state",
73 8 : nb::overload_cast<std::shared_ptr<const typename Basis<T>::ket_t>>(
74 : &Basis<T>::get_corresponding_state, nb::const_))
75 8 : .def("get_corresponding_state_index",
76 8 : nb::overload_cast<size_t>(&Basis<T>::get_corresponding_state_index, nb::const_))
77 8 : .def("get_corresponding_state_index",
78 8 : nb::overload_cast<std::shared_ptr<const typename Basis<T>::ket_t>>(
79 : &Basis<T>::get_corresponding_state_index, nb::const_))
80 8 : .def("get_corresponding_ket",
81 8 : nb::overload_cast<size_t>(&Basis<T>::get_corresponding_ket, nb::const_))
82 8 : .def("get_corresponding_ket",
83 8 : nb::overload_cast<std::shared_ptr<const T>>(&Basis<T>::get_corresponding_ket,
84 : nb::const_))
85 8 : .def("get_corresponding_ket_index",
86 8 : nb::overload_cast<size_t>(&Basis<T>::get_corresponding_ket_index, nb::const_))
87 8 : .def("get_corresponding_ket_index",
88 8 : nb::overload_cast<std::shared_ptr<const T>>(&Basis<T>::get_corresponding_ket_index,
89 : nb::const_))
90 11 : .def("copy", [](const Basis<T> &self) {
91 3 : return std::make_shared<T>(static_cast<const T &>(self));
92 : });
93 8 : }
94 :
95 : template <typename T>
96 4 : static void declare_basis_atom(nb::module_ &m, std::string const &type_name) {
97 4 : std::string pyclass_name = "BasisAtom" + type_name;
98 4 : nb::class_<BasisAtom<T>, Basis<BasisAtom<T>>> pyclass(m, pyclass_name.c_str());
99 4 : }
100 :
101 : template <typename T>
102 4 : static void declare_basis_atom_creator(nb::module_ &m, std::string const &type_name) {
103 4 : std::string pyclass_name = "BasisAtomCreator" + type_name;
104 8 : nb::class_<BasisAtomCreator<T>> pyclass(m, pyclass_name.c_str());
105 8 : pyclass.def(nb::init<>())
106 4 : .def("set_species", &BasisAtomCreator<T>::set_species)
107 4 : .def("restrict_energy", &BasisAtomCreator<T>::restrict_energy)
108 4 : .def("restrict_quantum_number_f", &BasisAtomCreator<T>::restrict_quantum_number_f)
109 4 : .def("restrict_quantum_number_m", &BasisAtomCreator<T>::restrict_quantum_number_m)
110 4 : .def("restrict_parity", &BasisAtomCreator<T>::restrict_parity)
111 4 : .def("restrict_quantum_number_n", &BasisAtomCreator<T>::restrict_quantum_number_n)
112 4 : .def("restrict_quantum_number_nu", &BasisAtomCreator<T>::restrict_quantum_number_nu)
113 4 : .def("restrict_quantum_number_nui", &BasisAtomCreator<T>::restrict_quantum_number_nui)
114 4 : .def("restrict_quantum_number_l", &BasisAtomCreator<T>::restrict_quantum_number_l)
115 4 : .def("restrict_quantum_number_s", &BasisAtomCreator<T>::restrict_quantum_number_s)
116 4 : .def("restrict_quantum_number_j", &BasisAtomCreator<T>::restrict_quantum_number_j)
117 4 : .def("restrict_quantum_number_l_ryd", &BasisAtomCreator<T>::restrict_quantum_number_l_ryd)
118 4 : .def("restrict_quantum_number_j_ryd", &BasisAtomCreator<T>::restrict_quantum_number_j_ryd)
119 4 : .def("append_ket", &BasisAtomCreator<T>::append_ket)
120 4 : .def("create", &BasisAtomCreator<T>::create);
121 4 : }
122 :
123 : template <typename T>
124 4 : static void declare_basis_pair(nb::module_ &m, std::string const &type_name) {
125 4 : std::string pyclass_name = "BasisPair" + type_name;
126 4 : nb::class_<BasisPair<T>, Basis<BasisPair<T>>> pyclass(m, pyclass_name.c_str());
127 : pyclass
128 4 : .def("get_amplitudes", // define again, so nanobind can resolve all overloads correctly
129 4 : nb::overload_cast<std::shared_ptr<const BasisPair<T>>>(
130 : &Basis<BasisPair<T>>::get_amplitudes, nb::const_))
131 4 : .def("get_amplitudes",
132 : nb::overload_cast<std::shared_ptr<const BasisAtom<T>>,
133 4 : std::shared_ptr<const BasisAtom<T>>>(&BasisPair<T>::get_amplitudes,
134 : nb::const_))
135 4 : .def("get_amplitudes", // define again, so nanobind can resolve all overloads correctly
136 4 : nb::overload_cast<std::shared_ptr<const typename BasisPair<T>::ket_t>>(
137 : &Basis<BasisPair<T>>::get_amplitudes, nb::const_))
138 4 : .def("get_amplitudes",
139 4 : nb::overload_cast<std::shared_ptr<const KetAtom>, std::shared_ptr<const KetAtom>>(
140 : &BasisPair<T>::get_amplitudes, nb::const_))
141 4 : .def("get_overlaps", // define again, so nanobind can resolve all overloads correctly
142 4 : nb::overload_cast<std::shared_ptr<const BasisPair<T>>>(
143 : &Basis<BasisPair<T>>::get_overlaps, nb::const_))
144 4 : .def("get_overlaps",
145 : nb::overload_cast<std::shared_ptr<const BasisAtom<T>>,
146 4 : std::shared_ptr<const BasisAtom<T>>>(&BasisPair<T>::get_overlaps,
147 : nb::const_))
148 4 : .def("get_overlaps", // define again, so nanobind can resolve all overloads correctly
149 4 : nb::overload_cast<std::shared_ptr<const typename BasisPair<T>::ket_t>>(
150 : &Basis<BasisPair<T>>::get_overlaps, nb::const_))
151 4 : .def("get_overlaps",
152 4 : nb::overload_cast<std::shared_ptr<const KetAtom>, std::shared_ptr<const KetAtom>>(
153 : &BasisPair<T>::get_overlaps, nb::const_))
154 4 : .def("get_matrix_elements",
155 : nb::overload_cast<std::shared_ptr<const BasisPair<T>>, OperatorType, OperatorType, int,
156 4 : int>(&BasisPair<T>::get_matrix_elements, nb::const_))
157 4 : .def("get_matrix_elements",
158 : nb::overload_cast<std::shared_ptr<const BasisAtom<T>>,
159 : std::shared_ptr<const BasisAtom<T>>, OperatorType, OperatorType, int,
160 4 : int>(&BasisPair<T>::get_matrix_elements, nb::const_))
161 4 : .def("get_matrix_elements",
162 : nb::overload_cast<std::shared_ptr<const typename BasisPair<T>::ket_t>, OperatorType,
163 4 : OperatorType, int, int>(&BasisPair<T>::get_matrix_elements,
164 : nb::const_))
165 : .def("get_matrix_elements",
166 : nb::overload_cast<std::shared_ptr<const KetAtom>, std::shared_ptr<const KetAtom>,
167 4 : OperatorType, OperatorType, int, int>(
168 : &BasisPair<T>::get_matrix_elements, nb::const_));
169 4 : }
170 :
171 : template <typename T>
172 4 : static void declare_basis_pair_creator(nb::module_ &m, std::string const &type_name) {
173 4 : std::string pyclass_name = "BasisPairCreator" + type_name;
174 8 : nb::class_<BasisPairCreator<T>> pyclass(m, pyclass_name.c_str());
175 8 : pyclass.def(nb::init<>())
176 4 : .def("add", &BasisPairCreator<T>::add)
177 4 : .def("restrict_energy", &BasisPairCreator<T>::restrict_energy)
178 4 : .def("restrict_quantum_number_m", &BasisPairCreator<T>::restrict_quantum_number_m)
179 4 : .def("restrict_product_of_parities", &BasisPairCreator<T>::restrict_product_of_parities)
180 4 : .def("create", &BasisPairCreator<T>::create);
181 4 : }
182 :
183 2 : void bind_basis(nb::module_ &m) {
184 2 : declare_basis<BasisAtom<double>>(m, "BasisAtomReal");
185 2 : declare_basis<BasisAtom<std::complex<double>>>(m, "BasisAtomComplex");
186 2 : declare_basis_atom<double>(m, "Real");
187 2 : declare_basis_atom<std::complex<double>>(m, "Complex");
188 2 : declare_basis_atom_creator<double>(m, "Real");
189 2 : declare_basis_atom_creator<std::complex<double>>(m, "Complex");
190 :
191 2 : declare_basis<BasisPair<double>>(m, "BasisPairReal");
192 2 : declare_basis<BasisPair<std::complex<double>>>(m, "BasisPairComplex");
193 2 : declare_basis_pair<double>(m, "Real");
194 2 : declare_basis_pair<std::complex<double>>(m, "Complex");
195 2 : declare_basis_pair_creator<double>(m, "Real");
196 2 : declare_basis_pair_creator<std::complex<double>>(m, "Complex");
197 2 : }
|
