Open In Colab

fastKDE and xarray

This notebook demonstrates basic usage of fastkde.plot with xarray.

import numpy as np

try:
    import fastkde
    import xarray
except:
    # install fastkde
    !pip install --upgrade fastkde
    import fastkde

import matplotlib.pyplot as plt

For this example, we will generate data with the following relationships:

\[x := \mathcal{N}(0,\pi)\]
\[y := \mathcal{N}(\sin(x), 1)\]
""" Sample the two variables """
N = int(1e5)
x = np.random.normal(size=N, scale=np.pi / 2)
y = np.sin(x) + np.random.normal(scale=1, size=N)

# plot the data
plt.scatter(x, y, s=1, alpha=0.1)
plt.xlabel("x")
plt.ylabel("y")
plt.show()
../_images/a0e19297e2e3f2e96dace75830959613224a1b778335d6c49241323b76da8f0b.png
""" Calculate and plot 1D PDFs. """

# calculate the PDFs of x and y
pdf_x = fastkde.pdf(x, var_names=["x"])
pdf_y = fastkde.pdf(y, var_names=["y"])

# plot PDF(x)
pdf_x.plot()
plt.show()

# plot PDF(y)
pdf_y.plot()
plt.show()
../_images/82680469c0e4a2fcdfc99a19fa89a55c6f4e20673c72e6c1f7e032569ee44d1f.png ../_images/464c121e440da596a490db0718a7af504209d2548f10aa78315594d9a16742a0.png
""" Compute the 2D PDF. """

pdf = fastkde.pdf(x, y, var_names=["x", "y"])
pdf
<xarray.DataArray (y: 129, x: 128)> Size: 132kB
array([[0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       ...,
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.],
       [0., 0., 0., ..., 0., 0., 0.]], shape=(129, 128))
Coordinates:
  * x        (x) float64 1kB -6.906 -6.802 -6.698 -6.594 ... 6.125 6.229 6.334
  * y        (y) float64 1kB -5.46 -5.381 -5.302 -5.222 ... 4.554 4.633 4.713
Attributes:
    long_name:  PDF(x,y)
""" Plot the PDF using xarray. """
# plot the 2D pdf
pdf.plot();
../_images/3aab6faa1dea2b1ccf2458766585c941e107f5e5f4c3789f412ed0c12dd0fae4.png
""" Compute and plot the conditional PDF using xarray. """

cpdf = fastkde.conditional(y, x, var_names=["x", "y"])

# plot the conditional
cpdf.plot()

# plot the true conditional mean
plt.plot(cpdf.x, np.sin(cpdf.x), color="white", alpha=0.5);
../_images/6db5ee2d24ec33cfc68578f7f22abad8316f2c1f8db396d93bbbaebc781537f8.png