def test_multinomial_coefficients():
assert multinomial_coefficients(1, 1) == {(1,): 1}
assert multinomial_coefficients(1, 2) == {(2,): 1}
assert multinomial_coefficients(1, 3) == {(3,): 1}
assert multinomial_coefficients(2, 0) == {(0, 0): 1}
assert multinomial_coefficients(2, 1) == {(0, 1): 1, (1, 0): 1}
assert multinomial_coefficients(2, 2) == {(2, 0): 1, (0, 2): 1, (1, 1): 2}
assert multinomial_coefficients(2, 3) == {(3, 0): 1, (1, 2): 3, (0, 3): 1,
(2, 1): 3}
assert multinomial_coefficients(3, 1) == {(1, 0, 0): 1, (0, 1, 0): 1,
(0, 0, 1): 1}
assert multinomial_coefficients(3, 2) == {(0, 1, 1): 2, (0, 0, 2): 1,
(1, 1, 0): 2, (0, 2, 0): 1, (1, 0, 1): 2, (2, 0, 0): 1}
mc = multinomial_coefficients(3, 3)
assert mc == {(2, 1, 0): 3, (0, 3, 0): 1,
(1, 0, 2): 3, (0, 2, 1): 3, (0, 1, 2): 3, (3, 0, 0): 1,
(2, 0, 1): 3, (1, 2, 0): 3, (1, 1, 1): 6, (0, 0, 3): 1}
assert dict(multinomial_coefficients_iterator(2, 0)) == {(0, 0): 1}
assert dict(
multinomial_coefficients_iterator(2, 1)) == {(0, 1): 1, (1, 0): 1}
assert dict(multinomial_coefficients_iterator(2, 2)) == \
{(2, 0): 1, (0, 2): 1, (1, 1): 2}
assert dict(multinomial_coefficients_iterator(3, 3)) == mc
it = multinomial_coefficients_iterator(7, 2)
assert [next(it) for i in range(4)] == \
[((2, 0, 0, 0, 0, 0, 0), 1), ((1, 1, 0, 0, 0, 0, 0), 2),
((0, 2, 0, 0, 0, 0, 0), 1), ((1, 0, 1, 0, 0, 0, 0), 2)]
def test_multinomial_coefficients():
assert multinomial_coefficients(1, 1) == {(1,): 1}
assert multinomial_coefficients(1, 2) == {(2,): 1}
assert multinomial_coefficients(1, 3) == {(3,): 1}
assert multinomial_coefficients(2, 0) == {(0, 0): 1}
assert multinomial_coefficients(2, 1) == {(0, 1): 1, (1, 0): 1}
assert multinomial_coefficients(2, 2) == {(2, 0): 1, (0, 2): 1, (1, 1): 2}
assert multinomial_coefficients(2, 3) == {(3, 0): 1, (1, 2): 3, (0, 3): 1,
(2, 1): 3}
assert multinomial_coefficients(3, 1) == {(1, 0, 0): 1, (0, 1, 0): 1,
(0, 0, 1): 1}
assert multinomial_coefficients(3, 2) == {(0, 1, 1): 2, (0, 0, 2): 1,
(1, 1, 0): 2, (0, 2, 0): 1, (1, 0, 1): 2, (2, 0, 0): 1}
mc = multinomial_coefficients(3, 3)
assert mc == {(2, 1, 0): 3, (0, 3, 0): 1,
(1, 0, 2): 3, (0, 2, 1): 3, (0, 1, 2): 3, (3, 0, 0): 1,
(2, 0, 1): 3, (1, 2, 0): 3, (1, 1, 1): 6, (0, 0, 3): 1}
assert dict(multinomial_coefficients_iterator(2, 0)) == {(0, 0): 1}
assert dict(
multinomial_coefficients_iterator(2, 1)) == {(0, 1): 1, (1, 0): 1}
assert dict(multinomial_coefficients_iterator(2, 2)) == \
{(2, 0): 1, (0, 2): 1, (1, 1): 2}
assert dict(multinomial_coefficients_iterator(3, 3)) == mc
it = multinomial_coefficients_iterator(7, 2)
assert [next(it) for i in range(4)] == \
[((2, 0, 0, 0, 0, 0, 0), 1), ((1, 1, 0, 0, 0, 0, 0), 2),
((0, 2, 0, 0, 0, 0, 0), 1), ((1, 0, 1, 0, 0, 0, 0), 2)]
def _eval_derivative_n_times(self, s, n):
# adapted from Mul._eval_derivative_n_times
from sympy import Integer
from sympy.ntheory.multinomial import multinomial_coefficients_iterator
args = self.args
m = len(args)
if isinstance(n, (int, Integer)):
# https://en.wikipedia.org/wiki/General_Leibniz_rule#More_than_two_factors
terms = []
for kvals, c in multinomial_coefficients_iterator(m, n):
p = VecMul(*[arg.diff((s, k)) for k, arg in zip(kvals, args)])
terms.append(c * p)
return VecAdd(*terms)
raise TypeError("Derivative order must be integer")
