SymPy quaternion composition verification and code generation

from sympy import *
init_printing()

def make_quaternion_expression(x: tuple, i: Symbol, j: Symbol, k: Symbol):
    return x[0] + x[1]*i + x[2]*j + x[3]*k

def quaternion_simplification(expr: Expr, i: Symbol, j: Symbol, k: Symbol):
    # Only does one round: if you know nothing about expr you must call this function until it doesn't change `expr`
    # (or do something more clever!)
    expr = expr.subs(i**2, -1)
    expr = expr.subs(j**2, -1)
    expr = expr.subs(k**2, -1)

    expr = expr.subs(i*j, k)
    expr = expr.subs(i*k, -j)

    expr = expr.subs(j*i, -k)
    expr = expr.subs(j*k, i)

    expr = expr.subs(k*i, j)
    expr = expr.subs(k*j, -i)
    return expr

def quaternion_expression_to_coefficient_tuple(expr, i, j, k):
    # Strip off each imaginary term one by one
    i_coeff = expr.coeff(i, 1)
    expr = simplify(expr - i_coeff * i)

    j_coeff = expr.coeff(j, 1)
    expr = simplify(expr - j_coeff * j)

    k_coeff = expr.coeff(k, 1)
    expr = simplify(expr - k_coeff * k)

    # Only the real part should be left... let's double check that
    real_part = expr
    assert real_part == real_part.coeff(i, 0).coeff(j, 0).coeff(k, 0)

    return (real_part, i_coeff, j_coeff, k_coeff)

def quaternion_composition(a: tuple, b: tuple):
    i, j, k = symbols("i, j, k", commutative=False)

    # These are our two quaternions
    # 0th index holds the real part
    A = make_quaternion_expression(a, i, j, k)
    B = make_quaternion_expression(b, i, j, k)

    # FOIL the product of two quaternions
    expr = expand( B * A )

    # Use quaternion algebra to simplify expression
    # we happen to know that no more than two imaginary terms are multiplied together,
    # so we can just do each substitution once
    expr = quaternion_simplification(expr, i, j, k)

    return quaternion_expression_to_coefficient_tuple(expr, i, j, k)

def quaternion_coefficient_conjugate(a: tuple):
    return (a[0], -a[1], -a[2], -a[3])

def quaternion_rotation(quat_coeff: tuple, vec3: ImmutableMatrix):
    i, j, k = symbols('i, j, k', commutative=False)
    q = make_quaternion(quat_coeff, i, j, k)

    conj_coeff = quaternion_coefficient_conjugate( quat_coeff )
    q_conj = make_quaternion(conj_coeff, i, j, k)

    expr = expand( q * vec3 * q_conj )
    expr = quaternion_simplification(expr, i, j, k)
    return expr

# Let us verify that rotating a vector by quaternions `a` then `b` is equal to
# composing `a` and `b` into one quaternion `ab` then rotating a vector by `ab`
vec3 = ImmutableMatrix( symbols("v[:3]") )
a = symbols("a[:4]")
b = symbols("b[:4]")

# Method 1
vec3_a = quaternion_rotation(a, vec3)
vec3_a_b = quaternion_rotation(b, vec3_a)

# Method 2
ab = quaternion_composition(a, b)
vec3_ab = quaternion_rotation(ab, vec3)

# We just verified that composition does indeed work the way we think it does
assert vec3_a_b == vec3_ab

# Let's print these functions out as C++ code
from sympy.printing.cxx import *
printer = CXX17CodePrinter()

print("template <typename T>")
print("void")
print("quaternion_composition(")
print("\t\tT const * a,")
print("\t\tT const * b,")
print("\t\tT * a_then_b)")
print("{")
for i in range(len(a_then_b)):
    print(f"\ta_then_b[{i}] = " + str(a_then_b[i]) + ";")
print("}")
print("")
# While we're at it, lets emit some code for rotating a vector by a quaternion
print("template <typename T>")
print("void")
print("quaternion_rotation(")
print("\t\tT const * quat,")
print("\t\tT const * vec3,")
print("\t\tT * out)")
print("{")
rotated_vector = quaternion_rotation(
        symbols("quat[:4]"),
         ImmutableMatrix( symbols("vec3[:3]") ))
for i in range(len(rotated_vector)):
    print(f"\tout[{i}] = " + printer.doprint(rotated_vector[i]) + ";")
print("}")

# template <typename T>
# void
# quaternion_composition(
#       T const * a,
#       T const * b,
#       T * a_then_b)
# {
#   a_then_b[0] = a[0]*b[0] - a[1]*b[1] - a[2]*b[2] - a[3]*b[3];
#   a_then_b[1] = a[0]*b[1] + a[1]*b[0] - a[2]*b[3] + a[3]*b[2];
#   a_then_b[2] = a[0]*b[2] + a[1]*b[3] + a[2]*b[0] - a[3]*b[1];
#   a_then_b[3] = a[0]*b[3] - a[1]*b[2] + a[2]*b[1] + a[3]*b[0];
# }

# template <typename T>
# void
# quaternion_rotation(
#       T const * quat,
#       T const * vec3,
#       T * out)
# {
#   out[0] = std::pow(quat[0], 2)*vec3[0] + std::pow(quat[1], 2)*vec3[0] + std::pow(quat[2], 2)*vec3[0] + std::pow(quat[3], 2)*vec3[0];
#   out[1] = std::pow(quat[0], 2)*vec3[1] + std::pow(quat[1], 2)*vec3[1] + std::pow(quat[2], 2)*vec3[1] + std::pow(quat[3], 2)*vec3[1];
#   out[2] = std::pow(quat[0], 2)*vec3[2] + std::pow(quat[1], 2)*vec3[2] + std::pow(quat[2], 2)*vec3[2] + std::pow(quat[3], 2)*vec3[2];
# }