abacus-mc/python/pyabacus/src/ModuleBase/py_base_math.cpp at 0be2efa68ed8decda667fa3f10d9158f6807a09d · mohanchen/abacus-mc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>

#include "source_base/math_sphbes.h"
#include "source_base/math_integral.h"
#include "source_base/sph_bessel_tf.h"

#include "../utils/pybind_utils.h"

namespace py = pybind11;
using namespace pybind11::literals;
using namespace pyabacus::utils;

template <typename... Args>
using overload_cast_ = pybind11::detail::overload_cast_impl<Args...>;

void bind_base_math(py::module& m)
{
    // python binding for class Sphbes
    py::class_<ModuleBase::Sphbes>(m, "Sphbes")
        .def(py::init<>())
        .def_static("sphbesj", overload_cast_<const int, const double>()(&ModuleBase::Sphbes::sphbesj), "l"_a, "x"_a)
        .def_static("dsphbesj", overload_cast_<const int, const double>()(&ModuleBase::Sphbes::dsphbesj), "l"_a, "x"_a)
        .def_static("sphbesj",
                    [](const int n, py::array_t<double> r, const double q, const int l, py::array_t<double> jl) {
                        check_1d_array(r, "r");
                        check_1d_array(jl, "jl");
                        ModuleBase::Sphbes::sphbesj(n,
                                                    get_array_ptr(r),
                                                    q,
                                                    l,
                                                    get_array_ptr(jl));
                    })
        .def_static("dsphbesj",
                    [](const int n, py::array_t<double> r, const double q, const int l, py::array_t<double> djl) {
                        check_1d_array(r, "r");
                        check_1d_array(djl, "djl");
                        ModuleBase::Sphbes::dsphbesj(n,
                                                     get_array_ptr(r),
                                                     q,
                                                     l,
                                                     get_array_ptr(djl));
                    })
        .def_static("sphbes_zeros", [](const int l, const int n, py::array_t<double> zeros) {
            check_1d_array(zeros, "zeros");
            ModuleBase::Sphbes::sphbes_zeros(l, n, get_array_ptr(zeros));
        });

    // python binding for class Integral
    py::class_<ModuleBase::Integral>(m, "Integral")
        .def(py::init<>())
        .def_static("Simpson_Integral", [](const int mesh, py::array_t<double> func, py::array_t<double> rab, double asum) {
            check_1d_array(func, "func");
            check_1d_array(rab, "rab");

            double isum = asum;
            ModuleBase::Integral::Simpson_Integral(mesh,
                                                    get_array_ptr(func),
                                                    get_array_ptr(rab),
                                                    isum);
            return isum;
        })
        .def_static("Simpson_Integral", [](const int mesh, py::array_t<double> func, const double dr, double asum){
            check_1d_array(func, "func");

            double isum = asum;
            ModuleBase::Integral::Simpson_Integral(mesh,
                                                    get_array_ptr(func),
                                                    dr,
                                                    isum);
            return isum;
        })
        .def_static("Simpson_Integral_0toall", [](const int mesh, py::array_t<double> func, py::array_t<double> rab, py::array_t<double> asum){
            check_1d_array(func, "func");
            check_1d_array(rab, "rab");
            check_1d_array(asum, "asum");
            ModuleBase::Integral::Simpson_Integral_0toall(mesh,
                                                            get_array_ptr(func),
                                                            get_array_ptr(rab),
                                                            get_array_ptr(asum));
        })
        .def_static("Simpson_Integral_alltoinf", [](const int mesh, py::array_t<double> func, py::array_t<double> rab, py::array_t<double> asum){
            check_1d_array(func, "func");
            check_1d_array(rab, "rab");
            check_1d_array(asum, "asum");
            ModuleBase::Integral::Simpson_Integral_alltoinf(mesh,
                                                            get_array_ptr(func),
                                                            get_array_ptr(rab),
                                                            get_array_ptr(asum));
        })
        .def_static("simpson", [](const int n, py::array_t<double> f, const double dx){
            check_1d_array(f, "f");
            return ModuleBase::Integral::simpson(n,
                                                    get_array_ptr(f),
                                                    dx);
        })
        .def_static("simpson", [](const int n, py::array_t<double> f, py::array_t<double> h){
            check_1d_array(f, "f");
            check_1d_array(h, "h");
            return ModuleBase::Integral::simpson(n,
                                                    get_array_ptr(f),
                                                    get_array_ptr(h));
        })
        .def_static("Gauss_Legendre_grid_and_weight", [](const int n, py::array_t<double> x, py::array_t<double> w){
            check_1d_array(x, "x");
            check_1d_array(w, "w");
            ModuleBase::Integral::Gauss_Legendre_grid_and_weight(n,
                                                                    get_array_ptr(x),
                                                                    get_array_ptr(w));
        })
        .def_static("Gauss_Legendre_grid_and_weight", [](const double xmin, const double xmax, const int n, py::array_t<double> x, py::array_t<double> w){
            check_1d_array(x, "x");
            check_1d_array(w, "w");
            ModuleBase::Integral::Gauss_Legendre_grid_and_weight(xmin,
                                                                    xmax,
                                                                    n,
                                                                    get_array_ptr(x),
                                                                    get_array_ptr(w));
        });
    py::class_<ModuleBase::SphericalBesselTransformer>(m, "SphericalBesselTransformer")
        .def(py::init<>());
}

PYBIND11_MODULE(_base_pack, m)
{
    m.doc() = "Submodule for pyabacus: ModuleBase";

    bind_base_math(m);
}