"""
reblocking.py
Module for reblocking operations on block models.
"""
from typing import TYPE_CHECKING
import json
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from parq_blockmodel.reblocking.conversion import tabular_to_3d_dict, dict_3d_to_tabular
from parq_blockmodel.reblocking.downsample import downsample_attributes
from parq_blockmodel.reblocking.upsample import upsample_attributes
from parq_blockmodel.utils.spatial_encoding import encode_world_coordinates, get_world_id_encoding_params
from parq_blockmodel.io.ingest_utils import build_world_id_encoding_from_xyz
if TYPE_CHECKING:
from parq_blockmodel import ParquetBlockModel
from parq_blockmodel.geometry import RegularGeometry, LocalGeometry, WorldFrame
def _validate_params(config, new_block_size):
if not isinstance(new_block_size, tuple) or len(new_block_size) != 3:
raise ValueError("new_block_size must be a tuple of three floats (dx, dy, dz).")
if any(dim <= 0 for dim in new_block_size):
raise ValueError("All dimensions in new_block_size must be positive.")
if not isinstance(config, dict):
raise ValueError("config must be a dictionary.")
def _persisted_attribute_columns(blockmodel) -> list[str]:
geometry_cols = {"block_id", "world_id", "x", "y", "z", "i", "j", "k"}
return [c for c in blockmodel.columns if c not in geometry_cols]
def _required_downsample_inputs(aggregation_config: dict[str, dict]) -> set[str]:
required: set[str] = set(aggregation_config)
for config in aggregation_config.values():
if not isinstance(config, dict):
continue
basis_name = config.get("basis", config.get("weight"))
if isinstance(basis_name, str) and basis_name:
required.add(basis_name)
fill_ratio_name = config.get("fill_ratio")
if isinstance(fill_ratio_name, str) and fill_ratio_name:
required.add(fill_ratio_name)
return required
def _prepare_arrays_and_index(
blockmodel,
requested_columns: set[str] | None = None,
) -> tuple[dict[str, np.ndarray], dict[str, pd.Index]]:
"""Prepare dense attribute arrays and category indices for reblocking.
This helper reads the blockmodel on a dense ijk grid and converts it
into 3D arrays suitable for up/downsampling. Geometry-derived columns
(coordinates and linear indices) are excluded and later recomputed from
the new :class:`RegularGeometry` to avoid inconsistencies.
"""
# Read a dense ijk-indexed view of the data. ``ParquetBlockModel.read``
# already returns a C-order (i, j, k) MultiIndex when ``index="ijk"``,
# which is exactly what ``tabular_to_3d_dict`` expects.
#
# Only load true attributes; drop geometry/identity columns which
# will be regenerated for the reblocked grid.
cols = _persisted_attribute_columns(blockmodel)
if requested_columns is not None:
cols = [c for c in cols if c in requested_columns]
extras = sorted(c for c in requested_columns if c not in cols)
cols.extend(extras)
df: pd.DataFrame = blockmodel.read(columns=cols, index="ijk", dense=True)
arrays, categories = tabular_to_3d_dict(df)
return arrays, categories
def _calculate_factors(blockmodel, new_block_size):
old_shape = blockmodel.geometry.local.shape
old_block_size = blockmodel.geometry.local.block_size
factors = []
for ob, nb in zip(old_block_size, new_block_size):
factor = nb / ob
if np.isclose(factor, round(factor)):
factors.append(round(factor))
elif np.isclose(1 / factor, round(1 / factor)):
factors.append(1 / round(1 / factor))
else:
raise ValueError("Block sizes must be integer multiples or divisors for up/downsampling.")
fx, fy, fz = factors
if all(f >= 1 for f in factors):
# Downsampling
new_shape = tuple(int(o // f) for o, f in zip(old_shape, factors))
elif all(f <= 1 for f in factors):
# Upsampling
up_factors = [int(round(1 / f)) for f in factors]
new_shape = tuple(int(o * uf) for o, uf in zip(old_shape, up_factors))
else:
raise ValueError("Cannot mix upsampling and downsampling in one operation.")
return fx, fy, fz, new_shape
def _build_new_geometry(blockmodel, new_block_size, new_shape) -> RegularGeometry:
return RegularGeometry(
local=LocalGeometry(
corner=blockmodel.geometry.local.corner,
block_size=new_block_size,
shape=new_shape,
),
world=WorldFrame(
origin=blockmodel.geometry.world.origin,
axis_u=blockmodel.geometry.world.axis_u,
axis_v=blockmodel.geometry.world.axis_v,
axis_w=blockmodel.geometry.world.axis_w,
srs=blockmodel.geometry.world.srs,
),
world_id_encoding=blockmodel.geometry.world_id_encoding,
)
[docs]
def downsample_blockmodel(blockmodel, new_block_size, aggregation_config) -> "ParquetBlockModel":
"""
Downsample a block model to a coarser grid with specified aggregation methods for each attribute.
This function supports downsampling of both categorical and numeric attributes.
Args:
blockmodel: ParquetBlockModel instance.
new_block_size: tuple of floats (dx, dy, dz) for the new block size.
aggregation_config: dict mapping attribute names to aggregation methods.
Use ``basis`` for ``weighted_mean`` configurations.
Example:
aggregation_config = {
'grade': {'method': 'weighted_mean', 'basis': 'dry_mass'},
'density': {'method': 'weighted_mean', 'basis': 'volume'},
'dry_mass': {'method': 'sum'},
'volume': {'method': 'sum'},
'rock_type': {'method': 'mode'}
}
Returns:
ParquetBlockModel: A new ParquetBlockModel instance with the downsampled grid.
"""
from parq_blockmodel import ParquetBlockModel
_validate_params(aggregation_config, new_block_size)
fx, fy, fz, new_shape = _calculate_factors(blockmodel, new_block_size)
if fx > 1 or fy > 1 or fz > 1:
# Downsampling: expect dict of dicts
if not all(isinstance(v, dict) for v in (aggregation_config or {}).values()):
raise ValueError("Downsampling config must be a dict of dicts per attribute.")
required_inputs = _required_downsample_inputs(aggregation_config or {})
arrays, categories = _prepare_arrays_and_index(blockmodel, requested_columns=required_inputs)
arrays = downsample_attributes(arrays, fx, fy, fz, aggregation_config or {})
transient_inputs = required_inputs - set(_persisted_attribute_columns(blockmodel)) - set(aggregation_config or {})
for column in transient_inputs:
arrays.pop(column, None)
elif fx < 1 or fy < 1 or fz < 1:
raise ValueError("reblocking factors imply upsampling, not downsampling.")
elif any(dim < 1 for dim in (fx, fy, fz)) and any(dim > 1 for dim in (fx, fy, fz)):
raise ValueError("Reblocking factors cannot specify upsample and downsample at the same time.")
# Convert back to DataFrame (ijk-indexed attributes only)
new_geometry = _build_new_geometry(blockmodel, new_block_size, new_shape)
# Use canonical .pbm suffix expected by ParquetBlockModel
new_path = blockmodel.blockmodel_path.with_name(f"{blockmodel.name}.reblocked.pbm")
reblocked_df = dict_3d_to_tabular(arrays, categories)
# assert the indexes are aligned in ijk space
if not reblocked_df.index.equals(new_geometry.to_multi_index_ijk()):
raise ValueError("Reblocking ijk index mismatch.")
# ------------------------------------------------------------------
# Reconstruct geometry-derived columns for the new grid
# ------------------------------------------------------------------
# x, y, z centroids in C-order matching the dense ijk index
reblocked_df["x"] = new_geometry.centroid_x
reblocked_df["y"] = new_geometry.centroid_y
reblocked_df["z"] = new_geometry.centroid_z
# Canonical linear id in C-order of local (i, j, k).
N = int(np.prod(new_geometry.local.shape))
rows = np.arange(N, dtype=np.uint32)
reblocked_df["block_id"] = rows
i, j, k = new_geometry.ijk_from_row_index(rows)
reblocked_df["i"] = i.astype(np.int32)
reblocked_df["j"] = j.astype(np.int32)
reblocked_df["k"] = k.astype(np.int32)
if new_geometry.world_id_encoding is None:
new_geometry.world_id_encoding = build_world_id_encoding_from_xyz(
reblocked_df["x"].to_numpy(dtype=float),
reblocked_df["y"].to_numpy(dtype=float),
reblocked_df["z"].to_numpy(dtype=float),
)
offset, scale, bits_per_axis = get_world_id_encoding_params(new_geometry.world_id_encoding)
reblocked_df["world_id"] = encode_world_coordinates(
reblocked_df["x"].to_numpy(dtype=float),
reblocked_df["y"].to_numpy(dtype=float),
reblocked_df["z"].to_numpy(dtype=float),
offset=offset,
scale=scale,
bits_per_axis=bits_per_axis,
).astype(np.int64)
ordered_cols = ["block_id", "world_id", "x", "y", "z", "i", "j", "k"]
ordered_cols.extend([c for c in reblocked_df.columns if c not in ordered_cols])
table = pa.Table.from_pandas(reblocked_df[ordered_cols].reset_index(drop=True), preserve_index=False)
meta = dict(table.schema.metadata or {})
meta[b"parq-blockmodel"] = json.dumps(new_geometry.to_metadata_dict()).encode("utf-8")
table = table.replace_schema_metadata(meta)
pq.write_table(table, new_path)
return ParquetBlockModel(blockmodel_path=new_path, geometry=new_geometry)
[docs]
def upsample_blockmodel(blockmodel, new_block_size, upsample_config=None, interpolation_config=None) -> "ParquetBlockModel":
"""
Upsample a block model to a finer grid with explicit per-attribute methods.
This function supports both continuous interpolation and class-preserving
assignment methods.
Args:
blockmodel: ParquetBlockModel instance.
new_block_size: tuple of floats (dx, dy, dz) for the new block size.
upsample_config: dict mapping each attribute name to one method
("linear", "nearest", "mode", or "parent").
interpolation_config: deprecated alias of ``upsample_config``.
Example:
upsample_config = {
'grade': 'linear',
'density': 'linear',
'dry_mass': 'linear',
'volume': 'linear',
'rock_type': 'mode'
}
Returns:
ParquetBlockModel: A new ParquetBlockModel instance with the upsampled grid.
"""
from parq_blockmodel import ParquetBlockModel
if upsample_config is not None and interpolation_config is not None:
raise ValueError("Provide only one of 'upsample_config' or deprecated 'interpolation_config'.")
effective_config = upsample_config if upsample_config is not None else interpolation_config
_validate_params(effective_config, new_block_size)
arrays, categories = _prepare_arrays_and_index(blockmodel)
fx, fy, fz, new_shape = _calculate_factors(blockmodel, new_block_size)
if fx > 1 or fy > 1 or fz > 1:
raise ValueError("reblocking factors imply downsampling, not upsampling.")
elif fx < 1 or fy < 1 or fz < 1:
# Upsampling: expect dict of strings
if not all(isinstance(v, str) for v in (effective_config or {}).values()):
raise ValueError("Upsampling config must be a dict of interpolation methods per attribute.")
missing = sorted(set(arrays) - set((effective_config or {})))
if missing:
raise ValueError(
"Upsampling config must specify a method for every attribute. "
f"Missing: {missing}"
)
arrays = upsample_attributes(arrays, int(1 // fx), int(1 // fy), int(1 // fz), effective_config or {})
elif any(dim < 1 for dim in (fx, fy, fz)) and any(dim > 1 for dim in (fx, fy, fz)):
raise ValueError("Reblocking factors cannot specify upsample and downsample at the same time.")
# Convert back to DataFrame (ijk-indexed attributes only)
new_geometry = _build_new_geometry(blockmodel, new_block_size, new_shape)
# Use canonical .pbm suffix expected by ParquetBlockModel
new_path = blockmodel.blockmodel_path.with_name(f"{blockmodel.name}.reblocked.pbm")
reblocked_df = dict_3d_to_tabular(arrays, categories)
# assert the indexes are aligned in ijk space
if not reblocked_df.index.equals(new_geometry.to_multi_index_ijk()):
raise ValueError("Upsampling ijk index mismatch.")
# ------------------------------------------------------------------
# Reconstruct geometry-derived columns for the new grid
# ------------------------------------------------------------------
reblocked_df["x"] = new_geometry.centroid_x
reblocked_df["y"] = new_geometry.centroid_y
reblocked_df["z"] = new_geometry.centroid_z
N = int(np.prod(new_geometry.local.shape))
rows = np.arange(N, dtype=np.uint32)
reblocked_df["block_id"] = rows
i, j, k = new_geometry.ijk_from_row_index(rows)
reblocked_df["i"] = i.astype(np.int32)
reblocked_df["j"] = j.astype(np.int32)
reblocked_df["k"] = k.astype(np.int32)
if new_geometry.world_id_encoding is None:
new_geometry.world_id_encoding = build_world_id_encoding_from_xyz(
reblocked_df["x"].to_numpy(dtype=float),
reblocked_df["y"].to_numpy(dtype=float),
reblocked_df["z"].to_numpy(dtype=float),
)
offset, scale, bits_per_axis = get_world_id_encoding_params(new_geometry.world_id_encoding)
reblocked_df["world_id"] = encode_world_coordinates(
reblocked_df["x"].to_numpy(dtype=float),
reblocked_df["y"].to_numpy(dtype=float),
reblocked_df["z"].to_numpy(dtype=float),
offset=offset,
scale=scale,
bits_per_axis=bits_per_axis,
).astype(np.int64)
ordered_cols = ["block_id", "world_id", "x", "y", "z", "i", "j", "k"]
ordered_cols.extend([c for c in reblocked_df.columns if c not in ordered_cols])
table = pa.Table.from_pandas(reblocked_df[ordered_cols].reset_index(drop=True), preserve_index=False)
meta = dict(table.schema.metadata or {})
meta[b"parq-blockmodel"] = json.dumps(new_geometry.to_metadata_dict()).encode("utf-8")
table = table.replace_schema_metadata(meta)
pq.write_table(table, new_path)
return ParquetBlockModel(blockmodel_path=new_path, geometry=new_geometry)
