File indexing completed on 2026-05-27 07:23:54
0001 import sys
0002 import sympy
0003
0004 from detray_sympy.common import name_expr
0005 from detray_sympy.output import write_out_file
0006 from detray_sympy.codegen import gen_cxx_code
0007 from detray_sympy.matrices import add_transport_jacobian_substructure
0008
0009
0010 def gen_code(gradient=True):
0011 transport_jacobian = (
0012 sympy.MatrixSymbol("transport_jacobian", 8, 8).as_explicit().as_mutable()
0013 )
0014 transport_jacobian = add_transport_jacobian_substructure(
0015 transport_jacobian, gradient=gradient
0016 )
0017
0018 b2f_dpos_dloc = sympy.MatrixSymbol("b2f_dpos_dloc", 3, 2).as_explicit().as_mutable()
0019 b2f_ddir_dangle = (
0020 sympy.MatrixSymbol("b2f_ddir_dangle", 3, 2).as_explicit().as_mutable()
0021 )
0022 b2f_ddir_dangle[2, 0] = 0
0023 b2f_dpos_dangle = (
0024 sympy.MatrixSymbol("b2f_dpos_dangle", 3, 2).as_explicit().as_mutable()
0025 )
0026
0027 path_to_free_derivative = (
0028 sympy.MatrixSymbol("path_to_free_derivative", 8, 1).as_explicit().as_mutable()
0029 )
0030 path_to_free_derivative[3, 0] = 0
0031
0032 free_to_path_derivative = (
0033 sympy.MatrixSymbol("free_to_path_derivative", 1, 8).as_explicit().as_mutable()
0034 )
0035 free_to_path_derivative[0, 3] = 0
0036 free_to_path_derivative[0, 7] = 0
0037
0038 f2b_dloc_dpos = sympy.MatrixSymbol("f2b_dloc_dpos", 2, 3).as_explicit().as_mutable()
0039 f2b_dangle_ddir = (
0040 sympy.MatrixSymbol("f2b_dangle_ddir", 2, 3).as_explicit().as_mutable()
0041 )
0042 f2b_dangle_ddir[0, 2] = 0
0043
0044 bound_to_free_jacobian = sympy.zeros(8, 6)
0045 bound_to_free_jacobian[0:3, 0:2] = b2f_dpos_dloc
0046 bound_to_free_jacobian[0:3, 2:4] = b2f_dpos_dangle
0047 bound_to_free_jacobian[4:7, 2:4] = b2f_ddir_dangle
0048 bound_to_free_jacobian[7, 4] = 1
0049 bound_to_free_jacobian[3, 5] = 1
0050
0051 free_to_bound_jacobian = sympy.zeros(6, 8)
0052 free_to_bound_jacobian[0:2, 0:3] = f2b_dloc_dpos
0053 free_to_bound_jacobian[2:4, 4:7] = f2b_dangle_ddir
0054 free_to_bound_jacobian[4, 7] = 1
0055 free_to_bound_jacobian[5, 3] = 1
0056
0057 full_jacobian = (
0058 free_to_bound_jacobian
0059 * (path_to_free_derivative * free_to_path_derivative + sympy.eye(8))
0060 * transport_jacobian
0061 * bound_to_free_jacobian
0062 )
0063
0064 input_name_exprs = [
0065 name_expr("transport_jacobian", transport_jacobian),
0066 name_expr("b2f_dpos_dloc", b2f_dpos_dloc),
0067 name_expr("b2f_ddir_dangle", b2f_ddir_dangle),
0068 name_expr("b2f_dpos_dangle", b2f_dpos_dangle),
0069 name_expr("path_to_free_derivative", path_to_free_derivative),
0070 name_expr("free_to_path_derivative", free_to_path_derivative),
0071 name_expr("f2b_dloc_dpos", f2b_dloc_dpos),
0072 name_expr("f2b_dangle_ddir", f2b_dangle_ddir),
0073 ]
0074 output_name_exprs = [name_expr("full_jacobian", full_jacobian)]
0075 code = gen_cxx_code(
0076 "update_full_jacobian_"
0077 + ("with_gradient" if gradient else "without_gradient")
0078 + "_impl",
0079 input_name_exprs,
0080 output_name_exprs,
0081 run_cse=True,
0082 )
0083 return code
0084
0085
0086 if __name__ == "__main__":
0087 if len(sys.argv) > 1:
0088 output = sys.argv[1]
0089 else:
0090 output = None
0091
0092 c1 = gen_code(gradient=True)
0093 c2 = gen_code(gradient=False)
0094
0095 write_out_file(c1 + c2, output)
