🤖 machine learning

MST-Direct at Scale: Multivariate and Conditional Geostatistical Simulation via Sinkhorn Optimal Transport

This paper extends the MST-Direct geostatistical simulation method to handle multivariate, conditional, and large-scale scenarios by employing a sparse Sinkhorn optimal transport approach that exactly preserves joint distributions and honors hard data while outperforming the approximate Projection Pursuit Multivariate Transform (PPMT).

Original authors: Tcharlies Bachmann Schmitz

Published 2026-06-09

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tcharlies Bachmann Schmitz

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are an artist trying to recreate a complex, chaotic painting of a geological landscape. You have a bucket of thousands of unique, multi-colored tiles (each tile represents a specific combination of rock properties, like porosity, density, and mineral content). Your goal is to arrange these tiles on a giant grid to create a new map that looks and feels exactly like the original, but with one catch: the tiles need to be arranged in a way that respects the "neighborly" rules of the landscape (e.g., if one spot is sandy, the spot next to it should probably be sandy too).

This paper introduces a new, super-efficient way to solve this puzzle, called MST-Direct. Here is how it works, broken down into simple concepts:

1. The Old Problem: The "Perfect Match" Puzzle

In the past, scientists tried to shuffle these tiles using complex math. However, they hit three major walls:

Too Slow: If the grid got too big (like a 200x200 map), the computer would crash trying to figure out which tile goes where.
Too Simple: They could only handle two types of data at a time (like just sand and clay), but real geology has many variables (sand, clay, gold, water, etc.).
Rigid Rules: If you had some real measurements from a drill hole (hard data), the old methods couldn't force the map to match those exact spots without messing up the rest of the picture.

2. The New Solution: The "Smart Shuffle"

The author, Tcharlies Bachmann Schmitz, upgraded the method to solve all three problems at once. Think of it as a three-step magic trick:

Step A: The "Sparse" Search (Solving the Speed Problem)
Imagine you have 40,000 tiles to place. A brute-force method would try to compare every single tile with every single spot on the map (like checking 1.6 billion pairs). That's too slow.

The Fix: The new method is like a smart shopper who only looks at the 50 closest items in the store aisle instead of the whole warehouse. By only checking a few "candidate" tiles for each spot, the computer finishes the job in under a minute, even on huge maps.

Step B: The "Invisible Backbone" (Solving the Multi-Variable Problem)
To handle many variables (like 6 different rock properties) at once, the method uses a trick.

The Analogy: Imagine a skeleton (a "backbone") made of invisible, perfectly smooth Gaussian clouds. This skeleton already has the right shape and neighborly rules built-in.
The Trick: The method takes your messy, complex pile of real-world tiles and "snaps" them onto this smooth skeleton. Because the skeleton is mathematically perfect, when you snap the tiles onto it, the complex relationships between the 6 variables stay exactly the same. It's like taking a messy pile of LEGOs and snapping them onto a pre-built frame; the frame ensures they fit together perfectly without breaking the design.

Step C: The "Pinning" Strategy (Solving the Hard Data Problem)
What if you have a few spots where you know the exact values (from a drill)?

The Fix: The method "pins" those specific tiles to their exact locations on the map, like taping a photo to a wall. Then, it uses a technique called "kriging" (a fancy way of guessing the neighbors based on the pinned photo) to fill in the rest of the map around it. This ensures the map honors the real data perfectly while still looking natural everywhere else.

3. The Results: A Perfect Copy

The author tested this new method against an older, popular method called PPMT.

The Test: They used a very difficult, non-linear 6-variable dataset (think of it as a very complex, wavy pattern that is hard to copy).
The Winner:
- MST-Direct: Produced a map that was a perfect copy of the original data. The histogram (the chart showing how many of each tile type exists) had zero error. It preserved every single detail of the complex relationships.
- PPMT: Got the general idea right but distorted the fine details. It was an "approximation," meaning it smoothed out the rough edges and lost some of the complex patterns.

Summary

In short, this paper presents a new tool that allows geologists to create massive, detailed maps of the underground. It is fast enough to handle huge grids, smart enough to juggle many different rock properties at once, and precise enough to lock in real-world measurements without breaking the pattern. It essentially allows computers to shuffle complex data perfectly, preserving the "personality" of the geological model better than previous methods.

Technical Summary: MST-Direct at Scale

Problem Statement
Geostatistical simulation aims to generate multiple realizations of geological models that replicate spatial features and support uncertainty quantification. A significant challenge arises when geological relationships are strongly non-linear and heteroscedastic, rendering classical multivariate methods that rely on linear correlation coefficients inadequate. Existing approaches, such as Stepwise Conditional Transformation (SCT), Projection Pursuit Multivariate Transform (PPMT), and Direct Multivariate Simulation (DMS), attempt to address this by transforming variables into Gaussian space and back. However, these transforms are only exact asymptotically, and their errors tend to grow as the number of variables increases.

