Add the ResFlow epoch-40 weights, config and model card.

README.md

@@ -0,0 +1,117 @@
+---
+license: mit
+library_name: resflow
+datasets:
+- SciLM/SiliciclasticReservoirs
+tags:
+- geology
+- reservoir
+- subsurface
+- 3d
+- flow-matching
+- diffusion-transformer
+- generative
+---
+
+# ResFlow
+
+ResFlow generates 3D reservoir facies (sand and mud) with one flow-matching model for eight simulated
+siliciclastic environments: deep-water lobes, a distributary delta, and six fluvial classes. It is conditioned on
+geological parameters and wells, and it generates single 64×64×32 volumes or whole fields of any size in one pass.
+These are the weights of the paper *SiliciclasticReservoirs: A Million-Reservoir Dataset and Flow-Matching
+Foundation Model for 3D Siliciclastic Reservoir Generation* (NeurIPS 2026, Track on Evaluations and Datasets).
+
+**Code:** [github.com/SciLM-ai/ResFlow](https://github.com/SciLM-ai/ResFlow) ·
+**Dataset:** [SciLM/SiliciclasticReservoirs](https://huggingface.co/datasets/SciLM/SiliciclasticReservoirs) ·
+**Benchmark:** [ResBench](https://github.com/SciLM-ai/ResBench) ·
+**Simulator:** [ResMill](https://github.com/SciLM-ai/ResMill)
+
+## Quick start
+
+```bash
+pip install git+https://github.com/SciLM-ai/ResFlow
+```
+
+```python
+import resflow
+
+model = resflow.load_pretrained()                     # this repository
+
+vols = model.generate('meander', n=8)                 # (8, 64, 64, 32) uint8, 1 = sand
+vols = model.generate('lobe', n=4, ntg=0.35, width_cells=30, azimuth=45)
+ens = model.generate('meander', n=32,                 # every realisation honours the well exactly
+                     wells=[resflow.Well(x=32, y=32, facies=column)])
+field = model.generate_field('lobe', shape=(512, 512, 32), ntg=0.5)
+```
+
+A walkthrough notebook is in the code repository at `examples/pretrained/quickstart.ipynb`. A GPU is recommended:
+on an NVIDIA GH200 a 64×64×32 volume takes about 1.6 s in batches of 128 and a 512×512×32 field about 4 minutes.
+A CPU works, but takes tens of minutes per volume.
+
+## Inputs
+
+| Input | Meaning |
+|---|---|
+| environment | `lobe`, `pv_shoestring`, `cb_labyrinth`, `cb_jigsaw`, `sh_distal`, `sh_proximal`, `meander`, `delta` |
+| `ntg` | sand fraction of the volume (0–1) |
+| `width_cells`, `depth_cells` | characteristic body width and thickness, in cells |
+| `azimuth` | flow direction in degrees, 0 = along +x (default 0) |
+| `asp` | lobe aspect ratio (lobe) |
+| `mCHsinu` | channel sinuosity (channels, delta) |
+| `mFFCHprop` | fraction of abandoned channels plugged with mud (channels, delta) |
+| `probAvulInside` | in-belt avulsion probability (channels, delta) |
+| `trunk_length_fraction` | share of the trunk kept free of bifurcations (delta) |
+| `wells` / `observed` | vertical `Well(x, y, facies)`, or an array of known cells: 1 sand, 0 mud, −1 unknown |
+
+Parameters left out take the environment's typical value (the median over its training data), and
+`model.parameters(env)` lists each one with its training range. In `generate_field` any parameter can also be an
+(X, Y) map, so trends can vary across one reservoir. Volumes have axes (x, y, z) with z = 0 at the base; cells are
+100 m laterally for lobes and 10 m for channels and deltas, and 1 m vertically.
+
+## Model
+
+A 32.6M-parameter 3D diffusion transformer (4×4×2 patches, 12 blocks, width 384, 6 heads, QK-norm, 3D rotary
+positions with θ = 100, adaptive layer norm) trained with conditional flow matching. The 18-number condition vector
+holds the environment one-hot, NTG, width, depth, the sine and cosine of the azimuth, and five environment-specific
+parameters; wells enter as two extra input channels. Sampling uses Heun's method with 100 steps and
+classifier-free guidance 3. Volumes use global attention as in training; fields switch to sliding-window attention
+of ±12 tokens per axis, so any extent is generated whole, without tiling.
+
+Training used 900,000 volumes of SiliciclasticReservoirs (all eight environments) for 40 epochs on 16 NVIDIA GH200
+GPUs (global batch 512, AdamW, peak learning rate 5.77×10⁻⁴, cosine schedule, about 7.5 hours). These weights are
+the final epoch-40 exponential moving average.
+
+## Evaluation
+
+On ResBench, which scores each check against the simulator's own run-to-run variability (1 means
+indistinguishable from the simulator, lower is better), ResFlow scores 2.56 overall: 2.33 on unconditional
+volumes, 3.12 on well-conditioned volumes, and 2.24 on fields. Given the same inputs, it is closer to the simulator
+than GANSim-3D, DiffSIM and a 3D latent diffusion model in every environment where they were compared, and it beats single-environment specialists with its own
+architecture and recipe in six of eight environments. Details are in the paper.
+
+## Limitations
+
+- Trained only on simulated reservoirs from the ResMill simulator, and not validated on real subsurface data.
+  Forecasting a specific real reservoir needs separate calibration.
+- Covers the eight training environments only; parameter values outside the training ranges are not supported
+  (the API warns).
+- Under a fixed input it underestimates the simulator's uncertainty in most channel environments, so its ensembles
+  understate uncertainty there.
+- Distal sheets contain isolated sand specks; at field scale, connectivity in shoestring and delta fields deviates
+  most from the simulator.
+
+## Citation
+
+```bibtex
+@inproceedings{baghishov2026siliciclastic,
+  title     = {SiliciclasticReservoirs: A Million-Reservoir Dataset and Flow-Matching Foundation Model
+               for 3D Siliciclastic Reservoir Generation},
+  author    = {Baghishov, Ilgar and Rustamzade, Elnara and Henkelman, Graeme and Foster, John T. and Pyrcz, Michael J.},
+  booktitle = {Advances in Neural Information Processing Systems (NeurIPS), Track on Evaluations and Datasets},
+  year      = {2026}
+}
+```
+
+## License
+
+MIT.

config.json

@@ -0,0 +1,360 @@
+{
+ "model_type": "resflow-dit3d",
+ "architecture": {
+  "in_channels": 3,
+  "out_channels": 1,
+  "volume_shape": [
+   64,
+   64,
+   32
+  ],
+  "patch_size": [
+   4,
+   4,
+   2
+  ],
+  "hidden": 384,
+  "depth": 12,
+  "num_heads": 6,
+  "num_cond": 18,
+  "num_time_embs": 1,
+  "expand_angle_idx": null,
+  "qk_norm": true,
+  "conv_io": false,
+  "pos_embed": "rope",
+  "rope_theta": 100.0,
+  "window": null,
+  "window_shift": false,
+  "token_conv": 0,
+  "token_conv_kernel": [
+   3,
+   3,
+   3
+  ],
+  "attn_mode": "window",
+  "attn_radius": [
+   8,
+   8,
+   8
+  ]
+ },
+ "layer_types": [
+  "lobe",
+  "channel:PV_SHOESTRING",
+  "channel:CB_LABYRINTH",
+  "channel:CB_JIGSAW",
+  "channel:SH_DISTAL",
+  "channel:SH_PROXIMAL",
+  "channel:MEANDER_OXBOW",
+  "delta"
+ ],
+ "cont_cols": [
+  "ntg",
+  "width_cells",
+  "depth_cells",
+  "asp",
+  "mCHsinu",
+  "mFFCHprop",
+  "probAvulInside",
+  "trunk_length_fraction"
+ ],
+ "n_universal": 3,
+ "cont_min": [
+  0.05118560791015625,
+  2.404618263244629,
+  3.000000238418579,
+  1.000009536743164,
+  1.0500000715255737,
+  0.0,
+  0.0,
+  0.10000130534172058
+ ],
+ "cont_max": [
+  0.9156951904296875,
+  79.82032775878906,
+  19.993122100830078,
+  2.499992847442627,
+  1.8499996662139893,
+  0.6999967098236084,
+  0.7999995946884155,
+  0.4999952018260956
+ ],
+ "env_defaults": {
+  "lobe": {
+   "ntg": 0.517,
+   "width_cells": 46.0001,
+   "depth_cells": 12.4999,
+   "asp": 1.75
+  },
+  "channel:PV_SHOESTRING": {
+   "ntg": 0.2077,
+   "width_cells": 7.5233,
+   "depth_cells": 6.9282,
+   "mCHsinu": 1.575,
+   "mFFCHprop": 0.15,
+   "probAvulInside": 0.06
+  },
+  "channel:CB_LABYRINTH": {
+   "ntg": 0.357,
+   "width_cells": 10.9679,
+   "depth_cells": 6.9282,
+   "mCHsinu": 1.275,
+   "mFFCHprop": 0.325,
+   "probAvulInside": 0.06
+  },
+  "channel:CB_JIGSAW": {
+   "ntg": 0.388,
+   "width_cells": 11.6044,
+   "depth_cells": 6.9282,
+   "mCHsinu": 1.25,
+   "mFFCHprop": 0.4,
+   "probAvulInside": 0.375
+  },
+  "channel:SH_DISTAL": {
+   "ntg": 0.6356,
+   "width_cells": 12.3673,
+   "depth_cells": 6.9282,
+   "mCHsinu": 1.275,
+   "mFFCHprop": 0.1,
+   "probAvulInside": 0.06
+  },
+  "channel:SH_PROXIMAL": {
+   "ntg": 0.5099,
+   "width_cells": 14.4069,
+   "depth_cells": 6.9282,
+   "mCHsinu": 1.225,
+   "mFFCHprop": 0.225,
+   "probAvulInside": 0.35
+  },
+  "channel:MEANDER_OXBOW": {
+   "ntg": 0.494,
+   "width_cells": 7.5244,
+   "depth_cells": 6.9282,
+   "mCHsinu": 1.625,
+   "mFFCHprop": 0.0,
+   "probAvulInside": 0.0
+  },
+  "delta": {
+   "ntg": 0.1934,
+   "width_cells": 10.1923,
+   "depth_cells": 6.9282,
+   "mCHsinu": 1.075,
+   "mFFCHprop": 0.2,
+   "probAvulInside": 0.65,
+   "trunk_length_fraction": 0.3
+  }
+ },
+ "env_ranges": {
+  "lobe": {
+   "ntg": [
+    0.0448,
+    0.9392
+   ],
+   "width_cells": [
+    12.1034,
+    79.9134
+   ],
+   "depth_cells": [
+    5.0161,
+    19.9931
+   ],
+   "asp": [
+    1.0,
+    2.5
+   ]
+  },
+  "channel:PV_SHOESTRING": {
+   "ntg": [
+    0.0,
+    0.6591
+   ],
+   "width_cells": [
+    2.4042,
+    22.3059
+   ],
+   "depth_cells": [
+    3.0,
+    15.9999
+   ],
+   "mCHsinu": [
+    1.3,
+    1.85
+   ],
+   "mFFCHprop": [
+    0.0,
+    0.3
+   ],
+   "probAvulInside": [
+    0.02,
+    0.1
+   ]
+  },
+  "channel:CB_LABYRINTH": {
+   "ntg": [
+    0.0,
+    0.9188
+   ],
+   "width_cells": [
+    3.6057,
+    31.8716
+   ],
+   "depth_cells": [
+    3.0,
+    15.9999
+   ],
+   "mCHsinu": [
+    1.1,
+    1.45
+   ],
+   "mFFCHprop": [
+    0.1,
+    0.55
+   ],
+   "probAvulInside": [
+    0.02,
+    0.1
+   ]
+  },
+  "channel:CB_JIGSAW": {
+   "ntg": [
+    0.0,
+    0.9547
+   ],
+   "width_cells": [
+    3.6068,
+    35.1057
+   ],
+   "depth_cells": [
+    3.0,
+    15.9999
+   ],
+   "mCHsinu": [
+    1.1,
+    1.4
+   ],
+   "mFFCHprop": [
+    0.1,
+    0.7
+   ],
+   "probAvulInside": [
+    0.2,
+    0.55
+   ]
+  },
+  "channel:SH_DISTAL": {
+   "ntg": [
+    0.0004,
+    0.9996
+   ],
+   "width_cells": [
+    4.2059,
+    35.0671
+   ],
+   "depth_cells": [
+    3.0,
+    15.9999
+   ],
+   "mCHsinu": [
+    1.1,
+    1.45
+   ],
+   "mFFCHprop": [
+    0.0,
+    0.2
+   ],
+   "probAvulInside": [
+    0.02,
+    0.1
+   ]
+  },
+  "channel:SH_PROXIMAL": {
+   "ntg": [
+    0.0,
+    0.9993
+   ],
+   "width_cells": [
+    4.8072,
+    41.4372
+   ],
+   "depth_cells": [
+    3.0,
+    15.9999
+   ],
+   "mCHsinu": [
+    1.05,
+    1.4
+   ],
+   "mFFCHprop": [
+    0.05,
+    0.4
+   ],
+   "probAvulInside": [
+    0.2,
+    0.5
+   ]
+  },
+  "channel:MEANDER_OXBOW": {
+   "ntg": [
+    0.0001,
+    0.9756
+   ],
+   "width_cells": [
+    2.4096,
+    22.3713
+   ],
+   "depth_cells": [
+    3.0,
+    15.9997
+   ],
+   "mCHsinu": [
+    1.4,
+    1.85
+   ],
+   "mFFCHprop": [
+    0.0,
+    0.0
+   ],
+   "probAvulInside": [
+    0.0,
+    0.0
+   ]
+  },
+  "delta": {
+   "ntg": [
+    0.0025,
+    0.8131
+   ],
+   "width_cells": [
+    3.0148,
+    31.9141
+   ],
+   "depth_cells": [
+    3.0,
+    15.9999
+   ],
+   "mCHsinu": [
+    1.05,
+    1.1
+   ],
+   "mFFCHprop": [
+    0.0,
+    0.4
+   ],
+   "probAvulInside": [
+    0.5,
+    0.8
+   ],
+   "trunk_length_fraction": [
+    0.1,
+    0.5
+   ]
+  }
+ },
+ "sampler": {
+  "solver": "heun",
+  "steps": 100,
+  "guidance": 3.0,
+  "field_radius": 12
+ },
+ "checkpoint": "inference_epoch040.pt"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95b3d751111999439136950e8ddfcb23baa9e2c4b6356eeb14481f04600a5618
+size 130361036