DryLab/drylab.cpp at main · farismuhammad17/DryLab · 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
#include "include/ElemsDB.h"
#include "include/system.h"
#include "include/atom.h"
#include "include/vector3.h"

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/operators.h>

#include <vector>

namespace py = pybind11;

PYBIND11_MODULE(drylab, m) {
    py::class_<ElementDB>(m, "ElementDB")
        .def_static("load", &ElementDB::load);

    py::class_<System>(m, "System")
        .def(py::init<std::vector<Atom*>, double, double, double, double>(),
			py::arg("atoms"),
			py::arg("dimensionX"),
			py::arg("dimensionY"),
			py::arg("dimensionZ"),
			py::arg("cellSize") = -1.0f)
        .def("run", &System::run)
        .def("run_steps", &System::run_steps)

        .def("makeBond", &System::makeBond)

        .def("getPositions", &System::getPositions)
        .def("getBondCoords", &System::getBondCoords)
        .def("getTotalKE", &System::getTotalKE)

        .def_readwrite("atoms", &System::atoms);

    py::class_<Atom>(m, "Atom")
        .def(py::init<int, Vector3, Vector3, double, double>(),
            py::arg("atomic_number"),
            py::arg("position"),
            py::arg("velocity"),
            py::arg("charge") = 0.0f,
            py::arg("velocityloss_pertick") = 0.0)
        .def("bondWith", &Atom::bondWith)
        .def("breakWith", &Atom::breakWith)

        .def("spawnLonePairs", &Atom::spawnLonePairs)
    	.def("getLPPositions", &Atom::getLPPositions)
        .def_readwrite("numLonePairs", &Atom::numLonePairs)

        .def_readwrite("pos", &Atom::pos)
        .def_readwrite("vel", &Atom::vel)

        .def_readwrite("bonds", &Atom::bonds)

        .def_readwrite("atomID", &Atom::atomID)
        .def_readwrite("moleculeID", &Atom::moleculeID)

        .def("__repr__", [](const Atom &p) {
            return "(" + std::to_string(p.AtomicNumber) + ")";
        });

    py::class_<Vector3>(m, "Vector3")
        .def(py::init<double, double, double>())

        .def("magSq", &Vector3::magSq)
        .def("normalized", &Vector3::normalized)
        .def("cross", &Vector3::cross)

        .def(py::self + py::self)
        .def(py::self - py::self)
        .def(py::self * py::self)
        .def(py::self * double())
        .def(py::self / double())

        .def_readwrite("x", &Vector3::x)
        .def_readwrite("y", &Vector3::y)
        .def_readwrite("z", &Vector3::z)

        .def("__repr__", [](const Vector3 &v) {
            return "(" + std::to_string(v.x) + ", " + std::to_string(v.y) + ", " + std::to_string(v.z) + ")";
        });
}