# 1. Machine Learning

## Qué es Machine Learning?

El término Machine Learning según la definición común es "el estudio de los algoritmos de computador que mejoran automáticamente a través de la experiencia".

## Aplicaciones recientes:
### Computación visual.
- ImageNet
- Vehículos conducidos autonomamente.
- Redes Neuronales Convolucionales
- Robots retroalimentados con CV.

### Reconocimiento de voz
- Asistentes por voz.
- Traductores automáticos en tiempo real.

### Traducción de idioma.
- Traducción de idioma no supervisada.

### Videojuegos.
- Aprendizaje por refuerzo para solucionar juegos.
- Aprendizaje para mejorar experiencia del jugador.

### Ranking:
- Sistemas de recomendación.
...

## Tipos de Aprendizaje

Hay tres grandes categorías en las que se puede enmarcar el Machine Learning. Es importante poder identificar el tipo de algoritmo de Machine Learning que se va a utilizar dependiendo del caso de estudio ya que esto permite planear el ambiente más apropiado y entender qué hacer para que funcione.

## Aprendizaje Supervisado.

El aprendizaje supervisado es el más común y popular de los paradigmas de Machine Learning. El aprendizaje supervisado logra, a partir de datos conocidos y respuestas conocidas a estos datos, aproximar una respuesta, de manera que ante la presencia de datos nuevos estima una nueva respuesta. (ie. Aprender $y=f(x)$ a partir de ${x_i,y_i}^{n}_{i=1}$)

Aplicaciones:
- Clasificación de Spam.
- Reconocimiento facial.

## Aprendizaje No Supervisado.

El aprendizaje no supervisado es opuesto al aprendizaje supervisado en varios sentidos. El objetivo del aprendizaje no supervisado es que, a partir de grandes cantidades de información no etiquetada, sea posible entender las propiedades de la información. A partir de aquí, es posible aprender a agrupar y organizar la información de manera que un humano (u otro algoritmo inteligente) pueda encontrar un sentido a partir de la información organizada. El aprendizaje no supervisado es interesante debido a que mucha información en el mundo real no está etiquetada. (ie. Aprender similaridades en ${x_i}^{n}_{i=1}$)

Aplicaciones:

- Sistemas generadores de recomendaciones.

- Estimar hábitos de compra.

Aprendizaje por Refuerzo.

El aprendizaje por refuerzo es diferente a los aprendizajes supervisado y no supervisado. El objetivo del aprendizaje por refuerzo es determinar las políticas (o acciones) que debe tener un agente en orden de maximizar una función de recompensa. En Machine Learning, es muy común que el ambiente del agente sea modelado como un proceso de decisión de Markov. En general, se busca encontrar un equilibrio entre la exploración y la explotación.

## Aplicaciones:

- Simulaciones industriales.
- Manejo de recursos en ambientes complejos.
- Videojuegos.

## Aspectos importantes de Machine Learning.

- Validación y selección de Modelo.
- Aprendizaje Estadístico.
- Machine Learning Computacional.
- Optimización.

# 2. Aprendizaje Supervisado: Regresión Lineal

A partir de un conjunto de variables de entrada, se busca predecir una característica objetivo de salida. Se denota una pareja $(x^{(i)},y^{(i)})$, la muestra de entrenamiento y el conjunto desde $i=1,...,n$ es denotado, conjunto de entrenamiento. El objetivo es a partir de nuestro conjunto deentrada $\mathcal{X}$ obtener una hipótesis $h:\mathcal{X} \rightarrow \mathcal{Y}$. Cuando $y$ únicamente puede tomar un pequeño número de valores discretos, se dice que este es un problema de clasificación.

## 2.1 Regresión Lineal

Para realizar aprendizaje supervisado, se necesita decidir cómo representar funciones hipótesis $h$. Inicialmente, es posible decidir representarlas como funciones lineales de x:

Denotaremos

\begin{equation}
h_{\theta}(x)=\theta_0 + \theta_1 x_1 +\theta_2 x_2
\end{equation}

Aquí, las $\theta_i$ son los parámetros, también llamados pesos (weights). También es posible escribir lo anterior de la siguiente manera.

\begin{equation}
h_{\theta}(x)=\sum_{i=0}^{d}\theta_i x_i=\theta^T x
\end{equation}

Ahora, es importante en un conjunto de entrenamiento, escoger una función de costo que permita hacer que $h(x)$ sea tan cercana a $y$ como sea posible para obtener unos parámetros $\theta$ apropiados. Definimos una función de costo:

\begin{equation}
J(\theta)=\frac{1}{2}\sum_{i=1}^{n}(h_{\theta}(x^{(i)})-y^{(i)})^2
\end{equation}

Y podemos recordar esta ecuación como la solución al modelo de mínimos cuadrados.

```python
import numpy as np
import matplotlib.pyplot as plt

def linReg(X, y, theta=0, alpha=0.01, num_iters=100):
    m = y.size
    costs=[]
    for i in range(num_iters):
        y_hat = np.dot(X, theta)
        costi =1.0/2.0 * sum((y_hat-y)**2)
        costs.append(costi)
        theta = theta - alpha * (1.0/m) * np.dot(X.T, y_hat-y)
    return theta, np.array(costs)

#Generamos un conjunto de datos sintético de prueba.
x=np.linspace(-1,5,10)
y=7*x
y=y+(np.random.rand(10)*2-1)*7

#Realizamos el ajuste
t,c=linReg(x,y)

#Graficamos el ajuste
plt.plot(x,t*x,color='k',label='fit')
plt.scatter(x,y,color='blue',label='datos')
plt.legend()
plt.title('Regresión lineal, descenso de gradiente')
plt.show()

#Graficamos la función de costo en función de las iteraciones
plt.plot(c)
plt.xlabel('Iteración')
plt.ylabel(r'$J(\theta)$')
plt.title('Función de costo')
plt.show()

```

# 3. Mínimos cuadrados
Es claro que a partir de la función de costo $J(\theta)$ que se definió previamente, el parámetro óptimo que se puede obtener utilizando mínimos cuadrados es $\hat{\theta}=(X^T X)^{-1}X^T \vec{y}$. Sin embargo, no siempre es factible llegar a esta solución debido a distintas ra

## 3.1 Interpretación probabilística

Cuando se tiene un problema de regresión es común pensar en por qué la escogencia de la función de costo $J$ es apropiada. Asumamos que el target y las entradas están relacionadas a través de

\begin{equation}
y=\theta^T x+ \epsilon
\end{equation}

donde $\epsilon$ es un término de error que contiene toda la información que no se ha modelado. Se puede asumir que este término es IID y se modela gaussiano estándar. Por lo tanto, se puede decir que

\begin{equation}
p(y|x;\theta)=\frac{1}{\sqrt{2\pi \sigma^2}} \exp{\frac{-(y-\theta^T x)^2}{2\sigma^2}}
\end{equation}

A partir de esta probabilidad se desea observar esta función explicitamente como una función de $\theta$, a esta función se la llamará función de likelihood.

\begin{equation}
L(\theta)=L(\theta;X,y)=p(y|X;\theta)
\end{equation}

Donde
\begin{equation}
L(\theta)=\prod{\frac{1}{\sqrt{2\pi \sigma^2}} \exp{\frac{-(y-\theta^T x)^2}{2\sigma^2}}}
\end{equation}

Ahora, dado este modelo probabilístico que relaciona $y$ con $x$, se desea cuál es la manera razonable de escoger el parámetro $\theta$. Como se ha visto previamente, en lugar de maximizar esta función, es posible maximizar la función log-llikelihood $l(\theta)$

Donde
\begin{equation}
l(\theta)=\log{\prod{\frac{1}{\sqrt{2\pi \sigma^2}} \exp{\frac{-(y-\theta^T x)^2}{2\sigma^2}}}}
\end{equation}

De manera que el máximo es

Donde
\begin{equation}
\frac{1}{2}\sum_{i=1}^{n}{(y^{(i)}-\theta^T x^{(i)})^2}
\end{equation}

Por lo que la escogenca de nuestra función de costo es justificada.

```python
import matplotlib.pyplot as plt
import numpy as np
from sklearn import datasets, linear_model
from sklearn.metrics import mean_squared_error, r2_score

diabetes = datasets.load_diabetes()

diabetes_X = diabetes.data[:, np.newaxis, 2]

diabetes_X_train = diabetes_X[:-20]
diabetes_X_test = diabetes_X[-20:]

diabetes_y_train = diabetes.target[:-20]
diabetes_y_test = diabetes.target[-20:]

regr = linear_model.LinearRegression()

regr.fit(diabetes_X_train, diabetes_y_train)

diabetes_y_pred = regr.predict(diabetes_X_test)

print('Coefficients: \n', regr.coef_)
print("Mean squared error: %.2f"
      % mean_squared_error(diabetes_y_test, diabetes_y_pred))
print('Variance score: %.2f' % r2_score(diabetes_y_test, diabetes_y_pred))

plt.scatter(diabetes_X_test, diabetes_y_test,  color='blue')
plt.plot(diabetes_X_test, diabetes_y_pred, color='black', linewidth=3)
plt.title('Regresión lineal')
plt.xlabel('x')
plt.ylabel('y')
plt.show()
```

```python

```


# Truncated and folded distributions

This tutorial will cover how to work with truncated and folded
distributions in NumPyro.
It is assumed that you're already familiar with the basics of NumPyro.
To get the most out of this tutorial you'll need some background in probability.

### Table of contents

