detray-sympy/tests/test_matrices.py

0001 import pytest
0002 import sympy
0003
0004 import detray_sympy.checks
0005 import detray_sympy.matrices
0006
0007
0008 # Test that the D-matrix getters with gradients turned on give matrices with
0009 # the same known substructure.
0010 def test_generic_D_matrix_with_gradient():
0011     dFdr = sympy.MatrixSymbol("dFdr", 3, 3).as_explicit().as_mutable()
0012     dGdr = sympy.MatrixSymbol("dGdr", 3, 3).as_explicit().as_mutable()
0013     dFdt = sympy.MatrixSymbol("dFdt", 3, 3).as_explicit().as_mutable()
0014     dGdt = sympy.MatrixSymbol("dGdt", 3, 3).as_explicit().as_mutable()
0015     dFdqop = sympy.MatrixSymbol("dFdqop", 3, 1).as_explicit().as_mutable()
0016     dGdqop = sympy.MatrixSymbol("dGdqop", 3, 1).as_explicit().as_mutable()
0017     dqopqop = sympy.Symbol("dqopqop")
0018
0019     D = detray_sympy.matrices.get_matrix_D(
0020         dFdr, dGdr, dFdqop, dGdqop, dFdt, dGdt, dqopqop, gradient=True
0021     )
0022     Dg = detray_sympy.matrices.get_generic_matrix_D(gradient=True)
0023     assert detray_sympy.checks.has_same_known_substructure(D, Dg)
0024
0025
0026 # Test that the D-matrix getters with gradients turned off give matrices with
0027 # the same known substructure.
0028 def test_generic_D_matrix_without_gradient():
0029     dFdr = sympy.MatrixSymbol("dFdr", 3, 3).as_explicit().as_mutable()
0030     dGdr = sympy.MatrixSymbol("dGdr", 3, 3).as_explicit().as_mutable()
0031     dFdt = sympy.MatrixSymbol("dFdt", 3, 3).as_explicit().as_mutable()
0032     dGdt = sympy.MatrixSymbol("dGdt", 3, 3).as_explicit().as_mutable()
0033     dFdqop = sympy.MatrixSymbol("dFdqop", 3, 1).as_explicit().as_mutable()
0034     dGdqop = sympy.MatrixSymbol("dGdqop", 3, 1).as_explicit().as_mutable()
0035     dqopqop = sympy.Symbol("dqopqop")
0036
0037     D = detray_sympy.matrices.get_matrix_D(
0038         dFdr, dGdr, dFdqop, dGdqop, dFdt, dGdt, dqopqop, gradient=False
0039     )
0040     Dg = detray_sympy.matrices.get_generic_matrix_D(gradient=False)
0041     assert detray_sympy.checks.has_same_known_substructure(D, Dg)