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 ourccode(*a, **kw):
 172   return ccode(*a, user_functions={'sinc':'sinc'}, **kw)
 173
 174 def cprintraw(*s):
 175   print(*s)
 176
 177 def cprint(s):
 178   for l in s.split('\n'):
 179     cprintraw(l, '\\')
 180
 181 def cse_prep_cprint(v, tmp_prefix):
 182   # => v, but also having cprint'd the common subexpression assignments
 183   sym_iter = map((lambda i: symbols('%s%d' % (tmp_prefix,i))),
 184                  itertools.count())
 185   (defs, vs) = cse(v, symbols=sym_iter)
 186   for defname, defval in defs:
 187     cprint('double '+ourccode(defval, assign_to=defname))
 188   return vs[0]
 189
 190 def cassign(v, assign_to, tmp_prefix):
 191   v = cse_prep_cprint(v, tmp_prefix)
 192   cprint(ourccode(v, assign_to=assign_to))
 193
 194 def gen_diff(current, smalls):
 195   global j
 196   if not smalls:
 197     j = zeros(len(params),0)
 198     for param in params:
 199       global d
 200       paramv = eval(param)
 201       d = diff(current, paramv)
 202       dbg('d')
 203       j = j.row_join(d)
 204     dbg('j')
 205     j = cse_prep_cprint(j, 'jtmp')
 206     for ix in range(0, j.cols):
 207       cprint(ourccode(j.col(ix), 'J_COL'))
 208       cprint('J_END_COL(%d)' % ix)
 209   else:
 210     small = smalls[0]
 211     smalls = smalls[1:]
 212     cprint('if (!IS_SMALL(' + ourccode(small) + ')) {')
 213     gen_diff(current, smalls)
 214     cprint('} else { /* %s small */' % small)
 215     gen_diff(current.replace(
 216       sinc(small),
 217       1 - small*small/factorial(3) - small**4/factorial(5),
 218       ),
 219       smalls
 220     )
 221     cprint('} /* %s small */' % small)
 222
 223 def gen_misc():
 224   cprintraw('// AUTOGENERATED - DO NOT EDIT\n')
 225   cprintraw('#define N %d\n' % len(params))
 226
 227 def gen_x_extract():
 228   cprint('#define X_EXTRACT')
 229   for ix in range(0, len(params)):
 230     cprint('double %s = X(%d);' % (eval(params[ix]), ix))
 231   cprintraw()
 232
 233 def gen_f_populate():
 234   cprint('#define F_POPULATE')
 235   cassign(result_dirnscaled,'F','ftmp')
 236   cprintraw('')
 237
 238 def gen_j_populate():
 239   cprint('#define J_POPULATE')
 240   gen_diff(result_dirnscaled, (sh*sh*la, th*th*la))
 241   cprintraw('')
 242
 243 def gen_C():
 244   gen_misc()
 245   gen_x_extract()
 246   gen_f_populate()
 247   gen_j_populate()
 248
 249 gen_C()
 250
 251 defs, vs = cse(q_sqparm)
 252 for dn, dv in defs:
 253   pp = python(dv.subs(dn, 'NONE'))
 254   print('\n', dn.name, pp)
 255 print_python(vs[0])