Add support for fit from list

Author: krzysztofrusek

src/gsd/fit.py

@@ -27,6 +27,8 @@ def fit_moments(data: ArrayLike) -> GSDParams:
     :param data: A 5d Array of counts of each response.
     :return: GSD Parameters
     """
+
+    data= jnp.asarray(data)
     psi = jnp.dot(data, jnp.arange(1, 6)) / jnp.sum(data)
     V = jnp.dot(data, jnp.arange(1, 6) ** 2) / jnp.sum(data) - psi ** 2
     return GSDParams(psi=psi, rho=(vmax(psi) - V) / (vmax(psi) - vmin(psi)))
@@ -65,6 +67,7 @@ def fit_mle(data: ArrayLike, max_iterations: int = 100, log_lr_min: ArrayLike =
     :return: An opt state whore params filed contains estimated values of GSD Parameters
     """
 
+    data = jnp.asarray(data)
     def ll(theta: GSDParams) -> Array:
         logits = jax.vmap(log_prob, (None, None, 0), (0))(theta.psi, theta.rho, jnp.arange(1, 6))
         return jnp.dot(data, logits) / jnp.sum(data)

tests/fit_test.py

@@ -13,6 +13,14 @@ def test_mle(self):
         self.assertAlmostEqual(os.params.rho, 1)
 
 
+    def test_list(self):
+        #     1  2  3 4 5
+        data=[0,10,10,0,0.]
+        _,os = gsd.fit.fit_mle(data)
+        self.assertAlmostEqual(os.params.psi, 2.5)
+        self.assertAlmostEqual(os.params.rho, 1)
+
+
 
 if __name__ == '__main__':
     unittest.main()