chiark / gitweb /
visdebug works now
[moebius3.git] / symbolic.py
1
2 from __future__ import print_function
3
4 from sympy import *
5 import itertools
6 from sympy.utilities.lambdify import lambdify, implemented_function
7 import numpy as np
8
9 from moedebug import *
10
11 r, theta, s, la, mu, kappa = symbols('r theta s lambda mu kappa')
12
13 # start      original formulation
14 # rightvars  replaces 
15
16 def dprint(*args):
17   if not dbg_enabled(): return
18   print(*args)
19
20 def dbg(*args):
21   if not dbg_enabled(): return
22   for vn in args:
23     print('\n    ' + vn + '\n')
24     pprint(eval(vn))
25     print('\n          =\n')
26     pprint(cse(eval(vn)))
27
28 calculated = False
29
30 def calculate():
31   global calculated
32   if calculated: return
33
34   p_start = Matrix([
35     r * (1 - cos(theta)),
36     r * sin(theta),
37     mu * s,
38   ])
39
40   global p_rightvars
41   p_rightvars = p_start.subs( theta, s/r ).subs( r, 1/la )
42   dbg('p_rightvars')
43
44   global p_dirn_rightvars
45   p_dirn_rightvars = diff(p_rightvars, s)
46   dbg('p_dirn_rightvars')
47
48   zeta = Wild('zeta')
49
50   global p_nosing
51   p_nosing = (p_rightvars
52               .replace( 1-cos(zeta)  ,   2*sin(zeta/2)**2          )
53               .replace( sin(zeta)**2 ,   zeta*sinc(zeta)*sin(zeta) )
54               )
55   p_nosing[1] = (p_nosing[1]
56               .replace( sin(zeta) , zeta * sinc(zeta)                )
57                  )
58
59   dbg('p_nosing')
60
61   global t
62   t = symbols('t')
63
64   global q_owncoords, q_dirn_owncoords
65   q_owncoords = p_nosing.replace(s,t).replace(la,-la)
66   q_dirn_owncoords = p_dirn_rightvars.replace(s,t).replace(la,-la)
67
68   dbg('q_owncoords','q_dirn_owncoords')
69   dbg('q_owncoords.replace(t,0)','q_dirn_owncoords.replace(t,0)')
70
71   global p2q_translate, p2q_rotate
72   p2q_translate = p_nosing
73   #p2q_rotate_2d = Matrix([ p_dirn_rightvars[0:2],
74
75   #p2q_rotate = eye(3)
76   #p2q_rotate[0:2, 0] = Matrix([ p_dirn_rightvars[1], -p_dirn_rightvars[0] ])
77   #p2q_rotate[0:2, 1] = p_dirn_rightvars[0:2]
78
79   p2q_rotate = Matrix([[  cos(theta), sin(theta), 0 ],
80                        [ -sin(theta), cos(theta), 0 ],
81                        [  0         , 0,          1 ]]).subs(theta,la*s)
82   #p2q_rotate.add_col([0,0])
83   #p2q_rotate.add_row([0,0,1])
84
85   dbg('p2q_rotate')
86
87   global q_dirn_maincoords, q_maincoords
88   q_dirn_maincoords = p2q_rotate * q_dirn_owncoords;
89   q_maincoords = p2q_rotate * q_owncoords + p2q_translate
90
91   dbg('diff(p_dirn_rightvars,s)')
92   dbg('diff(q_dirn_maincoords,t)')
93   dbg('diff(q_dirn_maincoords,t).replace(t,0)')
94
95   assert(Eq(p2q_rotate * Matrix([0,1,mu]), p_dirn_rightvars))
96
97   #for v in 's','t','la','mu':
98   #  dbg('diff(q_maincoords,%s)' % v)
99
100   #print('\n eye3 subs etc.\n')
101   #dbg('''Eq(eye(3) * Matrix([1,0,mu]),
102   #     p_dirn_rightvars .cross(Matrix([0,0,1]) .subs(s,0)))''')
103
104   #dbg('''Eq(p2q_rotate * Matrix([1,0,mu]),
105   #          p_dirn_rightvars .cross(Matrix([0,0,1])))''')
106
107   #eq = Eq(qmat * q_dirn_owncoords_0, p_dirn_rightvars)
108   #print
109   #pprint(eq)
110   #solve(eq, Q)
111
112   dbg('q_maincoords.replace(t,0)','q_dirn_maincoords.replace(t,0)')
113
114   dbg('q_maincoords','q_dirn_maincoords')
115
116   global sinof_mu, cosof_mu
117   sinof_mu = sin(atan(mu))
118   cosof_mu = cos(atan(mu))
119
120   dbg('cosof_mu','sinof_mu')
121
122   o2p_rotate1 = Matrix([[ 1,  0,         0        ],
123                         [ 0,  cosof_mu, +sinof_mu ],
124                         [ 0, -sinof_mu,  cosof_mu ]])
125
126   global check_dirn_p_s0
127   check_dirn_p_s0 = o2p_rotate1 * p_dirn_rightvars.replace(s,0)
128   check_dirn_p_s0.simplify()
129   dbg('check_dirn_p_s0')
130
131   o2p_rotate2 = Matrix([[  cos(kappa), 0, -sin(kappa) ],
132                         [  0,          1,  0          ],
133                         [ +sin(kappa), 0,  cos(kappa) ]])
134
135   p_dirn_orgcoords = o2p_rotate2 * o2p_rotate1 * p_dirn_rightvars
136
137   check_dirn_p_s0 = p_dirn_orgcoords.replace(s,0)
138   check_dirn_p_s0.simplify()
139   dbg('check_dirn_p_s0')
140
141   global check_accel_p_s0
142   check_accel_p_s0 = diff(p_dirn_orgcoords,s).replace(s,0)
143   check_accel_p_s0.simplify()
144   dbg('check_accel_p_s0')
145
146   global q_dirn_orgcoords, q_orgcoords
147   q_dirn_orgcoords = o2p_rotate2 * o2p_rotate1 * q_dirn_maincoords;
148   q_orgcoords = o2p_rotate2 * o2p_rotate1 * q_maincoords;
149   dbg('q_orgcoords','q_dirn_orgcoords')
150
151   global sh, th
152   sh, th = symbols('alpha beta')
153
154   global q_dirn_sqparm, q_sqparm
155   q_dirn_sqparm = q_dirn_orgcoords.replace(s, sh**2).replace(t, th**2)
156   q_sqparm      = q_orgcoords     .replace(s, sh**2).replace(t, th**2)
157
158   dprint('----------------------------------------')
159   dbg('q_sqparm', 'q_dirn_sqparm')
160   dprint('----------------------------------------')
161   for v in 'sh','th','la','mu':
162     dbg('diff(q_sqparm,%s)' % v)
163     dbg('diff(q_dirn_sqparm,%s)' % v)
164   dprint('----------------------------------------')
165
166   gamma = symbols('gamma')
167
168   q_dirn_dirnscaled = q_dirn_sqparm * gamma
169
170   global result_dirnscaled
171   result_dirnscaled = q_sqparm.col_join(q_dirn_dirnscaled)
172   dbg('result_dirnscaled')
173
174   calculated = True
175
176 params = ('sh','th','la','mu','gamma','kappa')
177
178 def ourccode(*a, **kw):
179   return ccode(*a, user_functions={'sinc':'sinc'}, **kw)
180
181 def cprintraw(*s):
182   print(*s)
183
184 def cprint(s):
185   for l in s.split('\n'):
186     cprintraw(l, '\\')
187
188 def cse_prep_cprint(v, tmp_prefix):
189   # => v, but also having cprint'd the common subexpression assignments
190   sym_iter = map((lambda i: symbols('%s%d' % (tmp_prefix,i))),
191                  itertools.count())
192   (defs, vs) = cse(v, symbols=sym_iter)
193   for defname, defval in defs:
194     cprint('double '+ourccode(defval, assign_to=defname))
195   return vs[0]
196
197 def cassign(v, assign_to, tmp_prefix):
198   v = cse_prep_cprint(v, tmp_prefix)
199   cprint(ourccode(v, assign_to=assign_to))
200
201 def gen_diff(current, smalls):
202   global j
203   if not smalls:
204     j = zeros(len(params),0)
205     for param in params:
206       global d
207       paramv = eval(param)
208       d = diff(current, paramv)
209       dbg('d')
210       j = j.row_join(d)
211     dbg('j')
212     j = cse_prep_cprint(j, 'jtmp')
213     for ix in range(0, j.cols):
214       cprint(ourccode(j.col(ix), 'J_COL'))
215       cprint('J_END_COL(%d)' % ix)
216   else:
217     small = smalls[0]
218     smalls = smalls[1:]
219     cprint('if (!IS_SMALL(' + ourccode(small) + ')) {')
220     gen_diff(current, smalls)
221     cprint('} else { /* %s small */' % small)
222     gen_diff(current.replace(
223       sinc(small),
224       1 - small*small/factorial(3) - small**4/factorial(5),
225       ),
226       smalls
227     )
228     cprint('} /* %s small */' % small)
229
230 def gen_misc():
231   cprintraw('// AUTOGENERATED - DO NOT EDIT\n')
232   cprintraw('#define N %d\n' % len(params))
233
234 def gen_x_extract():
235   cprint('#define X_EXTRACT')
236   for ix in range(0, len(params)):
237     cprint('double %s = X(%d);' % (eval(params[ix]), ix))
238   cprintraw()
239
240 def gen_f_populate():
241   cprint('#define F_POPULATE')
242   cassign(result_dirnscaled,'F','ftmp')
243   cprintraw('')
244
245 def gen_j_populate():
246   cprint('#define J_POPULATE')
247   gen_diff(result_dirnscaled, (sh*sh*la, th*th*la))
248   cprintraw('')
249
250 def gen_C():
251   gen_misc()
252   gen_x_extract()
253   gen_f_populate()
254   gen_j_populate()
255
256 def get_python():
257   # https://github.com/sympy/sympy/issues/13642
258   # "lambdify sinc gives wrong answer!"
259   out = q_sqparm
260   sinc_fixed = Function('sinc_fixed')
261   implemented_function(sinc_fixed, lambda x: np.sinc(x/np.pi))
262   out = out.subs(sinc,sinc_fixed)
263   p = list(map(eval,params))
264   return lambdify(p, out)