1 #include "python_ngcore.hpp"
2 #include "bitarray.hpp"
3 #include "taskmanager.hpp"
4
5 using namespace ngcore;
6 using namespace std;
7 using namespace pybind11::literals;
8
PYBIND11_MODULE(pyngcore,m)9 PYBIND11_MODULE(pyngcore, m) // NOLINT
10 {
11 try
12 {
13 auto numpy = py::module::import("numpy");
14 ngcore_have_numpy = !numpy.is_none();
15 }
16 catch(...) {}
17 ExportArray<int>(m);
18 ExportArray<unsigned>(m);
19 ExportArray<size_t>(m);
20 ExportArray<double>(m);
21
22 ExportTable<int>(m);
23
24 py::class_<BitArray, shared_ptr<BitArray>> (m, "BitArray")
25 .def(py::init([] (size_t n) { return make_shared<BitArray>(n); }),py::arg("n"))
26 .def(py::init([] (const BitArray& a) { return make_shared<BitArray>(a); } ), py::arg("ba"))
27 .def(py::init([] (const vector<bool> & a)
28 {
29 auto ba = make_shared<BitArray>(a.size());
30 ba->Clear();
31 for (size_t i = 0; i < a.size(); i++)
32 if (a[i]) ba->SetBit(i);
33 return ba;
34 } ), py::arg("vec"))
35 .def(NGSPickle<BitArray>())
36 .def("__str__", &ToString<BitArray>)
37 .def("__len__", &BitArray::Size)
38 .def("__getitem__", [] (BitArray & self, int i)
39 {
40 if (i < 0 || i >= self.Size())
41 throw py::index_error();
42 return self.Test(i);
43 }, py::arg("pos"), "Returns bit from given position")
44 .def("__setitem__", [] (BitArray & self, int i, bool b)
45 {
46 if (i < 0 || i >= self.Size())
47 throw py::index_error();
48 if (b) self.SetBit(i); else self.Clear(i);
49 }, py::arg("pos"), py::arg("value"), "Clear/Set bit at given position")
50
51 .def("__setitem__", [] (BitArray & self, py::slice inds, bool b)
52 {
53 size_t start, step, stop, n;
54 if (!inds.compute(self.Size(), &start, &stop, &step, &n))
55 throw py::error_already_set();
56
57 if (start == 0 && n == self.Size() && step == 1)
58 { // base branch
59 if (b)
60 self.Set();
61 else
62 self.Clear();
63 }
64 else
65 {
66 if (b)
67 for (size_t i=0; i<n; i++, start+=step)
68 self.SetBit(start);
69 else
70 for (size_t i=0; i<n; i++, start+=step)
71 self.Clear(start);
72 }
73 }, py::arg("inds"), py::arg("value"), "Clear/Set bit at given positions")
74
75 .def("__setitem__", [] (BitArray & self, py::slice inds, BitArray & ba)
76 {
77 size_t start, step, stop, n;
78 if (!inds.compute(self.Size(), &start, &stop, &step, &n))
79 throw py::error_already_set();
80
81 if (start == 0 && n == self.Size() && step == 1)
82 {
83 self = ba;
84 }
85 else
86 {
87 for (size_t i = 0; i < n; i++, start += step)
88 {
89 bool b = ba.Test(i);
90 if (b)
91 self.SetBit(start);
92 else
93 self.Clear(start);
94 }
95 }
96 }, py::arg("inds"), py::arg("ba"), "copy BitArray")
97
98 .def("__setitem__", [](BitArray & self, IntRange range, bool b)
99 {
100 if (b)
101 for (size_t i : range)
102 self.SetBit(i);
103 else
104 for (size_t i : range)
105 self.Clear(i);
106 }, py::arg("range"), py::arg("value"), "Set value for range of indices" )
107
108 .def("NumSet", &BitArray::NumSet)
109 .def("Set", [] (BitArray & self) { self.Set(); }, "Set all bits")
110 .def("Set", &BitArray::SetBit, py::arg("i"), "Set bit at given position")
111 .def("Clear", [] (BitArray & self) { self.Clear(); }, "Clear all bits")
112 .def("Clear", [] (BitArray & self, int i)
113 {
114 self.Clear(i);
115 }, py::arg("i"), "Clear bit at given position")
116
117
118 .def(py::self | py::self)
119 .def(py::self & py::self)
120 .def(py::self |= py::self)
121 .def(py::self &= py::self)
122 .def(~py::self)
123 ;
124
125 py::class_<Flags>(m, "Flags")
126 .def(py::init<>())
127 .def("__str__", &ToString<Flags>)
128 .def(py::init([](py::object & obj) {
129 Flags flags;
130 py::dict d(obj);
131 SetFlag (flags, "", d);
132 return flags;
133 }), py::arg("obj"), "Create Flags by given object")
134 .def(py::pickle([] (const Flags& self)
135 {
136 std::stringstream str;
137 self.SaveFlags(str);
138 return py::make_tuple(py::cast(str.str()));
139 },
140 [] (py::tuple state)
141 {
142 string s = state[0].cast<string>();
143 std::stringstream str(s);
144 Flags flags;
145 flags.LoadFlags(str);
146 return flags;
147 }
148 ))
149 .def("Set",[](Flags & self,const py::dict & aflags)->Flags&
150 {
151 SetFlag(self, "", aflags);
152 return self;
153 }, py::arg("aflag"), "Set the flags by given dict")
154
155 .def("Set",[](Flags & self, const char * akey, const py::object & value)->Flags&
156 {
157 SetFlag(self, akey, value);
158 return self;
159 }, py::arg("akey"), py::arg("value"), "Set flag by given value.")
160
161 .def("__getitem__", [](Flags & self, const string& name) -> py::object {
162
163 if(self.NumListFlagDefined(name))
164 return py::cast(self.GetNumListFlag(name));
165
166 if(self.StringListFlagDefined(name))
167 return py::cast(self.GetStringListFlag(name));
168
169 if(self.NumFlagDefined(name))
170 return py::cast(*self.GetNumFlagPtr(name));
171
172 if(self.StringFlagDefined(name))
173 return py::cast(self.GetStringFlag(name));
174
175 if(self.FlagsFlagDefined(name))
176 return py::cast(self.GetFlagsFlag(name));
177
178 return py::cast(self.GetDefineFlag(name));
179 }, py::arg("name"), "Return flag by given name")
180 .def("ToDict", [](const Flags& flags)
181 {
182 return CreateDictFromFlags(flags);
183 })
184 ;
185 py::implicitly_convertible<py::dict, Flags>();
186
187
188 py::enum_<level::level_enum>(m, "LOG_LEVEL", "Logging level")
189 .value("Trace", level::trace)
190 .value("Debug", level::debug)
191 .value("Info", level::info)
192 .value("Warn", level::warn)
193 .value("Error", level::err)
194 .value("Critical", level::critical)
195 .value("Off", level::off);
196
197 m.def("SetLoggingLevel", &SetLoggingLevel, py::arg("level"), py::arg("logger")="",
198 "Set logging level, if name is given only to the specific logger, else set the global logging level");
199 m.def("AddFileSink", &AddFileSink, py::arg("filename"), py::arg("level"), py::arg("logger")="",
200 "Add File sink, either only to logger specified or globally to all loggers");
201 m.def("AddConsoleSink", &AddConsoleSink, py::arg("level"), py::arg("logger")="",
202 "Add console output for specific logger or all if none given");
203 m.def("ClearLoggingSinks", &ClearLoggingSinks, py::arg("logger")="",
204 "Clear sinks of specific logger, or all if none given");
205 m.def("FlushOnLoggingLevel", &FlushOnLoggingLevel, py::arg("level"), py::arg("logger")="",
206 "Flush every message with level at least `level` for specific logger or all loggers if none given.");
207
208 m.def("RunWithTaskManager",
209 [](py::object lam)
210 {
211 GetLogger("TaskManager")->info("running Python function with task-manager");
212 RunWithTaskManager ([&] () { lam(); });
213 }, py::arg("lam"), R"raw_string(
214 Parameters:
215
216 lam : object
217 input function
218
219 )raw_string")
220 ;
221
222 m.def("SetNumThreads", &TaskManager::SetNumThreads, py::arg("threads"), R"raw_string(
223 Set number of threads
224
225 Parameters:
226
227 threads : int
228 input number of threads
229
230 )raw_string");
231
232 // local TaskManager class to be used as context manager in Python
233 class ParallelContextManager {
234 int num_threads;
235 public:
236 ParallelContextManager() : num_threads(0) {};
237 ParallelContextManager(size_t pajesize) : num_threads(0) {
238 TaskManager::SetPajeTrace(pajesize > 0);
239 PajeTrace::SetMaxTracefileSize(pajesize);
240 }
241 void Enter() {num_threads = EnterTaskManager(); }
242 void Exit(py::object exc_type, py::object exc_value, py::object traceback) {
243 ExitTaskManager(num_threads);
244 }
245 };
246
247 py::class_<ParallelContextManager>(m, "TaskManager")
248 .def(py::init<>())
249 .def(py::init<size_t>(), "pajetrace"_a, "Run paje-tracer, specify buffersize in bytes")
250 .def("__enter__", &ParallelContextManager::Enter)
251 .def("__exit__", &ParallelContextManager::Exit)
252 .def("__timing__", &TaskManager::Timing)
253 ;
254
255 py::class_<PajeTrace>(m, "PajeTrace")
256 .def(py::init( [] (string filename, size_t size_mb, bool threads, bool thread_counter, bool memory)
257 {
258 PajeTrace::SetMaxTracefileSize(size_mb*1014*1024);
259 PajeTrace::SetTraceThreads(threads);
260 PajeTrace::SetTraceMemory(memory);
261 PajeTrace::SetTraceThreadCounter(thread_counter);
262 trace = new PajeTrace(TaskManager::GetMaxThreads(), filename);
263 return trace;
264 }), py::arg("filename")="ng.trace", py::arg("size")=1000,
265 py::arg("threads")=true, py::arg("thread_counter")=false,
266 py::arg("memory")=true,
267 "size in Megabytes"
268 )
269 .def("__enter__", [](PajeTrace & self) { })
270 .def("__exit__", [](PajeTrace & self, py::args) { self.StopTracing(); })
271 .def("__del__", [](PajeTrace & self) { trace = nullptr; })
272 .def_static("SetTraceThreads", &PajeTrace::SetTraceThreads)
273 .def_static("SetTraceThreadCounter", &PajeTrace::SetTraceThreadCounter)
274 .def_static("SetMaxTracefileSize", &PajeTrace::SetMaxTracefileSize)
275 #ifdef NETGEN_TRACE_MEMORY
276 .def_static("WriteMemoryChart", [](string filename){ if(trace) trace->WriteMemoryChart(filename); }, py::arg("filename")="memory" )
277 #endif // NETGEN_TRACE_MEMORY
278 ;
279
280
281 }
282