atlas-python 0.9.0__cp310-abi3-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
atlas/xarray.py ADDED
@@ -0,0 +1,578 @@
1
+ """xarray integration for atlas.
2
+
3
+ `xarray` and `dask` are required dependencies.
4
+
5
+ Reads automatically return dask-backed variables when an array was stored with a
6
+ non-trivial chunk shape (`chunk_shape != shape`); full-shape arrays come back
7
+ eager as numpy. The dask chunks mirror the on-disk chunk grid one-to-one.
8
+
9
+ Bulk ingestion — Atlas itself batches; explicit flush on the atlas persists:
10
+ with atlas:
11
+ for nc_path in nc_paths:
12
+ ds = xr.open_dataset(nc_path)
13
+ atlas.add_xr_dataset(ds, name=Path(nc_path).stem)
14
+ # atlas.close() (== flush) runs on __exit__.
15
+ """
16
+ from __future__ import annotations
17
+
18
+ import itertools
19
+ import json
20
+ import time
21
+ from concurrent.futures import Future, ThreadPoolExecutor
22
+ from typing import TYPE_CHECKING, Any, Iterator, Optional, Sequence
23
+
24
+ import numpy as np
25
+
26
+ from ._atlas import log_chunk_event as _log_chunk_event
27
+
28
+ if TYPE_CHECKING:
29
+ import xarray as xr
30
+
31
+ from ._atlas import Atlas, DatasetView
32
+
33
+
34
+ _COORDS_ATTR = "_pyatlas_coords"
35
+ _JSON_PREFIX = "json:"
36
+
37
+ _NUMPY_TO_ATLAS = {
38
+ np.dtype("int8"): "int8",
39
+ np.dtype("int16"): "int16",
40
+ np.dtype("int32"): "int32",
41
+ np.dtype("int64"): "int64",
42
+ np.dtype("uint8"): "uint8",
43
+ np.dtype("uint16"): "uint16",
44
+ np.dtype("uint32"): "uint32",
45
+ np.dtype("uint64"): "uint64",
46
+ np.dtype("float32"): "float32",
47
+ np.dtype("float64"): "float64",
48
+ np.dtype("datetime64[ns]"): "timestamp_nanoseconds",
49
+ }
50
+
51
+
52
+ def _np_to_atlas_dtype(np_dtype: np.dtype) -> str:
53
+ if np_dtype in _NUMPY_TO_ATLAS:
54
+ return _NUMPY_TO_ATLAS[np_dtype]
55
+ # Object (Python str/bytes) and fixed-size byte/unicode strings all
56
+ # become variable-length atlas strings.
57
+ if np_dtype.kind in ("O", "S", "U"):
58
+ return "string"
59
+ supported = ", ".join(sorted(set(_NUMPY_TO_ATLAS.values())))
60
+ raise NotImplementedError(
61
+ f"numpy dtype {np_dtype!r} is not supported by atlas "
62
+ f"(supported: {supported}, plus object/bytes/unicode → string)"
63
+ )
64
+
65
+
66
+ def _sanitize_str(s: str) -> str:
67
+ """Strip lone Unicode surrogates from a Python str.
68
+
69
+ NetCDF backends often surface byte attrs as Python strs that were decoded
70
+ with ``errors='surrogateescape'``; the resulting strs hold pseudo-codepoints
71
+ in U+DC80..U+DCFF which Rust's UTF-8 strs can't represent. We try to recover
72
+ the original bytes via surrogateescape and re-decode as UTF-8 (the common
73
+ case: bytes that *were* valid UTF-8 all along but were treated as Latin-1
74
+ upstream); if that fails we fall back to lossy replacement.
75
+ """
76
+ try:
77
+ s.encode("utf-8")
78
+ return s
79
+ except UnicodeEncodeError:
80
+ pass
81
+ raw = s.encode("utf-8", errors="surrogateescape")
82
+ try:
83
+ return raw.decode("utf-8")
84
+ except UnicodeDecodeError:
85
+ return raw.decode("utf-8", errors="replace")
86
+
87
+
88
+ def _encode_attr_value(value: Any) -> Any:
89
+ """Coerce an xarray attr value to something atlas can store.
90
+
91
+ Primitives (bool/int/float/str) are returned as-is. Everything else is
92
+ JSON-encoded and prefixed with ``json:`` so it can be decoded losslessly
93
+ on read. Numpy scalars are unwrapped first. Numpy arrays are converted to
94
+ lists. Values that can't be JSON-serialised raise ``TypeError``.
95
+ """
96
+ # Unwrap numpy scalars
97
+ if isinstance(value, np.generic):
98
+ value = value.item()
99
+
100
+ if isinstance(value, str):
101
+ return _sanitize_str(value)
102
+ if isinstance(value, (bool, int, float)):
103
+ return value
104
+
105
+ # Convert numpy arrays to nested lists
106
+ if isinstance(value, np.ndarray):
107
+ value = value.tolist()
108
+
109
+ try:
110
+ return _JSON_PREFIX + json.dumps(value)
111
+ except (TypeError, ValueError) as exc:
112
+ raise TypeError(
113
+ f"attribute value of type {type(value).__name__} is not JSON-serialisable "
114
+ f"and cannot be stored: {value!r}"
115
+ ) from exc
116
+
117
+
118
+ def _decode_attr_value(value: Any) -> Any:
119
+ """Inverse of :func:`_encode_attr_value`."""
120
+ if isinstance(value, str) and value.startswith(_JSON_PREFIX):
121
+ return json.loads(value[len(_JSON_PREFIX):])
122
+ return value
123
+
124
+
125
+ def _normalize_fill_value(value: Any, np_dtype: np.dtype) -> Any:
126
+ """Coerce an xarray `_FillValue` attribute into a Python scalar matching the array dtype.
127
+
128
+ The binding's `define_array(fill_value=...)` expects a plain Python scalar
129
+ that is type-consistent with the array dtype. xarray often stores `_FillValue`
130
+ as a 0-D numpy array or numpy scalar; this unwraps that and (for datetime64
131
+ arrays) reinterprets the value as nanoseconds since the epoch. For
132
+ string-kind arrays (object/bytes/unicode), `bytes` are decoded to `str`
133
+ since NetCDF stores fixed-width string fill values as bytes.
134
+ """
135
+ if value is None:
136
+ return None
137
+ if isinstance(value, np.ndarray) and value.ndim == 0:
138
+ value = value.item() if np_dtype.kind != "M" else value.view(np.int64).item()
139
+ elif isinstance(value, np.generic):
140
+ if isinstance(value, np.datetime64):
141
+ return value.astype("datetime64[ns]").view(np.int64).item()
142
+ value = value.item()
143
+ if np_dtype.kind in ("O", "S", "U") and isinstance(value, bytes):
144
+ return value.decode("utf-8", errors="replace")
145
+ return value
146
+
147
+
148
+ def _is_dask_array(arr: Any) -> bool:
149
+ """Return True if `arr` is a `dask.array.Array`. False if dask isn't installed."""
150
+ try:
151
+ import dask.array as da
152
+ except ImportError:
153
+ return False
154
+ return isinstance(arr, da.Array)
155
+
156
+
157
+ def _dask_chunk_shape(arr: Any) -> list[int]:
158
+ """First chunk size along each dim of a dask array — used as the atlas chunk_shape."""
159
+ return [c[0] for c in arr.chunks]
160
+
161
+
162
+ def _contiguous(a: np.ndarray) -> np.ndarray:
163
+ # np.ascontiguousarray promotes 0-D arrays to 1-D (ndmin=1 default),
164
+ # which breaks scalar-array writes. Force-copy non-contiguous arrays
165
+ # via np.asarray + .copy(order='C') instead, which preserves rank.
166
+ if a.flags["C_CONTIGUOUS"]:
167
+ return a
168
+ return a.copy(order="C")
169
+
170
+
171
+ # Tuned for the "many small NetCDF chunks" case. batch_size controls how many
172
+ # dask blocks ride one scheduler invocation (lower dask plumbing overhead);
173
+ # prefetch_depth bounds peak memory at batch_size * prefetch_depth chunks.
174
+ _DEFAULT_BATCH_SIZE = 8
175
+ _DEFAULT_PREFETCH_DEPTH = 2
176
+
177
+
178
+ def _iter_blocks(
179
+ arr: Any,
180
+ var_name: str = "",
181
+ batch_size: int = _DEFAULT_BATCH_SIZE,
182
+ prefetch_depth: int = _DEFAULT_PREFETCH_DEPTH,
183
+ ) -> Iterator[tuple[list[int], np.ndarray]]:
184
+ """Yield ``(start_index, block_np)`` for each chunk in `arr`.
185
+
186
+ Numpy-backed inputs are yielded as one full-shape block. Dask-backed inputs
187
+ are prefetched in batches: a single background thread runs
188
+ ``dask.compute(*K)`` on the next ``batch_size`` blocks while the main thread
189
+ consumes already-materialised blocks, capping in-flight work at
190
+ ``prefetch_depth`` batches. This collapses N per-chunk dask scheduler
191
+ invocations into N/batch_size, and overlaps NetCDF I/O with atlas writes.
192
+
193
+ Emits one ``event=dask_compute`` debug event per fulfilled batch via the
194
+ Rust tracing subscriber (see ``log_chunk_event``).
195
+ """
196
+ if not _is_dask_array(arr):
197
+ a = np.asarray(arr)
198
+ yield [0] * a.ndim, _contiguous(a)
199
+ return
200
+
201
+ import dask
202
+
203
+ chunks = arr.chunks
204
+ offsets = [[0, *itertools.accumulate(c)][:-1] for c in chunks]
205
+ pairs: list[tuple[list[int], Any]] = []
206
+ for block_idx in itertools.product(*[range(len(c)) for c in chunks]):
207
+ start = [offsets[d][i] for d, i in enumerate(block_idx)]
208
+ pairs.append((start, arr.blocks[block_idx]))
209
+
210
+ def _fetch(batch: list[tuple[list[int], Any]]) -> list[tuple[list[int], np.ndarray]]:
211
+ starts = [s for s, _ in batch]
212
+ delayed = [b for _, b in batch]
213
+ t0 = time.perf_counter_ns()
214
+ materialised = dask.compute(*delayed)
215
+ elapsed_us = (time.perf_counter_ns() - t0) // 1000
216
+ _log_chunk_event("dask_compute", var_name, elapsed_us, chunks=len(batch))
217
+ return [(s, _contiguous(b)) for s, b in zip(starts, materialised)]
218
+
219
+ with ThreadPoolExecutor(max_workers=1) as executor:
220
+ in_flight: list[Future] = []
221
+ cursor = 0
222
+ while cursor < len(pairs) and len(in_flight) < prefetch_depth:
223
+ in_flight.append(executor.submit(_fetch, pairs[cursor : cursor + batch_size]))
224
+ cursor += batch_size
225
+ while in_flight:
226
+ future = in_flight.pop(0)
227
+ for start, block in future.result():
228
+ yield start, block
229
+ if cursor < len(pairs):
230
+ in_flight.append(
231
+ executor.submit(_fetch, pairs[cursor : cursor + batch_size])
232
+ )
233
+ cursor += batch_size
234
+
235
+
236
+ def _write_xarray_to_view(
237
+ view: "DatasetView",
238
+ ds: "xr.Dataset",
239
+ chunks: Optional[dict[str, Sequence[int]]] = None,
240
+ ) -> None:
241
+ """Populate an empty `DatasetView` with the contents of an xarray Dataset.
242
+
243
+ Writes every coordinate and data variable as an atlas array, the
244
+ coordinate names as ``_pyatlas_coords`` (JSON list), all dataset attrs,
245
+ and all per-variable attrs flattened as ``{var}.{attr}``.
246
+ """
247
+ coord_names = [str(n) for n in ds.coords.keys()]
248
+
249
+ # Write coords first, then data_vars. Order doesn't matter to atlas but
250
+ # makes the on-disk file layout predictable.
251
+ for var_name in coord_names + [str(n) for n in ds.data_vars.keys()]:
252
+ var = ds[var_name]
253
+ atlas_dtype = _np_to_atlas_dtype(np.dtype(var.dtype))
254
+ dims = [str(d) for d in var.dims]
255
+ shape = [int(s) for s in var.shape]
256
+
257
+ # Pick the atlas chunk_shape:
258
+ # 1. explicit user override via the `chunks=` kwarg, else
259
+ # 2. the dask chunk shape if the variable is dask-backed, else
260
+ # 3. None (atlas defaults to a single full-shape chunk).
261
+ if chunks is not None and var_name in chunks:
262
+ chunk_shape: Optional[list[int]] = [int(s) for s in chunks[var_name]]
263
+ elif _is_dask_array(var.data):
264
+ chunk_shape = _dask_chunk_shape(var.data)
265
+ else:
266
+ chunk_shape = None
267
+
268
+ # Extract the CF/netCDF `_FillValue` attribute (if any) so it's passed
269
+ # to define_array as a typed fill value, not stored as a flattened atlas
270
+ # attribute. Copy attrs first to avoid mutating the user's Dataset.
271
+ var_attrs = dict(var.attrs)
272
+ fill_value = _normalize_fill_value(var_attrs.pop("_FillValue", None), var.dtype)
273
+
274
+ view.define_array(
275
+ var_name,
276
+ dtype=atlas_dtype,
277
+ dims=dims,
278
+ shape=shape,
279
+ chunk_shape=chunk_shape,
280
+ fill_value=fill_value,
281
+ )
282
+
283
+ # Stream blocks: prefetched batches for dask-backed data, a single
284
+ # full-shape block for numpy-backed data.
285
+ for start, block in _iter_blocks(var.data, var_name=var_name):
286
+ # TimestampNs columns: the bindings accept np.int64 only; cast the
287
+ # numpy datetime64 view to int64 without copying.
288
+ if block.dtype.kind == "M":
289
+ block = block.view(np.int64)
290
+ t0 = time.perf_counter_ns()
291
+ view.write_array(var_name, start=start, data=block)
292
+ _log_chunk_event(
293
+ "write",
294
+ var_name,
295
+ (time.perf_counter_ns() - t0) // 1000,
296
+ bytes=int(block.nbytes),
297
+ )
298
+
299
+ # Per-variable attrs → flattened as `{var}.{attr}` (sans `_FillValue`).
300
+ for attr_key, attr_val in var_attrs.items():
301
+ encoded = _encode_attr_value(attr_val)
302
+ view.set_attribute(_sanitize_str(f"{var_name}.{attr_key}"), encoded)
303
+
304
+ # Dataset-level attrs
305
+ for attr_key, attr_val in ds.attrs.items():
306
+ encoded = _encode_attr_value(attr_val)
307
+ view.set_attribute(_sanitize_str(str(attr_key)), encoded)
308
+
309
+ # Marker so we can faithfully restore coord/var distinction on read.
310
+ view.set_attribute(_COORDS_ATTR, json.dumps(coord_names))
311
+
312
+
313
+ _ATLAS_TO_NUMPY = {atlas: np_dt for np_dt, atlas in _NUMPY_TO_ATLAS.items()}
314
+
315
+
316
+ def _atlas_to_numpy_dtype(atlas_dtype: str) -> np.dtype:
317
+ """Numpy dtype that `view.read_array` returns for a given atlas dtype string."""
318
+ if atlas_dtype in _ATLAS_TO_NUMPY:
319
+ return _ATLAS_TO_NUMPY[atlas_dtype]
320
+ if atlas_dtype == "string":
321
+ return np.dtype("object")
322
+ raise NotImplementedError(
323
+ f"atlas dtype {atlas_dtype!r} is not supported on the dask read path"
324
+ )
325
+
326
+
327
+ def _dask_chunks_for(shape: Sequence[int], chunk_shape: Sequence[int]) -> tuple:
328
+ """Per-dim chunk-length tuples in the form dask expects."""
329
+ chunks: list[tuple[int, ...]] = []
330
+ for dim_size, dim_chunk in zip(shape, chunk_shape):
331
+ if dim_chunk <= 0 or dim_size == 0:
332
+ chunks.append((dim_size,))
333
+ continue
334
+ full = dim_size // dim_chunk
335
+ rem = dim_size - full * dim_chunk
336
+ c = (dim_chunk,) * full + ((rem,) if rem else ())
337
+ chunks.append(c if c else (0,))
338
+ return tuple(chunks)
339
+
340
+
341
+ def _view_to_dask_array(view: "DatasetView", name: str) -> Any:
342
+ """Build a `dask.array.Array` that lazily reads `name` chunk-by-chunk.
343
+
344
+ Each on-disk chunk becomes one dask task; values are fetched via
345
+ `view.read_array(name, start, block_shape)` on demand. Used by
346
+ `_view_to_xarray` when an array's `chunk_shape != shape`.
347
+ """
348
+ import dask
349
+ import dask.array as da
350
+
351
+ meta = view.array_meta(name)
352
+ assert meta is not None, f"array {name!r} not found in view {view.name!r}"
353
+ shape: list[int] = list(meta["shape"])
354
+ chunk_shape: list[int] = list(meta["chunk_shape"])
355
+ np_dtype = _atlas_to_numpy_dtype(meta["dtype"])
356
+
357
+ per_dim_chunks = _dask_chunks_for(shape, chunk_shape)
358
+ offsets = [
359
+ [0, *itertools.accumulate(dim_chunks)][:-1] for dim_chunks in per_dim_chunks
360
+ ]
361
+
362
+ def _read_block(start: list[int], block_shape: list[int]) -> np.ndarray:
363
+ return view.read_array(name, start=start, shape=block_shape)
364
+
365
+ def _nested_blocks(axis: int, prefix_start: list[int], prefix_shape: list[int]):
366
+ if axis == len(shape):
367
+ block_shape = list(prefix_shape)
368
+ block_start = list(prefix_start)
369
+ delayed = dask.delayed(_read_block)(block_start, block_shape)
370
+ return da.from_delayed(delayed, shape=tuple(block_shape), dtype=np_dtype)
371
+ return [
372
+ _nested_blocks(
373
+ axis + 1,
374
+ prefix_start + [offsets[axis][i]],
375
+ prefix_shape + [per_dim_chunks[axis][i]],
376
+ )
377
+ for i in range(len(per_dim_chunks[axis]))
378
+ ]
379
+
380
+ nested = _nested_blocks(0, [], [])
381
+ return da.block(nested)
382
+
383
+
384
+ def _view_to_xarray(view: "DatasetView", force_lazy: bool = False) -> "xr.Dataset":
385
+ """Convert an atlas `DatasetView` into an xarray Dataset.
386
+
387
+ Variables stored with `chunk_shape != shape` come back dask-backed (one task
388
+ per on-disk chunk); full-shape arrays come back eager as numpy unless
389
+ ``force_lazy=True``, in which case they're wrapped in a single-chunk
390
+ `dask.array` so the returned Dataset is uniformly lazy (used by
391
+ `_atlas_to_xarray_many` so concat returns a lazy graph regardless of the
392
+ source chunk_shape).
393
+ """
394
+ import xarray as xr
395
+
396
+ array_names = list(view.list_arrays())
397
+ array_name_set = set(array_names)
398
+
399
+ # Coord/var assignment
400
+ coords_marker = view.get_attribute(_COORDS_ATTR)
401
+ if isinstance(coords_marker, str):
402
+ try:
403
+ coord_names = set(json.loads(coords_marker))
404
+ except (TypeError, ValueError):
405
+ coord_names = set()
406
+ else:
407
+ # Fallback heuristic: 1-D array whose single dim matches its name.
408
+ coord_names = set()
409
+ for name in array_names:
410
+ meta = view.array_meta(name)
411
+ if meta is None:
412
+ continue
413
+ if (
414
+ len(meta["dimension_names"]) == 1
415
+ and meta["dimension_names"][0] == name
416
+ ):
417
+ coord_names.add(name)
418
+
419
+ # Pull array data. Chunked arrays (chunk_shape != shape) come back as a
420
+ # lazy dask.array; full-shape and 0-D arrays come back eager.
421
+ data_vars: dict[str, tuple] = {}
422
+ coords: dict[str, tuple] = {}
423
+ for name in array_names:
424
+ meta = view.array_meta(name)
425
+ if meta is None:
426
+ continue
427
+ shape = list(meta["shape"])
428
+ chunk_shape = list(meta["chunk_shape"])
429
+ if not shape or chunk_shape == shape:
430
+ arr = view.read_array(name)
431
+ if arr is None:
432
+ continue
433
+ if force_lazy and shape:
434
+ import dask.array as da
435
+ arr = da.from_array(arr, chunks=tuple(shape))
436
+ else:
437
+ arr = _view_to_dask_array(view, name)
438
+ dims = list(meta["dimension_names"])
439
+ entry = (dims, arr, {}) # placeholder for per-var attrs; filled below
440
+ if name in coord_names:
441
+ coords[name] = entry
442
+ else:
443
+ data_vars[name] = entry
444
+
445
+ # Split dataset attrs vs flattened per-var attrs
446
+ raw_attrs = dict(view.attributes())
447
+ raw_attrs.pop(_COORDS_ATTR, None)
448
+
449
+ dataset_attrs: dict[str, Any] = {}
450
+ per_var_attrs: dict[str, dict[str, Any]] = {n: {} for n in array_names}
451
+ for key, value in raw_attrs.items():
452
+ if "." in key:
453
+ var, rest = key.split(".", 1)
454
+ if var in array_name_set:
455
+ per_var_attrs[var][rest] = _decode_attr_value(value)
456
+ continue
457
+ dataset_attrs[key] = _decode_attr_value(value)
458
+
459
+ # Restore _FillValue for any array that was defined with one.
460
+ for name in array_names:
461
+ fv = view.array_fill_value(name)
462
+ if fv is None:
463
+ continue
464
+ per_var_attrs.setdefault(name, {})
465
+ per_var_attrs[name]["_FillValue"] = fv
466
+
467
+ # Inject per-var attrs into the (dims, data, attrs) triples
468
+ def _with_attrs(name: str, triple: tuple) -> tuple:
469
+ dims, arr, _ = triple
470
+ return (dims, arr, per_var_attrs.get(name, {}))
471
+
472
+ data_vars = {n: _with_attrs(n, t) for n, t in data_vars.items()}
473
+ coords = {n: _with_attrs(n, t) for n, t in coords.items()}
474
+
475
+ return xr.Dataset(data_vars=data_vars, coords=coords, attrs=dataset_attrs)
476
+
477
+
478
+ def _atlas_to_xarray_many(
479
+ atlas: "Atlas",
480
+ names: list[str],
481
+ concat_dim: str = "dataset",
482
+ parallel: bool = True, # noqa: ARG001 — kept for API compat; ignored
483
+ ) -> "xr.Dataset":
484
+ """Open many atlas datasets and stack them into one xr.Dataset along
485
+ `concat_dim`. atlas-native equivalent of `xr.open_mfdataset(...)`.
486
+
487
+ Implementation: opens the first dataset to discover the schema (vars,
488
+ dims, dtypes, coords, per-var attrs), then for each data variable calls
489
+ `Atlas.read_array_across` to bulk-read across all `names` in one Rust
490
+ call. The N reads share one `RwLock::read` guard on the shared physical
491
+ file and dispatch concurrently on the tokio runtime — avoids the N
492
+ Python ↔ Rust round-trips the prior dask-delayed implementation paid.
493
+
494
+ Returns eager numpy-backed arrays of shape `(len(names), *original_shape)`.
495
+ Wrap with `.chunk(...)` downstream if you need dask laziness.
496
+
497
+ The `parallel` parameter is accepted for API compatibility but no longer
498
+ selects an implementation — the bulk path is always taken.
499
+ """
500
+ import numpy as np
501
+ import xarray as xr
502
+
503
+ if not names:
504
+ raise ValueError("to_xarray_many: `names` is empty")
505
+
506
+ # Schema discovery from the first dataset (cheap: in-memory meta lookup).
507
+ first_view = atlas.open_dataset(names[0])
508
+ template = _view_to_xarray(first_view, force_lazy=False)
509
+
510
+ # Bulk-read each data variable across all datasets in one Rust call,
511
+ # returning a pre-stacked (N, *shape) numpy array. Skips the Python-side
512
+ # `np.stack` copy that the list-returning `read_array_across` would
513
+ # require — significant on big workloads (a 1000-dataset gridded run
514
+ # saves several seconds of memory bandwidth).
515
+ #
516
+ # Variables are processed serially because each Rust call internally
517
+ # parallelises N reads up to num_cpus — running multiple vars in
518
+ # parallel would oversubscribe CPU.
519
+ data_vars: dict[str, xr.DataArray] = {}
520
+ for var in template.data_vars:
521
+ stacked = atlas.read_array_across_stacked(var, names)
522
+ original_dims = list(template[var].dims)
523
+ original_attrs = dict(template[var].attrs)
524
+ data_vars[var] = xr.DataArray(
525
+ stacked,
526
+ dims=[concat_dim, *original_dims],
527
+ attrs=original_attrs,
528
+ )
529
+
530
+ # Coords + dataset-level attrs come from the first dataset, matching
531
+ # xarray.open_mfdataset(coords="minimal", compat="override") semantics.
532
+ coords = {name: template.coords[name] for name in template.coords}
533
+ coords[concat_dim] = xr.DataArray(np.asarray(names), dims=[concat_dim])
534
+
535
+ return xr.Dataset(data_vars=data_vars, coords=coords, attrs=dict(template.attrs))
536
+
537
+
538
+ def _write_xarray_new_dataset(
539
+ atlas: "Atlas",
540
+ ds: "xr.Dataset",
541
+ name: str,
542
+ chunks: Optional[dict[str, Sequence[int]]] = None,
543
+ ) -> None:
544
+ """Rust-delegated helper: create a fresh atlas dataset and populate it.
545
+
546
+ Both `atlas.add_xr_dataset` and the `ds.atlas.write` accessor route through
547
+ this function.
548
+ """
549
+ view = atlas.create_dataset(name)
550
+ _write_xarray_to_view(view, ds, chunks=chunks)
551
+
552
+
553
+ # --- xarray accessor ----------------------------------------------------------
554
+ # Registered as `ds.atlas` once `atlas` is imported. The whole module is
555
+ # side-effect-imported from `atlas/__init__.py` so importing `atlas` is
556
+ # enough to activate this.
557
+
558
+ import xarray as _xr # noqa: E402
559
+
560
+
561
+ @_xr.register_dataset_accessor("atlas")
562
+ class _AtlasAccessor:
563
+ """Methods exposed on an `xr.Dataset` as `ds.atlas.*`."""
564
+
565
+ def __init__(self, ds: "_xr.Dataset") -> None:
566
+ self._ds = ds
567
+
568
+ def write(
569
+ self,
570
+ atlas: "Atlas",
571
+ name: str,
572
+ chunks: Optional[dict[str, Sequence[int]]] = None,
573
+ ) -> None:
574
+ """Append this Dataset to the open atlas store under `name`.
575
+
576
+ Equivalent to `atlas.add_xr_dataset(self_ds, name, chunks)`.
577
+ """
578
+ atlas.add_xr_dataset(self._ds, name, chunks)
@@ -0,0 +1,17 @@
1
+ Metadata-Version: 2.4
2
+ Name: atlas-python
3
+ Version: 0.9.0
4
+ Requires-Dist: numpy>=1.23
5
+ Requires-Dist: xarray>=2023.1
6
+ Requires-Dist: dask>=2023.1
7
+ Requires-Dist: zarr>=3 ; extra == 'bench'
8
+ Requires-Dist: numcodecs ; extra == 'bench'
9
+ Requires-Dist: netcdf4 ; extra == 'bench'
10
+ Requires-Dist: obstore>=0.9 ; extra == 'cloud'
11
+ Requires-Dist: pytest ; extra == 'test'
12
+ Requires-Dist: netcdf4 ; extra == 'test'
13
+ Provides-Extra: bench
14
+ Provides-Extra: cloud
15
+ Provides-Extra: test
16
+ Summary: Python bindings for the ATLAS array store
17
+ Requires-Python: >=3.10
@@ -0,0 +1,9 @@
1
+ atlas/__init__.py,sha256=Q2om8rCNbD_dYWw0C4Av_Jfvi7dazDqFxAbJ5dQRqtU,222
2
+ atlas/__init__.pyi,sha256=m5Hcnj7tlLii4B_UIToVQIndd74ke1qvMYG96-_t6Ss,12708
3
+ atlas/_atlas.pyd,sha256=xTxpL7r2rYy2a5s3gsT_PjukYK_GELjHThwdJXY576s,19614208
4
+ atlas/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
5
+ atlas/xarray.py,sha256=A80ty9shMNn5ynER48NA9StltSxAxPZ1L0-6zU1syiU,22826
6
+ atlas_python-0.9.0.dist-info/METADATA,sha256=aYbfSEtT97GKtS-zj9zyJJU5-WO3b2jucixZv8U5VBM,537
7
+ atlas_python-0.9.0.dist-info/WHEEL,sha256=OUT0XP5TL9Hq-6CIgsb5m6BAU8pfcNqYjx0xnFDWhNs,96
8
+ atlas_python-0.9.0.dist-info/sboms/atlas-python.cyclonedx.json,sha256=zJj1b8tP1-cpy0i50fjXrXZ7WP9QWEpZTMfqddsTVno,265661
9
+ atlas_python-0.9.0.dist-info/RECORD,,
@@ -0,0 +1,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: maturin (1.13.3)
3
+ Root-Is-Purelib: false
4
+ Tag: cp310-abi3-win_amd64