Bayesian_Switchpoint_Analysis.ipynb

Kernel: Python 3

Copyright 2019 The TensorFlow Probability Authors.

Licensed under the Apache License, Version 2.0 (the "License");

#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Análisis del punto de conmutación bayesiano

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

En este bloc de notas se vuelve a implementar y se amplía el ejemplo bayesiano de “Análisis de punto de conmutación” de la documentación de pymc3.

Requisitos previos

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (15,8)
%config InlineBackend.figure_format = 'retina'
import numpy as np
import pandas as pd

Conjunto de datos

El conjunto de datos proviene de aquí. Tenga en cuenta que hay otra versión de este ejemplo dando vueltas, pero tiene datos "faltantes", en cuyo caso deberá imputar los valores faltantes. (De lo contrario, su modelo nunca abandonará sus parámetros iniciales porque la función de probabilidad no estará definida).

disaster_data = np.array([ 4, 5, 4, 0, 1, 4, 3, 4, 0, 6, 3, 3, 4, 0, 2, 6,
                           3, 3, 5, 4, 5, 3, 1, 4, 4, 1, 5, 5, 3, 4, 2, 5,
                           2, 2, 3, 4, 2, 1, 3, 2, 2, 1, 1, 1, 1, 3, 0, 0,
                           1, 0, 1, 1, 0, 0, 3, 1, 0, 3, 2, 2, 0, 1, 1, 1,
                           0, 1, 0, 1, 0, 0, 0, 2, 1, 0, 0, 0, 1, 1, 0, 2,
                           3, 3, 1, 1, 2, 1, 1, 1, 1, 2, 4, 2, 0, 0, 1, 4,
                           0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1])
years = np.arange(1851, 1962)
plt.plot(years, disaster_data, 'o', markersize=8);
plt.ylabel('Disaster count')
plt.xlabel('Year')
plt.title('Mining disaster data set')
plt.show()

Modelo probabilístico

El modelo supone un “punto de conmutación” (por ejemplo, un año durante el cual cambiaron las normas de seguridad) y una tasa de catástrofes distribuida por Poisson con tasas constantes (pero potencialmente diferentes) antes y después de ese punto de conmutación.

El recuento real de catástrofes es fijo (observado); cualquier muestra de este modelo deberá especificar tanto el punto de conmutación como la tasa de catástrofes “tempranas” y “tardías”.

Modelo original del ejemplo de documentación de pymc3:

$$\begin{align*} (D_t|s,e,l)&\sim \text{Poisson}(r_t), \\ & \,\quad\text{with}\; r_t = \begin{cases}e & \text{if}\; t < s\\l &\text{if}\; t \ge s\end{cases} \\ s&\sim\text{Discrete Uniform}(t_l,\,t_h) \\ e&\sim\text{Exponential}(r_e)\\ l&\sim\text{Exponential}(r_l) \end{align*}$$

No obstante, la tasa media de catástrofes $r_t$ tiene una discontinuidad en el punto de conmutación $s$, lo que hace que no sea diferenciable. Por lo tanto, no proporciona ninguna señal de gradiente al algoritmo hamiltoniano de Monte Carlo (HMC), pero como la previa $s$ es continua, la alternativa del HMC a un paseo aleatorio es lo suficientemente buena como para encontrar las áreas de masa de alta probabilidad en este ejemplo.

Como segundo modelo, modificamos el modelo original con ayuda de una “conmutación” sigmoide entre e y l para que la transición sea diferenciable, y utilizamos una distribución uniforme continua para el punto de conmutación $s$. (Se podría argumentar que este modelo es más fiel a la realidad, ya que una “conmutación” en la tasa media probablemente se extendería a lo largo de varios años). El nuevo modelo es el siguiente:

$$\begin{align*} (D_t|s,e,l)&\sim\text{Poisson}(r_t), \\ & \,\quad \text{with}\; r_t = e + \frac{1}{1+\exp(s-t)}(l-e) \\ s&\sim\text{Uniform}(t_l,\,t_h) \\ e&\sim\text{Exponential}(r_e)\\ l&\sim\text{Exponential}(r_l) \end{align*}$$

Ante la falta de más información, asumimos $r_e = r_l = 1$ como parámetros para las previas. Ejecutaremos ambos modelos y compararemos sus resultados de inferencia.

def disaster_count_model(disaster_rate_fn):
  disaster_count = tfd.JointDistributionNamed(dict(
    e=tfd.Exponential(rate=1.),
    l=tfd.Exponential(rate=1.),
    s=tfd.Uniform(0., high=len(years)),
    d_t=lambda s, l, e: tfd.Independent(
        tfd.Poisson(rate=disaster_rate_fn(np.arange(len(years)), s, l, e)),
        reinterpreted_batch_ndims=1)
  ))
  return disaster_count

def disaster_rate_switch(ys, s, l, e):
  return tf.where(ys < s, e, l)

def disaster_rate_sigmoid(ys, s, l, e):
  return e + tf.sigmoid(ys - s) * (l - e)

