acts/codegen/generate_sympy_jac.py

0001 import sympy as sym
0002 from sympy import MatrixSymbol
0003
0004 from sympy_common import name_expr, find_by_name, cxx_printer, my_expression_print
0005
0006
0007 step_path_derivatives = (
0008     MatrixSymbol("step_path_derivatives", 8, 1).as_explicit().as_mutable()
0009 )
0010 step_path_derivatives[7, 0] = 0  # qop
0011
0012 surface_path_derivatives = (
0013     MatrixSymbol("surface_path_derivatives", 1, 8).as_explicit().as_mutable()
0014 )
0015 surface_path_derivatives[0, 3] = 0
0016 surface_path_derivatives[0, 7] = 0
0017
0018 J_bf = MatrixSymbol("J_bf", 8, 6).as_explicit().as_mutable()
0019 tmp = sym.zeros(8, 6)
0020 tmp[0:3, 0:2] = J_bf[0:3, 0:2]
0021 tmp[0:3, 2:4] = J_bf[0:3, 2:4]
0022 tmp[4:7, 2:4] = J_bf[4:7, 2:4]  # line surface
0023 tmp[3, 5] = 1
0024 tmp[7, 4] = 1
0025 J_bf = tmp
0026
0027 J_t = MatrixSymbol("J_t", 8, 8).as_explicit().as_mutable()
0028 tmp = sym.eye(8)
0029 tmp[0:3, 4:8] = J_t[0:3, 4:8]
0030 tmp[3, 7] = J_t[3, 7]
0031 tmp[4:7, 4:8] = J_t[4:7, 4:8]
0032 J_t = tmp
0033
0034 J_fb = MatrixSymbol("J_fb", 6, 8).as_explicit().as_mutable()
0035 tmp = sym.zeros(6, 8)
0036 tmp[0:2, 0:3] = J_fb[0:2, 0:3]
0037 tmp[2:4, 4:7] = J_fb[2:4, 4:7]
0038 tmp[5, 3] = 1
0039 tmp[4, 7] = 1
0040 J_fb = tmp
0041
0042
0043 def full_transport_jacobian_generic():
0044     J_full = name_expr(
0045         "J_full",
0046         J_fb
0047         * (sym.eye(8) + step_path_derivatives * surface_path_derivatives)
0048         * J_t
0049         * J_bf,
0050     )
0051
0052     return [J_full]
0053
0054
0055 def full_transport_jacobian_curvilinear(direction):
0056     surface_path_derivatives = (
0057         MatrixSymbol("surface_path_derivatives", 1, 8).as_explicit().as_mutable()
0058     )
0059     surface_path_derivatives[0, 0:3] = -direction.as_explicit().transpose()
0060     surface_path_derivatives[0, 3:8] = sym.zeros(1, 5)
0061
0062     J_full = name_expr(
0063         "J_full",
0064         J_fb
0065         * (sym.eye(8) + step_path_derivatives * surface_path_derivatives)
0066         * J_t
0067         * J_bf,
0068     )
0069
0070     return [J_full]
0071
0072
0073 def my_full_transport_jacobian_generic_function_print(name_exprs, run_cse=True):
0074     printer = cxx_printer
0075     outputs = [find_by_name(name_exprs, name)[0] for name in ["J_full"]]
0076
0077     lines = []
0078
0079     head = "template <typename T> void boundToBoundTransportJacobianImpl(const T* J_fb, const T* J_t, const T* J_bf, const T* step_path_derivatives, const T* surface_path_derivatives, T* J_full) {"
0080     lines.append(head)
0081
0082     code = my_expression_print(
0083         printer,
0084         name_exprs,
0085         outputs,
0086         run_cse=run_cse,
0087     )
0088     lines.extend([f"  {l}" for l in code.split("\n")])
0089
0090     lines.append("}")
0091
0092     return "\n".join(lines)
0093
0094
0095 def my_full_transport_jacobian_curvilinear_function_print(name_exprs, run_cse=True):
0096     printer = cxx_printer
0097     outputs = [find_by_name(name_exprs, name)[0] for name in ["J_full"]]
0098
0099     lines = []
0100
0101     head = "template <typename T> void boundToCurvilinearTransportJacobianImpl(const T* J_fb, const T* J_t, const T* J_bf, const T* step_path_derivatives, const T* dir, T* J_full) {"
0102     lines.append(head)
0103
0104     code = my_expression_print(
0105         printer,
0106         name_exprs,
0107         outputs,
0108         run_cse=run_cse,
0109     )
0110     lines.extend([f"  {l}" for l in code.split("\n")])
0111
0112     lines.append("}")
0113
0114     return "\n".join(lines)
0115
0116
0117 print(
0118     """// This file is part of the ACTS project.
0119 //
0120 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0121 //
0122 // This Source Code Form is subject to the terms of the Mozilla Public
0123 // License, v. 2.0. If a copy of the MPL was not distributed with this
0124 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0125
0126 // Note: This file is generated by generate_sympy_jac.py
0127 //       Do not modify it manually.
0128
0129 #pragma once
0130
0131 #include <cmath>
0132 """
0133 )
0134
0135 all_name_exprs = full_transport_jacobian_generic()
0136 code = my_full_transport_jacobian_generic_function_print(
0137     all_name_exprs,
0138     run_cse=True,
0139 )
0140 print(code)
0141 print()
0142
0143 all_name_exprs = full_transport_jacobian_curvilinear(MatrixSymbol("dir", 3, 1))
0144 code = my_full_transport_jacobian_curvilinear_function_print(
0145     all_name_exprs,
0146     run_cse=True,
0147 )
0148 print(code)