geoml.storage

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: 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.

class geoml.storage.ArrayStore(array, backend, store_path=None, _tempdir=None)

Bases: object

A single array backed by NumPy (in RAM) or Zarr (on disk, chunked).

classmethod from_numpy(values)

Wrap an existing array in a NumPy-backed store (no copy).

Return type:

ArrayStore

classmethod from_values(values, owner=None, threshold=None)

Store an existing array, spilling to disk when it is large.

Unlike from_numpy(), which always keeps the array in RAM, the backend is chosen by size as in allocate(). Use this for arrays a container owns for its whole life (coordinates, input variance) so a large one does not pin memory.

Return type:

ArrayStore

classmethod allocate(shape, dtype=<class 'float'>, fill_value=nan, chunks=None, backend='auto', store=None, threshold=None, owner=None)

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.

Return type:

ArrayStore

classmethod open(path, mode='r+')

Reopen an existing on-disk Zarr array.

Return type:

ArrayStore

classmethod wrap_zarr(zarr_array)

Wrap an already-open Zarr array (e.g. a child of a reopened group).

Return type:

ArrayStore

write_into(group, name)

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.

property shape

property dtype

property ndim

property size

copy()

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 type:

ndarray

ravel()

to_numpy()

Return type:

ndarray

as_dask()

A dask array view (lazy & chunked for Zarr, single-chunk for NumPy).

Return type:

Array

as_xarray(dims=None, coords=None, name=None)

A labelled xarray.DataArray over this store (dask-backed).

row_bands(rows=None)

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.

Return type:

list[slice]

row_quantiles(qs)

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.

Return type:

list[Array]

row_cdf(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 row_quantiles(). Returns an uncomputed dask array of shape (n_rows, len(cutoffs)) with values in [0, 1].

property backend

property store_path

close()

Release the array and delete the temp store if we created it.

geoml.storage.store_columns(columns, stores)

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 (Sequence[ArrayStore]) – One store per column, in the same order.