# geoML - machine learning models for geospatial data
# Copyright (C) 2026 Ítalo Gomes Gonçalves
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR a PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Array storage backend for data objects.
``ArrayStore`` wraps a single array that is held either in RAM (NumPy) or on
disk in chunks (Zarr), and exposes a small, NumPy-compatible surface so callers
can treat it like an ``ndarray`` regardless of where the data lives:
- region reads/writes: ``store[idx]`` and ``store[idx] = value`` touch only the
affected chunks, which is what the batched prediction write path needs;
- full materialization on demand: ``numpy.asarray(store)`` (used implicitly by
``numpy`` ufuncs, ``reshape``, filters, ...);
- a lazy labelled view: :meth:`ArrayStore.as_xarray` (dask-backed for Zarr) for
out-of-core reductions and export.
The backend is chosen by size: arrays whose in-RAM footprint would exceed
``DEFAULT_THRESHOLD`` spill to a chunked Zarr array, everything else stays in
NumPy. This module is deliberately independent of ``data.py`` so it can be
tested in isolation.
"""
__all__ = ["ArrayStore", "DEFAULT_THRESHOLD", "store_columns",
"set_realization_dtype", "realization_dtype"]
import os as _os
import shutil as _shutil
from collections.abc import Sequence
import tempfile as _tempfile
import weakref as _weakref
from typing import Any as _Any
import numpy as _np
import zarr as _zarr
from zarr.codecs import Shuffle as _Shuffle, ZstdCodec as _Zstd
import geoml._types as _types
import dask.array as _da
import xarray as _xr
# Arrays whose uncompressed in-RAM size would exceed this many bytes are stored
# on disk (Zarr) rather than in NumPy. Keeps ordinary point data in RAM while
# large grids / simulation cubes spill to disk.
DEFAULT_THRESHOLD = 50 * 1024 ** 2 # 50 MB
# Target uncompressed size of a single Zarr chunk.
_TARGET_CHUNK_BYTES = 8 * 1024 ** 2 # 8 MB
# Past this many realizations the trailing axis is split as well. Chunking
# the location axis alone puts every realization of a band of rows in one
# chunk, which is what the row-wise reductions want -- and it makes reading
# ONE realization read them all: `simulation(i)` on a (5 000 000, 100) store
# walked the whole 4 GB to hand back 40 MB. Below the threshold the store is
# small enough that the pass costs nothing worth a second chunk axis.
_MIN_SPLIT_COLUMNS = 32
# How many realizations a chunk holds once the axis is split: the share of
# the store one realization's read must visit is this over the total.
# Measured cold (page cache dropped) on a (2 000 000, 100) float64 store,
# 1.49 GB, at 100, 25 and 10 columns a chunk: reading one realization 0.76,
# 0.13 and 0.06 s, and the reductions no worse for it -- a quantile pass
# 1.34, 1.09 and 1.26 s, a pass in row bands 2.53, 1.19 and 1.07 s, since
# ten columns a chunk is ten times the rows and so fewer, longer reads.
# Ten wins on every measure; `docs/benchmarks/realization_chunks.py`.
_COLUMNS_PER_CHUNK = 10
# What realizations are stored as. The model computes in float64 and every
# read of a float32 store widens back to float64, so the narrower type is a
# matter of disk and memory only: seven significant digits are far finer
# than the spread between realizations, and a store of them takes half the
# room before compression and less after (GeoScape's item 40).
_REALIZATION_DTYPE = _np.dtype(_np.float32)
[docs]
def set_realization_dtype(dtype: _Any) -> None:
"""Set what realizations are stored as from now on.
Realizations are stored as float32 by default; the model computes in
float64 whatever they are stored as, and a store of float32 reads back
as float64. A store keeps the type it was made with.
Parameters
----------
dtype
`"float32"` (the default) or `"float64"`.
Raises
------
ValueError
For any other type.
"""
global _REALIZATION_DTYPE
dtype = _np.dtype(dtype)
if dtype not in (_np.dtype(_np.float32), _np.dtype(_np.float64)):
raise ValueError("realizations are stored as float32 or float64, "
"not %s" % dtype)
_REALIZATION_DTYPE = dtype
[docs]
def realization_dtype() -> _np.dtype:
"""What realizations are stored as now; see `set_realization_dtype`."""
return _REALIZATION_DTYPE
def _widened(values):
"""`values` as the model computes with them: float32 read as float64."""
if getattr(values, "dtype", None) == _np.float32:
return values.astype(_np.float64)
return values
def _compressors(dtype):
"""The codecs after `bytes`: a byte shuffle at the element's width, then
zstd. Shuffled, the bytes of equal significance sit together, which is
what lets floats compress; GeoScape's reader undoes `numcodecs.shuffle`
without a dependency, which Blosc's built-in shuffle would need. A
one-byte type has nothing to shuffle."""
dtype = _np.dtype(dtype)
if dtype.kind in "biufc" and dtype.itemsize > 1:
return [_Shuffle(elementsize=dtype.itemsize), _Zstd()]
return "auto"
def _leading_chunk(shape, dtype):
"""Chunk shape targeting ``_TARGET_CHUNK_BYTES``.
Splits the leading (data-location) axis, so batched writes align to whole
chunks and a band of rows is complete. A 2-D store holding more than
``_MIN_SPLIT_COLUMNS`` realizations splits the trailing axis as well, so
that reading one realization reads a share of the store rather than all
of it; the row-wise reductions rechunk a band back to whole rows, which
is one pass over the same bytes.
"""
itemsize = _np.dtype(dtype).itemsize
trailing = int(_np.prod(shape[1:])) if len(shape) > 1 else 1
columns = tuple(int(s) for s in shape[1:])
if len(shape) == 2 and shape[1] > _MIN_SPLIT_COLUMNS:
columns = (min(int(shape[1]), _COLUMNS_PER_CHUNK),)
trailing = columns[0]
row_bytes = max(trailing * itemsize, 1)
rows = max(1, _TARGET_CHUNK_BYTES // row_bytes)
if shape[0] > 0:
rows = min(rows, shape[0])
return (int(rows),) + columns
def _whole_rows(darr):
"""`darr` with its trailing axis in one chunk, for a row-wise reduction.
A reduction across realizations must see all of a row's at once, and
since 0.7.0 a wide store is chunked on that axis too. Rechunking reads
the same bytes in the same pass; it is a no-op on a store that was
never split.
"""
if darr.ndim == 2 and len(darr.chunks[1]) > 1:
return darr.rechunk({1: -1})
return darr
def _zarr_dtype(dtype):
"""`dtype` under the class NumPy names it by, which is the class Zarr
matches. A C `long long` array on Linux is 64-bit integers of another
dtype class, printing as int64 -- VTK 9.7 hands its integer arrays back
that way -- and Zarr refuses it ("no Zarr data type found that matches
dtype('int64')"). Rebuilt from its type string, a number's dtype keeps
its bytes and takes the canonical class."""
dtype = _np.dtype(dtype)
return _np.dtype(dtype.str) if dtype.kind in "biufc" else dtype
def _rows_of_masks(item):
"""A selection with every boolean mask in it turned into the indices it
keeps. Zarr takes a mask only of the array's whole shape, so a mask of
rows on a 2-D store -- what subsetting a container opened from its store
hands down -- was refused (`VindexInvalidSelectionError`); the indices
select the same rows."""
def one(part):
if isinstance(part, _np.ndarray) and part.dtype == bool \
and part.ndim == 1:
return _np.flatnonzero(part)
return part
if isinstance(item, tuple):
return tuple(one(part) for part in item)
return one(item)
def _use_zarr(shape, dtype, threshold):
"""Whether an array of this shape/dtype should live on disk."""
if _np.dtype(dtype) == object:
# object arrays (categorical labels, ...) stay in NumPy; Zarr would need
# a variable-length codec and they are small in practice.
return False
nbytes = int(_np.prod(shape)) * _np.dtype(dtype).itemsize
return nbytes > threshold
class _ScratchGroup:
"""An owner's consolidated scratch store.
One temporary directory holding a single Zarr group into which all of the
owner's large working arrays are allocated (instead of one temp directory
per array). The whole directory is removed when this object dies — which,
through the weak registry below, happens when the owning container is
garbage-collected. Arrays inside must not outlive their owner.
"""
def __init__(self):
self._tempdir = _tempfile.mkdtemp(prefix="geoml_scratch_")
self.path = _os.path.join(self._tempdir, "scratch.zarr")
self._group: _Any = _zarr.open_group(self.path, mode="w")
self._count = 0
def create_array(self, shape, dtype, fill_value, chunks):
name = "a%d" % self._count
self._count += 1
return self._group.create_array(
name=name, shape=shape, chunks=chunks, dtype=dtype,
fill_value=fill_value, compressors=_compressors(dtype))
def close(self):
self._group = None
if self._tempdir is not None and _os.path.isdir(self._tempdir):
_shutil.rmtree(self._tempdir, ignore_errors=True)
self._tempdir = None
def __del__(self):
try:
self.close()
except Exception:
pass
# owner (e.g. a data container) -> its _ScratchGroup. Weak keys: when the owner
# is collected the group is dropped and its directory deleted. Deep copies of
# an owner are not in the registry (and their stores were materialized to
# NumPy by ArrayStore.__deepcopy__), so no double-delete can occur.
_scratch_groups = _weakref.WeakKeyDictionary()
def _scratch_for(owner):
group = _scratch_groups.get(owner)
if group is None:
group = _ScratchGroup()
_scratch_groups[owner] = group
return group
[docs]
def store_columns(columns, stores: "Sequence[ArrayStore]") -> None:
"""Write each column of a lazy 2-D dask array into its target store.
All columns are computed in a single chunk-by-chunk pass over the source;
the targets may be NumPy- or Zarr-backed.
Parameters
----------
columns
A two-dimensional dask array, one column per target.
stores
One store per column, in the same order.
"""
_da.store([columns[:, i] for i in range(columns.shape[1])],
[s._array for s in stores], lock=False)
[docs]
class ArrayStore:
"""A single array backed by NumPy (in RAM) or Zarr (on disk, chunked)."""
# Always present -- a NumPy array or a Zarr one, as `_backend` says.
_array: _Any
def __init__(self, array, backend: str, store_path=None,
_tempdir=None):
# Low-level constructor; prefer the ``from_numpy`` / ``allocate`` /
# ``open`` factories.
self._array = array
self._backend = backend # "numpy" or "zarr"
self._store_path = store_path # on-disk location, if any
self._tempdir = _tempdir # root to clean up when we own it
# ------------------------------------------------------------------ #
# construction
# ------------------------------------------------------------------ #
[docs]
@classmethod
def from_numpy(cls, values: _types.ArrayLike) -> "ArrayStore":
"""Wrap an existing array in a NumPy-backed store (no copy)."""
return cls(_np.asarray(values), backend="numpy")
[docs]
@classmethod
def from_values(cls, values: _types.ArrayLike, owner=None,
threshold: int | None = None) -> "ArrayStore":
"""Store an existing array, spilling to disk when it is large.
Unlike :meth:`from_numpy`, which always keeps the array in RAM, the
backend is chosen by size as in :meth:`allocate`. Use this for arrays a
container owns for its whole life (coordinates, input variance) so a
large one does not pin memory.
"""
values = _np.asarray(values)
if threshold is None:
threshold = DEFAULT_THRESHOLD
if not _use_zarr(values.shape, values.dtype, threshold):
return cls(values, backend="numpy")
store = cls.allocate(values.shape, dtype=values.dtype, fill_value=0,
backend="zarr", owner=owner)
store[...] = values
return store
[docs]
@classmethod
def allocate(cls, shape, dtype: _Any = float, fill_value=_np.nan,
chunks=None, backend: str = "auto", store=None,
threshold: int | None = None, owner=None) -> "ArrayStore":
"""Create a new, filled array.
Parameters
----------
shape : tuple
Full array shape; axis 0 is the data-location axis.
dtype : data-type
fill_value : scalar
Initial value for every element (``nan`` by default).
chunks : tuple, optional
Zarr chunk shape. Defaults to splitting axis 0 only.
backend : {"auto", "numpy", "zarr"}
``"auto"`` picks Zarr past ``threshold`` bytes, NumPy otherwise.
store : str or zarr store, optional
Where a Zarr array lives. If omitted, a temporary location is
used (see ``owner``).
threshold : int
Size in bytes above which ``"auto"`` chooses Zarr.
owner : object, optional
Scratch-lifecycle owner (typically the data container). Temporary
Zarr arrays of the same owner are consolidated into one on-disk
store, deleted when the owner is garbage-collected. Without an
owner (and without ``store``) the array gets its own temporary
directory, cleaned up with this object.
"""
shape = tuple(int(s) for s in _np.atleast_1d(shape))
if threshold is None:
# Read at call time so the module-level default stays configurable.
threshold = DEFAULT_THRESHOLD
if backend == "auto":
backend = "zarr" if _use_zarr(shape, dtype, threshold) else "numpy"
if backend == "numpy":
return cls(_np.full(shape, fill_value, dtype=dtype), backend="numpy")
if backend != "zarr":
raise ValueError(f"unknown backend '{backend}'")
if chunks is None:
chunks = _leading_chunk(shape, dtype)
if store is None and owner is not None:
scratch = _scratch_for(owner)
array = scratch.create_array(
shape, _zarr_dtype(dtype), fill_value, chunks)
return cls(array, backend="zarr", store_path=scratch.path)
tempdir = None
if store is None:
tempdir = _tempfile.mkdtemp(prefix="geoml_zarr_")
store = _os.path.join(tempdir, "array.zarr")
store_path = store if isinstance(store, str) else getattr(store, "path", None)
array = _zarr.create_array(
store=store, shape=shape, chunks=chunks,
dtype=_zarr_dtype(dtype), fill_value=fill_value,
compressors=_compressors(dtype))
return cls(array, backend="zarr", store_path=store_path, _tempdir=tempdir)
[docs]
@classmethod
def open(cls, path: _types.PathLike,
mode: str = "r+") -> "ArrayStore":
"""Reopen an existing on-disk Zarr array."""
return cls(_zarr.open_array(path, mode=mode), backend="zarr",
store_path=path)
[docs]
@classmethod
def wrap_zarr(cls, zarr_array) -> "ArrayStore":
"""Wrap an already-open Zarr array (e.g. a child of a reopened group)."""
return cls(zarr_array, backend="zarr",
store_path=getattr(zarr_array, "store_path", None))
[docs]
def write_into(self, group, name: str) -> None:
"""Stream this store into a new array ``name`` of an open Zarr group.
The copy is chunk-by-chunk via dask, so a large on-disk source is never
fully materialized. Returns the created Zarr array.
"""
if self._backend == "zarr":
chunks = self._array.chunks
else:
chunks = _leading_chunk(self.shape, self.dtype)
fill = _np.nan if _np.issubdtype(_np.dtype(self.dtype), _np.floating) else 0
target = group.create_array(
name=name, shape=self.shape, chunks=chunks,
dtype=_zarr_dtype(self.dtype), fill_value=fill,
compressors=_compressors(self.dtype))
_da.store(self.as_dask(), target, lock=False)
return target
# ------------------------------------------------------------------ #
# ndarray-compatible surface
# ------------------------------------------------------------------ #
def __getitem__(self, item):
if self._backend == "zarr":
item = _rows_of_masks(item)
return _widened(self._array[item])
def __setitem__(self, item, value):
if self._backend == "zarr":
if self._array.read_only:
# zarr refuses too, but without saying what to do instead
raise ValueError(
"%s was opened read-only; open it with mode='r+' to "
"write into it, or write the container to another "
"store with to_zarr" % (self._store_path or "this store"))
item = _rows_of_masks(item)
self._array[item] = value
def _stored(self):
"""The values as stored, not widened: what a copy keeps."""
if self._backend == "zarr":
return _np.asarray(self._array[...])
return _np.asarray(self._array)
def __array__(self, dtype=None, copy=None):
array = self._stored()
made_copy = self._backend == "zarr"
if array.dtype == _np.float32:
array = _widened(array)
made_copy = True
if dtype is not None and array.dtype != _np.dtype(dtype):
array = array.astype(dtype)
made_copy = True
if copy and not made_copy:
array = array.copy()
return array
@property
def shape(self):
return tuple(self._array.shape)
@property
def dtype(self):
"""The type stored, which a store made like this one takes; a
float32 store reads back as float64."""
return self._array.dtype
@property
def ndim(self):
return len(self._array.shape)
@property
def size(self):
return int(_np.prod(self._array.shape))
def __len__(self):
return int(self._array.shape[0])
[docs]
def copy(self) -> _np.ndarray:
"""Materialize to a fresh NumPy array.
Returns an array rather than another store: the callers of this
want the values in hand, and `__copy__` is what makes an
independent store.
"""
return _np.array(self.__array__())
[docs]
def ravel(self):
return self.__array__().ravel()
[docs]
def to_numpy(self) -> _np.ndarray:
return self.__array__()
def __eq__(self, other):
return self.__array__() == other
def __ne__(self, other):
return self.__array__() != other
__hash__ = None # type: ignore[assignment] # unhashable, by design
def __repr__(self):
return f"ArrayStore(backend={self._backend!r}, shape={self.shape}, " \
f"dtype={self.dtype})"
def __copy__(self):
# Copies are always independent NumPy-backed stores: sharing a Zarr
# array (and its temp directory) across stores would risk double-free.
# They keep the type stored, not the one a read widens to.
return ArrayStore.from_numpy(_np.array(self._stored()))
def __deepcopy__(self, memo):
return ArrayStore.from_numpy(_np.array(self._stored()))
# ------------------------------------------------------------------ #
# labelled / lazy views
# ------------------------------------------------------------------ #
[docs]
def as_dask(self) -> "_da.Array":
"""A dask array view (lazy & chunked for Zarr, single-chunk for NumPy)."""
if self._backend == "zarr":
darr = _da.from_array(self._array, chunks=self._array.chunks)
else:
darr = _da.from_array(self._array, chunks=-1)
return _widened(darr)
[docs]
def as_xarray(self, dims=None, coords=None, name=None):
"""A labelled ``xarray.DataArray`` over this store (dask-backed)."""
return _xr.DataArray(self.as_dask(), dims=dims, coords=coords, name=name)
[docs]
def row_bands(self, rows: int | None = None) -> "list[slice]":
"""Slices covering axis 0, each one holding whole chunks.
Reading a store a band at a time is what keeps a reduction over
locations flat in memory. A band is a whole number of chunks along
axis 0 and every row in it is complete -- reading the band reads
each of its column chunks, where the trailing axis is split -- so a
reduction across simulations sees all of a location's at once, and
nothing has to be stitched back together afterwards.
A NumPy-backed store is already in RAM and comes back as a single band,
so a caller written this way costs nothing on small data.
"""
n_rows = int(self.shape[0])
if rows is None:
rows = int(self._array.chunks[0]) \
if self._backend == "zarr" else n_rows
band = max(1, int(rows))
return [slice(lo, min(lo + band, n_rows))
for lo in range(0, n_rows, band)]
[docs]
def row_quantiles(self, qs: _types.ArrayLike) -> "list[_da.Array]":
"""Lazy row-wise quantiles of a 2-D store.
Returns an uncomputed dask array of shape ``(n_rows, len(qs))``.
The realization axis is gathered into one chunk first, so every
block holds complete rows: the quantiles are exact and the full
store is never materialized, at one pass over the same bytes.
"""
darr = _whole_rows(self.as_dask())
qs = _np.atleast_1d(qs).astype(float)
def block_quantiles(block):
return _np.quantile(block, qs, axis=1).T
return darr.map_blocks(
block_quantiles, dtype=_np.float64,
chunks=(darr.chunks[0], len(qs)))
[docs]
def row_cdf(self, cutoffs):
"""Lazy row-wise empirical CDF of a 2-D store.
For each cutoff, the fraction of columns (simulations) at or below it
— the inverse view of :meth:`row_quantiles`. Returns an uncomputed
dask array of shape ``(n_rows, len(cutoffs))`` with values in [0, 1].
"""
darr = _whole_rows(self.as_dask())
cutoffs = _np.atleast_1d(cutoffs).astype(float)
def block_cdf(block):
return _np.mean(
block[:, :, None] <= cutoffs[None, None, :], axis=1)
return darr.map_blocks(
block_cdf, dtype=_np.float64,
chunks=(darr.chunks[0], len(cutoffs)))
# ------------------------------------------------------------------ #
# backend info / lifecycle
# ------------------------------------------------------------------ #
@property
def backend(self):
return self._backend
@property
def store_path(self):
return self._store_path
[docs]
def close(self):
"""Release the array and delete the temp store if we created it."""
self._array = None
if self._tempdir is not None and _os.path.isdir(self._tempdir):
_shutil.rmtree(self._tempdir, ignore_errors=True)
self._tempdir = None
def __del__(self):
try:
self.close()
except Exception:
pass