Add __main__ for mle

Author: krzysztofrusek

README.md

@@ -48,3 +48,13 @@ Gsd uses unitest for testing. To run the tests, use the following command:
 ```
 $ hatch run test 
 ```
+
+### Standalone estimator
+
+You can quickly estimate GSD parameters from a command line interface
+
+```shell
+python3 -m gsd 0 12 13 4 0
+```
+
+    GSDParams(psi=Array(2.6272388, dtype=float32), rho=Array(0.9041536, dtype=float32))

src/gsd/__main__.py

@@ -1,17 +1,15 @@
-import jax
+import argparse
 import jax.numpy as jnp
-
-import gsd
+from gsd import fit_mle
 
 if __name__ == '__main__':
-    gsd.log_prob(1., 0.5, 2)
-    m = gsd.mean(3., 0.7)
-    v = gsd.variance(3., 0.7)
-    k = jax.random.key(43)
-    s = gsd.sample(3., 0.7, (24,), k)
+    parser = argparse.ArgumentParser(description='GSD estimator')
 
-    jnp.mean(s), jnp.var(s)
+    parser.add_argument("response", nargs=5, type=int,
+                        metavar=("num1", "num2", "num3", "num4", "num5"),
+                        help="List of 5 counts")
 
-    # jax.vmap(gsd.log_prob, in_axes=(None,None,0))(3.,0.7,s)
+    args = parser.parse_args()
 
-    print('test')
+    hat,_ = fit_mle(data=jnp.asarray(args.response, dtype=jnp.float32))
+    print(hat)

src/gsd/fit.py

@@ -6,7 +6,7 @@
 from jax import Array
 from jax.typing import ArrayLike
 
-from gsd.gsd import vmax, vmin, log_prob
+from .gsd import vmax, vmin, log_prob
 
 
 class GSDParams(NamedTuple):
@@ -108,3 +108,4 @@ def cond_fun(state: OptState) -> bool:
                                    OptState(params=theta0, previous_params=jtu.tree_map(lambda _: jnp.inf, theta0),
                                             count=0))
     return opt_state.params, opt_state
+