Previous work introduced MST-Direct (Matching via Sinkhorn Transport), which frames simulation as an entropy-regularized optimal transport (OT) problem. While this approach achieved 100% shape preservation for bivariate relationships, the original formulation was limited to small grids (e.g., 25 × 25), unconditional settings, and bivariate data. It explicitly left three limitations unresolved: scalability beyond ~10,000 nodes, extension to more than two variables, and the ability to perform conditional simulation with hard data.

Methodology
This paper extends MST-Direct to address the three identified limitations through a scalable, multivariate, and conditional framework. The core principle remains casting simulation as an optimal transport problem where target value tuples are matched onto a spatial template.

Scalable Matching: To handle large grids (e.g., 40,000 nodes), the authors replace the dense $O(n^2)$ coupling with a sparse, candidate-restricted Sinkhorn matcher. Each grid location is restricted to its $C$ nearest target tuples (identified via a k-d tree), reducing memory complexity to $O(nC)$. The solution is found using the Sinkhorn algorithm in the log domain. A greedy assignment rounds the soft coupling into a permutation, with a repair step for any locations left without a free candidate.
Multivariate Backbone: To extend the method to $d$ variables, the target cloud is matched onto an independent FFT-MA (Fast Fourier Transform Moving Average) Gaussian backbone. This backbone carries the prescribed variogram (e.g., spherical). The matching process transfers the spatial structure of the backbone to every variable without imposing isotropy or stationarity requirements, though an isotropic spherical model is used here for comparability.
Conditional Simulation: Hard data is honored through a two-step process. First, the Gaussian backbone is conditioned to the hard data locations using simple kriging, ensuring the prescribed covariance is maintained elsewhere. Second, the hard data locations are "pinned": their specific tuples are fixed and removed from the matching pool, while the remaining locations are simulated. This ensures exact reproduction of hard data values.

The relational term, which augments the cost function to ensure spatially adjacent locations receive mutually similar tuples, is applied via refinement passes that pull query points toward the mean of their spatial neighbors.

Key Contributions
The paper presents four primary contributions:

A scalable MST-Direct matcher utilizing sparse, candidate-restricted Sinkhorn with greedy bijection completion, capable of processing 200 × 200 (40,000-node) and 100 × 100 grids in under a minute.
A multivariate formulation that matches target clouds onto an independent FFT-MA backbone, demonstrated successfully on six variables.
A conditional simulation mechanism that honors hard data exactly by pinning data tuples and conditioning the backbone via kriging.
A head-to-head validation against the Projection Pursuit Multivariate Transform (PPMT) on a 6-variate, heteroscedastic, non-linear reference distribution, covering both unconditional and conditional settings.

Results
The method was validated on the same 6-variate, heteroscedastic, strongly non-linear reference distribution used by DMS.

Unconditional Simulation (200 × 200 grid): MST-Direct produced realizations indistinguishable from the reference distribution, achieving zero histogram error (Mean Squared Error, MSE). In contrast, PPMT recovered the gross structure but distorted fine non-linear joints, resulting in MSEs on the order of $10^{-6}$ to $10^{-5}$ . Regarding spatial statistics, MST-Direct matched the variogram sill exactly, whereas PPMT overshot it by 10–20%.
Conditional Simulation (100 × 100 grid, 200 hard data): MST-Direct honored the hard data exactly (maximum absolute error of 0), while PPMT honored them to machine precision ( $\sim 10^{-14}$ ). MST-Direct again reproduced the target joint distribution with zero histogram error, while PPMT remained an approximation. Both methods reproduced the spatial variograms, with MST-Direct tracking the sill exactly and PPMT slightly undershooting it.
Efficiency: The sparse matcher successfully handled the 40,000-node grid, a scale where the dense matcher would be infeasible.

Significance and Claims
The paper claims to resolve the three open limitations of the original MST-Direct formulation. The authors assert that because the transport plan remains a permutation of the target tuples, the multivariate joint distribution is preserved exactly, avoiding the asymptotic errors and dimensionality-dependent factorization errors associated with SCT, PPMT, and DMS.

The method is presented as well-suited for applications involving complex, strongly non-linear multivariate dependencies. While the approach successfully reproduces target statistics and spatial models, the authors note a residual artifact: a mild short-lag component in unconditional variograms, which is controllable via the number of candidates and relational passes. The paper concludes that anisotropic variograms remain future work. An open-source implementation (v2.0.0) is available as the mst-direct Python package.

1. The Old Problem: The "Perfect Match" Puzzle

2. The New Solution: The "Smart Shuffle"

3. The Results: A Perfect Copy

Summary

Technical Summary: MST-Direct at Scale

More like this