EWMA alignement with pandas and speedup

docs/notebooks/01_demo_EWMA.md

@@ -9,7 +9,7 @@ jupyter:
       format_version: '1.3'
       jupytext_version: 1.13.3
   kernelspec:
-    display_name: Python 3
+    display_name: Python 3 (ipykernel)
     language: python
     name: python3
 ---
@@ -96,15 +96,30 @@ dataframe = dataset.air.to_series().unstack(["lon", "lat"])
 ### EWMA with pandas
 
 ```python
-air_temp_ewma = dataframe.ewm(alpha=1.0 / 10.0).mean()
-_ = air_temp_ewma.iloc[:, 0].plot()
+air_temp_ewma = dataframe.ewm(com=10).mean()
+_ = air_temp_ewma.mean(1).plot()
+```
+
+## wax numba ewma
+
+```python
+from wax.numba.ewma_numba import register_wax_numba
+```
+
+```python
+register_wax_numba()
+```
+
+```python
+air_temp_ewma = dataframe.wax_numba.ewm(com=10).mean()
+_ = air_temp_ewma.mean(1).plot()
 ```
 
 ### EWMA with WAX-ML
 
 ```python
-air_temp_ewma = dataframe.wax.ewm(alpha=1.0 / 10.0).mean()
-_ = air_temp_ewma.iloc[:, 0].plot()
+air_temp_ewma = dataframe.wax.ewm(com=10).mean()
+_ = air_temp_ewma.mean(1).plot()
 ```
 
 On small data, WAX-ML's EWMA is slower than Pandas' because of the expensive data conversion steps.
@@ -122,8 +137,8 @@ from wax.modules import EWMA
 
 def my_custom_function(dataset):
     return {
-        "air_10": EWMA(1.0 / 10.0)(dataset["air"]),
-        "air_100": EWMA(1.0 / 100.0)(dataset["air"]),
+        "air_10": EWMA(com=10)(dataset["air"]),
+        "air_100": EWMA(com=100)(dataset["air"]),
     }
 
 
@@ -132,5 +147,5 @@ output, state = dataset.wax.stream().apply(
     my_custom_function, format_dims=dataset.air.dims
 )
 
-_ = output.isel(lat=0, lon=0).drop(["lat", "lon"]).to_pandas().plot(figsize=(12, 8))
+_ = output.isel(lat=0, lon=0).drop(["lat", "lon"]).to_dataframe().plot(figsize=(12, 8))
 ```

docs/notebooks/01_demo_EWMA.py

@@ -9,7 +9,7 @@
 #       format_version: '1.5'
 #       jupytext_version: 1.13.3
 #   kernelspec:
-#     display_name: Python 3
+#     display_name: Python 3 (ipykernel)
 #     language: python
 #     name: python3
 # ---
@@ -84,13 +84,22 @@
 
 # ### EWMA with pandas
 
-air_temp_ewma = dataframe.ewm(alpha=1.0 / 10.0).mean()
-_ = air_temp_ewma.iloc[:, 0].plot()
+air_temp_ewma = dataframe.ewm(com=10).mean()
+_ = air_temp_ewma.mean(1).plot()
+
+# ## wax numba ewma
+
+from wax.numba.ewma_numba import register_wax_numba
+
+register_wax_numba()
+
+air_temp_ewma = dataframe.wax_numba.ewm(com=10).mean()
+_ = air_temp_ewma.mean(1).plot()
 
 # ### EWMA with WAX-ML
 
-air_temp_ewma = dataframe.wax.ewm(alpha=1.0 / 10.0).mean()
-_ = air_temp_ewma.iloc[:, 0].plot()
+air_temp_ewma = dataframe.wax.ewm(com=10).mean()
+_ = air_temp_ewma.mean(1).plot()
 
 # On small data, WAX-ML's EWMA is slower than Pandas' because of the expensive data conversion steps.
 # WAX-ML's accessors are interesting to use on large data loads
@@ -106,8 +115,8 @@
 
 def my_custom_function(dataset):
     return {
-        "air_10": EWMA(1.0 / 10.0)(dataset["air"]),
-        "air_100": EWMA(1.0 / 100.0)(dataset["air"]),
+        "air_10": EWMA(com=10)(dataset["air"]),
+        "air_100": EWMA(com=100)(dataset["air"]),
     }
 
 
@@ -116,4 +125,4 @@ def my_custom_function(dataset):
     my_custom_function, format_dims=dataset.air.dims
 )
 
-_ = output.isel(lat=0, lon=0).drop(["lat", "lon"]).to_pandas().plot(figsize=(12, 8))
+_ = output.isel(lat=0, lon=0).drop(["lat", "lon"]).to_dataframe().plot(figsize=(12, 8))

docs/notebooks/02_Synchronize_data_streams.md

@@ -9,7 +9,7 @@ jupyter:
       format_version: '1.3'
       jupytext_version: 1.13.3
   kernelspec:
-    display_name: Python 3
+    display_name: Python 3 (ipykernel)
     language: python
     name: python3
 ---
@@ -90,9 +90,9 @@ from wax.modules import EWMA
 
 def my_custom_function(dataset):
     return {
-        "air_10": EWMA(1.0 / 10.0)(dataset["air"]),
-        "air_100": EWMA(1.0 / 100.0)(dataset["air"]),
-        "ground_100": EWMA(1.0 / 100.0)(dataset["ground"]),
+        "air_10": EWMA(com=10)(dataset["air"]),
+        "air_100": EWMA(com=100.0)(dataset["air"]),
+        "ground_100": EWMA(com=100.0)(dataset["ground"]),
     }
 ```
 
