Covariate Effects

TLDR: We will incorporate covariates to try and make the data more IID. We will look at covariates such as the spatial coordinates, the temporal coordinates and any related variables of interest. Some variables include temperature and precipitation. We can also explore local and global covariates like “warming”. We will explore simple statistical models like linear functions and then more complex functions like basis functions or perhaps nonlinear functions. We will use Bayesian Hierarchical Models (BHMs) to separate the parameter uncertainty of the process with the hyper-parameters of the process Parameterization. To select the best parameterization, we can use some traditional model selection techniques like AIC or BIC.

Conditional Models¶

A second attempt would be to construct a conditionally parameterized distribution. In this case, we introduce some dependencies on the variable itself.

This changes this function from a generative parametric model to a likelihood model, i.e., a model that generates the parameters wrt some external parameter like time.

In general, most of the distributions that were mentioned above have a mean and standard deviation parameter. One could easily try to parameterize those parameters with functions that are dependent upon some space and time parameters. For example, the mean function could be depending upon space, time, and some control vector and likewise the standard deviation.

Data¶

For fitting this data, we assume that all elements within the time series are a sequential set of observations

where $u\in\mathbb{R}$.

Conditional Parametric Model¶

We could also include some conditioning variables, $\boldsymbol{u}$. The parameters, $\boldsymbol{\theta}$, will be functions of the conditioning variable and some hyperparameters, $\boldsymbol{\alpha}$.

Some ideas for conditioning variables, $\boldsymbol{u}$.

Standard Variables

• $\boldsymbol{u}$ - standard variables, e.g., temperature, precipitation, ENSO, etc. $\boldsymbol{\theta_\mu} = \exp \left( \mathbf{W}^\top\mathbf{u} + \mathbf{b}\right)$
• $t$ - temporal coordinate encoding, $\boldsymbol{\theta_\mu}=a\sin(2\pi t/\omega L)$
• $\mathbf{x}$ - spatial coordinates encoding

Temporal Encoding¶

Here, we assume that the parameters of the data are processes which are parameterized functions of time.

where $t$ is the time and $\boldsymbol{\alpha}$ are the parameters of the distribution. For example, the Gaussian distribution could have a parameterized mean and scale. Similarly, the GEVD could have a parameterized mean, scale and shape. Again, this function can be as simple or as complicated as necessary, e.g., a linear model, a basis function, a non-linear function or a stochastic process. So, we have a conditional data likelihood term.

The posterior is given by

This is known as a Bayesian Hierarchical Model (BHM) because we have hierarchical processes and priors which condition the data likelihood.

Spatial Encoding¶

where $t$ is the time and $\boldsymbol{\alpha}$ are the parameters of the distribution. For example, the Gaussian distribution could have a parameterized mean and scale. Similarly, the GEVD could have a parameterized mean, scale and shape. Again, this function can be as simple or as complicated as necessary, e.g., a linear model, a basis function, a non-linear function or a stochastic process. So, we have a conditional data likelihood term.

The posterior is given by

This is also a BHM because we have hierarchical processes and priors which condition the data likelihood.

Inference¶

We can use the standard posterior distribution to best fit the parameters of the distribution given the data.

Here, we have many examples of different inference methods to use. For example, we can use simple methods like Maximum Likelihood Estimation (MLE) or Maximum-A-Posterior (MAP). We can use approximate inference methods like Laplace or Variational Inference (VI). We could even use sampling methods like Markov Chain Monte-Carlo (MCMC).

Literature Review¶

PseudoCode¶

Model¶

Recall, we are using a Gaussian likelihood that describes our observations.

However, we recognize that a single parameter, $\boldsymbol{\mu}$, is too simple describe the complete complexity of the system. Let’s assume we have our covariate of interest, $\boldsymbol{x}$, which we think is related to our observations, $\boldsymbol{y}$. We can write a linear function which relates a parameter of our data likelihood to this covariate.

Where $\boldsymbol{\theta}=\{\mathbf{W},b\}$ are the learnable parameters. This mean parameter

Note: There are many things we can improve. For example, we can increase the model complexity, e.g., a basis function or a nonlinear model. We could also add more covariates.

def mean_function(
		x: Array["Dx"], weight: Array["Dy Dx"], bias: Array["Dy"]
	) -> Array["Dy"]:
    return weight @ x + bias

def model(x: Array["Dx"], y: Array["Dy"]=None):
    # prior parameters
    weight: Array["Dy Dx"] = sample("weight", dist.Normal(0.0, 1.0))
    bias: Array["Dy"] = sample("bias", dist.Normal(0.0, 1.0))
    # process parameterizations
    mu: Array["Dy"] = mean_function(x, weight, bias) 
    sigma: Array["Dy"] = sample("sigma", dist.Normal(0.0, 10.0))
    # data likelihood
    return sample("y", dist.Normal(mu, sigma), obs=y)

Prior Predictive Posterior¶

# initialize predictive utility
N = 1_000
predictive = Predictive(model, num_samples=N)
# make predictions
y_pred: Array["N"] = predictive(rng_key, X)["y"]

Inference¶

kernel = NUTS(model)
num_samples = 2_000
mcmc = MCMC(kernel, num_warmup=1_000, num_samples=num_samples)
mcmc.run(key, x=x, y=y)
mcmc.print_summary()

Posterior Samples¶

# extract posterior samples
posterior_samples: Dict = mcmc.get_samples()
# extract individual parameter samples
alpha: Array["Np Dalpha"] = posterior_samples["alpha"]
theta: Array["Np Dtheta"] = posterior_samples["theta"]

Posterior Predictive¶

# initialize predictive utility
predictive = Predictive(model, posterior_samples=posterior_samples, return_sites=["y"])
# make predictions
X_pred: Array["Nx"] = ...
y_pred: Array["Np"] = predictive(rng_key, X_pred)["y"]