forked from inaos/iron-array-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathconstructors.py
More file actions
321 lines (252 loc) · 9.32 KB
/
Copy pathconstructors.py
File metadata and controls
321 lines (252 loc) · 9.32 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
###########################################################################################
# Copyright ironArray SL 2021.
#
# All rights reserved.
#
# This software is the confidential and proprietary information of ironArray SL
# ("Confidential Information"). You shall not disclose such Confidential Information
# and shall use it only in accordance with the terms of the license agreement.
###########################################################################################
import numpy as np
import zarr
import s3fs
import iarray as ia
from iarray import iarray_ext as ext
from .utils import IllegalArgumentError, zarr_to_iarray_dtypes
from dataclasses import dataclass
from typing import Sequence
@dataclass
class DTShape:
"""Shape and data type dataclass.
Parameters
----------
shape: list, tuple
The shape of the array.
dtype: (np.float64, np.float32, np.int64, np.int32, np.int16, np.int8, np.uint64, np.uint32, np.uint16,
np.uint8, np.bool_)
The data type of the elements in the array. The default is np.float64.
"""
shape: Sequence
dtype: (
np.float64,
np.float32,
np.int64,
np.int32,
np.int16,
np.int8,
np.uint64,
np.uint32,
np.uint16,
np.uint8,
np.bool_,
) = np.float64
def __post_init__(self):
if self.shape is None:
raise ValueError("shape must be non-empty")
def empty(shape: Sequence, cfg: ia.Config = None, **kwargs) -> ia.IArray:
"""Return an empty array.
An empty array has no data and needs to be filled via a write iterator.
Parameters
----------
shape : tuple, list
The shape of the array to be created.
cfg : :class:`Config`
The configuration for running the expression.
If None (default), global defaults are used.
kwargs : dict
A dictionary for setting some or all of the fields in the :class:`Config`
dataclass that should override the current configuration.
Returns
-------
:ref:`IArray`
The new array.
"""
if cfg is None:
cfg = ia.get_config_defaults()
with ia.config(shape=shape, cfg=cfg, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
return ext.empty(cfg, dtshape)
def uninit(shape: Sequence, cfg: ia.Config = None, **kwargs) -> ia.IArray:
"""Return an uninitialized array.
An uninitialized array has no data and needs to be filled via a write iterator.
Parameters
----------
shape : tuple, list
The shape of the array to be created.
cfg : :class:`Config`
The configuration for running the expression.
If None (default), global defaults are used.
kwargs : dict
A dictionary for setting some or all of the fields in the :class:`Config`
dataclass that should override the current configuration.
Returns
-------
:ref:`IArray`
The new array.
"""
if cfg is None:
cfg = ia.get_config_defaults()
with ia.config(shape=shape, cfg=cfg, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
return ext.uninit(cfg, dtshape)
def arange(
shape: Sequence, start=None, stop=None, step=None, cfg: ia.Config = None, **kwargs
) -> ia.IArray:
"""Return evenly spaced values within a given interval.
`shape`, `cfg` and `kwargs` are the same than for :func:`empty`.
`start`, `stop`, `step` are the same as in `np.arange <https://numpy.org/doc/stable/reference/generated/numpy.arange.html>`_.
Returns
-------
:ref:`IArray`
The new array.
See Also
--------
empty : Create an empty array.
"""
if (start, stop, step) == (None, None, None):
stop = np.prod(shape)
start = 0
step = 1
elif (stop, step) == (None, None):
stop = start
start = 0
step = 1
elif step is None:
stop = stop
start = start
if shape is None:
step = 1
else:
step = (stop - start) / np.prod(shape)
slice_ = slice(start, stop, step)
if cfg is None:
cfg = ia.get_config_defaults()
with ia.config(shape=shape, cfg=cfg, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
return ext.arange(cfg, slice_, dtshape)
def linspace(
shape: Sequence, start: float, stop: float, cfg: ia.Config = None, **kwargs
) -> ia.IArray:
"""Return evenly spaced numbers over a specified interval.
`shape`, `cfg` and `kwargs` are the same than for :func:`empty`.
`start`, `stop` are the same as in `np.linspace <https://numpy.org/doc/stable/reference/generated/numpy.linspace.html>`_.
Returns
-------
:ref:`IArray`
The new array.
See Also
--------
empty : Create an empty array.
"""
if cfg is None:
cfg = ia.get_config_defaults()
with ia.config(shape=shape, cfg=cfg, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
return ext.linspace(cfg, start, stop, dtshape)
def zeros(shape: Sequence, cfg: ia.Config = None, **kwargs) -> ia.IArray:
"""Return a new array of given shape and type, filled with zeros.
`shape`, `cfg` and `kwargs` are the same than for :func:`empty`.
Returns
-------
:ref:`IArray`
The new array.
See Also
--------
empty : Create an empty array.
ones : Create an array filled with ones.
"""
if cfg is None:
cfg = ia.get_config_defaults()
with ia.config(shape=shape, cfg=cfg, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
return ext.zeros(cfg, dtshape)
def ones(shape: Sequence, cfg: ia.Config = None, **kwargs) -> ia.IArray:
"""Return a new array of given shape and type, filled with ones.
`shape`, `cfg` and `kwargs` are the same than for :func:`empty`.
Returns
-------
:ref:`IArray`
The new array.
See Also
--------
empty : Create an empty array.
zeros : Create an array filled with zeros.
"""
if cfg is None:
cfg = ia.get_config_defaults()
with ia.config(shape=shape, cfg=cfg, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
return ext.ones(cfg, dtshape)
def full(shape: Sequence, fill_value, cfg: ia.Config = None, **kwargs) -> ia.IArray:
"""Return a new array of given shape and type, filled with `fill_value`.
`shape`, `cfg` and `kwargs` are the same than for :func:`empty`.
Returns
-------
:ref:`IArray`
The new array.
See Also
--------
empty : Create an empty array.
zeros : Create an array filled with zeros.
"""
if cfg is None:
cfg = ia.get_config_defaults()
with ia.config(shape=shape, cfg=cfg, **kwargs) as cfg:
dtshape = ia.DTShape(shape, cfg.dtype)
return ext.full(cfg, fill_value, dtshape)
def zarr_proxy(zarr_urlpath, cfg: ia.Config = None, **kwargs) -> ia.IArray:
"""Return a read-only Zarr proxy array.
`cfg` and `kwargs` are the same than for :func:`empty` except by `nthreads`, which is
always set to 1 (multi-threading is not yet supported).
The data type and chunks must not differ from the original Zarr array.
A Zarr proxy is a regular IArray array but with a special attribute called `zproxy_urlpath`. This
attribute is protected when `attrs.clear()` is used; but can still be deleted with `del attrs["zproxy_urlpath"]`,
`attrs.popitem()` or `attrs.pop("zproxy_urlpath")`.
This IArray has an additional attribute called `proxy_attrs` which contains the Zarr attributes. The user can
get and set these attributes.
Parameters
----------
zarr_urlpath : str
The path to the Zarr array.
If it is stored in the cloud, the path must begin with ``s3://``.
Returns
-------
:ref:`IArray`
The zarr proxy array.
Notes
-----
As a proxy, this array does not contain the data from the original array, it only reads it when needed.
But if a :func:`save` is done, a copy of all the data will be made and assigned to a new
and usual on disk :ref:`IArray`. To create a persistent proxy on-disk,
you can specificy the :paramref:`urlpath` during :func:`zarr_proxy` execution time.
"""
z = ext._zarray_from_proxy(zarr_urlpath)
# Create iarray
dtype = zarr_to_iarray_dtypes[str(z.dtype)]
if cfg is None:
cfg = ia.get_config_defaults()
if kwargs != {}:
if "dtype" in kwargs:
if kwargs.pop("dtype") != dtype:
raise AttributeError("dtype cannot differ from the original array")
if "chunks" in kwargs:
if tuple(kwargs.pop("chunks")) != z.chunks:
raise AttributeError("chunks cannot differ from the original array")
if "blocks" in kwargs:
blocks = tuple(kwargs.pop("blocks"))
else:
blocks = z.chunks
if "nthreads" in kwargs:
if kwargs.pop("nthreads") != 1:
raise IllegalArgumentError("Cannot use parallelism when interacting with Zarr")
with ia.config(
cfg=cfg, dtype=dtype, chunks=z.chunks, blocks=blocks, nthreads=1, **kwargs
) as cfg:
a = uninit(shape=z.shape, cfg=cfg)
# Set special attr to identify zarr_proxy
a.attrs["zproxy_urlpath"] = zarr_urlpath
# Create reference to zarr.attrs
a.zarr_attrs = z.attrs
# Assign postfilter
ext.set_zproxy_postfilter(a)
return a