chiark / gitweb /
symbolic.py: make -q work
[moebius3.git] / symbolic.py
1 #!/usr/bin/python3
2
3 from sympy import *
4 import itertools
5
6 import sys, codecs
7 from optparse import OptionParser
8
9 opt_parser = OptionParser()
10 opt_parser.add_option('-q',dest='quiet',action='store_true',
11                       default=False,help='suppress diagnostic output')
12 opt_parser.add_option('-7',dest='ascii',action='store_true',
13                       default=False,help='use 7-bit output only')
14 opt_parser.add_option('-w',dest='width',action='store',type='int',
15                       default=80,help='width for printing')
16 (options, args) = opt_parser.parse_args()
17 assert(not len(args))
18
19 if (not options.ascii) and sys.stdout.encoding is None:
20   sys.stdout = codecs.open("/dev/stdout", "w", 'utf-8')
21
22 init_printing(use_unicode=(not options.ascii), num_columns=options.width)
23
24 r, theta, s, la, mu, kappa = symbols('r theta s lambda mu kappa')
25
26 # start      original formulation
27 # rightvars  replaces 
28
29 p_start = Matrix([
30   r * (1 - cos(theta)),
31   r * sin(theta),
32   mu * s,
33 ])
34
35 p_rightvars = p_start.subs( theta, s/r ).subs( r, 1/la )
36
37 def dprint(*args):
38   if options.quiet: return
39   print(*args)
40
41 def dbg(*args):
42   if options.quiet: return
43   for vn in args:
44     print('\n    ' + vn + '\n')
45     pprint(eval(vn))
46     print('\n          =\n')
47     pprint(cse(eval(vn)))
48
49 dbg('p_rightvars')
50
51 p_dirn_rightvars = diff(p_rightvars, s)
52
53 dbg('p_dirn_rightvars')
54
55 zeta = Wild('zeta')
56
57 p_nosing = (p_rightvars
58             .replace( 1-cos(zeta)  ,   2*sin(zeta/2)**2          )
59             .replace( sin(zeta)**2 ,   zeta*sinc(zeta)*sin(zeta) )
60             )
61 p_nosing[1] = (p_nosing[1]
62             .replace( sin(zeta) , zeta * sinc(zeta)                )
63                )
64
65 dbg('p_nosing')
66
67 t = symbols('t')
68
69 q_owncoords = p_nosing.replace(s,t).replace(la,-la)
70 q_dirn_owncoords = p_dirn_rightvars.replace(s,t).replace(la,-la)
71
72 dbg('q_owncoords','q_dirn_owncoords')
73 dbg('q_owncoords.replace(t,0)','q_dirn_owncoords.replace(t,0)')
74
75 p2q_translate = p_nosing
76 #p2q_rotate_2d = Matrix([ p_dirn_rightvars[0:2],
77                          
78 #p2q_rotate = eye(3)
79 #p2q_rotate[0:2, 0] = Matrix([ p_dirn_rightvars[1], -p_dirn_rightvars[0] ])
80 #p2q_rotate[0:2, 1] = p_dirn_rightvars[0:2]
81
82 p2q_rotate = Matrix([[  cos(theta), sin(theta), 0 ],
83                      [ -sin(theta), cos(theta), 0 ],
84                      [  0         , 0,          1 ]]).subs(theta,la*s)
85 #p2q_rotate.add_col([0,0])
86 #p2q_rotate.add_row([0,0,1])
87
88 dbg('p2q_rotate')
89
90 q_dirn_maincoords = p2q_rotate * q_dirn_owncoords;
91 q_maincoords = p2q_rotate * q_owncoords + p2q_translate
92
93 dbg('diff(p_dirn_rightvars,s)')
94 dbg('diff(q_dirn_maincoords,t)')
95 dbg('diff(q_dirn_maincoords,t).replace(t,0)')
96
97 assert(Eq(p2q_rotate * Matrix([0,1,mu]), p_dirn_rightvars))
98
99 #for v in 's','t','la','mu':
100 #  dbg('diff(q_maincoords,%s)' % v)
101
102 #print('\n eye3 subs etc.\n')
103 #dbg('''Eq(eye(3) * Matrix([1,0,mu]),
104 #     p_dirn_rightvars .cross(Matrix([0,0,1]) .subs(s,0)))''')
105
106 #dbg('''Eq(p2q_rotate * Matrix([1,0,mu]),
107 #          p_dirn_rightvars .cross(Matrix([0,0,1])))''')
108
109 #eq = Eq(qmat * q_dirn_owncoords_0, p_dirn_rightvars)
110 #print
111 #pprint(eq)
112 #solve(eq, Q)
113
114 dbg('q_maincoords.replace(t,0)','q_dirn_maincoords.replace(t,0)')
115
116 dbg('q_maincoords','q_dirn_maincoords')
117
118 sinof_mu = sin(atan(mu))
119 cosof_mu = cos(atan(mu))
120
121 dbg('cosof_mu','sinof_mu')
122
123 o2p_rotate1 = Matrix([[ 1,  0,         0        ],
124                       [ 0,  cosof_mu, +sinof_mu ],
125                       [ 0, -sinof_mu,  cosof_mu ]])
126
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 check_accel_p_s0 = diff(p_dirn_orgcoords,s).replace(s,0)
142 check_accel_p_s0.simplify()
143 dbg('check_accel_p_s0')
144
145 q_dirn_orgcoords = o2p_rotate2 * o2p_rotate1 * q_dirn_maincoords;
146 q_orgcoords = o2p_rotate2 * o2p_rotate1 * q_maincoords;
147 dbg('q_orgcoords','q_dirn_orgcoords')
148
149 sh, th = symbols('alpha beta')
150
151 q_dirn_sqparm = q_dirn_orgcoords.replace(s, sh**2).replace(t, th**2)
152 q_sqparm      = q_orgcoords     .replace(s, sh**2).replace(t, th**2)
153
154 dprint('----------------------------------------')
155 dbg('q_sqparm', 'q_dirn_sqparm')
156 dprint('----------------------------------------')
157 for v in 'sh','th','la','mu':
158   dbg('diff(q_sqparm,%s)' % v)
159   dbg('diff(q_dirn_sqparm,%s)' % v)
160 dprint('----------------------------------------')
161
162 gamma = symbols('gamma')
163
164 q_dirn_dirnscaled = q_dirn_sqparm * gamma
165
166 result_dirnscaled = q_sqparm.col_join(q_dirn_dirnscaled)
167 dbg('result_dirnscaled')
168
169 params = ('sh','th','la','mu','gamma','kappa')
170
171 def cprint(s):
172   print(s)
173
174 def cse_prep_cprint(v, tmp_prefix):
175   # => v, but also having cprint'd the common subexpression assignments
176   sym_iter = map((lambda i: symbols('%s%d' % (tmp_prefix,i))),
177                  itertools.count())
178   (defs, vs) = cse(v, symbols=sym_iter)
179   for defname, defval in defs:
180     cprint(ccode(defval, assign_to=defname))
181   return vs[0]
182
183 def cassign(v, assign_to, tmp_prefix):
184   v = cse_prep_cprint(v, tmp_prefix)
185   cprint(ccode(v, assign_to=assign_to))
186
187 def gen_diff(current, smalls):
188   global j
189   if not smalls:
190     j = zeros(len(params),0)
191     for param in params:
192       global d
193       paramv = eval(param)
194       d = diff(current, paramv)
195       dbg('d')
196       j = j.row_join(d)
197     dbg('j')
198     j = cse_prep_cprint(j, 'jtmp')
199     for ix in range(0, j.cols):
200       cprint(ccode(j.col(ix), 'J(%d)' % ix))
201       cprint('J_END_COL(%d)' % ix)
202   else:
203     small = smalls[0]
204     smalls = smalls[1:]
205     cprint('if (!IS_SMALL(' + ccode(small) + ')) {')
206     gen_diff(current, smalls)
207     cprint('} else { /* %s small */' % small)
208     gen_diff(current.replace(
209       sinc(small),
210       1 - small*small/factorial(3) - small**4/factorial(5),
211       ),
212       smalls
213     )
214     print('} /* %s small */' % small)
215
216 gen_diff(result_dirnscaled, (sh*sh*la, th*th*la))
217
218 #bad = q_orgcoords[0]
219 #badd = diff(bad, la)
220
221 #dbg('bad','badd')