Propagator/detail/generate_sympy_jac.py

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