From ddff666e7fdbe7f1ef65a8e94b517f66052edea4 Mon Sep 17 00:00:00 2001
From: Felipe Cybis Pereira <felipe.cybis-pereira@espci.psl.eu>
Date: Tue, 14 Mar 2023 19:17:35 +0100
Subject: [PATCH] [ENH]: allow axis and weights parameters in `vonmisesmle`

Add changes to PR 14533

Move PR changes to right folder

Use ``_length`` to compute the generalized sample length
- This change will allow weighted vonmisesmle and the use of the
  ``axis`` parameter in the correct way

Add test to use axis and weights parameters in `vonmisesmle`
---
 astropy/stats/circstats.py            | 34 +++++++++++++++++--------
 astropy/stats/tests/test_circstats.py | 36 +++++++++++++++++++++++++++
 docs/changes/stats/14533.feature.rst  |  3 +++
 3 files changed, 63 insertions(+), 10 deletions(-)
 create mode 100644 docs/changes/stats/14533.feature.rst

diff --git a/astropy/stats/circstats.py b/astropy/stats/circstats.py
index 7f4c262fd2ac..d31ab7ecb82f 100644
--- a/astropy/stats/circstats.py
+++ b/astropy/stats/circstats.py
@@ -494,15 +494,24 @@ def vtest(data, mu=0.0, axis=None, weights=None):
 def _A1inv(x):
     # Approximation for _A1inv(x) according R Package 'CircStats'
     # See http://www.scienceasia.org/2012.38.n1/scias38_118.pdf, equation (4)
-    if 0 <= x < 0.53:
-        return 2.0 * x + x * x * x + (5.0 * x**5) / 6.0
-    elif x < 0.85:
-        return -0.4 + 1.39 * x + 0.43 / (1.0 - x)
-    else:
-        return 1.0 / (x * x * x - 4.0 * x * x + 3.0 * x)
-
 
-def vonmisesmle(data, axis=None):
+    kappa1 = np.where(
+        np.logical_and(0 <= x, x < 0.53),
+        2.0 * x + x * x * x + (5.0 * x**5) / 6.0,
+        0)
+    kappa2 = np.where(
+        np.logical_and(0.53 <= x, x < 0.85),
+        -0.4 + 1.39 * x + 0.43 / (1.0 - x),
+        0)
+    kappa3 = np.where(
+        np.logical_or(x < 0, 0.85 < x),
+        1.0 / (x * x * x - 4.0 * x * x + 3.0 * x),
+        0)
+    
+    return kappa1 + kappa2 + kappa3
+
+
+def vonmisesmle(data, axis=None, weights=None):
     """Computes the Maximum Likelihood Estimator (MLE) for the parameters of
     the von Mises distribution.
 
@@ -513,6 +522,11 @@ def vonmisesmle(data, axis=None):
         radians whenever ``data`` is ``numpy.ndarray``.
     axis : int, optional
         Axis along which the mle will be computed.
+    weights : numpy.ndarray, optional
+        In case of grouped data, the i-th element of ``weights`` represents a
+        weighting factor for each group such that ``sum(weights, axis)``
+        equals the number of observations. See [1]_, remark 1.4, page 22,
+        for detailed explanation.
 
     Returns
     -------
@@ -538,7 +552,7 @@ def vonmisesmle(data, axis=None):
        Circular Statistics (2001)'". 2015.
        <https://cran.r-project.org/web/packages/CircStats/CircStats.pdf>
     """
-    mu = circmean(data, axis=None)
+    mu = circmean(data, axis=axis, weights=weights)
 
-    kappa = _A1inv(np.mean(np.cos(data - mu), axis))
+    kappa = _A1inv(_length(data, p=1, phi=0., axis=axis, weights=weights))
     return mu, kappa
diff --git a/astropy/stats/tests/test_circstats.py b/astropy/stats/tests/test_circstats.py
index dfceefbe4154..88448a403572 100644
--- a/astropy/stats/tests/test_circstats.py
+++ b/astropy/stats/tests/test_circstats.py
@@ -146,3 +146,39 @@ def test_vonmisesmle():
     data = np.rad2deg(data) * u.deg
     answer = np.rad2deg(3.006514) * u.deg
     assert_equal(np.around(answer, 3), np.around(vonmisesmle(data)[0], 3))
+
+    # testing for weighted vonmisesmle
+    data = np.array(
+        [
+            np.pi/2, np.pi, np.pi/2,
+        ]
+    ) # this data has twice more np.pi/2 than np.pi
+    # get answer using astropy vonmisesmle to test
+    answer = vonmisesmle(data)
+    data_to_weigh = np.array(
+        [np.pi/2, np.pi]
+    )
+    weights = [2, 1]
+    assert_allclose(answer[0], vonmisesmle(data_to_weigh, weights=weights)[0], atol=1e-5)
+    assert_allclose(answer[1], vonmisesmle(data_to_weigh, weights=weights)[1], atol=1e-5)
+
+    # testing for axis argument (stacking the data from the first test)
+    data = np.array(
+        [
+            [3.3699057, 4.0411630, 0.5014477, 2.6223103, 3.7336524,
+            1.8136389, 4.1566039, 2.7806317, 2.4672173, 2.8493644],
+            [3.3699057, 4.0411630, 0.5014477, 2.6223103, 3.7336524,
+            1.8136389, 4.1566039, 2.7806317, 2.4672173, 2.8493644]
+        ]
+    )
+    # answer should be duplicated
+    answer = (np.array([3.006514, 3.006514]), np.array([1.474132, 1.474132]))
+    assert_allclose(answer[0], vonmisesmle(data, axis=1)[0], atol=1e-5)
+    assert_allclose(answer[1], vonmisesmle(data, axis=1)[1], atol=1e-5)
+
+    # same test for Quantity
+    data = np.rad2deg(data) * u.deg
+    answer = (np.rad2deg(answer[0]) * u.deg, answer[1])
+    assert_allclose(answer[0], vonmisesmle(data, axis=1)[0], atol=1e-5)
+    assert_allclose(answer[1], vonmisesmle(data, axis=1)[1], atol=1e-5)
+
diff --git a/docs/changes/stats/14533.feature.rst b/docs/changes/stats/14533.feature.rst
new file mode 100644
index 000000000000..33e647dae251
--- /dev/null
+++ b/docs/changes/stats/14533.feature.rst
@@ -0,0 +1,3 @@
+This pull request is to address the lack of "weights" and "axis" parameters on the circstats function ``vonmisesmle``.
+The "axis" parameter that existed but was not actually being used allows for fast usage of 2D numpy arrays.
+The "weights" parameter is very convenient since it is very common to have data for binned angles.
\ No newline at end of file