def test_prepare_next_tree_regular_level(self, bivariate_mock, conditional_mock):
"""prepare_next_tree computes the conditional U matrices on its edges."""
# Setup
instance = Tree(TreeTypes.REGULAR)
instance.level = 2
edge = MagicMock(spec=Edge)
edge.parents = ['first_parent', 'second_parent']
edge.name = 'copula_type'
edge.theta = 'copula_theta'
instance.edges = [edge]
copula_mock = bivariate_mock.return_value
copula_mock.partial_derivative.return_value = np.array([0.0, 0.25, 0.5, 0.75, 1.0])
conditional_mock.return_value = (
['left_u_1', 'left_u_2'],
['right_u_1', 'right_u_2']
)
expected_univariate = np.array([
[EPSILON, 0.25, 0.50, 0.75, 1 - EPSILON],
[EPSILON, 0.25, 0.50, 0.75, 1 - EPSILON]
])
conditional_univariates = np.array([
['left_u_1', 'right_u_1'],
['left_u_2', 'right_u_2']
])
expected_partial_derivative_call_args = [
((conditional_univariates,), {}),
((conditional_univariates[:, np.argsort([1, 0])],), {})
]
# Run
instance.prepare_next_tree()
# Check
compare_nested_iterables(instance.edges[0].U, expected_univariate)
bivariate_mock.assert_called_once_with(copula_type='copula_type')
conditional_mock.assert_called_once_with('first_parent', 'second_parent')
assert copula_mock.theta == 'copula_theta'
compare_nested_iterables(
copula_mock.partial_derivative.call_args_list,
expected_partial_derivative_call_args
)
def test_prepare_next_tree_first_level(self, bivariate_mock):
"""prepare_next_tree computes the conditional U matrices on its edges."""
# Setup
instance = Tree(TreeTypes.REGULAR)
instance.level = 1
instance.u_matrix = np.array([
[0.1, 0.2],
[0.3, 0.4]
])
edge = MagicMock(spec=Edge)
edge.L = 0
edge.R = 1
edge.name = 'copula_type'
edge.theta = 'copula_theta'
instance.edges = [edge]
copula_mock = bivariate_mock.return_value
copula_mock.partial_derivative.return_value = np.array([0.0, 0.25, 0.5, 0.75, 1.0])
expected_univariate = np.array([
[EPSILON, 0.25, 0.50, 0.75, 1 - EPSILON],
[EPSILON, 0.25, 0.50, 0.75, 1 - EPSILON]
])
expected_partial_derivative_call_args = [
((instance.u_matrix,), {}),
((instance.u_matrix[:, np.argsort([1, 0])],), {})
]
# Run
instance.prepare_next_tree()
# Check
compare_nested_iterables(instance.edges[0].U, expected_univariate)
bivariate_mock.assert_called_once_with(copula_type='copula_type')
assert copula_mock.theta == 'copula_theta'
compare_nested_iterables(
copula_mock.partial_derivative.call_args_list,
expected_partial_derivative_call_args
)
