Predecir la calidad del aire a 24 h con DCRNN: redes de grafos espacio-temporales (con código)

La contaminación del aire no es un fenómeno local: es un fenómeno que se propaga. Si una estación detecta un pico de NO₂, la estación que está 2 km a favor del viento lo verá 30-60 minutos después. Cualquier modelo que trate cada estación de forma aislada está tirando a la basura la señal más valiosa que existe: la dinámica de dispersión atmosférica entre estaciones vecinas. Este artículo cuenta cómo aproveché esa señal.

El enfoque ingenuo (y por qué falla)

Lo primero que casi todo el mundo hace es entrenar un LSTM independiente por estación: metes la serie histórica de una estación y predices su futuro. Funciona regular, y tiene dos problemas de fondo:

• Ignora la correlación espacial: dos estaciones a 1 km se comportan casi igual, pero el modelo no lo sabe.
• Produce predicciones incoherentes: la estación A puede predecir "muy contaminado" y su vecina de al lado "excelente" a la misma hora, algo físicamente imposible.

La forma correcta de resolverlo es tratar la red de estaciones como lo que es: un grafo.

La idea: difusión sobre un grafo (DCRNN)

La Diffusion Convolutional Recurrent Neural Network (DCRNN) sustituye las multiplicaciones matriciales de una GRU normal por convoluciones de difusión sobre un grafo. En cristiano: en cada paso temporal, cada estación no solo mira su propio estado, sino que propaga y recibe información de sus vecinas siguiendo la matriz de adyacencia. Es exactamente el mecanismo físico de dispersión, pero aprendido.

El grafo se construye con un kernel gaussiano sobre la distancia euclídea entre estaciones, normalizado por filas para que se comporte como un random walk (las probabilidades de difusión suman 1):

import numpy as np

def build_adjacency(coords, sigma=0.1):
    # coords: (N, 2) lat/lng de las 24 estaciones
    d = np.linalg.norm(coords[:, None] - coords[None, :], axis=-1)
    A = np.exp(-(d ** 2) / (2 * sigma ** 2))   # kernel gaussiano
    np.fill_diagonal(A, 0)
    return A / A.sum(axis=1, keepdims=True)     # normalización random-walk

El corazón del modelo es una celda GRU en la que las tres puertas (reset, update, candidata) usan difusión de grafo en lugar de una capa lineal. Con K=1 salto de difusión y una hidden_dim=32, es una DCRNN-lite deliberadamente pequeña para que la inferencia sea instantánea:

import torch, torch.nn as nn

class DiffusionConv(nn.Module):
    def __init__(self, in_dim, out_dim, adj):
        super().__init__()
        self.register_buffer("adj", torch.tensor(adj, dtype=torch.float32))
        self.lin = nn.Linear(in_dim * 2, out_dim)   # x propia + x difundida

    def forward(self, x):                # x: (batch, N_nodos, in_dim)
        diff = torch.einsum("ij,bjf->bif", self.adj, x)  # difusión 1-hop
        return self.lin(torch.cat([x, diff], dim=-1))

Cada nodo entra con 10 features: NO₂, PM10, O₃, velocidad de viento, seno y coseno de la dirección del viento, temperatura, seno y coseno de la hora del día, y el ICA actual. Codificar la dirección del viento y la hora con seno/coseno es importante: son variables circulares (359° está al lado de 1°) y meterlas como número crudo confunde al modelo.

De 1 hora a 24: autoregresión

El modelo predice H=1 (la siguiente hora). Para llegar a 24 h, realimento la predicción como entrada del siguiente paso, de forma iterativa. Es la técnica clásica de seq2seq autoregresivo; acumula error con el horizonte, pero para 24 h con datos horarios el compromiso es razonable.

Resultados honestos

Entrenado con datos reales del servicio CAMS Reanalysis de Open-Meteo (Copernicus), 2022-2023, 24 estaciones de Madrid con coordenadas reales (8.760 horas):

Una advertencia de honestidad: la exactitud "bruta" es del 97%, pero no me fío de ese número y tú tampoco deberías. Las clases están muy desbalanceadas (93,5% de las horas son de calidad "Buena"), así que un modelo que dijera siempre "Buena" ya acertaría el 93,5%. Por eso la métrica que de verdad importa es el MCC = 0.74, que penaliza ese truco: un MCC ≥ 0.7 es excelente incluso con clases desbalanceadas, y significa que el modelo detecta de verdad los episodios de contaminación moderada, no solo la calma.

Qué mejoraría

El grafo es estático; un grafo dinámico que dependiera de la dirección del viento en cada hora (difusión anisótropa) capturaría mejor la física. Y sustituiría la autoregresión por un decodificador seq2seq entrenado con scheduled sampling para reducir el error acumulado a 24 h.

Pruébalo tú mismo

Es un proyecto personal y open source: no vendo nada, solo comparto cómo está construido. Puedes probar la demo en vivo y trastear con ella. El código está en GitHub. Si te resulta útil o tienes una idea para mejorarlo, encantado de leerla.