model_switch = disaster_count_model(disaster_rate_switch)
model_sigmoid = disaster_count_model(disaster_rate_sigmoid)

El código anterior define el modelo mediante distribuciones JointDistributionSequential. Las funciones disaster_rate se llaman con una matriz de [0, ..., len(years)-1] para producir un vector de variables aleatorias de len(years); los años anteriores al switchpoint son early_disaster_rate, los posteriores a late_disaster_rate (módulo de la transición sigmoide).

A continuación, se muestra una comprobación de que la función de probabilidad logarítmica objetivo es correcta:

def target_log_prob_fn(model, s, e, l):
  return model.log_prob(s=s, e=e, l=l, d_t=disaster_data)

models = [model_switch, model_sigmoid]
print([target_log_prob_fn(m, 40., 3., .9).numpy() for m in models])  # Somewhat likely result
print([target_log_prob_fn(m, 60., 1., 5.).numpy() for m in models])  # Rather unlikely result
print([target_log_prob_fn(m, -10., 1., 1.).numpy() for m in models]) # Impossible result

[-176.94559, -176.28717]
[-371.3125, -366.8816]
[-inf, -inf]

HMC para hacer inferencia bayesiana

Definimos la cantidad de resultados y pasos de grabación necesarios; el código se basa principalmente en la documentación de tfp.mcmc.HamiltonianMonteCarlo. Usa un tamaño de paso adaptable (de lo contrario, el resultado es muy sensible al valor del tamaño de paso elegido). Usamos valores de uno como estado inicial de la cadena.

Sin embargo, esta no es la historia completa. Si regresa a la definición del modelo anterior, notará que algunas de las distribuciones de probabilidad no están bien definidas en toda la recta numérica real. Por lo tanto, restringimos el espacio que el HMC examinará al envolver el núcleo de HMC con un TransformedTransitionKernel que especifica los biyectores directos para transformar los números reales en el dominio en el que se define la distribución de probabilidad (consulte los comentarios en el código a continuación).

num_results = 10000
num_burnin_steps = 3000

@tf.function(autograph=False, jit_compile=True)
def make_chain(target_log_prob_fn):
   kernel = tfp.mcmc.TransformedTransitionKernel(
       inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          step_size=0.05,
          num_leapfrog_steps=3),
       bijector=[
          # The switchpoint is constrained between zero and len(years).
          # Hence we supply a bijector that maps the real numbers (in a
          # differentiable way) to the interval (0;len(yers))
          tfb.Sigmoid(low=0., high=tf.cast(len(years), dtype=tf.float32)),
          # Early and late disaster rate: The exponential distribution is
          # defined on the positive real numbers
          tfb.Softplus(),
          tfb.Softplus(),
      ])
   kernel = tfp.mcmc.SimpleStepSizeAdaptation(
        inner_kernel=kernel,
        num_adaptation_steps=int(0.8*num_burnin_steps))

   states = tfp.mcmc.sample_chain(
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      current_state=[
          # The three latent variables
          tf.ones([], name='init_switchpoint'),
          tf.ones([], name='init_early_disaster_rate'),
          tf.ones([], name='init_late_disaster_rate'),
      ],
      trace_fn=None,
      kernel=kernel)
   return states

switch_samples = [s.numpy() for s in make_chain(
    lambda *args: target_log_prob_fn(model_switch, *args))]
sigmoid_samples = [s.numpy() for s in make_chain(
    lambda *args: target_log_prob_fn(model_sigmoid, *args))]

switchpoint, early_disaster_rate, late_disaster_rate = zip(
    switch_samples, sigmoid_samples)

Ejecute ambos modelos en paralelo:

Visualización del resultado

Visualizamos el resultado como histogramas de muestras de la distribución posterior para la tasa de catástrofes tempranas y tardías, así como el punto de conmutación. Los histogramas se superponen con una línea continua que representa la mediana de la muestra, así como los límites del intervalo de credibilidad del percentil 95 como líneas discontinuas.

def _desc(v):
  return '(median: {}; 95%ile CI: $[{}, {}]$)'.format(
      *np.round(np.percentile(v, [50, 2.5, 97.5]), 2))

for t, v in [
    ('Early disaster rate ($e$) posterior samples', early_disaster_rate),
    ('Late disaster rate ($l$) posterior samples', late_disaster_rate),
    ('Switch point ($s$) posterior samples', years[0] + switchpoint),
]:
  fig, ax = plt.subplots(nrows=1, ncols=2, sharex=True)
  for (m, i) in (('Switch', 0), ('Sigmoid', 1)):
    a = ax[i]
    a.hist(v[i], bins=50)
    a.axvline(x=np.percentile(v[i], 50), color='k')
    a.axvline(x=np.percentile(v[i], 2.5), color='k', ls='dashed', alpha=.5)
    a.axvline(x=np.percentile(v[i], 97.5), color='k', ls='dashed', alpha=.5)
    a.set_title(m + ' model ' + _desc(v[i]))
  fig.suptitle(t)
  plt.show()