@@ -103,5 +103,5 @@ results, state = dataset.wax.stream(
 ```
 
 ```python
-_ = results.isel(lat=0, lon=0).drop(["lat", "lon"]).to_pandas().plot(figsize=(12, 8))
+_ = results.isel(lat=0, lon=0).drop(["lat", "lon"]).to_dataframe().plot(figsize=(12, 8))
 ```

docs/notebooks/02_Synchronize_data_streams.py

@@ -9,7 +9,7 @@
 #       format_version: '1.5'
 #       jupytext_version: 1.13.3
 #   kernelspec:
-#     display_name: Python 3
+#     display_name: Python 3 (ipykernel)
 #     language: python
 #     name: python3
 # ---
@@ -82,9 +82,9 @@
 
 def my_custom_function(dataset):
     return {
-        "air_10": EWMA(1.0 / 10.0)(dataset["air"]),
-        "air_100": EWMA(1.0 / 100.0)(dataset["air"]),
-        "ground_100": EWMA(1.0 / 100.0)(dataset["ground"]),
+        "air_10": EWMA(com=10)(dataset["air"]),
+        "air_100": EWMA(com=100.0)(dataset["air"]),
+        "ground_100": EWMA(com=100.0)(dataset["ground"]),
     }
 
 
@@ -94,4 +94,4 @@ def my_custom_function(dataset):
     local_time="time", ffills={"day": 1}, pbar=True
 ).apply(my_custom_function, format_dims=dataset.air.dims)
 
-_ = results.isel(lat=0, lon=0).drop(["lat", "lon"]).to_pandas().plot(figsize=(12, 8))
+_ = results.isel(lat=0, lon=0).drop(["lat", "lon"]).to_dataframe().plot(figsize=(12, 8))

docs/requirements.txt

@@ -1,11 +1 @@
-# sphinx <4 required by myst-nb v0.12.0 (Feb 2021)
-# sphinx >=3 required by sphinx-autodoc-typehints v1.11.1 (Oct 2020)
-sphinx >=3, <4
-sphinx_rtd_theme
-sphinx-autodoc-typehints==1.11.1
-jupyter-sphinx>=0.3.2
-myst-nb
-# Packages used for notebook execution
-matplotlib
-sklearn
-.[dev,complete]
+.[dev,complete,docs]

setup.cfg

@@ -42,6 +42,7 @@ install_requires =
 	jaxlib>=0.1.67
 	jax<=0.2.21
 	dm-haiku >= 0.0.4
+	numba
 
 [options.extras_require]
 optional = 
@@ -81,13 +82,17 @@ docs =
 	sphinx
 	sphinxcontrib-napoleon
 	sphinx_rtd_theme
+	sphinx-autodoc-typehints
 	sphinx-autosummary-accessors
 	ipython
 	ipykernel
 	jupyter-client
+	jupyter-sphinx
+	myst-nb
 	nbsphinx
 	scanpydoc
-
+	matplotlib
+	sklearn
 [options.package_data]
 wax = 
 	py.typed
@@ -207,6 +212,9 @@ ignore_errors = True
 [mypy-numpy.*]
 ignore_missing_imports = True
 
+[mypy-numba.*]
+ignore_missing_imports = True
+
 [mypy-opt_einsum.*]
 ignore_missing_imports = True
 

wax/accessors.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Define accessors for xarray and pandas data containers."""
 from dataclasses import dataclass, field
-from typing import Any, Callable, Tuple, Union
+from typing import Any, Callable, Optional, Tuple, Union
 
 import jax.numpy as jnp
 import numpy as onp
@@ -315,22 +315,23 @@ def ewm(self, *args, **kwargs):
 @dataclass(frozen=True)
 class ExponentialMovingWindow:
     accessor: WaxAccessor
-    alpha: float
+    com: Optional[float] = None
+    alpha: Optional[float] = None
+    min_periods: int = 0
     adjust: bool = True
+    ignore_na: bool = False
+    initial_value: float = jnp.nan
     return_state: bool = False
     format_outputs: bool = True
 
     def mean(self):
         from wax.modules import EWMA
 
-        def _apply_ema(
-            accessor, alpha, adjust, params=None, state=None, *args, **kwargs
-        ):
-            return accessor.stream(*args, **kwargs).apply(
-                lambda x: EWMA(alpha, adjust)(x),
-                params=params,
-                state=state,
-                rng=None,
+        def _apply_ema(*, accessor, return_state, format_outputs, **kwargs):
+            return accessor.stream(
+                return_state=return_state, format_outputs=format_outputs
+            ).apply(
+                lambda x: EWMA(**kwargs)(x),
             )
 
         return _apply_ema(**self.__dict__)

wax/accessors_test.py

@@ -41,7 +41,7 @@ def module_map(x):
 
 
 def check_ema_state(state, ref_count=124):
-    assert (state["ewma"]["count"] == ref_count).all()
+    assert (state["ewma"]["nobs"] == ref_count).all()
 
 
 def prepare_format_data(format):