detray-sympy/tests/test_assumptions_D.py

0001 import pytest
0002 import sympy
0003
0004 import detray_sympy.checks
0005 import detray_sympy.matrices
0006
0007
0008 # Check that matrix multiplication with matrix D, computed with the gradient,
0009 # is idempotent in terms of the shape of the output. This is important because
0010 # the transport Jacobian matrix $J_T$ is only ever updated in two ways:
0011 #
0012 # 1. $J_T = I$
0013 # 2. $J_T = D J_T$
0014 #
0015 # This gives us a natural induction that the Jacobian matrix is always a
0016 # product of D-shaped matrices, and if the shape of $D_1 D_2$ is the same as
0017 # the shape of $D_1$, then the transport Jacobian matrix always has the exact
0018 # same known substructure.
0019 def test_shape_idempotence_D_with_gradient():
0020     dFdr = sympy.MatrixSymbol("dFdr", 3, 3).as_explicit().as_mutable()
0021     dGdr = sympy.MatrixSymbol("dGdr", 3, 3).as_explicit().as_mutable()
0022     dFdt = sympy.MatrixSymbol("dFdt", 3, 3).as_explicit().as_mutable()
0023     dGdt = sympy.MatrixSymbol("dGdt", 3, 3).as_explicit().as_mutable()
0024     dFdqop = sympy.MatrixSymbol("dFdqop", 3, 1).as_explicit().as_mutable()
0025     dGdqop = sympy.MatrixSymbol("dGdqop", 3, 1).as_explicit().as_mutable()
0026     dqopqop = sympy.Symbol("dqopqop")
0027
0028     D = detray_sympy.matrices.get_matrix_D(
0029         dFdr, dGdr, dFdqop, dGdqop, dFdt, dGdt, dqopqop, gradient=True
0030     )
0031     assert detray_sympy.checks.has_same_known_substructure(D, D * D)
0032     assert detray_sympy.checks.has_same_known_substructure(D, D * D * D)
0033     assert detray_sympy.checks.has_same_known_substructure(D, D * D * D * D)
0034
0035
0036 # Test that multiplication with D is idempotent even with the gradient
0037 # disabled.
0038 def test_shape_idempotence_D_without_gradient():
0039     dFdr = sympy.MatrixSymbol("dFdr", 3, 3).as_explicit().as_mutable()
0040     dGdr = sympy.MatrixSymbol("dGdr", 3, 3).as_explicit().as_mutable()
0041     dFdt = sympy.MatrixSymbol("dFdt", 3, 3).as_explicit().as_mutable()
0042     dGdt = sympy.MatrixSymbol("dGdt", 3, 3).as_explicit().as_mutable()
0043     dFdqop = sympy.MatrixSymbol("dFdqop", 3, 1).as_explicit().as_mutable()
0044     dGdqop = sympy.MatrixSymbol("dGdqop", 3, 1).as_explicit().as_mutable()
0045     dqopqop = sympy.Symbol("dqopqop")
0046
0047     D = detray_sympy.matrices.get_matrix_D(
0048         dFdr, dGdr, dFdqop, dGdqop, dFdt, dGdt, dqopqop, gradient=False
0049     )
0050     assert detray_sympy.checks.has_same_known_substructure(D, D * D)
0051     assert detray_sympy.checks.has_same_known_substructure(D, D * D * D)
0052     assert detray_sympy.checks.has_same_known_substructure(D, D * D * D * D)