Skip to content

Commit 95bd5e5

Browse files
Vonmises moments (#5232)
* Added moment to vonmises distribution Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
1 parent a16ec4a commit 95bd5e5

File tree

2 files changed

+22
-0
lines changed

2 files changed

+22
-0
lines changed

pymc/distributions/continuous.py

+6
Original file line numberDiff line numberDiff line change
@@ -2924,6 +2924,12 @@ def dist(cls, mu=0.0, kappa=None, *args, **kwargs):
29242924
assert_negative_support(kappa, "kappa", "VonMises")
29252925
return super().dist([mu, kappa], *args, **kwargs)
29262926

2927+
def get_moment(rv, size, mu, kappa):
2928+
mu, _ = at.broadcast_arrays(mu, kappa)
2929+
if not rv_size_is_none(size):
2930+
mu = at.full(size, mu)
2931+
return mu
2932+
29272933

29282934
class SkewNormalRV(RandomVariable):
29292935
name = "skewnormal"

pymc/tests/test_distributions_moments.py

+16
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
Triangular,
5555
TruncatedNormal,
5656
Uniform,
57+
VonMises,
5758
Wald,
5859
Weibull,
5960
ZeroInflatedBinomial,
@@ -438,6 +439,21 @@ def test_pareto_moment(alpha, m, size, expected):
438439
assert_moment_is_expected(model, expected)
439440

440441

442+
@pytest.mark.parametrize(
443+
"mu, kappa, size, expected",
444+
[
445+
(0, 1, None, 0),
446+
(0, np.ones(4), None, np.zeros(4)),
447+
(np.arange(4), 0.5, None, np.arange(4)),
448+
(np.arange(4), np.arange(1, 5), (2, 4), np.full((2, 4), np.arange(4))),
449+
],
450+
)
451+
def test_vonmises_moment(mu, kappa, size, expected):
452+
with Model() as model:
453+
VonMises("x", mu=mu, kappa=kappa, size=size)
454+
assert_moment_is_expected(model, expected)
455+
456+
441457
@pytest.mark.parametrize(
442458
"mu, lam, phi, size, expected",
443459
[

0 commit comments

Comments
 (0)