Changeset 0d19f42 in sasmodels


Ignore:
Timestamp:
Oct 17, 2017 4:29:18 PM (7 years ago)
Author:
Paul Kienzle <pkienzle@…>
Branches:
master, core_shell_microgels, magnetic_model, ticket-1257-vesicle-product, ticket_1156, ticket_1265_superball, ticket_822_more_unit_tests
Children:
57eb6a4
Parents:
6ab64c9
Message:

update angle generator to spit out C and python code for transformations

File:
1 edited

Legend:

Unmodified
Added
Removed
  • explore/angles.py

    rbf09f55 r0d19f42  
     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 
    122from __future__ import print_function 
    223 
     24import codecs 
     25import sys 
     26import re 
     27 
    328import sympy as sp 
    4 sp.init_printing() 
    5 dphi, dpsi, dtheta = sp.var("phi_d psi_d theta_d") 
     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 
     40sys.stdout = codecs.getwriter('utf8')(sys.stdout) 
     41sp.init_printing(use_unicode=True) 
     42 
     43def 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 
     54def 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 
     65def 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 
     109def 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 
     116def 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 
     120def 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 
     124def 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. 
     133dphi, dpsi, dtheta = sp.var("phi^+ psi^+ theta^+") 
    6134phi, 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") 
    7137x, y, z = sp.var("x y z") 
     138q = sp.var("q") 
    8139qx, qy, qz = sp.var("qx qy qz") 
     140dqx, dqy, dqz = sp.var("qx^+ qy^+ qz^+") 
    9141qa, 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) 
     142dqa, dqb, dqc = sp.var("qa^+ qb^+ qc^+") 
     143qab = sp.var("qab") 
     144 
     145# 3x3 matrix M 
     146J = Matrix([sp.var("J(1:4)(1:4)")]).reshape(3,3) 
     147V = Matrix([sp.var("V(1:4)(1:4)")]).reshape(3,3) 
     148R = Matrix([sp.var("R(1:4)(1:4)")]).reshape(3,3) 
     149 
     150# various vectors 
     151q_xy = Matrix([[qx], [qy], [0]]) 
     152q_abc = Matrix([[qa], [qb], [qc]]) 
     153q_xyz = Matrix([[qx], [qy], [qz]]) 
     154dq_abc = Matrix([[dqa], [dqb], [dqc]]) 
     155dq_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 
     161phi = -phi 
     162dphi = -dphi 
     163 
     164def 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 
     197comment("==== asymmetric ====") 
     198print_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 
     206comment("\n\n==== symmetric ====") 
     207print_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 
     215comment("\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") 
     221vprint(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.