Changeset 0d19f42 in sasmodels
 Timestamp:
 Oct 17, 2017 4:29:18 PM (5 years ago)
 Branches:
 master, core_shell_microgels, magnetic_model, ticket1257vesicleproduct, ticket_1156, ticket_1265_superball, ticket_822_more_unit_tests
 Children:
 57eb6a4
 Parents:
 6ab64c9
 File:

 1 edited
Legend:
 Unmodified
 Added
 Removed

explore/angles.py
rbf09f55 r0d19f42 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 1 22 from __future__ import print_function 2 23 24 import codecs 25 import sys 26 import re 27 3 28 import sympy as sp 4 sp.init_printing() 5 dphi, dpsi, dtheta = sp.var("phi_d psi_d theta_d") 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'(phipsitheta)\^\+', r'\1', s) # jitter rep: V^+ => V 49 s = re.sub(r'([az]*)\^\+', r'd\1', s) # jitter rep: V^+ => dV 50 s = re.sub(r'(cossin)\(([az]*)\)', r'\1_\2', s) # cos(V) => cos_V 51 s = re.sub(r'pow\(([az]*), 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 matrixvector 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 116 def Rx(a): 117 """Rotate y and z about x""" 118 return Matrix([[1, 0, 0], [0, cos(a), sin(a)], [0, sin(a), cos(a)]]) 119 120 def Ry(a): 121 """Rotate x and z about y""" 122 return Matrix([[cos(a), 0, sin(a)], [0, 1, 0], [sin(a), 0, cos(a)]]) 123 124 def Rz(a): 125 """Rotate x and y about z""" 126 return Matrix([[cos(a), sin(a), 0], [sin(a), cos(a), 0], [0, 0, 1]]) 127 128 129 ## =============== Describe the transforms ==================== 130 131 # Define symbols used. Note that if you change the symbols for the jitter 132 # angles, you will need to update the subs() function accordingly. 133 dphi, dpsi, dtheta = sp.var("phi^+ psi^+ theta^+") 6 134 phi, psi, theta = sp.var("phi psi theta") 135 #dphi, dpsi, dtheta = sp.var("beta^+ gamma^+ alpha^+") 136 #phi, psi, theta = sp.var("beta gamma alpha") 7 137 x, y, z = sp.var("x y z") 138 q = sp.var("q") 8 139 qx, qy, qz = sp.var("qx qy qz") 140 dqx, dqy, dqz = sp.var("qx^+ qy^+ qz^+") 9 141 qa, qb, qc = sp.var("qa qb qc") 10 11 def Rx(a): 12 return sp.Matrix([[1, 0, 0], [0, sp.cos(a), sp.sin(a)], [0, sp.sin(a), sp.cos(a)]]) 13 def Ry(a): 14 return sp.Matrix([[sp.cos(a), 0, sp.sin(a)], [0, 1, 0], [sp.sin(a), 0, sp.cos(a)]]) 15 def Rz(a): 16 return sp.Matrix([[sp.cos(a), sp.sin(a), 0], [sp.sin(a), sp.cos(a), 0], [0, 0, 1]]) 17 18 print("==== asymmetric ====") 19 q_xy = sp.Matrix([[qx], [qy], [0]]) 20 q_abc = sp.Matrix([[qa], [qb], [qc]]) 21 p_xyz = sp.Matrix([[x], [y], [z]]) 22 jitter = Rx(dphi)*Ry(dtheta)*Rz(dpsi) 23 view = Rz(phi)*Ry(theta)*Rz(psi) 24 view_inv = Rz(psi)*Ry(theta)*Rz(phi) 25 print(">> jitter rotation") 26 for row in jitter*p_xyz: sp.pprint(row) 27 #print("\n>> jitter plus view") 28 #for row in view*jitter*p_xyz: sp.pprint(row) 29 print("\n>> view reverse") 30 for row in view_inv*q_xy: sp.pprint(row) 31 print("\n>> jitter reverse") 32 jitter_inv = Rz(psi)*Ry(theta)*Rx(phi) 33 for row in jitter_inv*q_abc: sp.pprint(row) 34 print("\n>> jitter view reverse") 35 jitter_inv = Rz(dpsi)*Ry(dtheta)*Rx(dphi) 36 for row in jitter_inv*view_inv*q_xy: sp.pprint(row) 37 38 print("\n\n==== symmetric ====") 39 q_x = sp.Matrix([[qx], [0], [0]]) 40 q_ac = sp.Matrix([[qa], [0], [qc]]) 41 p_xyz = sp.Matrix([[x], [y], [z]]) 42 jitter = Rx(dphi)*Ry(dtheta) 43 view = Rz(phi)*Ry(theta) 44 jitter_inv = Ry(dtheta)*Rx(dphi) 45 view_inv = Ry(theta)*Rz(phi) 46 print(">> jitter rotation") 47 for row in jitter*p_xyz: sp.pprint(row) 48 #print("\n>> jitter plus view") 49 #for row in view*jitter*p_xyz: sp.pprint(row) 50 print("\n>> view reverse") 51 for row in view_inv*q_x: sp.pprint(row) 52 print("\n>> jitter reverse") 53 for row in jitter_inv*q_ac: sp.pprint(row) 142 dqa, dqb, dqc = sp.var("qa^+ qb^+ qc^+") 143 qab = sp.var("qab") 144 145 # 3x3 matrix M 146 J = Matrix([sp.var("J(1:4)(1:4)")]).reshape(3,3) 147 V = Matrix([sp.var("V(1:4)(1:4)")]).reshape(3,3) 148 R = Matrix([sp.var("R(1:4)(1:4)")]).reshape(3,3) 149 150 # various vectors 151 q_xy = Matrix([[qx], [qy], [0]]) 152 q_abc = Matrix([[qa], [qb], [qc]]) 153 q_xyz = Matrix([[qx], [qy], [qz]]) 154 dq_abc = Matrix([[dqa], [dqb], [dqc]]) 155 dq_xyz = Matrix([[dqx], [dqy], [dqz]]) 156 157 # By comparing the code we generate to sasmodel 4.x orientation code, we 158 # are apparently using the opposite convention for phi than I expected. 159 # Theta 160 #theta = pi/2  theta 161 phi = phi 162 dphi = dphi 163 164 def print_steps(jitter, jitter_inv, view, view_inv, qc_only): 165 """ 166 Show the forward/reverse transform code for view and jitter. 167 """ 168 if 0: # forward calculations 169 vprint(q_xyz, jitter*q_abc, "apply jitter") 170 vprint(dq_xyz, view*q_xyz, "apply view after jitter") 171 172 #vprint(dq_xyz, view*jitter*q_abc, "combine view and jitter") 173 mprint(M, view*jitter, "forward matrix") 174 175 if 1: # reverse calculations 176 pre_view = "set angles from view" if REUSE_SINCOS else None 177 pre_jitter = "set angles from jitter" if REUSE_SINCOS else None 178 index = slice(2,3) if qc_only else slice(None,None) 179 180 comment("\n**** direct ****") 181 vprint(q_abc, view_inv*q_xy, "reverse view", post=pre_view) 182 vprint(dq_abc[index,:], (jitter_inv*q_abc)[index,:], 183 "reverse jitter after view", post=pre_jitter) 184 185 comment("\n\n**** precalc ****") 186 #vprint(q_abc, jitter_inv*view_inv*q_xy, "combine jitter and view reverse") 187 mprint(V[:,:2], view_inv[:,:2], "reverse view matrix", post=pre_view) 188 mprint(J[index,:], jitter_inv[index,:], "reverse jitter matrix", post=pre_jitter) 189 mprint(R[index,:2], (J*V)[index,:2], "reverse matrix") 190 comment("\n**** per point ****") 191 mprint(q_abc[index,:], (R*q_xy)[index,:], "applied reverse matrix") 192 #mprint(q_abc, J*V*q_xy, "applied reverse matrix") 193 #mprint(M, jitter_inv*view_inv, "reverse matrix direct") 194 195 #vprint(q_abc, M*q_xy, "matrix application") 196 197 comment("==== asymmetric ====") 198 print_steps( 199 jitter=Rx(dphi)*Ry(dtheta)*Rz(dpsi), 200 jitter_inv=Rz(dpsi)*Ry(dtheta)*Rx(dphi), 201 view=Rz(phi)*Ry(theta)*Rz(psi), 202 view_inv=Rz(psi)*Ry(theta)*Rz(phi), 203 qc_only=False, 204 ) 205 206 comment("\n\n==== symmetric ====") 207 print_steps( 208 jitter=Rx(dphi)*Ry(dtheta), 209 jitter_inv=Ry(dtheta)*Rx(dphi), 210 view=Rz(phi)*Ry(theta), 211 view_inv=Ry(theta)*Rz(phi), 212 qc_only=QC_ONLY, 213 ) 214 215 comment("\n**** qab from qc ****") 216 # The indirect calculation of qab is better than directly c 217 # alculating qab^2 = qa^2 + qb^2 since qc can be computed 218 # as qc = M31*qx + M32*qy, thus requiring only two elements 219 # of the rotation matrix. 220 #vprint(qab, sqrt(qa**2 + qb**2), "Direct calculation of qab") 221 vprint(dqa, sqrt((qx**2+qy**2)  dqc**2), 222 "Indirect calculation of qab, from qab^2 = q^2  qc^2")
Note: See TracChangeset
for help on using the changeset viewer.