use super::frame::ValueRange;
use super::wire::{self, Block, BlockData, Head};
use super::SCHEMA_VERSION;
use crate::error::StaticError;
use crate::grid::{Dims, Ijk, Point3};
use crate::model::model::StaticModel;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{self, Write};
const FACES: [[usize; 4]; 6] = [
[0, 1, 3, 2], [4, 5, 7, 6], [0, 1, 5, 4], [2, 3, 7, 6], [0, 2, 6, 4], [1, 3, 7, 5], ];
const NEIGHBORS: [(isize, isize, isize); 6] = [
(0, 0, -1), (0, 0, 1), (0, -1, 0), (0, 1, 0), (-1, 0, 0), (1, 0, 0), ];
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct VolumeBundle {
pub schema_version: u32,
pub inputs_ref: String,
pub property: String,
pub cell_count: usize,
pub positions: Vec<f32>,
pub indices: Vec<u32>,
pub tri_cell: Vec<u32>,
pub cell_values: Vec<f32>,
pub zone_ids: Vec<u16>,
pub zone_names: Vec<String>,
pub value_range: ValueRange,
}
impl StaticModel {
pub fn volume_bundle(&self, property: &str) -> Result<VolumeBundle, StaticError> {
self.extract_shell(property, |v| v.is_finite())
}
pub fn volume_bundle_thresholded(
&self,
property: &str,
cutoff: f64,
keep_above: bool,
) -> Result<VolumeBundle, StaticError> {
self.extract_shell(property, |v| {
v.is_finite() && if keep_above { v >= cutoff } else { v <= cutoff }
})
}
fn extract_shell(
&self,
property: &str,
keep: impl Fn(f64) -> bool,
) -> Result<VolumeBundle, StaticError> {
let grid = self.view_grid()?;
let dims = grid.dims();
let cell_count = dims.cell_count();
let cube = &grid
.properties()
.get(property)
.ok_or_else(|| {
StaticError::InvalidInput(format!("volume_bundle: no property '{property}'"))
})?
.values;
let zones = self.zones().zones();
let zone_names: Vec<String> = zones.iter().map(|z| z.name.clone()).collect();
let mut zone_of_k = vec![0u16; dims.nk];
for (zid, z) in zones.iter().enumerate() {
for k in z.k_range.clone() {
zone_of_k[k] = zid as u16;
}
}
let present: Vec<bool> = (0..cell_count)
.map(|c| keep(cube[c]) && grid.cell_volume_at(c) > 0.0)
.collect();
let shell = extract_faces(
dims,
&present,
|c| grid.cell(ijk_of(c, dims)).corners,
|c| cube[c] as f32,
|c| zone_of_k[ijk_of(c, dims).k],
);
let value_range = ValueRange::of(shell.cell_values.iter().map(|&v| v as f64));
Ok(VolumeBundle {
schema_version: SCHEMA_VERSION,
inputs_ref: self.provenance().inputs_ref.clone(),
property: property.to_string(),
cell_count,
positions: shell.positions,
indices: shell.indices,
tri_cell: shell.tri_cell,
cell_values: shell.cell_values,
zone_ids: shell.zone_ids,
zone_names,
value_range,
})
}
}
#[inline]
fn ijk_of(c: usize, dims: Dims) -> Ijk {
let per_layer = dims.ni * dims.nj;
let k = c / per_layer;
let rem = c % per_layer;
Ijk::new(rem % dims.ni, rem / dims.ni, k)
}
struct Shell {
positions: Vec<f32>,
indices: Vec<u32>,
tri_cell: Vec<u32>,
cell_values: Vec<f32>,
zone_ids: Vec<u16>,
}
fn extract_faces(
dims: Dims,
present: &[bool],
corners: impl Fn(usize) -> [Point3; 8],
value: impl Fn(usize) -> f32,
zone: impl Fn(usize) -> u16,
) -> Shell {
let (ni, nj, nk) = (dims.ni, dims.nj, dims.nk);
let mut positions: Vec<f32> = Vec::new();
let mut indices: Vec<u32> = Vec::new();
let mut tri_cell: Vec<u32> = Vec::new();
let mut cell_values: Vec<f32> = Vec::new();
let mut zone_ids: Vec<u16> = Vec::new();
let mut vert_map: HashMap<[u32; 3], u32> = HashMap::new();
for k in 0..nk {
for j in 0..nj {
for i in 0..ni {
let c = (k * nj + j) * ni + i;
if !present[c] {
continue;
}
let mut cid = u32::MAX; let cell = corners(c);
for (f, face) in FACES.iter().enumerate() {
if neighbor_present(present, dims, (i, j, k), NEIGHBORS[f]) {
continue; }
if cid == u32::MAX {
cid = cell_values.len() as u32;
cell_values.push(value(c));
zone_ids.push(zone(c));
}
let mut v = [0u32; 4];
for (n, &corner) in face.iter().enumerate() {
let p = cell[corner];
let (x, y, z) = (p.x as f32, p.y as f32, p.z as f32);
let key = [x.to_bits(), y.to_bits(), z.to_bits()];
v[n] = *vert_map.entry(key).or_insert_with(|| {
let id = (positions.len() / 3) as u32;
positions.extend_from_slice(&[x, y, z]);
id
});
}
indices.extend_from_slice(&[v[0], v[1], v[2], v[0], v[2], v[3]]);
tri_cell.push(cid);
tri_cell.push(cid);
}
}
}
}
Shell {
positions,
indices,
tri_cell,
cell_values,
zone_ids,
}
}
#[inline]
fn neighbor_present(
present: &[bool],
dims: Dims,
(i, j, k): (usize, usize, usize),
(di, dj, dk): (isize, isize, isize),
) -> bool {
let (ni, nj, nk) = (dims.ni as isize, dims.nj as isize, dims.nk as isize);
let (ii, jj, kk) = (i as isize + di, j as isize + dj, k as isize + dk);
if ii < 0 || ii >= ni || jj < 0 || jj >= nj || kk < 0 || kk >= nk {
return false;
}
let c = (kk * nj + jj) * ni + ii;
present[c as usize]
}
impl VolumeBundle {
fn blocks(&self) -> [Block<'_>; 5] {
let vcount = self.positions.len() / 3;
let tcount = self.indices.len() / 3;
let scount = self.cell_values.len();
[
Block {
name: "positions",
shape: vec![vcount, 3],
data: BlockData::F32(&self.positions),
},
Block {
name: "indices",
shape: vec![tcount, 3],
data: BlockData::U32(&self.indices),
},
Block {
name: "tri_cell",
shape: vec![tcount],
data: BlockData::U32(&self.tri_cell),
},
Block {
name: "cell_values",
shape: vec![scount],
data: BlockData::F32(&self.cell_values),
},
Block {
name: "zone_ids",
shape: vec![scount],
data: BlockData::U16(&self.zone_ids),
},
]
}
fn head(&self) -> Head<'_> {
Head {
schema_version: self.schema_version,
inputs_ref: &self.inputs_ref,
property: &self.property,
cell_count: self.cell_count,
shell_cell_count: self.cell_values.len(),
vertex_count: self.positions.len() / 3,
triangle_count: self.indices.len() / 3,
zone_names: &self.zone_names,
value_range: &self.value_range,
}
}
pub fn write_self_contained<W: Write>(&self, w: &mut W) -> io::Result<()> {
wire::write_self_contained(&self.head(), &self.blocks(), w)
}
pub fn write_sidecar<W1: Write, W2: Write>(
&self,
json: &mut W1,
bin: &mut W2,
) -> io::Result<()> {
wire::write_sidecar(&self.head(), &self.blocks(), json, bin)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::gridder::{Conformity, SolveOpts};
use crate::model::{BuildOpts, ConstantPriors, StaticModelBuilder};
use base64::engine::general_purpose::STANDARD;
use base64::Engine;
fn model_nk(ni: usize, nj: usize, nk: usize) -> StaticModel {
let opts = BuildOpts {
area_m2: 10_000.0,
gross_height_m: 40.0,
nk,
conformity: Conformity::Proportional,
solve_opts: SolveOpts::default(),
priors: ConstantPriors {
porosity: 0.25,
net_to_gross: 0.8,
water_saturation: 0.3,
},
};
StaticModelBuilder::flat(ni, nj, 5000.0, 5100.0, opts)
.unwrap()
.build()
.unwrap()
}
fn model() -> StaticModel {
model_nk(3, 2, 4)
}
fn box_shell_tris(ni: usize, nj: usize, nk: usize) -> usize {
2 * (ni * nj + ni * nk + nj * nk) * 2
}
#[test]
fn full_box_is_exactly_the_outer_shell() {
let (ni, nj, nk) = (3, 2, 4);
let b = model_nk(ni, nj, nk).volume_bundle("PORO").unwrap();
assert_eq!(b.cell_count, ni * nj * nk);
assert_eq!(b.indices.len() / 3, box_shell_tris(ni, nj, nk));
assert_eq!(b.tri_cell.len(), b.indices.len() / 3);
let vcount = b.positions.len() / 3;
assert!(
vcount < b.indices.len(),
"vertices deduplicated across faces"
);
assert!(b.indices.iter().all(|&i| (i as usize) < vcount));
}
#[test]
fn shell_cells_carry_their_value_and_zone() {
let b = model().volume_bundle("PORO").unwrap();
assert!(b.cell_values.iter().all(|&v| (v - 0.25).abs() < 1e-6));
assert_eq!(b.zone_ids.len(), b.cell_values.len());
assert!(b.zone_ids.iter().all(|&z| z == 0));
assert_eq!(b.zone_names, vec!["RESERVOIR".to_string()]);
assert!((b.value_range.min - 0.25).abs() < 1e-6);
assert!((b.value_range.max - 0.25).abs() < 1e-6);
assert!(b
.tri_cell
.iter()
.all(|&c| (c as usize) < b.cell_values.len()));
}
fn unit_cube(dims: Dims) -> impl Fn(usize) -> [Point3; 8] {
move |c| {
let ijk = ijk_of(c, dims);
let mut cn = [Point3::new(0.0, 0.0, 0.0); 8];
for (idx, slot) in cn.iter_mut().enumerate() {
let (di, dj, dk) = (idx & 1, (idx >> 1) & 1, (idx >> 2) & 1);
*slot = Point3::new(
(ijk.i + di) as f64,
(ijk.j + dj) as f64,
(ijk.k + dk) as f64,
);
}
cn
}
}
#[test]
fn extract_faces_full_box_is_the_outer_shell_exactly() {
let (ni, nj, nk) = (3, 3, 3);
let dims = Dims::new(ni, nj, nk).unwrap();
let present = vec![true; ni * nj * nk];
let s = extract_faces(dims, &present, unit_cube(dims), |_| 1.0, |_| 0);
assert_eq!(s.indices.len() / 3, box_shell_tris(ni, nj, nk)); assert_eq!(s.positions.len() / 3, 56);
assert_eq!(s.cell_values.len(), 26); }
#[test]
fn collapsed_interior_cell_exposes_its_neighbours_faces() {
let (ni, nj, nk) = (3, 3, 3);
let dims = Dims::new(ni, nj, nk).unwrap();
let mut present = vec![true; ni * nj * nk];
let centre = (nj + 1) * ni + 1; present[centre] = false;
let s = extract_faces(dims, &present, unit_cube(dims), |_| 1.0, |_| 0);
assert_eq!(s.indices.len() / 3, box_shell_tris(ni, nj, nk) + 6 * 2); assert_eq!(s.cell_values.len(), 26);
assert_eq!(s.tri_cell.len(), s.indices.len() / 3);
assert!(s
.tri_cell
.iter()
.all(|&c| (c as usize) < s.cell_values.len()));
}
#[test]
fn thresholded_recut_gates_cells_by_cutoff() {
let m = model_nk(4, 4, 4);
let full = m.volume_bundle("PORO").unwrap();
let all = m.volume_bundle_thresholded("PORO", 0.20, true).unwrap();
assert_eq!(all.indices.len(), full.indices.len());
assert_eq!(all.cell_values.len(), full.cell_values.len());
let none = m.volume_bundle_thresholded("PORO", 0.30, true).unwrap();
assert!(none.indices.is_empty());
assert!(none.cell_values.is_empty());
let below = m.volume_bundle_thresholded("PORO", 0.30, false).unwrap();
assert_eq!(below.cell_values.len(), full.cell_values.len());
}
#[test]
fn missing_property_errors() {
assert!(model().volume_bundle("NOPE").is_err());
assert!(model()
.volume_bundle_thresholded("NOPE", 0.0, true)
.is_err());
}
fn read_f32_le(bytes: &[u8]) -> Vec<f32> {
bytes
.chunks_exact(4)
.map(|c| f32::from_le_bytes(c.try_into().unwrap()))
.collect()
}
fn read_u32_le(bytes: &[u8]) -> Vec<u32> {
bytes
.chunks_exact(4)
.map(|c| u32::from_le_bytes(c.try_into().unwrap()))
.collect()
}
fn read_u16_le(bytes: &[u8]) -> Vec<u16> {
bytes
.chunks_exact(2)
.map(|c| u16::from_le_bytes(c.try_into().unwrap()))
.collect()
}
#[test]
fn self_contained_base64_round_trips_bit_exact() {
let b = model_nk(4, 3, 5).volume_bundle("PORO").unwrap();
let mut out = Vec::new();
b.write_self_contained(&mut out).unwrap();
let env: serde_json::Value = serde_json::from_slice(&out).unwrap();
assert_eq!(env["schema_version"], serde_json::json!(5));
assert_eq!(env["encoding"], serde_json::json!("base64"));
let blk = &env["blocks"];
let dec = |name: &str| {
STANDARD
.decode(blk[name]["data"].as_str().unwrap())
.unwrap()
};
assert_eq!(read_f32_le(&dec("positions")), b.positions);
assert_eq!(read_u32_le(&dec("indices")), b.indices);
assert_eq!(read_u32_le(&dec("tri_cell")), b.tri_cell);
assert_eq!(read_f32_le(&dec("cell_values")), b.cell_values);
assert_eq!(read_u16_le(&dec("zone_ids")), b.zone_ids);
}
#[test]
fn sidecar_manifest_round_trips_bit_exact() {
let b = model_nk(4, 3, 5).volume_bundle("PORO").unwrap();
let (mut json, mut bin) = (Vec::new(), Vec::new());
b.write_sidecar(&mut json, &mut bin).unwrap();
let env: serde_json::Value = serde_json::from_slice(&json).unwrap();
assert_eq!(env["encoding"], serde_json::json!("sidecar"));
let blk = &env["blocks"];
let slice = |name: &str| {
let off = blk[name]["offset"].as_u64().unwrap() as usize;
let len = blk[name]["length"].as_u64().unwrap() as usize;
&bin[off..off + len]
};
assert_eq!(read_f32_le(slice("positions")), b.positions);
assert_eq!(read_u32_le(slice("indices")), b.indices);
assert_eq!(read_u32_le(slice("tri_cell")), b.tri_cell);
assert_eq!(read_f32_le(slice("cell_values")), b.cell_values);
assert_eq!(read_u16_le(slice("zone_ids")), b.zone_ids);
let total: usize = [
"positions",
"indices",
"tri_cell",
"cell_values",
"zone_ids",
]
.iter()
.map(|n| blk[n]["length"].as_u64().unwrap() as usize)
.sum();
assert_eq!(total, bin.len());
}
#[test]
fn schema_snapshot_v4_envelope() {
let b = model().volume_bundle("PORO").unwrap();
let mut out = Vec::new();
b.write_self_contained(&mut out).unwrap();
let env: serde_json::Value = serde_json::from_slice(&out).unwrap();
let mut keys: Vec<&str> = env
.as_object()
.unwrap()
.keys()
.map(String::as_str)
.collect();
keys.sort_unstable();
assert_eq!(
keys,
[
"blocks",
"cell_count",
"encoding",
"inputs_ref",
"kind",
"property",
"schema_version",
"shell_cell_count",
"triangle_count",
"value_range",
"vertex_count",
"zone_names",
]
);
assert_eq!(env["kind"], serde_json::json!("volume"));
assert_eq!(env["schema_version"], serde_json::json!(5));
let blocks = env["blocks"].as_object().unwrap();
let mut bnames: Vec<&str> = blocks.keys().map(String::as_str).collect();
bnames.sort_unstable();
assert_eq!(
bnames,
[
"cell_values",
"indices",
"positions",
"tri_cell",
"zone_ids"
]
);
assert_eq!(
env["blocks"]["positions"]["dtype"],
serde_json::json!("f32")
);
assert_eq!(env["blocks"]["indices"]["dtype"], serde_json::json!("u32"));
assert_eq!(env["blocks"]["zone_ids"]["dtype"], serde_json::json!("u16"));
}
}