Output Area Foundation Model (OAFM) prototype

This repository presents a prototype foundation model capturing a range of socio-demographic factors in England and Wales at the 2021 Output Area and Lower layer Super Output Area levels, as a first step towards a population dynamics foundation model for the UK. A masked heterogeneous graph autoencoder is used to create embeddings based on five open datasets available from the Geographic Data Service, as described in more detail below.

The architecture lives in code/ and its heterogeneous graph in data/. Trained-model artifacts (info, logs, checkpoint) are written to model/ and the embeddings are exported to embeddings/. The eval/ folder contains three small evaluation experiments. See also the respective README.md for further details.

Background and scope

Foundation models can generate general-purpose, rich and concise representations of complex information or systems. In geography, creating general, digital fingerprints integrating and compressing a broad spectrum of data on population, health, housing, mobility, the environment, and the economy of a place can drive useful insights while limiting the amount of preprocessing and complexity of models used in each specific downstream task. For instance, Agarwal et al. (2025) developed Google's Population Dynamics Foundation Model (PDFM), which incorporates a wide range of environmental and population data, as well as Google's own search trends across many countries, and ESRI developed a Geodemographic Foundation Model.

The aim of this prototype is deliberately narrow: to establish a starting point for the development of an open-source population dynamics foundation model for the UK. As such, the model mirrors many key aspects of PDFM, while also integrating aspects of my previous work on developing deep embedding geodemographics for the UK (De Sabbata and Liu, 2019, De Sabbata and Liu, 2023, De Sabbata et al. 2026). For instance, the prototype architecture uses single-layer heads reconstructing the input data from separate sections of the embeddings (as in PDFM), but the graph is heterogeneous, including separate nodes for Output Areas and Lower layer Super Output Areas, with edges indicating spatial proximity and hierarchy, as well as commuting patterns.

The development of this prototype posed a wide range of questions regarding the selection of the data, their preprocessing, the model architecture and hyperparameters, and the final evaluation of the outputs. Many choices were driven by data availability, which affected both the scale and the type of data used, and therefore the architecture necessary to achieve satisfactory results. Other choices were guided or assessed using intrinsic (i.e., validation and test losses) and extrinsic (i.e., downstream tasks) evaluation. Nevertheless, the limited time available to complete this prototype didn't allow for an extensive exploration of the opportunities and challenges, nor a thorough assessment of the model. As such, this model is provided "as is", and should be regarded as a demonstrator and a starting point for further development, rather than as a final product.

Data

