-
Notifications
You must be signed in to change notification settings - Fork 4
Expand file tree
/
Copy pathkepler.py
More file actions
135 lines (104 loc) · 3.94 KB
/
kepler.py
File metadata and controls
135 lines (104 loc) · 3.94 KB
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
130
131
132
133
134
135
# -*- coding: utf-8 -*-
"""
Implementation of a Kepler solver from from: https://github.com/dfm/theano_ops
"""
from __future__ import division, print_function
__all__ = ["KeplerOp", "get_eccentric_anomaly"]
import numpy as np
import theano
from theano import gof
import theano.tensor as tt
class KeplerOp(gof.Op):
__props__ = ("tol", "maxiter")
def __init__(self, tol=1e-8, maxiter=2000, **kwargs):
self.tol = tol
self.maxiter = maxiter
super(KeplerOp, self).__init__(**kwargs)
def make_node(self, mean_anom, eccen):
output_var = tt.TensorType(
dtype=theano.scalar.upcast(mean_anom.dtype, eccen.dtype),
broadcastable=[False] * mean_anom.ndim)()
return gof.Apply(self, [mean_anom, eccen], [output_var])
def c_code_cache_version(self):
return (0, 0, 1)
def grad(self, inputs, gradients):
M, e = inputs
E = self(M, e)
bM = gradients[0] / (1.0 - e * tt.cos(E))
be = tt.sin(E) * bM
return [bM, be]
def c_support_code_apply(self, node, name):
dtype_mean_anom = node.inputs[0].dtype
dtype_eccen = node.inputs[1].dtype
dtype_eccen_anom = node.outputs[0].dtype
c_support_code = """
inline npy_%(dtype_eccen_anom)s solve_kepler_%(name)s (
npy_%(dtype_mean_anom)s M, npy_%(dtype_eccen)s e,
int maxiter, float tol
) {
typedef npy_%(dtype_eccen_anom)s T;
T E0 = M, E = M;
if (fabs(e) < tol) return E;
for (int i = 0; i < maxiter; ++i) {
T g = E0 - e * sin(E0) - M, gp = 1.0 - e * cos(E0);
T delta = g / (gp + tol);
delta = (fabs(delta) < T(1)) ? delta : delta / fabs(delta);
E = E0 - delta;
if (fabs(E - E0) <= T(tol)) {
return E;
}
E0 = E;
}
return E;
}
"""
return c_support_code % locals()
def c_code(self, node, name, inp, out, sub):
tol = self.tol
maxiter = self.maxiter
mean_anom, eccen = inp
eccen_anom, = out
dtype_mean_anom = node.inputs[0].dtype
dtype_eccen = node.inputs[1].dtype
dtype_eccen_anom = node.outputs[0].dtype
itemsize_mean_anom = np.dtype(dtype_mean_anom).itemsize
itemsize_eccen = np.dtype(dtype_eccen).itemsize
itemsize_eccen_anom = np.dtype(dtype_eccen_anom).itemsize
typenum_eccen_anom = np.dtype(dtype_eccen_anom).num
fail = sub['fail']
c_code = """
npy_intp size = PyArray_SIZE(%(mean_anom)s);
npy_%(dtype_mean_anom)s* mean_anom;
npy_%(dtype_eccen)s* eccen;
npy_%(dtype_eccen_anom)s* eccen_anom;
// Validate that the inputs have the same shape
if ( !PyArray_SAMESHAPE(%(mean_anom)s, %(eccen)s) )
{
PyErr_Format(PyExc_ValueError, "shape mismatch");
%(fail)s;
}
// Validate that the output storage exists and has the same
// shape.
if (NULL == %(eccen_anom)s ||
!PyArray_SAMESHAPE(%(mean_anom)s, %(eccen_anom)s))
{
Py_XDECREF(%(eccen_anom)s);
%(eccen_anom)s = (PyArrayObject*)PyArray_EMPTY(
PyArray_NDIM(%(mean_anom)s),
PyArray_DIMS(%(mean_anom)s),
%(typenum_eccen_anom)s,
0);
if (!%(eccen_anom)s) {
%(fail)s;
}
}
mean_anom = (npy_%(dtype_mean_anom)s*)PyArray_DATA(%(mean_anom)s);
eccen = (npy_%(dtype_eccen)s*)PyArray_DATA(%(eccen)s);
eccen_anom = (npy_%(dtype_eccen_anom)s*)PyArray_DATA(%(eccen_anom)s);
for (npy_intp i = 0; i < size; ++i) {
eccen_anom[i] = solve_kepler_%(name)s (mean_anom[i], eccen[i],
%(maxiter)s, %(tol)s);
}
"""
return c_code % locals()
get_eccentric_anomaly = KeplerOp()