# geoML - machine learning models for geospatial data
# Copyright (C) 2019 Í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"]
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
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. Chunking is done along the
# leading (data-location) axis only, so batched writes align to whole chunks.
_TARGET_CHUNK_BYTES = 8 * 1024 ** 2 # 8 MB
def _leading_chunk(shape, dtype):
"""Chunk shape that splits only axis 0, targeting ``_TARGET_CHUNK_BYTES``."""
itemsize = _np.dtype(dtype).itemsize
trailing = int(_np.prod(shape[1:])) if len(shape) > 1 else 1
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),) + tuple(int(s) for s in shape[1:])
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)
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, _np.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=_np.dtype(dtype), fill_value=fill_value)
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=_np.dtype(self.dtype), fill_value=fill)
_da.store(self.as_dask(), target, lock=False)
return target
# ------------------------------------------------------------------ #
# ndarray-compatible surface
# ------------------------------------------------------------------ #
def __getitem__(self, item):
return self._array[item]
def __setitem__(self, item, value):
self._array[item] = value
def __array__(self, dtype=None, copy=None):
if self._backend == "zarr":
array = _np.asarray(self._array[...])
made_copy = True
else:
array = _np.asarray(self._array)
made_copy = False
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):
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.
return ArrayStore.from_numpy(self.__array__().copy())
def __deepcopy__(self, memo):
return ArrayStore.from_numpy(self.__array__().copy())
# ------------------------------------------------------------------ #
# 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":
return _da.from_array(self._array, chunks=self._array.chunks)
return _da.from_array(self._array, chunks=-1)
[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. Chunking splits axis 0 only, so a band is a
whole number of chunks and every row in it is complete: 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))``.
Because chunking splits only axis 0, every chunk holds complete rows,
so the quantiles are exact and the full store is never materialized.
"""
darr = 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 = 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