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
|
#
# Copyright 2014 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 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),
("solution_branch", ctypes.c_uint),
]
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')
|