* [0. Setup](#0)
* [1. What is a truncated distribution?](#1)
* [2. What is a folded distribution?](#2)
* [3. Sampling from truncated and folded distributions](#3)
* [4. Ready-to-use truncated and folded distributions](#4)
* [5. Building your own truncanted distributions](#5)
    * [5.1 Recap of NumPyro distributions](#5.1)
    * [5.2 Right-truncated normal](#5.2)
    * [5.3 Left-truncated Poisson](#5.3)
* [6. References and related material](#references)

### Setup
To run this notebook, we are going to need the following imports

```python
!pip install -q git+https://github.com/pyro-ppl/numpyro.git
```

```python
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import numpyro.distributions as dist
from jax import lax, random
from jax.scipy.special import ndtr, ndtri
from jax.scipy.stats import poisson, norm
from numpyro.distributions import (
    constraints,
    Distribution,
    FoldedDistribution,
    SoftLaplace,
    StudentT,
    TruncatedDistribution,
    TruncatedNormal,
)
from numpyro.distributions.util import promote_shapes
from numpyro.infer import DiscreteHMCGibbs, MCMC, NUTS, Predictive
from scipy.stats import poisson as sp_poisson

numpyro.enable_x64()
RNG = random.PRNGKey(0)
PRIOR_RNG, MCMC_RNG, PRED_RNG = random.split(RNG, 3)
MCMC_KWARGS = dict(
    num_warmup=2000,
    num_samples=2000,
    num_chains=4,
    chain_method="sequential",
)
```

### 1. What are truncated distributions?

The **support** of a probability distribution is the set of values
in the domain with **non-zero probability**. For example, the
support of the normal distribution is the whole real line (even if
the density gets very small as we move away from the mean, technically
speaking, it is never quite zero). The support of the uniform distribution,
as coded in `jax.random.uniform` with the default arguments, is the interval $\left[0, 1)\right.$, because any
value outside of that interval has zero probability. The support of the Poisson distribution is the set of non-negative integers, etc.

**Truncating** a distribution makes its support smaller
so that any value outside our desired domain has zero probability. In practice, this can be useful
for modelling situations in which certain biases are introduced during data collection.
For example, some physical detectors only get triggered when the signal is above some
minimum threshold, or sometimes the detectors fail if the signal exceeds a certain value.
As a result, the **observed values are constrained to be within a limited range of values**,
even though the true signal does not have the same constraints.
See, for example, section 3.1 of _Information Theory and Learning Algorithms_ by David Mackay.
Naively, if $S$ is the support of the original density $p_Y(y)$, then by truncating to a new support
$T\subset S$ we are effectively defining a new random variable $Z$ for which the density is

$$
\begin{align}
  p_Z(z) \propto
    \begin{cases}
      p_Y(z) & \text{if $z$ is in $T$}\\
      0 & \text{if $z$ is outside $T$}\\
    \end{cases}
\end{align}
$$

The reason for writing a $\propto$ (proportional to) sign instead of a strict equation is that,
defined in the above way, the resulting function does not integrate to $1$ and so it cannot be strictly considered a probability density. To make it into a probability density **we need to re-distribute the truncated mass**
among the part of the distribution that remains. To do this, we simply re-weight every point by the same constant:

$$
\begin{align}
  p_Z(z) =
    \begin{cases}
      \frac{1}{M}p_Y(z) & \text{if $z$ is in $T$}\\
      0 & \text{if $z$ is outside $T$}\\
    \end{cases}
\end{align}
$$

where $M = \int_T p_Y(y)\mathrm{d}y$.

In practice, the truncation is often one-sided. This means that if, for example, the support before truncation is the interval $(a, b)$, then the support after truncation is of the form $(a, c)$ or $(c, b)$, with $a < c < b$. The figure below illustrates a left-sided truncation at zero of a normal distribution $N(1, 1)$.

The original distribution (left side) is truncated at the vertical dotted line. The truncated mass (orange region) is redistributed in the new support (right side image) so that the total area under the curve remains equal to 1 even after truncation. This method of re-weighting ensures that the density ratio between any two points, $p(a)/p(b)$ remains the same before and after the reweighting is done (as long as the points are inside the new support, of course).

**Note**: Truncated data is different from _censored_ data. Censoring also hides values that are outside some desired support but, contrary to truncated data, we know when a value has been censored. The typical example is the household scale which does not report values above 300 pounds. Censored data will not be covered in this tutorial.

### 2. What is a folded distribution?

**Folding** is achieved by taking the absolute value of a random variable, $Z = \lvert Y \rvert$. This obviously modifies the support of the original distribution since negative values now have zero
probability:

$$
\begin{align}
  p_Z(z) =
    \begin{cases}
      p_Y(z) + p_Y(-z) & \text{if $z\ge 0$}\\
      0 & \text{if $z\lt 0$}\\
    \end{cases}
\end{align}
$$

The figure below illustrates a folded normal distribution $N(1, 1)$.

As you can see, the resulting distribution is different from the truncated case. In particular, the density ratio between points, $p(a)/p(b)$, is in general not the same after folding. For some examples in which folding is relevant see [references 3 and 4](#references)

If the original distribution is symmetric around zero, then folding and truncating at zero have the same effect.

### 3. Sampling from truncated and folded distributions

**Truncated distributions**

Usually, we already have a sampler for the pre-truncated distribution (e.g. `np.random.normal`).
So, a seemingly simple way of generating samples from the truncated distribution would be to
sample from the original distribution, and then discard the samples that are outside the
desired support. For example, if we wanted samples from a normal distribution truncated to the
support $(-\infty, 1)$, we'd simply do:

```python
upper = 1
samples = np.random.normal(size=1000)
truncated_samples = samples[samples < upper]
```

This is called **_rejection sampling_ but it is not very efficient**.
If the region we truncated had a sufficiently high probability mass, then we'd be discarding a lot of samples and it might be a while before we accumulate sufficient samples for the truncated distribution. For example, the above snippet would only result in approximately 840 truncated samples even though we initially drew 1000. This can easily get a lot worse for other combinations of parameters.
A **more efficient** approach is to use a method known as [inverse transform sampling](https://en.wikipedia.org/wiki/Inverse_transform_sampling).
In this method, we first sample from a uniform distribution in (0, 1) and then transform those samples with the inverse cumulative distribution of our truncated distribution.
This method ensures that no samples are wasted in the process, though it does have the slight complication that
**we need to calculate the inverse CDF (ICDF)** of our truncated distribution. This might sound too complicated at first but, with a bit of algebra, we can often calculate the truncated ICDF in terms of the untruncated ICDF. The untruncated ICDF for many distributions is already available.

**Folded distributions**

This case is a lot simpler. Since we already have a sampler for the pre-folded distribution, all we need to do is to take the absolute value of those samples:

```python
samples = np.random.normal(size=1000)
folded_samples = np.abs(samples)
```

### 4. Ready to use truncated and folded distributions

The later sections in this tutorial will show you how to construct your own truncated and folded distributions, but you don't have to reinvent the wheel. NumPyro has [a bunch of truncated distributions](https://num.pyro.ai/en/stable/distributions.html#truncated-distributions) already implemented.

Suppose, for example, that you want a normal distribution truncated on the right.
For that purpose, we use the [TruncatedNormal](https://num.pyro.ai/en/stable/distributions.html#truncatednormal) distribution. The parameters of this distribution are `loc` and `scale`, corresponding to the `loc` and `scale` of the _untruncated_ normal, and `low` and/or `high` corresponding to the truncation points. Importantly, the `low` and `high` are **keyword only** arguments, only `loc` and `scale` are valid as positional arguments.
This is how you can use this class in a model:

```python
def truncated_normal_model(num_observations, high, x=None):
    loc = numpyro.sample("loc", dist.Normal())
    scale = numpyro.sample("scale", dist.LogNormal())
    with numpyro.plate("observations", num_observations):
        numpyro.sample("x", TruncatedNormal(loc, scale, high=high), obs=x)
```

Let's now check that we can use this model in a typical MCMC workflow.

**Prior simulation**

```python
high = 1.2
num_observations = 250
num_prior_samples = 100

prior = Predictive(truncated_normal_model, num_samples=num_prior_samples)
prior_samples = prior(PRIOR_RNG, num_observations, high)
```

**Inference**

To test our model, we run mcmc against some synthetic data.
The synthetic data can be any arbitrary sample from the prior simulation.

```python
# -- select an arbitrary prior sample as true data
true_idx = 0
true_loc = prior_samples["loc"][true_idx]
true_scale = prior_samples["scale"][true_idx]
true_x = prior_samples["x"][true_idx]
```

```python
plt.hist(true_x.copy(), bins=20)
plt.axvline(high, linestyle=":", color="k")
plt.xlabel("x")
plt.show()
```

```python
# --- Run MCMC and check estimates and diagnostics
mcmc = MCMC(NUTS(truncated_normal_model), **MCMC_KWARGS)
mcmc.run(MCMC_RNG, num_observations, high, true_x)
mcmc.print_summary()

# --- Compare to ground truth
print(f"True loc  : {true_loc:3.2}")
print(f"True scale: {true_scale:3.2}")
```

    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:02<00:00, 1909.24it/s, 1 steps of size 5.65e-01. acc. prob=0.93]
    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 10214.14it/s, 3 steps of size 5.16e-01. acc. prob=0.95]
    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 15102.95it/s, 1 steps of size 6.42e-01. acc. prob=0.90]
    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 16522.03it/s, 3 steps of size 6.39e-01. acc. prob=0.90]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
           loc     -0.58      0.15     -0.59     -0.82     -0.35   2883.69      1.00
         scale      1.49      0.11      1.48      1.32      1.66   3037.78      1.00

    Number of divergences: 0
    True loc  : -0.56
    True scale: 1.4

**Removing the truncation**

Once we have inferred the parameters of our model, a common task is to understand what the data would look like _without_ the truncation. In this example, this is easily done by simply "pushing" the value of `high` to infinity.

```python
pred = Predictive(truncated_normal_model, posterior_samples=mcmc.get_samples())
pred_samples = pred(PRED_RNG, num_observations, high=float("inf"))
```

Let's finally plot these samples and compare them to the original, observed data.

```python
# thin the samples to not saturate matplotlib
samples_thinned = pred_samples["x"].ravel()[::1000]
```

```python
f, axes = plt.subplots(1, 2, figsize=(15, 5), sharex=True)

axes[0].hist(
    samples_thinned.copy(), label="Untruncated posterior", bins=20, density=True
)
axes[0].set_title("Untruncated posterior")

vals, bins, _ = axes[1].hist(
    samples_thinned[samples_thinned < high].copy(),
    label="Tail of untruncated posterior",
    bins=10,
    density=True,
)
axes[1].hist(
    true_x.copy(), bins=bins, label="Observed, truncated data", density=True, alpha=0.5
)
axes[1].set_title("Comparison to observed data")

for ax in axes:
    ax.axvline(high, linestyle=":", color="k", label="Truncation point")
    ax.legend()

plt.show()
```

The plot on the left shows data simulated from the posterior distribution with the truncation removed, so we are able to see how the data would look like if it were not truncated. To sense check this, we discard the simulated samples that are above the truncation point and make histogram of those and compare it to a histogram of the true data (right plot).

**The TruncatedDistribution class**

The source code for the [TruncatedNormal](https://num.pyro.ai/en/stable/distributions.html#truncatednormal) in NumPyro uses a class called [TruncatedDistribution](https://num.pyro.ai/en/stable/distributions.html#truncateddistribution) which abstracts away the logic for `sample` and `log_prob` that
we will discuss in the next sections. At the moment, though, this logic only works continuous, symmetric distributions with _real_ support.

We can use this class to quickly construct other truncated distributions. For example, if we need a truncated [SoftLaplace](https://num.pyro.ai/en/stable/distributions.html#softlaplace) we can use the following pattern:

```python
def TruncatedSoftLaplace(
    loc=0.0, scale=1.0, *, low=None, high=None, validate_args=None
):
    return TruncatedDistribution(
        base_dist=SoftLaplace(loc, scale),
        low=low,
        high=high,
        validate_args=validate_args,
    )
```

```python
def truncated_soft_laplace_model(num_observations, high, x=None):
    loc = numpyro.sample("loc", dist.Normal())
    scale = numpyro.sample("scale", dist.LogNormal())
    with numpyro.plate("obs", num_observations):
        numpyro.sample("x", TruncatedSoftLaplace(loc, scale, high=high), obs=x)
```

And, as before, we check that we can use this model in the steps of a typical workflow:

```python
high = 2.3
num_observations = 200
num_prior_samples = 100

prior = Predictive(truncated_soft_laplace_model, num_samples=num_prior_samples)
prior_samples = prior(PRIOR_RNG, num_observations, high)

true_idx = 0
true_x = prior_samples["x"][true_idx]
true_loc = prior_samples["loc"][true_idx]
true_scale = prior_samples["scale"][true_idx]

mcmc = MCMC(
    NUTS(truncated_soft_laplace_model),
    **MCMC_KWARGS,
)

mcmc.run(
    MCMC_RNG,
    num_observations,
    high,
    true_x,
)

mcmc.print_summary()

print(f"True loc  : {true_loc:3.2}")
print(f"True scale: {true_scale:3.2}")
```

    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:02<00:00, 1745.70it/s, 1 steps of size 6.78e-01. acc. prob=0.93]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 9294.56it/s, 1 steps of size 7.02e-01. acc. prob=0.93]
    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 10412.30it/s, 1 steps of size 7.20e-01. acc. prob=0.92]
    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 10583.85it/s, 3 steps of size 7.01e-01. acc. prob=0.93]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
           loc     -0.37      0.17     -0.38     -0.65     -0.10   4034.96      1.00
         scale      1.46      0.12      1.45      1.27      1.65   3618.77      1.00

    Number of divergences: 0
    True loc  : -0.56
    True scale: 1.4

**Important**

The `sample` method of the [TruncatedDistribution](https://num.pyro.ai/en/stable/distributions.html#truncateddistribution) class relies on inverse-transform sampling.
This has the implicit requirement that the base distribution should have an `icdf` method already available.
If this is not the case, we will not be able to call the `sample` method on any instances of our distribution, nor use it with the `Predictive` class.
However, the `log_prob` method only depends on the `cdf` method (which is more frequently available than the `icdf`). If the `log_prob` method is available, then we _can_ use our distribution as prior/likelihood in a model.

**The FoldedDistribution class**

Similar to truncated distributions, NumPyro has the [FoldedDistribution](https://num.pyro.ai/en/stable/distributions.html#foldeddistribution) class to help you quickly construct folded distributions. Popular examples of folded distributions are the so-called "half-normal", "half-student" or "half-cauchy". As the name suggests, these distributions keep only (the positive) _half_ of the distribution. Implicit in the name of these "half" distributions is that they are centered at zero before folding. But, of course, you can fold a distribution even if its not centered at zero. For instance, this is how you would define a folded student-t distribution.

```python
def FoldedStudentT(df, loc=0.0, scale=1.0):
    return FoldedDistribution(StudentT(df, loc=loc, scale=scale))
```

```python
def folded_student_model(num_observations, x=None):
    df = numpyro.sample("df", dist.Gamma(6, 2))
    loc = numpyro.sample("loc", dist.Normal())
    scale = numpyro.sample("scale", dist.LogNormal())
    with numpyro.plate("obs", num_observations):
        numpyro.sample("x", FoldedStudentT(df, loc, scale), obs=x)
```

And we check that we can use our distribution in a typical workflow:

```python
# --- prior sampling
num_observations = 500
num_prior_samples = 100
prior = Predictive(folded_student_model, num_samples=num_prior_samples)
prior_samples = prior(PRIOR_RNG, num_observations)

# --- choose any prior sample as the ground truth
true_idx = 0
true_df = prior_samples["df"][true_idx]
true_loc = prior_samples["loc"][true_idx]
true_scale = prior_samples["scale"][true_idx]
true_x = prior_samples["x"][true_idx]

# --- do inference with MCMC
mcmc = MCMC(
    NUTS(folded_student_model),
    **MCMC_KWARGS,
)
mcmc.run(MCMC_RNG, num_observations, true_x)

# --- Check diagostics
mcmc.print_summary()

# --- Compare to ground truth:
print(f"True df   : {true_df:3.2f}")
print(f"True loc  : {true_loc:3.2f}")
print(f"True scale: {true_scale:3.2f}")
```

    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:02<00:00, 1343.54it/s, 7 steps of size 3.51e-01. acc. prob=0.75]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:01<00:00, 3644.99it/s, 7 steps of size 3.56e-01. acc. prob=0.73]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:01<00:00, 3137.13it/s, 7 steps of size 2.62e-01. acc. prob=0.91]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:01<00:00, 3028.93it/s, 7 steps of size 1.85e-01. acc. prob=0.96]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
            df      3.12      0.52      3.07      2.30      3.97   2057.60      1.00
           loc     -0.02      0.88     -0.03     -1.28      1.34    925.84      1.01
         scale      2.23      0.21      2.25      1.89      2.57   1677.38      1.00

    Number of divergences: 33
    True df   : 3.01
    True loc  : 0.37
    True scale: 2.41

### 5. Building your own truncated distribution

If the
[TruncatedDistribution](https://num.pyro.ai/en/stable/distributions.html#truncateddistribution) and
[FoldedDistribution](https://num.pyro.ai/en/stable/distributions.html#foldeddistribution)
classes are not sufficient to solve your problem,
you might want to look into writing your own truncated distribution from the ground up.
This can be a tedious process, so this section will give you some guidance and examples to help you with it.

#### 5.1 Recap of NumPyro distributions

A NumPyro distribution should subclass [Distribution](https://num.pyro.ai/en/stable/distributions.html#distribution) and implement a few basic ingredients:

**Class attributes**

The class attributes serve a few different purposes. Here we will mainly care about two:
1. `arg_constraints`: Impose some requirements on the parameters of the distribution. Errors are raised at instantiation time if the parameters passed do not satisfy the constraints.
2. `support`: It is used in some inference algorithms like MCMC and SVI with auto-guides, where we need to perform the algorithm in the unconstrained space. Knowing the support, we can automatically reparametrize things under the hood.

We'll explain other class attributes as we go.

**The** `__init__` **method**

This is where we define the parameters of the distribution.
We also use `jax` and `lax` to promote the parameters to shapes that are valid for broadcasting.
The `__init__` method of the parent class is also required because that's where the validation of our parameters is done.

**The** `log_prob` **method**

Implementing the `log_prob` method ensures that we can do inference. As the name suggests, this method returns the logarithm of the density evaluated at the argument.

**The** `sample` **method**

This method is used for drawing independent samples from our distribution. It is particularly useful for doing prior and posterior predictive checks. Note, in particular, that this method is not needed if you only need to use your distribution as prior in a model - the `log_prob` method will suffice.

The place-holder code for any of our implementations can be written as
```python
class MyDistribution(Distribution):
    # class attributes
    arg_constraints = {}
    support = None
    def __init__(self):
        pass

    def log_prob(self, value):
        pass

    def sample(self, key, sample_shape=()):
        pass
```

#### 5.2 Example: Right-truncated normal

We are going to modify a normal distribution so that its new support is
of the form `(-inf, high)`, with `high` a real number. This could be done with the `TruncatedNormal` distribution but, for the sake of illustration, we are not going to rely on it.
We'll call our distribution `RightTruncatedNormal`. Let's write the skeleton code and then proceed to fill in the blanks.

```python
class RightTruncatedNormal(Distribution):
    # <class attributes>
    def __init__(self):
        pass

    def log_prob(self, value):
        pass

    def sample(self, key, sample_shape=()):
        pass
```

**Class attributes**

Remember that a non-truncated normal distribution is specified in NumPyro by two parameters, `loc` and `scale`,
which correspond to the mean and standard deviation.
Looking at the [source code](https://github.com/pyro-ppl/numpyro/blob/0664c2d2dd1eb5f41ea6a0bcef91e5fa2a417ce5/numpyro/distributions/continuous.py#L1337) for the `Normal` distribution we see the following lines:

```python
arg_constraints = {"loc": constraints.real, "scale": constraints.positive}
support = constraints.real
reparametrized_params = ["loc", "scale"]
```

The `reparametrized_params` attribute is used by variational inference algorithms when constructing gradient estimators. The parameters of many common distributions with continuous support (e.g. the Normal distribution) are reparameterizable, while the parameters of discrete distributions are not. Note that `reparametrized_params` is irrelevant for MCMC algorithms like HMC. See [SVI Part III](https://pyro.ai/examples/svi_part_iii.html#Tricky-Case:-Non-reparameterizable-Random-Variables) for more details.

We must adapt these attributes to our case by including the `"high"` parameter, but there are two issues we need to deal with:

1. `constraints.real` is a bit too restrictive. We'd like `jnp.inf` to be a valid value for `high` (equivalent to no truncation), but at the moment infinity is not a valid real number. We deal with this situation by defining our own constraint. The source code for `constraints.real` is easy to imitate:

```python
class _RightExtendedReal(constraints.Constraint):
    """
    Any number in the interval (-inf, inf].
    """
    def __call__(self, x):
        return (x == x) & (x != float("-inf"))

    def feasible_like(self, prototype):
        return jnp.zeros_like(prototype)

right_extended_real = _RightExtendedReal()
```

2. `support` can no longer be a class attribute as it will depend on the value of `high`. So instead we implement it as a dependent property.

Our distribution then looks as follows:
```python
class RightTruncatedNormal(Distribution):
    arg_constraints = {
        "loc": constraints.real,
        "scale": constraints.positive,
        "high": right_extended_real,
    }
    reparametrized_params = ["loc", "scale", "high"]

    # ...

    @constraints.dependent_property
    def support(self):
        return constraints.lower_than(self.high)
```

**The** `__init__` **method**

Once again we take inspiration from the [source code](https://github.com/pyro-ppl/numpyro/blob/0664c2d2dd1eb5f41ea6a0bcef91e5fa2a417ce5/numpyro/distributions/continuous.py#L1342) for the normal distribution. The key point is the use of `lax` and `jax` to check the shapes of the arguments passed and make sure that such shapes are consistent for broadcasting. We follow the same pattern for our use case -- all we need to do is include the `high` parameter.

In the source implementation of `Normal`, both parameters `loc` and `scale` are given defaults so that one recovers a standard normal distribution if no arguments are specified. In the same spirit, we choose `float("inf")` as a default for `high` which would be equivalent to no truncation.

```python
# ...
    def __init__(self, loc=0.0, scale=1.0, high=float("inf"), validate_args=None):
        batch_shape = lax.broadcast_shapes(
            jnp.shape(loc),
            jnp.shape(scale),
            jnp.shape(high),
        )
        self.loc, self.scale, self.high = promote_shapes(loc, scale, high)
        super().__init__(batch_shape, validate_args=validate_args)
# ...
```

**The** `log_prob` **method**

For a truncated distribution, the log density is given by

$$
\begin{align}
  \log p_Z(z) =
    \begin{cases}
      \log p_Y(z) - \log M & \text{if $z$ is in $T$}\\
      -\infty & \text{if $z$ is outside $T$}\\
    \end{cases}
\end{align}
$$

where, again, $p_Z$ is the density of the truncated distribution, $p_Y$ is the density before truncation, and $M = \int_T p_Y(y)\mathrm{d}y$. For the specific case of truncating the normal distribution to the interval `(-inf, high)`, the constant $M$ is equal to the cumulative density evaluated at the truncation point. We can easily implement this log-density method because `jax.scipy.stats` already has a `norm` module that we can use.

```python
# ...
    def log_prob(self, value):
        log_m = norm.logcdf(self.high, self.loc, self.scale)
        log_p = norm.logpdf(value, self.loc, self.scale)
        return jnp.where(value < self.high, log_p - log_m, -jnp.inf)
# ...
```

**The** `sample` **method**

To implement the sample method using inverse-transform sampling, we need to also implement the inverse cumulative distribution function. For this, we can use the `ndtri` function that lives inside `jax.scipy.special`. This function returns the inverse cdf for the standard normal distribution. We can do a bit of algebra to obtain the inverse cdf of the truncated, non-standard normal. First recall that if $X\sim Normal(0, 1)$ and $Y = \mu + \sigma X$, then $Y\sim Normal(\mu, \sigma)$. Then if $Z$ is the truncated $Y$, its cumulative density is given by:

$$
\begin{align}
F_Z(y) &= \int_{-\infty}^{y}p_Z(r)dr\newline
       &= \frac{1}{M}\int_{-\infty}^{y}p_Y(s)ds \quad\text{if $y < high$} \newline
       &= \frac{1}{M}F_Y(y)
\end{align}
$$

And so its inverse is

$$
\begin{align}
F_Z^{-1}(u) = \left(\frac{1}{M}F_Y\right)^{-1}(u)
            = F_Y^{-1}(M u)
            = F_{\mu + \sigma X}^{-1}(Mu)
            = \mu + \sigma F_X^{-1}(Mu)
\end{align}
$$

The translation of the above math into code is

```python
# ...
    def sample(self, key, sample_shape=()):
        shape = sample_shape + self.batch_shape
        minval = jnp.finfo(jnp.result_type(float)).tiny
        u = random.uniform(key, shape, minval=minval)
        return self.icdf(u)

    def icdf(self, u):
        m = norm.cdf(self.high, self.loc, self.scale)
        return self.loc + self.scale * ndtri(m * u)

```

With everything in place, the final implementation is as below.

```python
class _RightExtendedReal(constraints.Constraint):
    """
    Any number in the interval (-inf, inf].
    """

    def __call__(self, x):
        return (x == x) & (x != float("-inf"))

    def feasible_like(self, prototype):
        return jnp.zeros_like(prototype)

right_extended_real = _RightExtendedReal()

class RightTruncatedNormal(Distribution):
    """
    A truncated Normal distribution.
    :param numpy.ndarray loc: location parameter of the untruncated normal
    :param numpy.ndarray scale: scale parameter of the untruncated normal
    :param numpy.ndarray high: point at which the truncation happens
    """

    arg_constraints = {
        "loc": constraints.real,
        "scale": constraints.positive,
        "high": right_extended_real,
    }
    reparametrized_params = ["loc", "scale", "high"]

    def __init__(self, loc=0.0, scale=1.0, high=float("inf"), validate_args=True):
        batch_shape = lax.broadcast_shapes(
            jnp.shape(loc),
            jnp.shape(scale),
            jnp.shape(high),
        )
        self.loc, self.scale, self.high = promote_shapes(loc, scale, high)
        super().__init__(batch_shape, validate_args=validate_args)

    def log_prob(self, value):
        log_m = norm.logcdf(self.high, self.loc, self.scale)
        log_p = norm.logpdf(value, self.loc, self.scale)
        return jnp.where(value < self.high, log_p - log_m, -jnp.inf)

    def sample(self, key, sample_shape=()):
        shape = sample_shape + self.batch_shape
        minval = jnp.finfo(jnp.result_type(float)).tiny
        u = random.uniform(key, shape, minval=minval)
        return self.icdf(u)

    def icdf(self, u):
        m = norm.cdf(self.high, self.loc, self.scale)
        return self.loc + self.scale * ndtri(m * u)

    @constraints.dependent_property
    def support(self):
        return constraints.less_than(self.high)
```

Let's try it out!

```python
def truncated_normal_model(num_observations, x=None):
    loc = numpyro.sample("loc", dist.Normal())
    scale = numpyro.sample("scale", dist.LogNormal())
    high = numpyro.sample("high", dist.Normal())
    with numpyro.plate("observations", num_observations):
        numpyro.sample("x", RightTruncatedNormal(loc, scale, high), obs=x)
```

```python
num_observations = 1000
num_prior_samples = 100
prior = Predictive(truncated_normal_model, num_samples=num_prior_samples)
prior_samples = prior(PRIOR_RNG, num_observations)
```

As before, we run mcmc against some synthetic data.
We select any random sample from the prior as the ground truth:

```python
true_idx = 0
true_loc = prior_samples["loc"][true_idx]
true_scale = prior_samples["scale"][true_idx]
true_high = prior_samples["high"][true_idx]
true_x = prior_samples["x"][true_idx]
```

```python
plt.hist(true_x.copy())
plt.axvline(true_high, linestyle=":", color="k")
plt.xlabel("x")
plt.show()
```

Run MCMC and check the estimates:

```python
mcmc = MCMC(NUTS(truncated_normal_model), **MCMC_KWARGS)
mcmc.run(MCMC_RNG, num_observations, true_x)
mcmc.print_summary()
```

    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:02<00:00, 1850.91it/s, 15 steps of size 8.88e-02. acc. prob=0.88]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 7434.51it/s, 5 steps of size 1.56e-01. acc. prob=0.78]
    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 7792.94it/s, 54 steps of size 5.41e-02. acc. prob=0.91]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 7404.07it/s, 9 steps of size 1.77e-01. acc. prob=0.78]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
          high      0.88      0.01      0.88      0.88      0.89    590.13      1.01
           loc     -0.58      0.07     -0.58     -0.70     -0.46    671.04      1.01
         scale      1.40      0.05      1.40      1.32      1.48    678.30      1.01

    Number of divergences: 6310

Compare estimates against the ground truth:

```python
print(f"True high : {true_high:3.2f}")
print(f"True loc  : {true_loc:3.2f}")
print(f"True scale: {true_scale:3.2f}")
```

    True high : 0.88
    True loc  : -0.56
    True scale: 1.45

Note that, even though we can recover good estimates for the true values,
we had a very high number of divergences. These divergences happen because
the data can be outside of the support that we are allowing with our priors.
To fix this, we can change the prior on `high` so that it depends on the observations:

```python
def truncated_normal_model_2(num_observations, x=None):
    loc = numpyro.sample("loc", dist.Normal())
    scale = numpyro.sample("scale", dist.LogNormal())
    if x is None:
        high = numpyro.sample("high", dist.Normal())
    else:
        # high is greater or equal to the max value in x:
        delta = numpyro.sample("delta", dist.HalfNormal())
        high = numpyro.deterministic("high", delta + x.max())

    with numpyro.plate("observations", num_observations):
        numpyro.sample("x", RightTruncatedNormal(loc, scale, high), obs=x)
```

```python
mcmc = MCMC(NUTS(truncated_normal_model_2), **MCMC_KWARGS)
mcmc.run(MCMC_RNG, num_observations, true_x)
mcmc.print_summary(exclude_deterministic=False)
```

    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:03<00:00, 1089.76it/s, 15 steps of size 4.85e-01. acc. prob=0.93]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 8802.95it/s, 7 steps of size 5.19e-01. acc. prob=0.92]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 8975.35it/s, 3 steps of size 5.72e-01. acc. prob=0.89]
    sample: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 8471.94it/s, 15 steps of size 3.76e-01. acc. prob=0.96]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
         delta      0.01      0.01      0.00      0.00      0.01   6104.22      1.00
          high      0.88      0.01      0.88      0.88      0.89   6104.22      1.00
           loc     -0.58      0.08     -0.58     -0.71     -0.46   3319.65      1.00
         scale      1.40      0.06      1.40      1.31      1.49   3377.38      1.00

    Number of divergences: 0

And the divergences are gone.

In practice, we usually want to understand how the data
would look like without the truncation. To do that in NumPyro,
there is no need of writing a separate model, we can simply
rely on the `condition` handler to push the truncation point to infinity:

```python
model_without_truncation = numpyro.handlers.condition(
    truncated_normal_model,
    {"high": float("inf")},
)
estimates = mcmc.get_samples().copy()
estimates.pop("high")  # Drop to make sure these are not used
pred = Predictive(
    model_without_truncation,
    posterior_samples=estimates,
)
pred_samples = pred(PRED_RNG, num_observations=1000)
```

```python
# thin the samples for a faster histogram
samples_thinned = pred_samples["x"].ravel()[::1000]
```

```python
f, axes = plt.subplots(1, 2, figsize=(15, 5))

axes[0].hist(
    samples_thinned.copy(), label="Untruncated posterior", bins=20, density=True
)
axes[0].axvline(true_high, linestyle=":", color="k", label="Truncation point")
axes[0].set_title("Untruncated posterior")
axes[0].legend()

axes[1].hist(
    samples_thinned[samples_thinned < true_high].copy(),
    label="Tail of untruncated posterior",
    bins=20,
    density=True,
)
axes[1].hist(true_x.copy(), label="Observed, truncated data", density=True, alpha=0.5)
axes[1].axvline(true_high, linestyle=":", color="k", label="Truncation point")
axes[1].set_title("Comparison to observed data")
axes[1].legend()
plt.show()
```

#### 5.3 Example: Left-truncated Poisson

As a final example, we now implement a left-truncated Poisson distribution.
Note that a right-truncated Poisson could be reformulated as a particular
case of a categorical distribution, so we focus on the less trivial case.

**Class attributes**

For a truncated Poisson we need two parameters, the `rate` of the original Poisson
distribution and a `low` parameter to indicate the truncation point.
As this is a discrete distribution, we need to clarify whether or not the truncation point is included
in the support. In this tutorial, we'll take the convention that the truncation point `low`
_is_ part of the support.

The `low` parameter has to be given a 'non-negative integer' constraint. As it is a discrete parameter, it will not be possible to do inference for this parameter using [NUTS](https://num.pyro.ai/en/stable/mcmc.html#nuts). This is likely not a problem since the truncation point is often known in advance. However, if we really must infer the `low` parameter, it is possible to do so with [DiscreteHMCGibbs](https://num.pyro.ai/en/stable/mcmc.html#discretehmcgibbs) though one is limited to using priors with enumerate support.

Like in the case of the truncated normal, the support of this distribution will be defined as a property and not as a class attribute because it depends on the specific value of the `low` parameter.
```python
class LeftTruncatedPoisson:
    arg_constraints = {
        "low": constraints.nonnegative_integer,
        "rate": constraints.positive,
    }

    # ...
    @constraints.dependent_property(is_discrete=True)
    def support(self):
        return constraints.integer_greater_than(self.low - 1)
```

The `is_discrete` argument passed in the `dependent_property` decorator is used to tell the inference algorithms which variables are discrete latent variables.

**The** `__init__` **method**

Here we just follow the same pattern as in the previous example.
```python
    # ...
    def __init__(self, rate=1.0, low=0, validate_args=None):
        batch_shape = lax.broadcast_shapes(
            jnp.shape(low), jnp.shape(rate)
        )
        self.low, self.rate = promote_shapes(low, rate)
        super().__init__(batch_shape, validate_args=validate_args)
    # ...
```

**The** `log_prob` **method**

The logic is very similar to the truncated normal case. But this time we are truncating on the left, so the correct normalization is the complementary cumulative density:

$$
\begin{align}
M = \sum_{n=L}^{\infty} p_Y(n) = 1 - \sum_{n=0}^{L - 1} p_Y(n) = 1 - F_Y(L - 1)
\end{align}
$$

For the code, we can rely on the `poisson` module that lives inside `jax.scipy.stats`.

```python
    # ...
    def log_prob(self, value):
        m = 1 - poisson.cdf(self.low - 1, self.rate)
        log_p = poisson.logpmf(value, self.rate)
        return jnp.where(value >= self.low, log_p - jnp.log(m), -jnp.inf)
    # ...
```

**The** `sample` **method**

Inverse-transform sampling also works for discrete distributions. The "inverse" cdf of a discrete distribution being defined as:

$$
\begin{align}
F^{-1}(u) = \max\left\{n\in \mathbb{N} \rvert F(n) \lt u\right\}
\end{align}
$$

Or, in plain English, $F^{-1}(u)$ is the highest number for which the cumulative density is less than $u$.
However, there's currently no implementation of $F^{-1}$ for the Poisson distribution in Jax (at least, at the moment of writing this tutorial). We have to rely on our own implementation. Fortunately, we can take advantage of the discrete nature of the distribution and easily implement a "brute-force" version that will work for most cases. The brute force approach consists of simply scanning all non-negative integers in order, one by one, until the value of the cumulative density exceeds the argument $u$. The implicit requirement is that we need a way to evaluate the cumulative density for the truncated distribution, but we can calculate that:

$$
\begin{align}
F_Z(z) &= \sum_{n=0}^z p_z(n)\newline
       &= \frac{1}{M}\sum_{n=L}^z p_Y(n)\quad \text{assuming $z >= L$}\newline
       &= \frac{1}{M}\left(\sum_{n=0}^z p_Y(n) - \sum_{n=0}^{L-1}p_Y(n)\right)\newline
       &= \frac{1}{M}\left(F_Y(z) - F_Y (L-1)\right)
\end{align}
$$

And, of course, the value of $F_Z(z)$ is equal to zero if $z < L$.
(As in the previous example, we are using $Y$ to denote the original, un-truncated variable, and we are using $Z$ to denote the truncated variable)

```python
    # ...
    def sample(self, key, sample_shape=()):
        shape = sample_shape + self.batch_shape
        minval = jnp.finfo(jnp.result_type(float)).tiny
        u = random.uniform(key, shape, minval=minval)
        return self.icdf(u)

    def icdf(self, u):
        def cond_fn(val):
            n, cdf = val
            return jnp.any(cdf < u)

        def body_fn(val):
            n, cdf = val
            n_new = jnp.where(cdf < u, n + 1, n)
            return n_new, self.cdf(n_new)

        low = self.low * jnp.ones_like(u)
        cdf = self.cdf(low)
        n, _ = lax.while_loop(cond_fn, body_fn, (low, cdf))
        return n.astype(jnp.result_type(int))

    def cdf(self, value):
        m = 1 - poisson.cdf(self.low - 1, self.rate)
        f = poisson.cdf(value, self.rate) - poisson.cdf(self.low - 1, self.rate)
        return jnp.where(k >= self.low, f / m, 0)
```

A few comments with respect to the above implementation:
* Even with double precision, if `rate` is much less than `low`, the above code will not work. Due to numerical limitations, one obtains that `poisson.cdf(low - 1, rate)` is equal (or very close) to `1.0`. This makes it impossible to re-weight the distribution accurately because the normalization constant would be `0.0`.
* The brute-force `icdf` is of course very slow, particularly when `rate` is high. If you need faster sampling, one option would be to rely on a faster search algorithm. For example:

```python
def icdf_faster(self, u):
    num_bins = 200 # Choose a reasonably large value
    bins = jnp.arange(num_bins)
    cdf = self.cdf(bins)
    indices = jnp.searchsorted(cdf, u)
    return bins[indices]
```

The obvious limitation here is that the number of bins has to be fixed a priori (jax does not allow for dynamically sized arrays). Another option would be to rely on an _approximate_ implementation, as proposed in [this article](https://people.maths.ox.ac.uk/gilesm/codes/poissinv/paper.pdf).

* Yet another alternative for the `icdf` is to rely on `scipy`'s implementation and make use of Jax's `host_callback` module. This feature allows you to use Python functions without having to code them in `Jax`. This means that we can simply make use of `scipy`'s implementation of the Poisson ICDF! From the last equation, we can write the _truncated_ icdf as:

$$
\begin{align}
F_Z^{-1}(u) = F_Y^{-1}(Mu + F_Y(L-1))
\end{align}
$$

And in python:

```python
    def scipy_truncated_poisson_icdf(args): # Note: all arguments are passed inside a tuple
        rate, low, u = args
        rate = np.asarray(rate)
        low = np.asarray(low)
        u = np.asarray(u)
        density = sp_poisson(rate)
        low_cdf = density.cdf(low - 1)
        normalizer = 1.0 - low_cdf
        x = normalizer * u + low_cdf
        return density.ppf(x)
```

In principle, it wouldn't be possible to use the above function in our NumPyro distribution because it is not  coded in Jax. The `jax.experimental.host_callback.call` function solves precisely that problem. The code below shows you how to use it, but keep in mind that this is currently an experimental feature so you should expect changes to the module. See the `host_callback` [docs](https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html) for more details.

```python
    # ...
    def icdf_scipy(self, u):
        result_shape = jax.ShapeDtypeStruct(
            u.shape,
            jnp.result_type(float) # int type not currently supported
        )
        result = jax.experimental.host_callback.call(
            scipy_truncated_poisson_icdf,
            (self.rate, self.low, u),
            result_shape=result_shape
        )
        return result.astype(jnp.result_type(int))
    # ...
```

Putting it all together, the implementation is as below:

```python
def scipy_truncated_poisson_icdf(args):  # Note: all arguments are passed inside a tuple
    rate, low, u = args
    rate = np.asarray(rate)
    low = np.asarray(low)
    u = np.asarray(u)
    density = sp_poisson(rate)
    low_cdf = density.cdf(low - 1)
    normalizer = 1.0 - low_cdf
    x = normalizer * u + low_cdf
    return density.ppf(x)

class LeftTruncatedPoisson(Distribution):
    """
    A truncated Poisson distribution.
    :param numpy.ndarray low: lower bound at which truncation happens
    :param numpy.ndarray rate: rate of the Poisson distribution.
    """

    arg_constraints = {
        "low": constraints.nonnegative_integer,
        "rate": constraints.positive,
    }

    def __init__(self, rate=1.0, low=0, validate_args=None):
        batch_shape = lax.broadcast_shapes(jnp.shape(low), jnp.shape(rate))
        self.low, self.rate = promote_shapes(low, rate)
        super().__init__(batch_shape, validate_args=validate_args)

    def log_prob(self, value):
        m = 1 - poisson.cdf(self.low - 1, self.rate)
        log_p = poisson.logpmf(value, self.rate)
        return jnp.where(value >= self.low, log_p - jnp.log(m), -jnp.inf)

    def sample(self, key, sample_shape=()):
        shape = sample_shape + self.batch_shape
        float_type = jnp.result_type(float)
        minval = jnp.finfo(float_type).tiny
        u = random.uniform(key, shape, minval=minval)
        # return self.icdf(u)        # Brute force
        # return self.icdf_faster(u) # For faster sampling.
        return self.icdf_scipy(u)  # Using `host_callback`

    def icdf(self, u):
        def cond_fn(val):
            n, cdf = val
            return jnp.any(cdf < u)

        def body_fn(val):
            n, cdf = val
            n_new = jnp.where(cdf < u, n + 1, n)
            return n_new, self.cdf(n_new)

        low = self.low * jnp.ones_like(u)
        cdf = self.cdf(low)
        n, _ = lax.while_loop(cond_fn, body_fn, (low, cdf))
        return n.astype(jnp.result_type(int))

    def icdf_faster(self, u):
        num_bins = 200  # Choose a reasonably large value
        bins = jnp.arange(num_bins)
        cdf = self.cdf(bins)
        indices = jnp.searchsorted(cdf, u)
        return bins[indices]

    def icdf_scipy(self, u):
        result_shape = jax.ShapeDtypeStruct(u.shape, jnp.result_type(float))
        result = jax.experimental.host_callback.call(
            scipy_truncated_poisson_icdf,
            (self.rate, self.low, u),
            result_shape=result_shape,
        )
        return result.astype(jnp.result_type(int))

    def cdf(self, value):
        m = 1 - poisson.cdf(self.low - 1, self.rate)
        f = poisson.cdf(value, self.rate) - poisson.cdf(self.low - 1, self.rate)
        return jnp.where(value >= self.low, f / m, 0)

    @constraints.dependent_property(is_discrete=True)
    def support(self):
        return constraints.integer_greater_than(self.low - 1)
```

Let's try it out!

```python
def discrete_distplot(samples, ax=None, **kwargs):
    """
    Utility function for plotting the samples as a barplot.
    """
    x, y = np.unique(samples, return_counts=True)
    y = y / sum(y)
    if ax is None:
        ax = plt.gca()

    ax.bar(x, y, **kwargs)
    return ax
```

```python
def truncated_poisson_model(num_observations, x=None):
    low = numpyro.sample("low", dist.Categorical(0.2 * jnp.ones((5,))))
    rate = numpyro.sample("rate", dist.LogNormal(1, 1))
    with numpyro.plate("observations", num_observations):
        numpyro.sample("x", LeftTruncatedPoisson(rate, low), obs=x)
```

**Prior samples**

```python
# -- prior samples
num_observations = 1000
num_prior_samples = 100
prior = Predictive(truncated_poisson_model, num_samples=num_prior_samples)
prior_samples = prior(PRIOR_RNG, num_observations)
```

**Inference**

As in the case for the truncated normal, here it is better to replace
the prior on the `low` parameter so that it is consistent with the observed data.
We'd like to have a categorical prior on `low` (so that we can use [DiscreteHMCGibbs](https://num.pyro.ai/en/stable/mcmc.html#discretehmcgibbs))
whose highest category is equal to the minimum value of `x` (so that prior and data are consistent).
However, we have to be careful in the way we write such model because Jax does not allow for dynamically sized arrays. A simple way of coding this model is to simply specify the number of categories as an argument:

```python
def truncated_poisson_model(num_observations, x=None, k=5):
    zeros = jnp.zeros((k,))
    low = numpyro.sample("low", dist.Categorical(logits=zeros))
    rate = numpyro.sample("rate", dist.LogNormal(1, 1))
    with numpyro.plate("observations", num_observations):
        numpyro.sample("x", LeftTruncatedPoisson(rate, low), obs=x)
```

```python
# Take any prior sample as the true process.
true_idx = 6
true_low = prior_samples["low"][true_idx]
true_rate = prior_samples["rate"][true_idx]
true_x = prior_samples["x"][true_idx]
discrete_distplot(true_x.copy());
```

To do inference, we set `k = x.min() + 1`. Note also the use of [DiscreteHMCGibbs](https://num.pyro.ai/en/stable/mcmc.html#discretehmcgibbs):

```python
mcmc = MCMC(DiscreteHMCGibbs(NUTS(truncated_poisson_model)), **MCMC_KWARGS)
mcmc.run(MCMC_RNG, num_observations, true_x, k=true_x.min() + 1)
mcmc.print_summary()
```

    sample: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:04<00:00, 808.70it/s, 3 steps of size 9.58e-01. acc. prob=0.93]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 5916.30it/s, 3 steps of size 9.14e-01. acc. prob=0.93]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 5082.16it/s, 3 steps of size 9.91e-01. acc. prob=0.92]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 6511.68it/s, 3 steps of size 8.66e-01. acc. prob=0.94]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
           low      4.13      2.43      4.00      0.00      7.00   7433.79      1.00
          rate     18.16      0.14     18.16     17.96     18.40   3074.46      1.00

```python
true_rate
```

    DeviceArray(18.2091848, dtype=float64)

As before, one needs to be extra careful when estimating the truncation point.
If the truncation point is known is best to provide it.

```python
model_with_known_low = numpyro.handlers.condition(
    truncated_poisson_model, {"low": true_low}
)
```

And note we can use [NUTS](https://num.pyro.ai/en/stable/mcmc.html#nuts) directly because there's no need to infer any discrete parameters.

```python
mcmc = MCMC(
    NUTS(model_with_known_low),
    **MCMC_KWARGS,
)
```

```python
mcmc.run(MCMC_RNG, num_observations, true_x)
mcmc.print_summary()
```

    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:03<00:00, 1185.13it/s, 1 steps of size 9.18e-01. acc. prob=0.93]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 5786.32it/s, 3 steps of size 1.00e+00. acc. prob=0.92]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 5919.13it/s, 1 steps of size 8.62e-01. acc. prob=0.94]
    sample: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4000/4000 [00:00<00:00, 7562.36it/s, 3 steps of size 9.01e-01. acc. prob=0.93]

                    mean       std    median      5.0%     95.0%     n_eff     r_hat
          rate     18.17      0.13     18.17     17.95     18.39   3406.81      1.00

    Number of divergences: 0

