Repository: PSLmodels/microdf

Issue:
MicroSeries(list) causes TypeError: __init__() got an unexpected keyword argument 'column'
This crashes my ipython session:
```
import microdf as mdf
mdf.Series([1, 2, 3])
```
Error is long, includes `TypeError: __init__() got an unexpected keyword argument 'column'`

Comments:

Patch:
diff --git a/microdf/generic.py b/microdf/generic.py
index 2b72bbd..654cba5 100644
--- a/microdf/generic.py
+++ b/microdf/generic.py
@@ -12,21 +12,29 @@ class MicroSeries(pd.Series):
         :type weights: np.array
         """
         super().__init__(*args, **kwargs)
-        self.weights = weights
+        self.set_weights(weights)
 
     def _init_micro(self, weights=None):
         self.weights = weights
 
+    def handles_zero_weights(fn):
+        def safe_fn(*args, **kwargs):
+            try:
+                return fn(*args, **kwargs)
+            except ZeroDivisionError:
+                return np.NaN
+
+        return safe_fn
+
     def set_weights(self, weights):
         """Sets the weight values.
 
         :param weights: Array of weights.
         :type weights: np.array.
-
-        :returns: A Pandas Series multiplying the MicroSeries by its weight.
         """
-        self.weights = weights
+        self.weights = pd.Series(weights)
 
+    @handles_zero_weights
     def weight(self):
         """Calculates the weighted value of the MicroSeries.
 
@@ -34,13 +42,15 @@ class MicroSeries(pd.Series):
         """
         return self.multiply(self.weights)
 
+    @handles_zero_weights
     def sum(self):
         """Calculates the weighted sum of the MicroSeries.
 
         :returns: The weighted sum.
         """
-        return self.weight().sum()
+        return self.multiply(self.weights).sum()
 
+    @handles_zero_weights
     def mean(self):
         """Calculates the weighted mean of the MicroSeries
 
@@ -48,6 +58,7 @@ class MicroSeries(pd.Series):
         """
         return np.average(self.values, weights=self.weights)
 
+    @handles_zero_weights
     def quantile(self, quantiles):
         """Calculates weighted quantiles of the MicroSeries.
 
@@ -76,6 +87,7 @@ class MicroSeries(pd.Series):
         weighted_quantiles /= np.sum(sample_weight)
         return np.interp(quantiles, weighted_quantiles, values)
 
+    @handles_zero_weights
     def median(self):
         """Calculates the weighted median of the MicroSeries.
 
@@ -83,6 +95,63 @@ class MicroSeries(pd.Series):
         """
         return self.quantile(0.5)
 
+    def groupby(self, *args, **kwargs):
+        gb = super().groupby(*args, **kwargs)
+        gb.__class__ = MicroSeriesGroupBy
+        gb.weights = pd.Series(self.weights).groupby(*args, **kwargs)
+        return gb
+
+    def __getitem__(self, key):
+        result = super().__getitem__(key)
+        if isinstance(result, pd.Series):
+            weights = self.weights.__getitem__(key)
+            return MicroSeries(result, weights=weights)
+        return result
+
+
+class MicroSeriesGroupBy(pd.core.groupby.generic.SeriesGroupBy):
+    def __init__(self, weights=None, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.weights = weights
+
+    def _weighted_agg(func):
+        def via_micro_series(row, fn, *args, **kwargs):
+            return getattr(MicroSeries(row.a, weights=row.w), fn.__name__)(
+                *args, **kwargs
+            )
+
+        def _weighted_agg_fn(self, *args, **kwargs):
+            arrays = self.apply(np.array)
+            weights = self.weights.apply(np.array)
+            df = pd.DataFrame(dict(a=arrays, w=weights))
+            result = df.agg(
+                lambda row: via_micro_series(row, func, *args, **kwargs),
+                axis=1,
+            )
+            return result
+
+        return _weighted_agg_fn
+
+    @_weighted_agg
+    def weight(self):
+        return MicroSeries.weight(self)
+
+    @_weighted_agg
+    def sum(self):
+        return MicroSeries.sum(self)
+
+    @_weighted_agg
+    def mean(self):
+        return MicroSeries.mean(self)
+
+    @_weighted_agg
+    def quantile(self, quantiles):
+        return MicroSeries.quantile(self, quantiles)
+
+    @_weighted_agg
+    def median(self):
+        return MicroSeries.median(self)
+
 
 class MicroDataFrame(pd.DataFrame):
     def __init__(self, *args, weights=None, **kwargs):
@@ -96,6 +165,11 @@ class MicroDataFrame(pd.DataFrame):
         super().__init__(*args, **kwargs)
         self.weights = weights
         self.weight_col = None
+        self._link_all_weights()
+
+    def __setitem__(self, *args, **kwargs):
+        super().__setitem__(*args, **kwargs)
+        self._link_all_weights()
 
     def _link_weights(self, column):
         # self[column] = ... triggers __setitem__, which forces pd.Series
@@ -132,3 +206,22 @@ class MicroDataFrame(pd.DataFrame):
         self.weights = np.array(self[column])
         self.weight_col = column
         self._link_all_weights()
+
+    def groupby(self, by, *args, **kwargs):
+        """Returns a GroupBy object with MicroSeriesGroupBy objects for each column
+
+        :param by: column to group by
+        :type by: str
+
+        return: DataFrameGroupBy object with columns using weights
+        rtype: DataFrameGroupBy
+        """
+        gb = super().groupby(by, *args, **kwargs)
+        weights = pd.Series(self.weights).groupby(self[by], *args, **kwargs)
+        for col in self.columns:  # df.groupby(...)[col]s use weights
+            if col != by:
+                res = gb[col]
+                res.__class__ = MicroSeriesGroupBy
+                res.weights = weights
+                setattr(gb, col, res)
+        return gb


Test patch:
diff --git a/microdf/tests/test_generic.py b/microdf/tests/test_generic.py
index 063e679..0cd6c57 100644
--- a/microdf/tests/test_generic.py
+++ b/microdf/tests/test_generic.py
@@ -2,6 +2,33 @@ import numpy as np
 import microdf as mdf
 
 
+def test_df_init():
+    arr = np.array([0, 1, 1])
+    w = np.array([3, 0, 9])
+    df = mdf.MicroDataFrame({"a": arr}, weights=w)
+    assert df.a.mean() == np.average(arr, weights=w)
+
+    df = mdf.MicroDataFrame()
+    df["a"] = arr
+    df.set_weights(w)
+    assert df.a.mean() == np.average(arr, weights=w)
+
+    df = mdf.MicroDataFrame()
+    df["a"] = arr
+    df["w"] = w
+    df.set_weight_col("w")
+    assert df.a.mean() == np.average(arr, weights=w)
+
+
+def test_series_getitem():
+    arr = np.array([0, 1, 1])
+    w = np.array([3, 0, 9])
+    s = mdf.MicroSeries(arr, weights=w)
+    assert s[[1, 2]].sum() == np.sum(arr[[1, 2]] * w[[1, 2]])
+
+    assert s[1:3].sum() == np.sum(arr[1:3] * w[1:3])
+
+
 def test_sum():
     arr = np.array([0, 1, 1])
     w = np.array([3, 0, 9])