[0d19f42] | 1 | #!/usr/bin/env python |
---|
| 2 | """ |
---|
| 3 | Generate code for orientation transforms using symbolic algebra. |
---|
| 4 | |
---|
| 5 | To make it easier to generate correct transforms for oriented shapes, we |
---|
| 6 | use the sympy symbolic alegbra package to do the matrix multiplication. |
---|
| 7 | The transforms are displayed both using an ascii math notation, and as |
---|
| 8 | C or python code which can be pasted directly into the kernel driver. |
---|
| 9 | |
---|
| 10 | If ever we decide to change conventions, we simply need to adjust the |
---|
| 11 | order and parameters to the rotation matrices. For display we want to |
---|
| 12 | use forward transforms for the mesh describing the shape, first applying |
---|
| 13 | jitter, then adjusting the view. For calculation we know the effective q |
---|
| 14 | so we instead need to first unwind the view, using the inverse rotation, |
---|
| 15 | then undo the jitter to get the q to calculate for the shape in its |
---|
| 16 | canonical orientation. |
---|
| 17 | |
---|
| 18 | Set *OUTPUT* to the type of code you want to see: ccode, python, math |
---|
| 19 | or any combination. |
---|
| 20 | """ |
---|
| 21 | |
---|
[bf09f55] | 22 | from __future__ import print_function |
---|
| 23 | |
---|
[0d19f42] | 24 | import codecs |
---|
| 25 | import sys |
---|
| 26 | import re |
---|
| 27 | |
---|
[e755d8a] | 28 | import sympy as sp |
---|
[0d19f42] | 29 | from sympy import pi, sqrt, sin, cos, Matrix, Eq |
---|
| 30 | |
---|
| 31 | # Select output |
---|
| 32 | OUTPUT = "" |
---|
| 33 | OUTPUT = OUTPUT + "ccode" |
---|
| 34 | #OUTPUT = OUTPUT + "python " |
---|
| 35 | OUTPUT = OUTPUT + "math " |
---|
| 36 | REUSE_SINCOS = True |
---|
| 37 | QC_ONLY = True # show only what is needed for dqc in the symmetric case |
---|
| 38 | |
---|
| 39 | # include unicode symbols in output, even if piping to a pager |
---|
| 40 | sys.stdout = codecs.getwriter('utf8')(sys.stdout) |
---|
| 41 | sp.init_printing(use_unicode=True) |
---|
| 42 | |
---|
| 43 | def subs(s): |
---|
| 44 | """ |
---|
| 45 | Transform sympy generated code to follow sasmodels naming conventions. |
---|
| 46 | """ |
---|
| 47 | if REUSE_SINCOS: |
---|
| 48 | s = re.sub(r'(phi|psi|theta)\^\+', r'\1', s) # jitter rep: V^+ => V |
---|
| 49 | s = re.sub(r'([a-z]*)\^\+', r'd\1', s) # jitter rep: V^+ => dV |
---|
| 50 | s = re.sub(r'(cos|sin)\(([a-z]*)\)', r'\1_\2', s) # cos(V) => cos_V |
---|
| 51 | s = re.sub(r'pow\(([a-z]*), 2\)', r'\1*\1', s) # pow(V, 2) => V*V |
---|
| 52 | return s |
---|
| 53 | |
---|
| 54 | def comment(s): |
---|
| 55 | r""" |
---|
| 56 | Add a comment to the generated code. Use '\n' to separate lines. |
---|
| 57 | """ |
---|
| 58 | if 'ccode' in OUTPUT: |
---|
| 59 | for line in s.split("\n"): |
---|
| 60 | print("// " + line if line else "") |
---|
| 61 | if 'python' in OUTPUT: |
---|
| 62 | for line in s.split("\n"): |
---|
| 63 | print(" ## " + line if line else "") |
---|
| 64 | |
---|
| 65 | def vprint(var, vec, comment=None, post=None): |
---|
| 66 | """ |
---|
| 67 | Generate assignment statements. |
---|
| 68 | |
---|
| 69 | *var* could be a single sympy symbol or a 1xN vector of symbols. |
---|
| 70 | |
---|
| 71 | *vec* could be a single sympy expression or a 1xN vector of expressions |
---|
| 72 | such as results from a matrix-vector multiplication. |
---|
| 73 | |
---|
| 74 | *comment* if present is added to the start of the block as documentation. |
---|
| 75 | """ |
---|
| 76 | #for v, row in zip(var, vec): sp.pprint(Eq(v, row)) |
---|
| 77 | desc = sp.pretty(Eq(var, vec), wrap_line=False) |
---|
| 78 | if not isinstance(var, Matrix): |
---|
| 79 | var, vec = [var], [vec] |
---|
| 80 | if 'ccode' in OUTPUT: |
---|
| 81 | if 'math' in OUTPUT: |
---|
| 82 | print("\n// " + comment if comment else "") |
---|
| 83 | print("/*") |
---|
| 84 | for line in desc.split("\n"): |
---|
| 85 | print(" * "+line) |
---|
| 86 | print(" *\n */") |
---|
| 87 | else: |
---|
| 88 | print("\n // " + comment if comment else "") |
---|
| 89 | if post: |
---|
| 90 | print(" // " + post) |
---|
| 91 | for v, row in zip(var, vec): |
---|
| 92 | print(subs(" const double " + sp.ccode(row, assign_to=v))) |
---|
| 93 | if 'python' in OUTPUT: |
---|
| 94 | if comment: |
---|
| 95 | print("\n ## " + comment) |
---|
| 96 | if 'math' in OUTPUT: |
---|
| 97 | for line in desc.split("\n"): |
---|
| 98 | print(" # " + line) |
---|
| 99 | if post: |
---|
| 100 | print(" ## " + post) |
---|
| 101 | for v, row in zip(var, vec): |
---|
| 102 | print(subs(" " + sp.ccode(row, assign_to=v)[:-1])) |
---|
| 103 | |
---|
| 104 | if OUTPUT == 'math ': |
---|
| 105 | print("\n// " + comment if comment else "") |
---|
| 106 | if post: print("// " + post) |
---|
| 107 | print(desc) |
---|
| 108 | |
---|
| 109 | def mprint(var, mat, comment=None, post=None): |
---|
| 110 | """ |
---|
| 111 | Generate assignment statements for matrix elements. |
---|
| 112 | """ |
---|
| 113 | n = sp.prod(var.shape) |
---|
| 114 | vprint(var.reshape(n, 1), mat.reshape(n, 1), comment=comment, post=post) |
---|
| 115 | |
---|
[ef8e68c] | 116 | # From wikipedia: |
---|
| 117 | # https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations |
---|
[0d19f42] | 118 | def Rx(a): |
---|
| 119 | """Rotate y and z about x""" |
---|
[ef8e68c] | 120 | R = [[1, 0, 0], |
---|
| 121 | [0, +cos(a), -sin(a)], |
---|
| 122 | [0, +sin(a), +cos(a)]] |
---|
| 123 | return Matrix(R) |
---|
[0d19f42] | 124 | |
---|
| 125 | def Ry(a): |
---|
| 126 | """Rotate x and z about y""" |
---|
[ef8e68c] | 127 | R = [[+cos(a), 0, +sin(a)], |
---|
| 128 | [0, 1, 0], |
---|
| 129 | [-sin(a), 0, +cos(a)]] |
---|
| 130 | return Matrix(R) |
---|
[0d19f42] | 131 | |
---|
| 132 | def Rz(a): |
---|
| 133 | """Rotate x and y about z""" |
---|
[ef8e68c] | 134 | R = [[+cos(a), -sin(a), 0], |
---|
| 135 | [+sin(a), +cos(a), 0], |
---|
| 136 | [0, 0, 1]] |
---|
| 137 | return Matrix(R) |
---|
[0d19f42] | 138 | |
---|
| 139 | |
---|
| 140 | ## =============== Describe the transforms ==================== |
---|
| 141 | |
---|
| 142 | # Define symbols used. Note that if you change the symbols for the jitter |
---|
| 143 | # angles, you will need to update the subs() function accordingly. |
---|
| 144 | dphi, dpsi, dtheta = sp.var("phi^+ psi^+ theta^+") |
---|
[e755d8a] | 145 | phi, psi, theta = sp.var("phi psi theta") |
---|
[0d19f42] | 146 | #dphi, dpsi, dtheta = sp.var("beta^+ gamma^+ alpha^+") |
---|
| 147 | #phi, psi, theta = sp.var("beta gamma alpha") |
---|
[e755d8a] | 148 | x, y, z = sp.var("x y z") |
---|
[0d19f42] | 149 | q = sp.var("q") |
---|
[e755d8a] | 150 | qx, qy, qz = sp.var("qx qy qz") |
---|
[0d19f42] | 151 | dqx, dqy, dqz = sp.var("qx^+ qy^+ qz^+") |
---|
[bf09f55] | 152 | qa, qb, qc = sp.var("qa qb qc") |
---|
[0d19f42] | 153 | dqa, dqb, dqc = sp.var("qa^+ qb^+ qc^+") |
---|
| 154 | qab = sp.var("qab") |
---|
[e755d8a] | 155 | |
---|
[0d19f42] | 156 | # 3x3 matrix M |
---|
| 157 | J = Matrix([sp.var("J(1:4)(1:4)")]).reshape(3,3) |
---|
| 158 | V = Matrix([sp.var("V(1:4)(1:4)")]).reshape(3,3) |
---|
| 159 | R = Matrix([sp.var("R(1:4)(1:4)")]).reshape(3,3) |
---|
| 160 | |
---|
| 161 | # various vectors |
---|
| 162 | q_xy = Matrix([[qx], [qy], [0]]) |
---|
| 163 | q_abc = Matrix([[qa], [qb], [qc]]) |
---|
| 164 | q_xyz = Matrix([[qx], [qy], [qz]]) |
---|
| 165 | dq_abc = Matrix([[dqa], [dqb], [dqc]]) |
---|
| 166 | dq_xyz = Matrix([[dqx], [dqy], [dqz]]) |
---|
| 167 | |
---|
| 168 | def print_steps(jitter, jitter_inv, view, view_inv, qc_only): |
---|
| 169 | """ |
---|
| 170 | Show the forward/reverse transform code for view and jitter. |
---|
| 171 | """ |
---|
| 172 | if 0: # forward calculations |
---|
| 173 | vprint(q_xyz, jitter*q_abc, "apply jitter") |
---|
| 174 | vprint(dq_xyz, view*q_xyz, "apply view after jitter") |
---|
| 175 | |
---|
| 176 | #vprint(dq_xyz, view*jitter*q_abc, "combine view and jitter") |
---|
| 177 | mprint(M, view*jitter, "forward matrix") |
---|
| 178 | |
---|
| 179 | if 1: # reverse calculations |
---|
| 180 | pre_view = "set angles from view" if REUSE_SINCOS else None |
---|
| 181 | pre_jitter = "set angles from jitter" if REUSE_SINCOS else None |
---|
| 182 | index = slice(2,3) if qc_only else slice(None,None) |
---|
| 183 | |
---|
| 184 | comment("\n**** direct ****") |
---|
| 185 | vprint(q_abc, view_inv*q_xy, "reverse view", post=pre_view) |
---|
| 186 | vprint(dq_abc[index,:], (jitter_inv*q_abc)[index,:], |
---|
| 187 | "reverse jitter after view", post=pre_jitter) |
---|
| 188 | |
---|
| 189 | comment("\n\n**** precalc ****") |
---|
| 190 | #vprint(q_abc, jitter_inv*view_inv*q_xy, "combine jitter and view reverse") |
---|
| 191 | mprint(V[:,:2], view_inv[:,:2], "reverse view matrix", post=pre_view) |
---|
| 192 | mprint(J[index,:], jitter_inv[index,:], "reverse jitter matrix", post=pre_jitter) |
---|
| 193 | mprint(R[index,:2], (J*V)[index,:2], "reverse matrix") |
---|
| 194 | comment("\n**** per point ****") |
---|
| 195 | mprint(q_abc[index,:], (R*q_xy)[index,:], "applied reverse matrix") |
---|
| 196 | #mprint(q_abc, J*V*q_xy, "applied reverse matrix") |
---|
| 197 | #mprint(M, jitter_inv*view_inv, "reverse matrix direct") |
---|
| 198 | |
---|
| 199 | #vprint(q_abc, M*q_xy, "matrix application") |
---|
| 200 | |
---|
| 201 | comment("==== asymmetric ====") |
---|
| 202 | print_steps( |
---|
| 203 | jitter=Rx(dphi)*Ry(dtheta)*Rz(dpsi), |
---|
| 204 | jitter_inv=Rz(-dpsi)*Ry(-dtheta)*Rx(-dphi), |
---|
| 205 | view=Rz(phi)*Ry(theta)*Rz(psi), |
---|
| 206 | view_inv=Rz(-psi)*Ry(-theta)*Rz(-phi), |
---|
| 207 | qc_only=False, |
---|
| 208 | ) |
---|
| 209 | |
---|
| 210 | comment("\n\n==== symmetric ====") |
---|
| 211 | print_steps( |
---|
| 212 | jitter=Rx(dphi)*Ry(dtheta), |
---|
| 213 | jitter_inv=Ry(-dtheta)*Rx(-dphi), |
---|
| 214 | view=Rz(phi)*Ry(theta), |
---|
| 215 | view_inv=Ry(-theta)*Rz(-phi), |
---|
| 216 | qc_only=QC_ONLY, |
---|
| 217 | ) |
---|
| 218 | |
---|
| 219 | comment("\n**** qab from qc ****") |
---|
| 220 | # The indirect calculation of qab is better than directly c |
---|
| 221 | # alculating qab^2 = qa^2 + qb^2 since qc can be computed |
---|
| 222 | # as qc = M31*qx + M32*qy, thus requiring only two elements |
---|
| 223 | # of the rotation matrix. |
---|
| 224 | #vprint(qab, sqrt(qa**2 + qb**2), "Direct calculation of qab") |
---|
| 225 | vprint(dqa, sqrt((qx**2+qy**2) - dqc**2), |
---|
| 226 | "Indirect calculation of qab, from qab^2 = |q|^2 - qc^2") |
---|