**Removing the truncation**

```python
model_without_truncation = numpyro.handlers.condition(
    truncated_poisson_model,
    {"low": 0},
)
pred = Predictive(model_without_truncation, posterior_samples=mcmc.get_samples())
pred_samples = pred(PRED_RNG, num_observations)
thinned_samples = pred_samples["x"][::500]
```

```python
discrete_distplot(thinned_samples.copy());
```

### References and related material

1. [Wikipedia page on inverse transform sampling](https://en.wikipedia.org/wiki/Inverse_transform_sampling)
2. [David Mackay's book on information theory](http://www.inference.org.uk/itprnn/book.pdf)
3. [Composite models with underlying folded distributions](https://www.sciencedirect.com/science/article/pii/S0377042720306427)
4. [Application of the generalized folded-normal distribution to the process capability measures](https://link.springer.com/article/10.1007/s00170-003-2043-x)
4. [Pyro SVI tutorial part 3](https://pyro.ai/examples/svi_part_iii.html)
5. [Approximation of the inverse Poisson cumulative distribution function](https://people.maths.ox.ac.uk/gilesm/codes/poissinv/paper.pdf)


```python
%matplotlib inline
import sympy as sympy
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sbn
from scipy import *
```

Numpy allows us to create vectors rather than lists.

```python
#Take list 1, 2, 3 and make it a vector
x_vector = np.array([1,2,3])
```

Vectors have *shape*, or *direction.*

```python
x_vector.shape
```

    (3,)

Additionally, we can create matrices.

```python
matrix = np.array([[1, 2, 3],
                   [4, 5, 6]])
matrix
```

    array([[1, 2, 3],
           [4, 5, 6]])

#### Addition/Subtraction

\begin{equation}
    A+3=\begin{bmatrix}
      a_{11} & a_{12} \\
      a_{21} & a_{22}
    \end{bmatrix}+3
    =\begin{bmatrix}
      a_{11}+3 & a_{12}+3 \\
      a_{21}+3 & a_{22}+3
    \end{bmatrix}
\end{equation}
\begin{equation}
    A-3=\begin{bmatrix}
      a_{11} & a_{12} \\
      a_{21} & a_{22}
    \end{bmatrix}-3
    =\begin{bmatrix}
      a_{11}-3 & a_{12}-3 \\
      a_{21}-3 & a_{22}-3
    \end{bmatrix}
\end{equation}

Match the rows/columns and add the scalar.

And similarly for matrices:

\begin{equation}
A_{2 \times 2} + B_{2 \times 2}= \begin{bmatrix}
  a_{11}+b_{11} & a_{12}+b_{12} \\
  a_{21}+b_{21} & a_{22}+b_{22}
\end{bmatrix}_{2 \times 2}
\end{equation}

Here's an example with `numpy`

```python
#Initial matrix has dimensions 2 x 3, this one should be 3 x 2
matrix2 = np.random.rand(3,2)

matrix*matrix2
```

#### Scalar Multiplication

#### Matrix Multiplicataion

\begin{align}
    A_{3 \times 2} \times C_{2 \times 3}=&
    \begin{bmatrix}
      a_{11} & a_{12} \\
      a_{21} & a_{22} \\
      a_{31} & a_{32}
    \end{bmatrix}_{3 \times 2}
    \times
    \begin{bmatrix}
      c_{11} & c_{12} & c_{13} \\
      c_{21} & c_{22} & c_{23}
    \end{bmatrix}_{2 \times 3} \\
    =&
    \begin{bmatrix}
      a_{11} c_{11}+a_{12} c_{21} & a_{11} c_{12}+a_{12} c_{22} & a_{11} c_{13}+a_{12} c_{23} \\
      a_{21} c_{11}+a_{22} c_{21} & a_{21} c_{12}+a_{22} c_{22} & a_{21} c_{13}+a_{22} c_{23} \\
      a_{31} c_{11}+a_{32} c_{21} & a_{31} c_{12}+a_{32} c_{22} & a_{31} c_{13}+a_{32} c_{23}
    \end{bmatrix}_{3 \times 3}
\end{align}

Requirements:
* Columns of first matrix = Rows of second matrix\
$$c_x=r_y$$

Result:
* A matrix of with the same amount of rows as the first matrix, and the columns of the second.
$$r_x \times c_y$$

Memory Devices
* C for condition, c for columns first
* R for result, r for rows first

#### Terminology

**Basis vectors** the $\color{red}{columns}$ of the matrix, which are the $\color{red}{\mbox{input dimensions}}$ mapping to a certain output.

**Linear Combination**:

**Span** - all possible linear combinations.  So if I have $a_{1}\vec{v_{1}} + a_{1}\vec{v_{2}}$, then the span is the entire 2D coordinate plane, *unless they are linearly dependent*

**Linear Dependence** - one of the vectors can be defined as a linear combination of the others.
* Consider a third vector, $\vec{u}$.  It would be linear dependent if: $$\vec{u} = a\vec{v} + b\vec{v}$$
* Basically, this would be having a vector that doesn't really do anything.  If I said, go 3 miles east, then four miles north, those vectors would be **linearly independent**.  But if I added on 'go five miles northeast,' that would be useless information, denoting **linear dependency**.

A **rank** is the dimensions of the output of a linear transformation.  It is equal to the number of columns in the column space.  When the rank is the same as the number of columns in the matrix, then it is spanning all possible dimensions, and it can be considered **full rank**

A **column space** is the set of all possible outputs of a linear transformation, also known as the *span of the columns in the matrix*

The **Null space**, or the **kernel**, of a vector space is all vector mappings $\vec{v}$ in  $$A\vec{x}=\vec{v}$$

where the mapping is the 0 vector.  Everything in the column space that is *not in the kernel* is known as the image of $A$

#### Linear Transformations

Essentially a function, taking in a vector as a parameter and outputting another vector.

Two important properties:
* Lines remain lines
* Origin remains fixed

Try to keep in mind that with any linear transformation in 2D, the *gridlines stay parallel and evenly spaced.*

Looking at this, where the vector started at $(-1, 2)$, we can see that this is essentially a linear algebra distributive property.

Commonly, this linear transformed $\hat{i} and \hat{j}$ is represented as follows:

And finally, we can represent the linear transformation as:

Common transformation is a 90° counterclockwise
* $\hat{i}$ = \begin{bmatrix}
0\\
1
\end{bmatrix}

* $\hat{j}$ = \begin{bmatrix}
-1\\
0
\end{bmatrix}

##### The big idea

**All matrices can be thought of as linear transformations.**

#### Matrix Composition

Multiple transformations can be thought of as one transformation, a **composition**, the product of two transformation matrices

Note here that **order matters**.  The right-hand matrix (rotation) is the *first* transformation, and the left-hand (shear) is the *second* transformation.  It is similar to how in the equation $f(g(x))$, $g(x)$ is evaluated first.

With that said, matrix multiplication *is* associative.  That is, if we have matrices $A, B \mbox{ and } C$
$$A(BC) = (AB)C$$

because we are just applying three transformations in the same order

#### Determinants

The **scaled area** enclosed by vector components $\hat{i}\mbox{ and }\hat{j}$ is called the **determinant**

* For 3D planes, the **scaled volume** is the determinant, a **parallelepiped**

Determinants of *linearly dependent* vectors will be zero.  Thus, solving the determinant is a way to solve for the linear dependence of two vectors.

When the **orientation** of space has been **inverted** ($\hat{i}\mbox{ now on the right of }\hat{j}$, the **determinant is negative**

#### Systems of Equations

Commonly, the matrix of the coefficients is called $A$, the variable vector is referred to as $\vec{x}$, and the vector of the constants is known as $\vec{v}$.  This gives us the equation:  $$A\vec{x}=\vec{v}$$

In words:  Applying the transformation $A$ to vector $\vec{x}$ makes $\vec{x}$ land on $\vec{v}$

If we want to go from $\vec{x}$ to $\vec{v}$, we have to divide the matrix $A$, or **multiply by the inverse matrix,** $A^{-1}$

#### Identity Matrix

Multiplying the transformation by the inverse transformation, $A\times A^{-1}$, results in no change at all.  This is called the identity matrix, and has form:
\begin{equation*}
\mbox{$A\times A^{-1}$} =
\begin{bmatrix}
1 & 0 \\
0 & 1
\end{bmatrix}
\end{equation*}

If the determinant is zero, than there **is no inverse matrix**, $A^{-1}$.  In 2D, this is when the output is squished on a line.  In 3D, it is squished onto a plane.  These have **ranks** of 1, and 2 respectively.

A **rank** is therefore the dimensions of the output of a linear transformation.
A **column space** is the set of all possible outputs of a linear transformation.

#### Nonsquare Matrices

Consider a $2 \times 3$ matrix, $B$

\begin{equation*}
B =
\begin{bmatrix}
a & b \\
d & e \\
g & h
\end{bmatrix}
\end{equation*}

This would mean we have a *2D input* mapping to a *3d output*.  The number of columns map to the input, and the number of rows to the output.

#### Dot Product

##### Properties

Thinking in terms of projections help explain these.

* Perpendicular:  Dot product is *zero*
* Similar direction: Dot product is *positive*
* Opposite direction: Dot product is *negative*

#### Cross Product

Sign changes, order matters.  Remember it with $\hat{i}\mbox{ and }\hat{j}$.  Their dot product, with $\hat{i}$ on the left of the cross product multiplication, would be positive.

Big idea with sign? - If the first vector is on the left of the second, then the result will be positive.

##### Properties:
* As vectors get closer to being perpendicular, the result increases
* The result is a vector
* Perpendicular to the parallelogram formed, considering the right hand rule

##### Calculation

It is the determinant, but with the first column being the basis vectors.

Remember that this involves striking through the rows and then carrying out the determinant like usual.

#### Duality

Any transformation from a 2D plane to the 1D number line can be equivalently thought of as a linear transformation (a $1\times 2$ matrix) or a vector (a $2 \times1$ cell)

The vector equivalent to this linear transformation is called the **dual vector**

#### Cramer's Rule

The "logic" behind this is that areas get scaled by the same transformation, $A$.  See [this video](https://www.youtube.com/watch?v=jBsC34PxzoM&list=PLZHQObOWTQDPD3MizzM2xVFitgF8hE_ab&index=12) for more information.

#### Change of Basis

We use the basis vectors $\hat{i}\mbox{ and }\hat{j}$, but we can change these and then scale them with vectors.

Translating between coordinate systems uses a matrix translator of sorts, called the **change of basis matrix**.

Multiplying the change of basis matrix written in our coordinates by the vector in the other coordinates gives the same vector in our coordinates.  If we used the inverse of the change of basis matrix instead and multiplied by our coordinates, we could translate our coordinates to the other coordinates.

$$A^{-1}MA$$

The above equation generally suggests a coordinate system shift.  Computing this, where $A$ is the change of basis matrix and $M$ is the transformation will result in the matrix that can translate a vector in the other coordinate system in the **same way** as it translates one in our coordinate system.

#### Eigenvectors

Vectors that remain on their same span when transformed.  They are only stretched or shrunk, but **not** knocked off of their span path.

To turn both sides into matrix-vector multiplication, we multiply $\lambda$ by the identity matrix, $I$, which brings us to this:

Often, it turns out this can be calculated by calculating the roots of a polynomial.  Here, we can see an example of this, showing that a matrix \begin{equation*}
B =
\begin{bmatrix}
0 & -1\\
1 & 0 \\
\end{bmatrix}
\end{equation*}

(the rotation matrix) has no eigenvectors.

On the other hand, we can also have times where the diagonals of the matrix are the eigenvalues themselves.  This happens with diagonal matrices.

Change of basis is used in conjunction with eigenvectors and diagonal matrices to simplify computations, rerouting the coordinate system to one with eigenvectors.

#### Linearity

```python
from IPython.core.display import HTML
def css_styling():
    styles = open("../../jupyter-styles/custom.css", "r").read()
    return HTML(styles)
css_styling()
```

<style>
    @font-face {
        font-family: "Computer Modern";
        src: url('http://9dbb143991406a7c655e-aa5fcb0a5a4ec34cff238a2d56ca4144.r56.cf5.rackcdn.com/cmunss.otf');
    }
    @font-face {
        font-family: "Computer Modern";
        font-weight: bold;
        src: url('http://9dbb143991406a7c655e-aa5fcb0a5a4ec34cff238a2d56ca4144.r56.cf5.rackcdn.com/cmunsx.otf');
    }
    @font-face {
        font-family: "Computer Modern";
        font-style: oblique;
        src: url('http://9dbb143991406a7c655e-aa5fcb0a5a4ec34cff238a2d56ca4144.r56.cf5.rackcdn.com/cmunsi.otf');
    }
    @font-face {
        font-family: "Computer Modern";
        font-weight: bold;
        font-style: oblique;
        src: url('http://9dbb143991406a7c655e-aa5fcb0a5a4ec34cff238a2d56ca4144.r56.cf5.rackcdn.com/cmunso.otf');
    }
    div.cell{
        width:800px;
        margin-left:16% !important;
        margin-right:auto;
    }
    h1 {
        font-family: Helvetica, serif;
    }
    h4{
        margin-top:12px;
        margin-bottom: 3px;
       }
    div.text_cell_render{
        font-family: Computer Modern, "Helvetica Neue", Arial, Helvetica, Geneva, sans-serif;
        line-height: 145%;
        font-size: 130%;
        width:800px;
        margin-left:auto;
        margin-right:auto;
    }
    .CodeMirror{
            font-family: "Source Code Pro", source-code-pro,Consolas, monospace;
    }
    .prompt{
        display: None;
    }
    .text_cell_render h5 {
        font-weight: 300;
        font-size: 22pt;
        color: #4057A1;
        font-style: italic;
        margin-bottom: .5em;
        margin-top: 0.5em;
        display: block;
    }

    .warning{
        color: rgb( 240, 20, 20 )
        }
</style>


# ETS models

The ETS models are a family of time series models with an underlying state space model consisting of a level component, a trend component (T), a seasonal component (S), and an error term (E).

This notebook shows how they can be used with `statsmodels`. For a more thorough treatment we refer to [1], chapter 8 (free online resource), on which the implementation in statsmodels and the examples used in this notebook are based.

`statsmodels` implements all combinations of:
- additive and multiplicative error model
- additive and multiplicative trend, possibly dampened
- additive and multiplicative seasonality

However, not all of these methods are stable. Refer to [1] and references therein for more info about model stability.

[1] Hyndman, Rob J., and Athanasopoulos, George. *Forecasting: principles and practice*, 3rd edition, OTexts, 2021. https://otexts.com/fpp3/expsmooth.html

```python
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

%matplotlib inline
from statsmodels.tsa.exponential_smoothing.ets import ETSModel
```

```python
plt.rcParams["figure.figsize"] = (12, 8)
```

## Simple exponential smoothing

The simplest of the ETS models is also known as *simple exponential smoothing*. In ETS terms, it corresponds to the (A, N, N) model, that is, a model with additive errors, no trend, and no seasonality. The state space formulation of Holt's method is:

\begin{align}
y_{t} &= y_{t-1} + e_t\\
l_{t} &= l_{t-1} + \alpha e_t\\
\end{align}

This state space formulation can be turned into a different formulation, a forecast and a smoothing equation (as can be done with all ETS models):

\begin{align}
\hat{y}_{t|t-1} &= l_{t-1}\\
l_{t} &= \alpha y_{t-1} + (1 - \alpha) l_{t-1}
\end{align}

Here, $\hat{y}_{t|t-1}$ is the forecast/expectation of $y_t$ given the information of the previous step. In the simple exponential smoothing model, the forecast corresponds to the previous level. The second equation (smoothing equation) calculates the next level as weighted average of the previous level and the previous observation.

```python
oildata = [
    111.0091,
    130.8284,
    141.2871,
    154.2278,
    162.7409,
    192.1665,
    240.7997,
    304.2174,
    384.0046,
    429.6622,
    359.3169,
    437.2519,
    468.4008,
    424.4353,
    487.9794,
    509.8284,
    506.3473,
    340.1842,
    240.2589,
    219.0328,
    172.0747,
    252.5901,
    221.0711,
    276.5188,
    271.1480,
    342.6186,
    428.3558,
    442.3946,
    432.7851,
    437.2497,
    437.2092,
    445.3641,
    453.1950,
    454.4096,
    422.3789,
    456.0371,
    440.3866,
    425.1944,
    486.2052,
    500.4291,
    521.2759,
    508.9476,
    488.8889,
    509.8706,
    456.7229,
    473.8166,
    525.9509,
    549.8338,
    542.3405,
]
oil = pd.Series(oildata, index=pd.date_range("1965", "2013", freq="AS"))
oil.plot()
plt.ylabel("Annual oil production in Saudi Arabia (Mt)")
```

The plot above shows annual oil production in Saudi Arabia in million tonnes. The data are taken from the R package `fpp2` (companion package to prior version [1]).
Below you can see how to fit a simple exponential smoothing model using statsmodels's ETS implementation to this data. Additionally, the fit using `forecast` in R is shown as comparison.

```python
model = ETSModel(oil)
fit = model.fit(maxiter=10000)
oil.plot(label="data")
fit.fittedvalues.plot(label="statsmodels fit")
plt.ylabel("Annual oil production in Saudi Arabia (Mt)")

# obtained from R
params_R = [0.99989969, 0.11888177503085334, 0.80000197, 36.46466837, 34.72584983]
yhat = model.smooth(params_R).fittedvalues
yhat.plot(label="R fit", linestyle="--")

plt.legend()
```

By default the initial states are considered to be fitting parameters and are estimated by maximizing log-likelihood. It is possible to only use a heuristic for the initial values:

```python
model_heuristic = ETSModel(oil, initialization_method="heuristic")
fit_heuristic = model_heuristic.fit()
oil.plot(label="data")
fit.fittedvalues.plot(label="estimated")
fit_heuristic.fittedvalues.plot(label="heuristic", linestyle="--")
plt.ylabel("Annual oil production in Saudi Arabia (Mt)")

# obtained from R
params = [0.99989969, 0.11888177503085334, 0.80000197, 36.46466837, 34.72584983]
yhat = model.smooth(params).fittedvalues
yhat.plot(label="with R params", linestyle=":")

plt.legend()
```

The fitted parameters and some other measures are shown using `fit.summary()`. Here we can see that the log-likelihood of the model using fitted initial states is fractionally lower than the one using a heuristic for the initial states.

```python
print(fit.summary())
```

```python
print(fit_heuristic.summary())
```

## Holt-Winters' seasonal method

The exponential smoothing method can be modified to incorporate a trend and a seasonal component. In the additive Holt-Winters' method, the seasonal component is added to the rest. This model corresponds to the ETS(A, A, A) model, and has the following state space formulation:

\begin{align}
y_t &= l_{t-1} + b_{t-1} + s_{t-m} + e_t\\
l_{t} &= l_{t-1} + b_{t-1} + \alpha e_t\\
b_{t} &= b_{t-1} + \beta e_t\\
s_{t} &= s_{t-m} + \gamma e_t
\end{align}

```python
austourists_data = [
    30.05251300,
    19.14849600,
    25.31769200,
    27.59143700,
    32.07645600,
    23.48796100,
    28.47594000,
    35.12375300,
    36.83848500,
    25.00701700,
    30.72223000,
    28.69375900,
    36.64098600,
    23.82460900,
    29.31168300,
    31.77030900,
    35.17787700,
    19.77524400,
    29.60175000,
    34.53884200,
    41.27359900,
    26.65586200,
    28.27985900,
    35.19115300,
    42.20566386,
    24.64917133,
    32.66733514,
    37.25735401,
    45.24246027,
    29.35048127,
    36.34420728,
    41.78208136,
    49.27659843,
    31.27540139,
    37.85062549,
    38.83704413,
    51.23690034,
    31.83855162,
    41.32342126,
    42.79900337,
    55.70835836,
    33.40714492,
    42.31663797,
    45.15712257,
    59.57607996,
    34.83733016,
    44.84168072,
    46.97124960,
    60.01903094,
    38.37117851,
    46.97586413,
    50.73379646,
    61.64687319,
    39.29956937,
    52.67120908,
    54.33231689,
    66.83435838,
    40.87118847,
    51.82853579,
    57.49190993,
    65.25146985,
    43.06120822,
    54.76075713,
    59.83447494,
    73.25702747,
    47.69662373,
    61.09776802,
    66.05576122,
]
index = pd.date_range("1999-03-01", "2015-12-01", freq="3MS")
austourists = pd.Series(austourists_data, index=index)
austourists.plot()
plt.ylabel("Australian Tourists")
```

```python
# fit in statsmodels
model = ETSModel(
    austourists,
    error="add",
    trend="add",
    seasonal="add",
    damped_trend=True,
    seasonal_periods=4,
)
fit = model.fit()

# fit with R params
params_R = [
    0.35445427,
    0.03200749,
    0.39993387,
    0.97999997,
    24.01278357,
    0.97770147,
    1.76951063,
    -0.50735902,
    -6.61171798,
    5.34956637,
]
fit_R = model.smooth(params_R)

austourists.plot(label="data")
plt.ylabel("Australian Tourists")

fit.fittedvalues.plot(label="statsmodels fit")
fit_R.fittedvalues.plot(label="R fit", linestyle="--")
plt.legend()
```

```python
print(fit.summary())
```

## Predictions

The ETS model can also be used for predicting. There are several different methods available:
- `forecast`: makes out of sample predictions
- `predict`: in sample and out of sample predictions
- `simulate`: runs simulations of the statespace model
- `get_prediction`: in sample and out of sample predictions, as well as prediction intervals

We can use them on our previously fitted model to predict from 2014 to 2020.

```python
pred = fit.get_prediction(start="2014", end="2020")
```

```python
df = pred.summary_frame(alpha=0.05)
df
```

In this case the prediction intervals were calculated using an analytical formula. This is not available for all models. For these other models, prediction intervals are calculated by performing multiple simulations (1000 by default) and using the percentiles of the simulation results. This is done internally by the `get_prediction` method.

We can also manually run simulations, e.g. to plot them. Since the data ranges until end of 2015, we have to simulate from the first quarter of 2016 to the first quarter of 2020, which means 17 steps.

```python
simulated = fit.simulate(anchor="end", nsimulations=17, repetitions=100)
```

```python
for i in range(simulated.shape[1]):
    simulated.iloc[:, i].plot(label="_", color="gray", alpha=0.1)
df["mean"].plot(label="mean prediction")
df["pi_lower"].plot(linestyle="--", color="tab:blue", label="95% interval")
df["pi_upper"].plot(linestyle="--", color="tab:blue", label="_")
pred.endog.plot(label="data")
plt.legend()
```

In this case, we chose "end" as simulation anchor, which means that the first simulated value will be the first out of sample value. It is also possible to choose other anchor inside the sample.