text: # -*- coding: utf-8 -*-
  # ---
  # jupyter:
  #   jupytext:
  #     formats: ipynb,py,md
  #     text_representation:
  #       extension: .py
  #       format_name: light
  #       format_version: '1.5'
  #       jupytext_version: 1.13.3
  #   kernelspec:
  #     display_name: Python 3
  #     language: python
  #     name: python3
  # ---
  
  # +
  # Uncomment to run the notebook in Colab
  # # ! pip install -q "wax-ml[complete]@git+https://github.com/eserie/wax-ml.git"
  # # ! pip install -q --upgrade jax jaxlib==0.1.70+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
  # -
  
  # check available devices
  import jax
  
  print("jax backend {}".format(jax.lib.xla_bridge.get_backend().platform))
  jax.devices()
  
  # # 〰 Compute exponential moving averages with xarray and pandas accessors 〰
  #
  # [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/eserie/wax-ml/blob/main/docs/notebooks/01_demo_EWMA.ipynb)
  
  # WAX-ML implements pandas and xarray accessors to ease the usage of machine-learning algorithms with
  # high-level data APIs :
  # - pandas's `DataFrame` and `Series`
  # - xarray's `Dataset` and `DataArray`.
  #
  # These accessors allow to easily execute any function using Haiku modules
  # on these data containers.
  #
  # For instance, WAX-ML propose an implementation of the exponential moving average realized
  # with this mechanism.
  #
  # Let's show how it works.
  
  # ## Load accessors
  
  # First you need to load accessors:
  
  # + tags=[]
  from wax.accessors import register_wax_accessors
  
  register_wax_accessors()
  # -
  
  # ## EWMA on dataframes
  
  # Let's look at a simple example: The exponential moving average (EWMA).
  #
  # Let's apply the EWMA algorithm to the [NCEP/NCAR 's Air temperature data](http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanalysis.html).
  
  # ### 🌡 Load temperature dataset 🌡
  
  # + tags=[]
  import xarray as xr
  
  dataset = xr.tutorial.open_dataset("air_temperature")
  # -
  
  # Let's see what this dataset looks like:
  
  # + tags=[]
  dataset
  # -
  
  # To compute a EWMA on some variables of a dataset, we usually need to convert data
  # in pandas
  # [series](https://pandas.pydata.org/docs/reference/api/pandas.Series.html) or
  # [dataframe](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html).
  #
  # So, let's convert the dataset into a dataframe to illustrate `accessors` on a dataframe:
  
  dataframe = dataset.air.to_series().unstack(["lon", "lat"])
  
  # ### EWMA with pandas
  
  air_temp_ewma = dataframe.ewm(alpha=1.0 / 10.0).mean()
  _ = air_temp_ewma.iloc[:, 0].plot()
  
  # ### EWMA with WAX-ML
  
  air_temp_ewma = dataframe.wax.ewm(alpha=1.0 / 10.0).mean()
  _ = air_temp_ewma.iloc[:, 0].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
  # (See our [three-steps_workflow](https://wax-ml.readthedocs.io/en/latest/notebooks/04_The_three_steps_workflow.html))
  #
  # ## Apply a custom function to a Dataset
  
  # Now let's illustrate how WAX-ML accessors work on [xarray datasets](http://xarray.pydata.org/en/stable/generated/xarray.Dataset.html).
  
  # +
  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"]),
      }
  
  
  dataset = xr.tutorial.open_dataset("air_temperature")
  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))
  