Add a python interface.python

1 files changed, 265 insertions, 0 deletions

diff --git a/mg2d.py b/mg2d.py
new file mode 100644
index 0000000..6419a88
--- /dev/null
+++ b/mg2d.py
@@ -0,0 +1,265 @@
+#
+# Copyright 2019 Anton Khirnov <anton@khirnov.net>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+
+
+import os
+import ctypes
+import numpy as np
+
+DIFF_COEFF_00 = 0
+DIFF_COEFF_01 = 1
+DIFF_COEFF_10 = 2
+DIFF_COEFF_11 = 3
+DIFF_COEFF_02 = 4
+DIFF_COEFF_20 = 5
+
+BC_TYPE_FIXVAL  = 0
+BC_TYPE_REFLECT = 1
+BC_TYPE_FALLOFF = 2
+
+BOUNDARY_0L = 0
+BOUNDARY_0U = 1
+BOUNDARY_1L = 2
+BOUNDARY_1U = 3
+
+ERR_DIVERGE = 0xff00
+
+_doubleptr = ctypes.POINTER(ctypes.c_double)
+
+if ctypes.sizeof(ctypes.c_void_p) == 8:
+    _ptrdiff_t = ctypes.c_int64
+elif ctypes.sizeof(ctypes.c_void_p) == 4:
+    _ptrdiff_t = ctypes.c_int32
+else:
+    raise NotImplementedError('Unsupported pointer size: %d' %\
+                              ctypes.sizeof(ctypes.c_void_p))
+
+def _wrap_array(obj, stride, len):
+    arr_np = np.ctypeslib.as_array(obj, (len, stride))
+    return arr_np[:, :len]
+
+class SolveError(Exception):
+    errno = None
+    def __init__(self, errno):
+        super().__init__(self, 'Error solving the equation: %s (%d) ' % (os.strerror(errno), errno))
+        self.errno = errno
+
+class _MG2DBoundary(ctypes.Structure):
+    _fields_ = [
+        ("type",        ctypes.c_int),
+        ("val",         _doubleptr),
+        ("val_stride",  ctypes.c_size_t),
+    ]
+
+class _MG2DDCBoundarySpec(ctypes.Structure):
+    _fields_ = [
+        ("flags",        ctypes.c_int),
+        ("val",          _doubleptr),
+    ]
+
+class _MG2DDiffCoeffs(ctypes.Structure):
+    _fields_ = [
+        ("data",        _doubleptr),
+        ("stride",      _ptrdiff_t),
+        ("boundaries",  _MG2DDCBoundarySpec * 4),
+    ]
+
+class _MG2DContext(ctypes.Structure):
+    _fields_ = [
+        ("priv",                   ctypes.c_void_p),
+        ("log_callback",           ctypes.c_void_p),
+        ("opaque",                 ctypes.c_void_p),
+        ("domain_size",            ctypes.c_size_t),
+        ("step",                   ctypes.c_double * 2),
+        ("fd_stencil",             ctypes.c_size_t),
+        ("boundaries",             ctypes.POINTER(_MG2DBoundary) * 4),
+        ("maxiter",                ctypes.c_uint),
+        ("tol",                    ctypes.c_double),
+        ("nb_cycles",              ctypes.c_uint),
+        ("nb_relax_pre",           ctypes.c_uint),
+        ("nb_relax_post",          ctypes.c_uint),
+        ("u",                      _doubleptr),
+        ("u_stride",               _ptrdiff_t),
+        ("rhs",                    _doubleptr),
+        ("rhs_stride",             _ptrdiff_t),
+        #("diff_coeffs",            ctypes.POINTER(_MG2DDiffCoeffs) * 6),
+        ("diff_coeffs",            _doubleptr * 6),
+        ("diff_coeffs_stride",     _ptrdiff_t),
+        ("nb_threads",             ctypes.c_uint),
+        ("cfl_factor",             ctypes.c_double),
+        ("log_level",              ctypes.c_int),
+        ("max_levels",             ctypes.c_uint),
+        ("max_exact_size",         ctypes.c_size_t),
+    ]
+
+class Boundary:
+    type   = None
+    lines  = None
+
+    _arr = None
+
+    def __init__(self, val, val_stride, max_stencil, domain_size):
+        self._arr = np.ctypeslib.as_array(val, (val_stride * max_stencil,))
+
+        self.lines = []
+        for i in range(max_stencil):
+            line_offset = i * (val_stride - 1)
+            line_len    = domain_size + 2 * i
+            self.lines.append(self._arr[line_offset:line_offset + line_len])
+
+class DiffCoeffsBoundarySpec:
+    _mg2dctx = None
+    _dc_idx  = None
+    _bnd_idx = None
+
+    def __init__(self, mg2dctx, dc_idx, bnd_idx):
+        self._mg2dctx = mg2dctx
+        self._dc_idx  = dc_idx
+        self._bnd_idx = bnd_idx
+
+    @property
+    def val(self):
+        dc = self._mg2dctx.contents.diff_coeffs[self._dc_idx].contents
+        return np.ctypeslib.as_array(dc.boundaries[self._bnd_idx].val, (self._mg2dctx.contents.domain_size,))
+
+    @property
+    def flags(self):
+        dc = self._mg2dctx.contents.diff_coeffs[self._dc_idx].contents
+        return dc.boundaries[self._bnd_idx].flags
+    @flags.setter
+    def flags(self, val):
+        dc = self._mg2dctx.contents.diff_coeffs[self._dc_idx].contents
+        dc.boundaries[self._bnd_idx].flags = val
+
+class DiffCoeffs:
+
+    boundaries = None
+
+    _mg2dctx = None
+    _idx     = None
+
+    def __init__(self, mg2dctx, idx):
+        self._mg2dctx = mg2dctx
+        self._idx     = idx
+
+        #self.boundaries = []
+        #for i in range(4):
+        #    self.boundaries.append(DiffCoeffsBoundarySpec(mg2dctx, idx, i))
+
+    @property
+    def data(self):
+        #dc = self._mg2dctx.contents.diff_coeffs[self._idx].contents
+        #return _wrap_array(dc.data, dc.stride, self._mg2dctx.contents.domain_size)
+        data = self._mg2dctx.contents.diff_coeffs[self._idx]
+        stride = self._mg2dctx.contents.diff_coeffs_stride
+        return _wrap_array(data, stride, self._mg2dctx.contents.domain_size)
+
+    def __getitem__(self, key):
+        return self.data[key]
+
+    def __setitem__(self, key, val):
+        self.data[key] = val
+
+class Solver:
+    diff_coeffs = None
+    boundaries  = None
+
+    _libmg2d        = None
+    _mg2dctx        = None
+
+    def __init__(self, domain_size):
+        self._lib_init()
+
+        self._mg2dctx = self._libmg2d.mg2d_solver_alloc(domain_size)
+        if self._mg2dctx is None:
+            raise MemoryError('Error allocating the solver')
+
+        self.diff_coeffs     = []
+        for i in range(len(self._mg2dctx.contents.diff_coeffs)):
+            self.diff_coeffs.append(DiffCoeffs(self._mg2dctx, i))
+
+        self.boundaries = []
+        for i, bnd in enumerate(self._mg2dctx.contents.boundaries):
+            boundary = Boundary(bnd.contents.val, bnd.contents.val_stride,
+                                self._libmg2d.mg2d_max_fd_stencil(), domain_size)
+            self.boundaries.append(boundary)
+
+    def _lib_init(self):
+        self._libmg2d = ctypes.CDLL('libmg2d.so')
+
+        solver_alloc = self._libmg2d.mg2d_solver_alloc
+        solver_alloc.argtypes = [ctypes.c_size_t]
+        solver_alloc.restype  = ctypes.POINTER(_MG2DContext)
+
+        solver_free = self._libmg2d.mg2d_solver_free
+        solver_free.argtypes  = [ctypes.POINTER(ctypes.POINTER(_MG2DContext))]
+        solver_free.restype   = None
+
+        solver_solve = self._libmg2d.mg2d_solve
+        solver_solve.argtypes = [ctypes.POINTER(_MG2DContext)]
+        solver_solve.restype  = ctypes.c_int
+
+        solver_print_stats = self._libmg2d.mg2d_print_stats
+        solver_print_stats.argtypes = [ctypes.POINTER(_MG2DContext), ctypes.c_char_p]
+        solver_print_stats.restype  = None
+
+        max_fd_stencil = self._libmg2d.mg2d_max_fd_stencil
+        max_fd_stencil.restype = ctypes.c_uint
+
+    def __del__(self):
+        if self._mg2dctx:
+            free = self._libmg2d.mg2d_solver_free
+            addr_mg2dctx = ctypes.cast((ctypes.addressof(self._mg2dctx)), free.argtypes[0])
+            free(addr_mg2dctx)
+            self._mg2dctx = None
+
+    @property
+    def u(self):
+        return _wrap_array(self._mg2dctx.contents.u, self._mg2dctx.contents.u_stride,
+                           self._mg2dctx.contents.domain_size)
+    @property
+    def rhs(self):
+        return _wrap_array(self._mg2dctx.contents.rhs, self._mg2dctx.contents.rhs_stride,
+                           self._mg2dctx.contents.domain_size)
+
+    def solve(self, **kwargs):
+        for arg, value in kwargs.items():
+            try:
+                self._mg2dctx.contents.__setattr__(arg, value)
+            except TypeError as e:
+                # try assigning items of an iterable
+                try:
+                    for i, it in enumerate(value):
+                        self._mg2dctx.contents.__getattribute__(arg)[i] = it
+                except:
+                    raise e
+
+        for i, bnd in enumerate(self.boundaries):
+            if bnd.type is None:
+                raise ValueError('Boundary %d not set' % i)
+
+            self._mg2dctx.contents.boundaries[i].contents.type = bnd.type
+
+        ret = self._libmg2d.mg2d_solve(self._mg2dctx)
+        if ret < 0:
+            raise SolveError(-ret)
+
+    def print_stats(self):
+        log_level = self._mg2dctx.contents.log_level
+        self._mg2dctx.contents.log_level = 0x40
+        self._libmg2d.mg2d_print_stats(self._mg2dctx, ctypes.c_char_p(None))
+        self._mg2dctx.contents.log_level = log_level