Line data Source code
1 : // SPDX-FileCopyrightText: 2024 PairInteraction Developers
2 : // SPDX-License-Identifier: LGPL-3.0-or-later
3 :
4 : #include "./System.py.hpp"
5 :
6 : #include "pairinteraction/basis/BasisAtom.hpp"
7 : #include "pairinteraction/basis/BasisPair.hpp"
8 : #include "pairinteraction/interfaces/DiagonalizerInterface.hpp"
9 : #include "pairinteraction/system/GreenTensorInterpolator.hpp"
10 : #include "pairinteraction/system/System.hpp"
11 : #include "pairinteraction/system/SystemAtom.hpp"
12 : #include "pairinteraction/system/SystemPair.hpp"
13 :
14 : #include <nanobind/eigen/dense.h>
15 : #include <nanobind/eigen/sparse.h>
16 : #include <nanobind/nanobind.h>
17 : #include <nanobind/stl/array.h>
18 : #include <nanobind/stl/complex.h>
19 : #include <nanobind/stl/optional.h>
20 : #include <nanobind/stl/shared_ptr.h>
21 : #include <nanobind/stl/variant.h>
22 : #include <nanobind/stl/vector.h>
23 :
24 : namespace nb = nanobind;
25 : using namespace nb::literals;
26 : using namespace pairinteraction;
27 :
28 : template <typename T>
29 8 : static void declare_system(nb::module_ &m, const std::string &type_name) {
30 : using S = System<T>;
31 : using scalar_t = typename System<T>::scalar_t;
32 : using real_t = typename System<T>::real_t;
33 :
34 8 : std::string pyclass_name = "System" + type_name;
35 :
36 8 : nb::class_<System<T>, SorterBuilderInterface> pyclass(m, pyclass_name.c_str());
37 8 : pyclass.def("get_basis", &S::get_basis, nb::call_guard<nb::gil_scoped_release>())
38 8 : .def("get_eigenbasis", &S::get_eigenbasis, nb::call_guard<nb::gil_scoped_release>())
39 8 : .def("get_eigenenergies", &S::get_eigenenergies, nb::call_guard<nb::gil_scoped_release>())
40 8 : .def("get_matrix", &S::get_matrix, nb::call_guard<nb::gil_scoped_release>())
41 8 : .def("get_sorter", &S::get_sorter, nb::call_guard<nb::gil_scoped_release>())
42 16 : .def("get_indices_of_blocks", &S::get_indices_of_blocks,
43 0 : nb::call_guard<nb::gil_scoped_release>())
44 8 : .def(
45 : "transform",
46 0 : [](S &self, const Eigen::SparseMatrix<scalar_t, Eigen::RowMajor> &transformation)
47 0 : -> T & { return static_cast<T &>(self.transform(transformation)); },
48 0 : nb::call_guard<nb::gil_scoped_release>())
49 8 : .def(
50 : "transform",
51 2 : [](S &self, const Eigen::PermutationMatrix<Eigen::Dynamic, Eigen::Dynamic> &sorter)
52 2 : -> T & { return static_cast<T &>(self.transform(sorter)); },
53 0 : nb::call_guard<nb::gil_scoped_release>())
54 0 : .def(
55 : "diagonalize",
56 8 : [](S &self, const DiagonalizerInterface<scalar_t> &diagonalizer,
57 : std::optional<real_t> min_eigenenergy, std::optional<real_t> max_eigenenergy,
58 : double rtol, bool sort_by_energy) -> T & {
59 0 : return static_cast<T &>(self.diagonalize(diagonalizer, min_eigenenergy,
60 0 : max_eigenenergy, rtol, sort_by_energy));
61 : },
62 32 : "diagonalizer"_a, "min_eigenenergy"_a = nb::none(), "max_eigenenergy"_a = nb::none(),
63 24 : "rtol"_a = 1e-6, "sort_by_energy"_a = true, nb::call_guard<nb::gil_scoped_release>())
64 16 : .def("is_diagonal", &S::is_diagonal, nb::call_guard<nb::gil_scoped_release>());
65 8 : }
66 :
67 : template <typename T>
68 4 : static void declare_system_atom(nb::module_ &m, const std::string &type_name) {
69 : using S = SystemAtom<T>;
70 : using basis_t = typename SystemAtom<T>::basis_t;
71 :
72 4 : std::string pyclass_name = "SystemAtom" + type_name;
73 :
74 8 : nb::class_<S, System<S>> pyclass(m, pyclass_name.c_str());
75 12 : pyclass.def(nb::init<std::shared_ptr<const basis_t>>())
76 4 : .def("set_electric_field", &S::set_electric_field)
77 4 : .def("set_magnetic_field", &S::set_magnetic_field)
78 4 : .def("set_diamagnetism_enabled", &S::set_diamagnetism_enabled)
79 4 : .def("set_ion_distance_vector", &S::set_ion_distance_vector)
80 4 : .def("set_ion_charge", &S::set_ion_charge)
81 4 : .def("set_ion_interaction_order", &S::set_ion_interaction_order)
82 8 : .def("set_green_tensor_interpolator", &S::set_green_tensor_interpolator);
83 4 : }
84 :
85 : template <typename T>
86 4 : static void declare_system_pair(nb::module_ &m, const std::string &type_name) {
87 : using S = SystemPair<T>;
88 : using basis_t = typename SystemPair<T>::basis_t;
89 :
90 4 : std::string pyclass_name = "SystemPair" + type_name;
91 :
92 8 : nb::class_<S, System<S>> pyclass(m, pyclass_name.c_str());
93 12 : pyclass.def(nb::init<std::shared_ptr<const basis_t>>())
94 4 : .def("set_interaction_order", &S::set_interaction_order)
95 4 : .def("set_distance_vector", &S::set_distance_vector)
96 8 : .def("set_green_tensor_interpolator", &S::set_green_tensor_interpolator);
97 4 : }
98 :
99 : template <typename T>
100 4 : static void declare_green_tensor_interpolator(nb::module_ &m, const std::string &type_name) {
101 : using CE = typename GreenTensorInterpolator<T>::ConstantEntry;
102 : using OE = typename GreenTensorInterpolator<T>::OmegaDependentEntry;
103 : using GT = GreenTensorInterpolator<T>;
104 :
105 4 : std::string ce_name = "ConstantEntry" + type_name;
106 4 : std::string oe_name = "OmegaDependentEntry" + type_name;
107 4 : std::string gt_name = "GreenTensorInterpolator" + type_name;
108 :
109 4 : nb::class_<CE>(m, ce_name.c_str())
110 4 : .def("row", &CE::row)
111 4 : .def("col", &CE::col)
112 8 : .def("val", &CE::val);
113 :
114 4 : nb::class_<OE>(m, oe_name.c_str())
115 4 : .def("row", &OE::row)
116 4 : .def("col", &OE::col)
117 8 : .def("val", &OE::val);
118 :
119 8 : nb::class_<GT>(m, gt_name.c_str())
120 4 : .def(nb::init<>())
121 4 : .def("create_entries_from_cartesian",
122 4 : nb::overload_cast<int, int, const Eigen::MatrixX<T> &>(
123 : >::create_entries_from_cartesian))
124 4 : .def("create_entries_from_cartesian",
125 : nb::overload_cast<int, int, const std::vector<Eigen::MatrixX<T>> &,
126 4 : const std::vector<double> &>(>::create_entries_from_cartesian))
127 : .def("get_spherical_entries",
128 8 : nb::overload_cast<int, int>(>::get_spherical_entries, nb::const_));
129 4 : }
130 :
131 2 : void bind_system(nb::module_ &m) {
132 2 : declare_system<SystemAtom<double>>(m, "SystemAtomReal");
133 2 : declare_system<SystemAtom<std::complex<double>>>(m, "SystemAtomComplex");
134 2 : declare_system_atom<double>(m, "Real");
135 2 : declare_system_atom<std::complex<double>>(m, "Complex");
136 :
137 2 : declare_system<SystemPair<double>>(m, "SystemPairReal");
138 2 : declare_system<SystemPair<std::complex<double>>>(m, "SystemPairComplex");
139 2 : declare_system_pair<double>(m, "Real");
140 2 : declare_system_pair<std::complex<double>>(m, "Complex");
141 :
142 2 : declare_green_tensor_interpolator<double>(m, "Real");
143 2 : declare_green_tensor_interpolator<std::complex<double>>(m, "Complex");
144 2 : }
|
