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