https://github.com/google/bayesnf

Bayesian Neural Field models for prediction in large-scale spatiotemporal datasets
https://github.com/google/bayesnf

bayesian-inference data-science machine-learning spatiotemporal-data-analysis statistics

Last synced: 3 days ago
JSON representation

Bayesian Neural Field models for prediction in large-scale spatiotemporal datasets

• Host: GitHub
• URL: https://github.com/google/bayesnf
• Owner: google
• License: apache-2.0
• Created: 2023-12-13T15:30:55.000Z (9 months ago)
• Default Branch: main
• Last Pushed: 2024-09-12T19:15:53.000Z (8 days ago)
• Last Synced: 2024-09-14T03:57:13.926Z (7 days ago)
• Topics: bayesian-inference, data-science, machine-learning, spatiotemporal-data-analysis, statistics
• Language: Python
• Homepage: https://google.github.io/bayesnf/
• Size: 14.3 MB
• Stars: 44
• Watchers: 6
• Forks: 3
• Open Issues: 10
• Metadata Files:
  • Readme: README.md
  • Contributing: CONTRIBUTING.md
  • License: LICENSE
  • Code of conduct: CODE_OF_CONDUCT.md

Awesome Lists containing this project

README

        
# Bayesian Neural Fields for Spatiotemporal Prediction

[![Unittests](https://github.com/google/bayesnf/actions/workflows/pytest_and_autopublish.yml/badge.svg)](https://github.com/google/bayesnf/actions/workflows/pytest_and_autopublish.yml)

[![PyPI version](https://badge.fury.io/py/bayesnf.svg)](https://badge.fury.io/py/bayesnf)

*This is not an officially supported Google product.*

Spatially referenced time series (i.e., spatiotemporal) datasets are

ubiquitous in scientific, engineering, and business-intelligence

applications. This repository contains an implementation of the _Bayesian

Neural Field_ (BayesNF), a spatiotemporal modeling method that

integrates hierarchical Bayesian inference for accurate uncertainty

estimation with deep neural networks for high-capacity function

approximation.

BayesNF infer joint probability distributions over field values at

arbitrary points in time and space, which makes the model suitable for many

data-analysis tasks including spatial interpolation, temporal forecasting,

and variography. Posterior inference is conducted using variationally

learned surrogates trained via mini-batch stochastic gradient descent for

handling large-scale data. The system is build on the

[JAX](https://jax.readthedocs.io/en/latest/) machine learning platform.

The probabilistic model and inference algorithm are described in the

following [paper](https://www.nature.com/articles/s41467-024-51477-5):

_Scalable spatiotemporal prediction with Bayesian neural fields_. Feras Saad,

Jacob Burnim, Colin Carroll, Brian Patton, Urs Köster, Rif A. Saurous,

Matthew Hoffman. Nature Communications **15**, 7942 (2024).

https://doi.org/10.1038/s41467-024-51477-5

```bibtex

@article{

title     = {Scalable Spatiotemporal Prediction with {Bayesian} Neural Fields},

authors   = {Saad, Feras and Burnim, Jacob and Carroll, Colin and Patton, Brian and Köster, Urs  and Saurous, Rif A. and Hoffman, Matthew}

journal   = {Nature Communications},

volume    = {15},

issue     = {7942},

year      = {2024},

doi       = {10.1038/s41467-024-51477-5},

publisher = {Springer Nature},

}

```

## Installation

`bayesnf` can be installed from the Python Package Index

([PyPI](https://pypi.org/project/bayesnf/)) using:

```

$ python -m pip install bayesnf

```

The typical install time is 1 minute. This software is tested on Python 3.10

with a standard Debian GNU/Linux setup. The large-scale experiments in

`scripts/` were run using [TPU v3-8 accelerators](https://cloud.google.com/tpu/docs/supported-tpu-configurations#tpu-v3-config).

To run BayesNF locally on medium to large-scale data, a GPU is

required at minimum.

Installation into a virtual environment is highly recommended, using the

following steps:

```

$ python -m venv pyenv

$ source pyenv/bin/activate

$ python -m pip install -U bayesnf

```

The versions of dependencies will depend on the Python version.

Github Actions tests the software using Python 3.10.

If encountering any version issues, please refer to the following file

for the versions of libraries used in the test suite:

[requirements.Python3.10.14.txt](https://github.com/google/bayesnf/blob/main/requirements.Python3.10.14.txt).

These specific versions can be installed into the virtual environment

using the following command:

```

$ python -m pip install -r requirements.Python3.10.14.txt

```

## Documentation and Tutorials

Please visit 

## Quick start

```python

# Load a dataframe with "long" format spatiotemporal data.

df_train = pd.read_csv('chickenpox.5.train.csv',

  index_col=0, parse_dates=['datetime'])

# Build a BayesianNeuralFieldEstimator

model = BayesianNeuralFieldMAP(

  width=256,

  depth=2,

  freq='W',

  seasonality_periods=['M', 'Y'],

  num_seasonal_harmonics=[2, 10],

  feature_cols=['datetime', 'latitude', 'longitude'],

  target_col='chickenpox',

  observation_model='NORMAL',

  timetype='index',

  standardize=['latitude', 'longitude'],

  interactions=[(0, 1), (0, 2), (1, 2)])

# Fit the model.

model = model.fit(

  df_train,

  seed=jax.random.PRNGKey(0),

  ensemble_size=ensemble_size,

  num_epochs=num_epochs)

# Make predictions of means and quantiles on test data.

df_test = pd.read_csv('chickenpox.5.test.csv',

  index_col=0, parse_dates=['datetime'])

yhat, yhat_quantiles = model.predict(df_test, quantiles=(0.025, 0.5, 0.975))

```