The model spans two spatial scales across England and Wales: the 188,880 2021 Output Areas (OAs) and the 35,672 2021 Lower layer Super Output Areas (LSOAs) (see also the Office for National Statistics' page on the Census 2021 geographies for further details). The two scales become the two node types of the heterogeneous graph, akin to the PDFM's use of two nested geographies (postcodes and counties). The embeddings encapsulate data from five different datasets as outlined below (more details are provided in data/README.md). Four datasets are available natively at one scale and harmonised to the other: census and broadband are available at OA and aggregated up to LSOA; AHAH and dwellings are available at LSOA and copied down to every constituent OA. A distance-to-nearest-k field is computed for POIs at both scales. The sources are listed here in the order they occupy in the 267-column feature matrix, which is also the order of the embedding partitions.

Source	Native geography	Input dims	At OA	At LSOA
Access to healthy assets and hazards (AHAH)	LSOA	15	copied down from LSOA	native
Broadband speed (Ofcom)	OA (2011, best fit to OA21)	3	native	aggregated up (median of the constituent OAs, re-standardised)
Unified UK Census (2021)	OA	182	native	aggregated up from raw OA counts
Dwelling ages and prices (VOA/ONS)	LSOA (2011, best fit to LSOA21)	6	copied down from LSOA	native
Points of interest (Overture)	Overture POIs, Basic Level Category	61	mean distance from the OA centroid to the k=10 nearest POIs of each category	same, from the LSOA centroid

The heterogeneous graph is encoded as a PyTorch Geometric HeteroData object: node types oa (188,880) and lsoa (35,672), with three kinds of edge (six edge types, one per relation per scale or direction):

1. Proximity (near) within each scale: the 16 nearest neighbours (KNN16 on area centroids), 3,022,080 edges at OA and 570,752 at LSOA. KNN guarantees every area has neighbours (islands, sparsely bounded areas) and gives a constant in-degree of 16.
2. Containment (to_iswithin / to_contains) across scales: each OA linked to its LSOA and back (188,880 edges each way), providing multi-resolution coupling.
3. Commuting (commute) within each scale: origin-destination flows from the 2011 and 2021 censuses, kept where a residence-workplace pair carries more than three commuters in either census, with self-flows dropped. Both the 2011 and 2021 census data are used to mitigate for the 2021 data having been collected the covid-19 pandemic. Each pair becomes an unweighted edge in both directions, so a node is linked to the areas its residents work in and to the areas its workers live in, and the OA pairs are mapped up to their LSOAs: 2,370,662 OA edges (median in-degree 7, and 3.3% of OAs have no in-edge) and 781,688 LSOA edges (median 14, one LSOA with none).

The graph also carries an 80 / 10 / 10 train / validation / test split, drawn at random over the LSOAs and inherited by their constituent OAs, so an OA and its LSOA never fall in different splits. The split only determines which nodes act as seeds during training, as all the edges and all the node features are retained in the graph. The same split is reused by the evaluation experiments below.

Architecture

The model is defined as a masked heterogeneous graph autoencoder, where two different node types are used to encode two different geographic scales and masking is used to foster robust learning (see also code/README.md). A single round of message passing over the graph is based on GraphSAGE convolutions (Hamilton et al. 2017), one per edge type. The same computation runs on each node type, but the approach to source partitioning and embedding split are shared across node types. The model is self supervised: the input features act as their own reconstruction targets. The exported artifact is the model trained on the 80% training split, rather than a model fully re-trained on the whole dataset.

raw features per node  (267, concatenated across sources, clipped + standardised)
concatenated with the per-feature mask indicator (267)
        |
    [ prep MLP (per node type): 534 -> 256 -> 192, GeLU ] 
        |
    [ HeteroConv: SAGEConv (per edge type), mean neighbour reduction, 128 dims ] 
        |
    [ post MLP (per node type): 128 -> 192 -> 256, GeLU ] -> embedding (256 dims)
        |
    embedding partitioned by source, each source slice split into self | near | commute
    then, per source:
        self head    -> reconstruct the area's own <source> features
        near head    -> reconstruct its neighbours' mean <source> features
        commute head -> reconstruct its neighbours' mean <source> features

The 256-dimensional embedding is allocated per source roughly in proportion to input dimensionality, and each source slice is then split between the three reconstruction targets, self taking the largest share (roughly half) and the remainder divided equally between the two same-scale relations. Allocations are adjusted to ensure each slice carries sufficient information to achieve similar validation loss.

Partition	Source	Input dims	Embedding dims	self	near	commute
0	AHAH	15	20	8	6	6
1	Broadband	3	6	2	2	2
2	Census	182	180	96	42	42
3	Dwellings	6	10	4	3	3
4	POIs	61	40	20	10	10
	total	267	256	130	63	63

The resulting model has 1,112,001 trainable parameters. The released model was trained for the full budget of 3,000 epochs, reaching the best validation loss at epoch 1,896.

Masking

Training is designed as a masked denoising objective for the seed node attributes reconstruction, applied every step. A random fraction of each seed's own features is replaced by a learned per-feature mask token and flagged with a per-feature mask-indicator channel. The seed must then reconstruct its clean features from the surviving ones and its neighbours, forcing the embedding to encode neighbourhood-predictive structure rather than a near-copy of its own inputs.

Dual-target reconstruction

Only part of each source's slice reconstructs the area's own values, as a standard source-partitioned autoencoder would. The rest reconstructs, separately for each relation that links a node to others of its own type (near and commute), the average of that relation's neighbours' values for that source. The aim in designing an architecture where a slice of the embeddings is used to predict the neighbourhood context is to train embeddings that can be predictive of the neighbourhood structure (near and commute) rather than only a copy of the area itself, complementing the masking objective.

Evaluation

The generated embeddings are used as input to simple downstream models on three tasks, one per notebook in eval/:

Notebook	Downstream task	Scale	Models
eval_regression_policeknn.ipynb	crime exposure: mean distance from each area to the 10 nearest recorded crimes, one target per crime type (police.uk, 2025)	OA	ridge regression (penalty chosen by RidgeCV)
eval_classif_gci.ipynb	Geodemographic Classification of Interactivity group	OA	logistic regression, gradient boosted trees
eval_classif_lnwc.ipynb	London Night-time Workers Classification group	LSOA	logistic regression, gradient boosted trees

Both geodemographic classifications cover Greater London only, that is 26,369 OAs and 4,994 LSOAs, 14% of the areas at each scale, so those two tasks are evaluated on London alone. Only the crime exposure task covers the whole of England and Wales.

Each experiment fits every model on the same feature sets, so both the embedding as a whole and its internal structure are evaluated:

• a dummy predictor (the mean, or the most frequent class), as the floor;
• the 267 raw input variables;
• the full 256-dimension embedding;
• subsets of the embedding by reconstruction target: self plus near, self plus commute, and each of self, near and commute alone, which test whether the neighbourhood-context dimensions carry signal of their own.

All models are trained on the train split and scored on the test split, reusing the model's own split: classification by accuracy, balanced accuracy, macro F1 and Cohen's kappa, regression by R², MAE, MSE and RMSE.

Results

The maps below provide an initial overview of the learned embeddings. The 256 dimensions are standardised and reduced to the first nine principal components, which are shown in groups of three as RGB channels (England and Wales on the top row, Greater London on the bottom row). While the maps are rather coarse projections, they exhibit coherent spatial structure at both scales: the urban-rural divide, broad regional contrasts, and finer neighbourhood patterns within London suggest that the embeddings capture meaningful geographic variation.

The three downstream experiments (see notebooks in eval/ for further details) compare the embeddings against the 267 normalised input variables. On the crime-exposure regression, which covers the whole of England and Wales, the full embedding clearly and consistently outperforms the raw inputs, improving $R^2$ for all fourteen crime types. However, crime exposure is a spatially-smooth target, and the embeddings carry neighbourhood context that the per-area raw inputs lack, so a spatial model applied to the raw inputs might improve the quality of the estimation. Indeed, most of the advantage seems to come from the neighbourhood-context slices: the self-only subset barely exceeds the raw inputs, while adding the near and commute slices substantially lifts $R^2$ across all targets. On the two Greater London geodemographic classifications, the two feature sets are close in macro F1, as the embeddings match or slightly exceed the raw inputs for the GCI and for the LNWC under logistic regression, and fall marginally behind only for the LNWC under gradient boosted trees. While the margins are narrow, the self half of the embedding (130 of the 256 dimensions, a roughly two-fold compression, excluding the near and commute slices) is competitive with the raw inputs across both classifications. The neighbourhood-context slices also carry signal of their own: used alone, the near and commute subsets score well above the dummy floor, and on the GCI the self slice combined with near gives the single best result. Moreover, the embeddings provide a non-linear representation, the importance of which is shown on the LNWC, where the raw inputs themselves score higher with gradient boosted trees than with logistic regression.

Overall, these results are an encouraging first indication that the approach is a promising starting point for a fuller population dynamics foundation model, with the clearest gains where spatial context matters most. At the same time, further development and testing will be necessary to achieve any robust conclusions.

Reproducing the prototype

Setting up the environment

The trained model checkpoint and the exported embeddings are stored using Git LFS, which needs to be installed and initialised before cloning this repository:

git lfs install

If the repository has already been cloned without Git LFS, installing it and then running git lfs pull replaces the pointer files with the actual content.

The Python environment is managed with pixi and pinned in pixi.toml and pixi.lock, resolved for linux-64 with CUDA 12.9 and Python 3.12. From the repository root:

pixi install

Every command in this repository is then run through pixi, from the repository root, for instance:

pixi run python code/oafm_train.py

The notebooks in data/ and eval/ are executed through a Jupyter kernel, which needs to be registered once:

pixi run python -m ipykernel install --user --name oafm-pixi --display-name "OAFM (pixi)"

A CUDA GPU is required to prepare the points of interest tables, which use FAISS on GPU, and to train the model in a reasonable time, although the training, test and embedding scripts fall back to CPU if no GPU is available.

Data preparation, model training and evaluation

The steps below assume that the environment is set up as above, that the source datasets have been downloaded into data/source, data/spatial and data/eval as detailed in data/README.md, and that the outputs of the first three data preparation steps are available (the identifier tables in data/ids and the LSOA best-fit lookup in data/spatial). All the commands are run from the repository root.

1. Prepare the model inputs. The script below clears and re-runs the remaining data preparation notebooks (steps 4 to 12 in data/README.md) in a single pass, regenerating the per-source tables in data/tables and the model inputs in data/input, including the heterogeneous graph.

pixi run python data/dataprep_run_totables_toinput.py

2. Train the model. The training script trains the model for the full epoch budget, then automatically runs the test step and exports the embeddings from the best checkpoint, writing them to model/{model_name} and embeddings/{model_name}, where {model_name} records the architecture, hyperparameter and dataset versions (see code/README.md). The random seed is fixed by default, and can be changed using the --set_seed and --randomseed options. Training for the full budget takes several hours on a GPU.

pixi run python code/oafm_train.py

3. Run the evaluation. The script below clears and re-runs the three evaluation notebooks in a single pass, preceded by the notebook that prepares the crime exposure targets from the police data, unless those targets are already available. The emb_path variable at the top of each evaluation notebook points to the released embeddings, and needs to be pointed at the folder created in step 2 to evaluate a newly trained model.

pixi run python eval/eval_run_all.py

References

• Agarwal, M. et al. (2025). General geospatial inference with a population dynamics foundation model. arXiv preprint arXiv:2411.07207.
• De Sabbata, S. and Liu, P. (2019). Deep learning geodemographics with autoencoders and geographic convolution. In Proceedings of the 22nd AGILE conference on geographic information science, Limassol, Greece.
• De Sabbata, S. and Liu, P. (2023). A graph neural network framework for spatial geodemographic classification. International Journal of Geographical Information Science, 37(12), pp.2464-2486.
• De Sabbata, S., Liu, P., Goodwin, O. and Singleton, A. (2026). AE-MoSE: an autoencoder mixture of spatial experts for geodemographic classification. GeoAI 2026.
• Hamilton, W., Ying, Z. and Leskovec, J. (2017). Inductive representation learning on large graphs. Advances in neural information processing systems, 30.

How to cite

Please cite this work as follows.

@misc{desabbata2026oafmprototype,
	title = {Output {Area} {Foundation} {Model} ({OAFM}) prototype},
	rights = {{MIT} Licence},
	url = {https://github.com/stefdesabbata/oafm-prototype},
	doi = {10.25392/leicester.data.33071534},
	author = {De Sabbata, Stef},
	year = {2026},
}

Acknowledgements

This work was supported by the Economic and Social Research Council, Smart Data Research UK Fellowship, grant number UKRI4013.

The data for this research have been provided by the Geographic Data Service (geods.ac.uk), a Smart Data Research UK Investment: ES/Z504464/1.

Licence

The code and output data in this repository are licensed under the MIT License (see LICENSE) and contain public sector information licensed under the Open Government Licence v3.0. Contains OS data © Crown copyright and database right 2023 and 2024. The images are licensed under the Creative Commons Attribution 4.0 License. The lookup tables included in data/ids/ and data/spatial/ are distributed under Open Government Licence (OGL). The maps use digital boundaries from the Office for National Statistics licensed under the Open Government Licence v3.0, contains OS data © Crown copyright and database right 2023.