source: sasmodels/explore/angles.py @ f39d4a3

core_shell_microgelsmagnetic_modelticket-1257-vesicle-productticket_1156ticket_1265_superballticket_822_more_unit_tests
Last change on this file since f39d4a3 was f39d4a3, checked in by Paul Kienzle <pkienzle@…>, 6 years ago

twiddle with explore/angles.py; no change in angle defn

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