From 234e722c5abfb15c9850b4224cabbb3e2d0c3657 Mon Sep 17 00:00:00 2001 From: Anton Khirnov Date: Fri, 9 Feb 2018 15:34:25 +0100 Subject: Add the remaining parts for teukolsky waves initial data --- teukolsky_data.py | 128 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 teukolsky_data.py (limited to 'teukolsky_data.py') diff --git a/teukolsky_data.py b/teukolsky_data.py new file mode 100644 index 0000000..5bc9364 --- /dev/null +++ b/teukolsky_data.py @@ -0,0 +1,128 @@ +# +# Copyright 2014 Anton Khirnov +# +# 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 . +# + + +import ctypes +import numpy as np + +class TeukolskyData(object): + + coeffs = None + + _libtd = None + _tdctx = None + + class _TDContext(ctypes.Structure): + _fields_ = [("priv", ctypes.c_void_p), + ("log_callback", ctypes.c_void_p), + ("opaque", ctypes.c_void_p), + ("amplitude", ctypes.c_double), + ("nb_coeffs", ctypes.c_uint * 2), + ("basis_scale_factor", ctypes.c_double * 2), + ("max_iter", ctypes.c_uint), + ("atol", ctypes.c_double), + ("nb_threads", ctypes.c_uint), + ("coeffs", ctypes.POINTER(ctypes.c_double) * 3), + ] + + def __init__(self, **kwargs): + self._libtd = ctypes.CDLL('libteukolskydata.so') + tdctx_alloc = self._libtd.td_context_alloc + tdctx_alloc.restype = ctypes.POINTER(self._TDContext) + self._tdctx = self._libtd.td_context_alloc() + + coeffs_init = ctypes.c_void_p() + + for arg, value in kwargs.iteritems(): + if arg == 'coeffs_init': + coeffs_init = (ctypes.POINTER(ctypes.c_double) * 3)() + coeffs_init[0] = ctypes.cast(np.ctypeslib.as_ctypes(value[0]), ctypes.POINTER(ctypes.c_double)) + coeffs_init[1] = ctypes.cast(np.ctypeslib.as_ctypes(value[1]), ctypes.POINTER(ctypes.c_double)) + coeffs_init[2] = ctypes.cast(np.ctypeslib.as_ctypes(value[2]), ctypes.POINTER(ctypes.c_double)) + continue + + try: + self._tdctx.contents.__setattr__(arg, value) + except TypeError as e: + # try assigning items of an iterable + try: + for i, it in enumerate(value): + self._tdctx.contents.__getattribute__(arg)[i] = it + except: + raise e + + ret = self._libtd.td_solve(self._tdctx, coeffs_init) + if ret < 0: + raise RuntimeError('Error solving the equation') + + self.coeffs = [None] * 3 + for i in xrange(3): + self.coeffs[i] = np.copy(np.ctypeslib.as_array(self._tdctx.contents.coeffs[i], (self._tdctx.contents.nb_coeffs[1], self._tdctx.contents.nb_coeffs[0]))) + + def __del__(self): + if self._tdctx: + addr_tdctx = ctypes.c_void_p(ctypes.addressof(self._tdctx)) + self._libtd.td_context_free(addr_tdctx) + self._tdctx = None + + def _eval_var(self, eval_func, r, theta, diff_order = None): + if diff_order is None: + diff_order = [0, 0] + + c_diff_order = (ctypes.c_uint * 2)() + c_diff_order[0] = diff_order[0] + c_diff_order[1] = diff_order[1] + + if r.ndim == 2: + if r.shape != theta.shape: + raise TypeError('r and theta must be identically-shaped 2-dimensional arrays') + R, Theta = r.view(), theta.view() + elif r.ndim == 1: + if theta.ndim != 1: + raise TypeError('r and theta must both be 1-dimensional NumPy arrays') + R, Theta = np.meshgrid(r, theta) + else: + raise TypeError('invalid r/theta parameters') + + out = np.empty(R.shape[0] * R.shape[1]) + + R.shape = out.shape + Theta.shape = out.shape + + c_out = np.ctypeslib.as_ctypes(out) + c_r = np.ctypeslib.as_ctypes(R) + c_theta = np.ctypeslib.as_ctypes(Theta) + + ret = eval_func(self._tdctx, out.shape[0], c_r, c_theta, + c_diff_order, c_out, ctypes.c_long(r.shape[0])) + if ret < 0: + raise RuntimeError('Error evaluating the variable') + + out.shape = (theta.shape[0], r.shape[0]) + return out + + def eval_psi(self, r, theta, diff_order = None): + return self._eval_var(self._libtd.td_eval_psi, r, theta, diff_order) + def eval_krr(self, r, theta, diff_order = None): + return self._eval_var(self._libtd.td_eval_krr, r, theta, diff_order) + def eval_kpp(self, r, theta, diff_order = None): + return self._eval_var(self._libtd.td_eval_kpp, r, theta, diff_order) + + + @property + def amplitude(self): + return self._tdctx.contents.__getattribute__('amplitude') -- cgit v1.2.3