174 lines
7.1 KiB
Python
174 lines
7.1 KiB
Python
#!/usr/bin/env python3
|
|
|
|
## Copyright 2009-2021 Intel Corporation
|
|
## SPDX-License-Identifier: Apache-2.0
|
|
|
|
## Used to generate code in kernels/common/motion_derivative.h
|
|
|
|
|
|
import sympy as sp
|
|
import numpy as np
|
|
import math
|
|
|
|
|
|
################################################################################
|
|
#### Utils
|
|
################################################################################
|
|
|
|
def getTerms(map, key):
|
|
if key in map.keys():
|
|
return map[key]
|
|
return 0
|
|
|
|
# simple linear interpolation wrapper
|
|
def lerp(v0,v1,t):
|
|
return v0*(1-t)+v1*t
|
|
|
|
# custom quaternion to matrix conversion
|
|
def to_rotation_matrix(q):
|
|
return sp.Matrix([[q.a*q.a + q.b*q.b - q.c*q.c - q.d*q.d, 2*(q.b*q.c - q.a*q.d), 2*(q.b*q.d + q.a*q.c), 0],
|
|
[2*(q.b*q.c + q.a*q.d), q.a*q.a - q.b*q.b + q.c*q.c - q.d*q.d, 2*(q.c*q.d - q.a*q.b), 0],
|
|
[2*(q.b*q.d - q.a*q.c), 2*(q.c*q.d + q.a*q.b), q.a*q.a - q.b*q.b - q.c*q.c + q.d*q.d, 0],
|
|
[0, 0, 0, 1]])
|
|
|
|
|
|
################################################################################
|
|
#### Set up symbolic objects
|
|
################################################################################
|
|
|
|
t, theta = sp.symbols("t, theta", real = True)
|
|
|
|
px0, py0, pz0 = sp.symbols("px0, py0, pz0", real=True) # vertex position at t=0
|
|
px1, py1, pz1 = sp.symbols("px1, py1, pz1", real=True) # vertex position at t=1
|
|
|
|
tx0, ty0, tz0 = sp.symbols("tx0, ty0, tz0", real=True) # translation at t=0
|
|
tx1, ty1, tz1 = sp.symbols("tx1, ty1, tz1", real=True) # translation at t=1
|
|
|
|
qx0, qy0, qz0, qw0 = sp.symbols("qx0, qy0, qz0, qw0", real=True) # quaternion at t=0
|
|
qx1, qy1, qz1, qw1 = sp.symbols("qx1, qy1, qz1, qw1", real=True) # quaternion at t=1
|
|
|
|
# coefficients for upper triangular matrices
|
|
s000, s001, s002, s003, s011, s012, s013, s022, s023 = sp.symbols("s000, s001, s002, s003, s011, s012, s013, s022, s023", real=True)
|
|
s100, s101, s102, s103, s111, s112, s113, s122, s123 = sp.symbols("s100, s101, s102, s103, s111, s112, s113, s122, s123", real=True)
|
|
|
|
q0 = sp.Quaternion(qw0, qx0, qy0, qz0)
|
|
q1 = sp.Quaternion(qw1, qx1, qy1, qz1)
|
|
|
|
# assuming that q1 is qperp = normalize(q1-q0*cosTheta), where cosTheta=dot(q0, q1) and theta = acos(cosTheta).
|
|
# this simplifies the terms of the symbolic expressions later
|
|
qt = q0 * sp.cos(t*theta) + q1 * sp.sin(t*theta)
|
|
|
|
S0 = sp.Matrix([[s000, s001, s002, s003],
|
|
[ 0, s011, s012, s013],
|
|
[ 0, 0, s022, s023],
|
|
[ 0, 0, 0, 1]])
|
|
S1 = sp.Matrix([[s100, s101, s102, s103],
|
|
[ 0, s111, s112, s113],
|
|
[ 0, 0, s122, s123],
|
|
[ 0, 0, 0, 1]])
|
|
D0 = sp.Matrix([[1, 0, 0, tx0],
|
|
[0, 1, 0, ty0],
|
|
[0, 0, 1, tz0],
|
|
[0, 0, 0, 1]])
|
|
D1 = sp.Matrix([[1, 0, 0, tx1],
|
|
[0, 1, 0, ty1],
|
|
[0, 0, 1, tz1],
|
|
[0, 0, 0, 1]])
|
|
p0 = sp.Matrix([px0, py0, pz0, 1])
|
|
p1 = sp.Matrix([px1, py1, pz1, 1])
|
|
|
|
Gamma = lerp(D0, D1, t)*to_rotation_matrix(qt)*lerp(S0, S1, t)*lerp(p0, p1, t)
|
|
C = sp.Matrix(np.empty(8)) # 8 coefficients
|
|
K = sp.Matrix(np.empty(7)) # 7 inputs
|
|
A = sp.Matrix(np.empty(8*7*3)) # 8 coefficients, 7 inputs (1, px0, py0, pz0, px1, py1, pz1), 3 dimensions (x, y, z)
|
|
dGamma = sp.diff(Gamma, t)
|
|
|
|
|
|
################################################################################
|
|
#### Group the coefficients (this might time a couple of seconds)
|
|
################################################################################
|
|
|
|
# loop over dimensions (x, y, z)
|
|
for dim in range(3):
|
|
dm = sp.expand(dGamma[dim])
|
|
dm = dm.subs(sp.sin(t*theta)*sp.sin(t*theta),(1-sp.cos(2*t*theta))/2) # remove sin(t*theta)^2
|
|
dm = dm.subs(sp.cos(t*theta)*sp.cos(t*theta),(1+sp.cos(2*t*theta))/2) # remove cos(t*theta)^2
|
|
dm = dm.subs(sp.sin(t*theta)*sp.cos(t*theta),sp.sin(2*t*theta)/2) # remove sin(t*theta)*cos(t*theta)
|
|
dm = sp.expand(dm)
|
|
|
|
# group all terms in the form a + b * cos(2*t*theta) + c * sin(2*t*theta)
|
|
dm_cos_sin = sp.collect(dm, (sp.cos(2*t*theta), sp.sin(2*t*theta)), evaluate=False)
|
|
|
|
# get the terms
|
|
coeff_cos = getTerms(dm_cos_sin, sp.cos(2*t*theta))
|
|
coeff_sin = getTerms(dm_cos_sin, sp.sin(2*t*theta))
|
|
coeff_const = getTerms(dm_cos_sin, 1)
|
|
|
|
# group the term in the form a + b * t
|
|
coeff_const_t = sp.collect(coeff_const, t, evaluate=False)
|
|
C[0] = getTerms(coeff_const_t, 1)
|
|
C[1] = getTerms(coeff_const_t, t)
|
|
|
|
# group the term in the form a + b * t + c * t^2
|
|
coeff_cos_t = sp.collect(coeff_cos, t, evaluate=False)
|
|
C[2] = getTerms(coeff_cos_t, 1)
|
|
C[3] = getTerms(coeff_cos_t, t)
|
|
C[4] = getTerms(coeff_cos_t, t*t)
|
|
|
|
# group the term in the form a + b * t + c * t^2
|
|
coeff_sin_t = sp.collect(coeff_sin, t, evaluate=False)
|
|
C[5] = getTerms(coeff_sin_t, 1)
|
|
C[6] = getTerms(coeff_sin_t, t)
|
|
C[7] = getTerms(coeff_sin_t, t*t)
|
|
|
|
for c in range(8):
|
|
kc = sp.collect(C[c], (px0, py0, pz0, px1, py1, pz1), evaluate=False)
|
|
K[0] = getTerms(kc, 1)
|
|
K[1] = getTerms(kc, px0)
|
|
K[2] = getTerms(kc, py0)
|
|
K[3] = getTerms(kc, pz0)
|
|
K[4] = getTerms(kc, px1)
|
|
K[5] = getTerms(kc, py1)
|
|
K[6] = getTerms(kc, pz1)
|
|
|
|
for k in range(7):
|
|
K[k] = sp.expand(K[k])
|
|
K[k] = K[k].subs(qw0*qw0, 1-qx0*qx0-qy0*qy0-qz0*qz0) # clean up substitutions
|
|
K[k] = K[k].subs(qw1*qw1, 1-qx1*qx1-qy1*qy1-qz1*qz1) # clean up substitutions
|
|
K[k] = sp.simplify(K[k])
|
|
A[8*7*dim + c*7 + k] = K[k]
|
|
|
|
|
|
################################################################################
|
|
#### Write code to file
|
|
################################################################################
|
|
|
|
from sympy.utilities.codegen import codegen, default_datatypes
|
|
from sympy.codegen.ast import real, float32
|
|
from sympy.printing.ccode import C99CodePrinter
|
|
printer = C99CodePrinter()
|
|
|
|
# custom code printer that will not generate such nonesene as x^2 -> pow(x, 2)
|
|
class CustomCodePrinter(C99CodePrinter):
|
|
def _print_Pow(self, expr):
|
|
if expr.exp.is_integer and expr.exp > 0 and expr.exp < 5:
|
|
return '*'.join([self._print(expr.base) for i in range(expr.exp)])
|
|
else:
|
|
return super()._print_Pow(expr)
|
|
|
|
customprinter = CustomCodePrinter()
|
|
customprinter.type_aliases[real] = float32 # cosf instead of cos
|
|
default_datatypes["float"].cname = "float" # float instead of double
|
|
params = [
|
|
theta,
|
|
tx0, ty0, tz0,
|
|
tx1, ty1, tz1,
|
|
qw0, qx0, qy0, qz0,
|
|
qw1, qx1, qy1, qz1,
|
|
s000, s001, s002, s003, s011, s012, s013, s022, s023,
|
|
s100, s101, s102, s103, s111, s112, s113, s122, s123]
|
|
R = sp.MatrixSymbol("coeff", A.shape[0], A.shape[1])
|
|
P = sp.MatrixSymbol('p', len(params), 1)
|
|
param_map = dict(zip(params, P))
|
|
B = A.xreplace(param_map)
|
|
codegen(('motion_derivative_coefficients', sp.Eq(R,B)), language='c', printer=customprinter, prefix='motion_derivative_coefficients', to_files=True)
|