forked from AntSimi/py-eddy-tracker
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathobservation.py
More file actions
2604 lines (2340 loc) · 92.3 KB
/
observation.py
File metadata and controls
2604 lines (2340 loc) · 92.3 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
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# -*- coding: utf-8 -*-
"""
Base class to manage eddy observation
"""
import logging
from datetime import datetime
from io import BufferedReader, BytesIO
from tarfile import ExFileObject
from tokenize import TokenError
import packaging.version
import zarr
from matplotlib.cm import get_cmap
from matplotlib.collections import PolyCollection
from matplotlib.colors import Normalize
from netCDF4 import Dataset
from numba import njit
from numba import types as numba_types
from numpy import (
absolute,
arange,
array,
array_equal,
ceil,
concatenate,
cos,
digitize,
empty,
errstate,
floor,
histogram,
histogram2d,
in1d,
isnan,
linspace,
ma,
nan,
ndarray,
ones,
percentile,
radians,
sin,
unique,
where,
zeros,
)
from pint import UnitRegistry
from pint.errors import UndefinedUnitError
from Polygon import Polygon
from .. import VAR_DESCR, VAR_DESCR_inv, __version__
from ..generic import (
bbox_indice_regular,
build_index,
distance,
distance_grid,
flatten_line_matrix,
hist_numba,
local_to_coordinates,
reverse_index,
wrap_longitude,
)
from ..poly import (
bbox_intersection,
close_center,
convexs,
create_vertice,
get_pixel_in_regular,
insidepoly,
poly_indexs,
reduce_size,
vertice_overlap,
winding_number_poly,
)
logger = logging.getLogger("pet")
# keep only major and minor version number
_software_version_reduced = packaging.version.Version(
"{v.major}.{v.minor}".format(v=packaging.version.parse(__version__))
)
def _check_versions(version):
"""Check if version of py_eddy_tracker used to create the file is compatible with software version
if not, warn user with both versions
:param version: string version of software used to create the file. If None, version was not provided
:type version: str, None
"""
file_version = packaging.version.parse(version) if version is not None else None
if file_version is None or file_version < _software_version_reduced:
logger.warning(
"File was created with py-eddy-tracker version '%s' but software version is '%s'",
file_version,
_software_version_reduced,
)
@njit(cache=True, fastmath=True)
def shifted_ellipsoid_degrees_mask2(lon0, lat0, lon1, lat1, minor=1.5, major=1.5):
"""
Work only if major is an array but faster * 6
"""
# c = (major ** 2 - minor ** 2) ** .5 + major
c = major
major = minor + 0.5 * (major - minor)
# r=.5*(c-c0)
# a=c0+r
# Focal
f_right = lon0
f_left = f_right - (c - minor)
# Ellipse center
x_c = (f_left + f_right) * 0.5
nb_0, nb_1 = lat0.shape[0], lat1.shape[0]
m = empty((nb_0, nb_1), dtype=numba_types.bool_)
for j in range(nb_0):
for i in range(nb_1):
dy = absolute(lat1[i] - lat0[j])
if dy > minor:
m[j, i] = False
continue
dx = absolute(lon1[i] - x_c[j])
if dx > 180:
dx = absolute((dx + 180) % 360 - 180)
if dx > major[j]:
m[j, i] = False
continue
d_normalize = dx ** 2 / major[j] ** 2 + dy ** 2 / minor ** 2
m[j, i] = d_normalize < 1.0
return m
class Table(object):
def __init__(self, values):
self.values = values
def _repr_html_(self):
rows = list()
if isinstance(self.values, ndarray):
row = "\n".join([f"<td>{v}</td >" for v in self.values.dtype.names])
rows.append(f"<tr>{row}</tr>")
for row in self.values:
row = "\n".join([f"<td>{v}</td >" for v in row])
rows.append(f"<tr>{row}</tr>")
rows = "\n".join(rows)
return (
f'<font size="2">'
f'<table class="docutils align-default">'
f"{rows}"
f"</table>"
f"</font>"
)
class EddiesObservations(object):
"""
Class to store eddy observations.
"""
__slots__ = (
"track_extra_variables",
"track_array_variables",
"array_variables",
"only_variables",
"observations",
"sign_type",
"raw_data",
"period_",
)
ELEMENTS = [
"lon",
"lat",
"radius_s",
"radius_e",
"amplitude",
"speed_average",
"time",
"shape_error_e",
"shape_error_s",
"speed_area",
"effective_area",
"nb_contour_selected",
"num_point_e",
"num_point_s",
"height_max_speed_contour",
"height_external_contour",
"height_inner_contour",
]
COLORS = [
"sienna",
"red",
"darkorange",
"gold",
"palegreen",
"limegreen",
"forestgreen",
"mediumblue",
"dodgerblue",
"lightskyblue",
"violet",
"blueviolet",
"darkmagenta",
"darkgrey",
"dimgrey",
"steelblue",
]
NB_COLORS = len(COLORS)
def __init__(
self,
size=0,
track_extra_variables=None,
track_array_variables=0,
array_variables=None,
only_variables=None,
raw_data=False,
):
self.period_ = None
self.only_variables = only_variables
self.raw_data = raw_data
self.track_extra_variables = (
track_extra_variables if track_extra_variables is not None else []
)
self.track_array_variables = track_array_variables
self.array_variables = array_variables if array_variables is not None else []
for elt in self.elements:
if elt not in VAR_DESCR:
raise Exception("Unknown element : %s" % elt)
self.observations = zeros(size, dtype=self.dtype)
self.sign_type = None
@property
def tracks(self):
return self.track
def __eq__(self, other):
if self.sign_type != other.sign_type:
return False
if self.dtype != other.dtype:
return False
return array_equal(self.obs, other.obs)
def get_color(self, i):
"""Return colors as a cyclic list"""
return self.COLORS[i % self.NB_COLORS]
@property
def sign_legend(self):
return "Cyclonic" if self.sign_type != 1 else "Anticyclonic"
@property
def shape(self):
return self.observations.shape
def get_infos(self):
infos = dict(
bins_lat=(-90, -60, -15, 15, 60, 90),
bins_amplitude=array((0, 1, 2, 3, 4, 5, 10, 500)),
bins_radius=array((0, 15, 30, 45, 60, 75, 100, 200, 2000)),
nb_obs=self.observations.shape[0],
)
t0, t1 = self.period
infos["t0"], infos["t1"] = t0, t1
infos["period"] = t1 - t0 + 1
return infos
def _repr_html_(self):
infos = self.get_infos()
return f"""<b>{infos['nb_obs']} observations from {infos['t0']} to {infos['t1']} </b>"""
def parse_varname(self, name):
return self[name] if isinstance(name, str) else name
def hist(self, varname, x, bins, percent=False, mean=False, nb=False):
"""Build histograms.
:param str,array varname: variable to use to compute stat
:param str,array x: variable to use to know in which bins
:param array bins:
:param bool percent: normalized by sum of all bins
:param bool mean: compute mean by bins
:param bool nb: only count by bins
:return: value by bins
:rtype: array
"""
x = self.parse_varname(x)
if nb:
v = hist_numba(x, bins=bins)[0]
else:
v = histogram(x, bins=bins, weights=self.parse_varname(varname))[0]
if percent:
v = v.astype("f4") / v.sum() * 100
elif mean:
v /= hist_numba(x, bins=bins)[0]
return v
@staticmethod
def box_display(value):
"""Return values evenly spaced with few numbers"""
return "".join([f"{v_:10.2f}" for v_ in value])
def field_table(self):
"""
Produce description table of the fields available in this object
"""
rows = [("Name (Unit)", "Long name", "Scale factor", "Offset")]
names = list(self.obs.dtype.names)
names.sort()
for field in names:
infos = VAR_DESCR[field]
rows.append(
(
f"{infos.get('nc_name', field)} ({infos['nc_attr'].get('units', '')})",
infos["nc_attr"].get("long_name", "").capitalize(),
infos.get("scale_factor", ""),
infos.get("add_offset", ""),
)
)
return Table(rows)
def __repr__(self):
"""
Return general informations on dataset as strings.
:return: informations on datasets
:rtype: str
"""
t0, t1 = self.period
period = t1 - t0 + 1
bins_lat = (-90, -60, -15, 15, 60, 90)
bins_amplitude = array((0, 1, 2, 3, 4, 5, 10, 500))
bins_radius = array((0, 15, 30, 45, 60, 75, 100, 200, 2000))
nb_obs = self.observations.shape[0]
return f""" | {nb_obs} observations from {t0} to {t1} ({period} days, ~{nb_obs / period:.0f} obs/day)
| Speed area : {self.speed_area.sum() / period / 1e12:.2f} Mkm²/day
| Effective area : {self.effective_area.sum() / period / 1e12:.2f} Mkm²/day
----Distribution in Amplitude:
| Amplitude bounds (cm) {self.box_display(bins_amplitude)}
| Percent of eddies : {
self.box_display(self.hist('time', 'amplitude', bins_amplitude / 100., percent=True, nb=True))}
----Distribution in Radius:
| Speed radius (km) {self.box_display(bins_radius)}
| Percent of eddies : {
self.box_display(self.hist('time', 'radius_s', bins_radius * 1000., percent=True, nb=True))}
| Effective radius (km) {self.box_display(bins_radius)}
| Percent of eddies : {
self.box_display(self.hist('time', 'radius_e', bins_radius * 1000., percent=True, nb=True))}
----Distribution in Latitude
Latitude bounds {self.box_display(bins_lat)}
Percent of eddies : {self.box_display(self.hist('time', 'lat', bins_lat, percent=True, nb=True))}
Percent of speed area : {self.box_display(self.hist('speed_area', 'lat', bins_lat, percent=True))}
Percent of effective area : {self.box_display(self.hist('effective_area', 'lat', bins_lat, percent=True))}
Mean speed radius (km) : {self.box_display(self.hist('radius_s', 'lat', bins_lat, mean=True) / 1000.)}
Mean effective radius (km): {self.box_display(self.hist('radius_e', 'lat', bins_lat, mean=True) / 1000.)}
Mean amplitude (cm) : {self.box_display(self.hist('amplitude', 'lat', bins_lat, mean=True) * 100.)}"""
def __dir__(self):
"""Provide method name lookup and completion."""
base = set(dir(type(self)))
intern_name = set(self.elements)
extern_name = set([VAR_DESCR[k]["nc_name"] for k in intern_name])
# Must be check in init not here
if base & intern_name:
logger.warning(
"Some variable name have a common name with class attrs: %s",
base & intern_name,
)
if base & extern_name:
logger.warning(
"Some variable name have a common name with class attrs: %s",
base & extern_name,
)
return sorted(base.union(intern_name).union(extern_name))
def __getitem__(self, attr: str):
if attr in self.elements:
return self.obs[attr]
elif attr in VAR_DESCR_inv:
return self.obs[VAR_DESCR_inv[attr]]
elif attr in ("lifetime", "age"):
return getattr(self, attr)
raise KeyError("%s unknown" % attr)
def __getattr__(self, attr):
if attr in self.elements:
return self.obs[attr]
elif attr in VAR_DESCR_inv:
return self.obs[VAR_DESCR_inv[attr]]
raise AttributeError(
"{!r} object has no attribute {!r}".format(type(self).__name__, attr)
)
@classmethod
def needed_variable(cls):
return None
@classmethod
def obs_dimension(cls, handler):
for candidate in ("obs", "Nobs", "observation", "i"):
if candidate in handler.dimensions.keys():
return candidate
def remove_fields(self, *fields):
"""
Copy with fields listed remove
"""
nb_obs = self.obs.shape[0]
fields = set(fields)
only_variables = set(self.obs.dtype.names) - fields
track_extra_variables = set(self.track_extra_variables) - fields
array_variables = set(self.array_variables) - fields
new = self.__class__(
size=nb_obs,
track_extra_variables=track_extra_variables,
track_array_variables=self.track_array_variables,
array_variables=array_variables,
only_variables=only_variables,
raw_data=self.raw_data,
)
new.sign_type = self.sign_type
for name in new.obs.dtype.names:
logger.debug("Copy of field %s ...", name)
new.obs[name] = self.obs[name]
return new
def add_fields(self, fields=list(), array_fields=list()):
"""
Add a new field.
"""
nb_obs = self.obs.shape[0]
new = self.__class__(
size=nb_obs,
track_extra_variables=list(
concatenate((self.track_extra_variables, fields))
),
track_array_variables=self.track_array_variables,
array_variables=list(concatenate((self.array_variables, array_fields))),
only_variables=list(
concatenate((self.obs.dtype.names, fields, array_fields))
),
raw_data=self.raw_data,
)
new.sign_type = self.sign_type
for name in self.obs.dtype.names:
logger.debug("Copy of field %s ...", name)
new.obs[name] = self.obs[name]
return new
def add_rotation_type(self):
new = self.add_fields(("type_cyc",))
new.type_cyc[:] = self.sign_type
return new
def circle_contour(self, only_virtual=False, factor=1):
"""
Set contours as circles from radius and center data.
.. minigallery:: py_eddy_tracker.EddiesObservations.circle_contour
"""
angle = radians(linspace(0, 360, self.track_array_variables))
x_norm, y_norm = cos(angle), sin(angle)
radius_s = "contour_lon_s" in self.obs.dtype.names
radius_e = "contour_lon_e" in self.obs.dtype.names
for i, obs in enumerate(self):
if only_virtual and not obs["virtual"]:
continue
x, y = obs["lon"], obs["lat"]
if radius_s:
r_s = obs["radius_s"] * factor
obs["contour_lon_s"], obs["contour_lat_s"] = local_to_coordinates(
x_norm * r_s, y_norm * r_s, x, y
)
if radius_e:
r_e = obs["radius_e"] * factor
obs["contour_lon_e"], obs["contour_lat_e"] = local_to_coordinates(
x_norm * r_e, y_norm * r_e, x, y
)
@property
def dtype(self):
"""Return dtype to build numpy array."""
dtype = list()
for elt in self.elements:
data_type = (
VAR_DESCR[elt].get("compute_type", VAR_DESCR[elt].get("nc_type"))
if not self.raw_data
else VAR_DESCR[elt]["output_type"]
)
if elt in self.array_variables:
dtype.append((elt, data_type, (self.track_array_variables,)))
else:
dtype.append((elt, data_type))
return dtype
@property
def elements(self):
"""Return all the names of the variables."""
elements = [i for i in self.ELEMENTS]
if self.track_array_variables > 0:
elements += self.array_variables
if len(self.track_extra_variables):
elements += self.track_extra_variables
if self.only_variables is not None:
elements = [i for i in elements if i in self.only_variables]
return list(set(elements))
def coherence(self, other):
"""Check coherence between two datasets."""
test = self.track_extra_variables == other.track_extra_variables
test *= self.track_array_variables == other.track_array_variables
test *= self.array_variables == other.array_variables
return test
@classmethod
def concatenate(cls, observations):
nb_obs = 0
ref_obs = observations[0]
for obs in observations:
if not ref_obs.coherence(obs):
raise Exception("Merge of different type of observations")
nb_obs += len(obs)
eddies = cls.new_like(ref_obs, nb_obs)
i = 0
for obs in observations:
nb_obs = len(obs)
eddies.obs[i : i + nb_obs] = obs.obs
i += nb_obs
eddies.sign_type = ref_obs.sign_type
return eddies
def merge(self, other):
"""Merge two datasets."""
nb_obs_self = len(self)
nb_obs = nb_obs_self + len(other)
eddies = self.new_like(self, nb_obs)
other_keys = other.obs.dtype.fields.keys()
self_keys = self.obs.dtype.fields.keys()
for key in eddies.obs.dtype.fields.keys():
eddies.obs[key][:nb_obs_self] = self.obs[key][:]
if key in other_keys:
eddies.obs[key][nb_obs_self:] = other.obs[key][:]
if "track" in other_keys and "track" in self_keys:
last_track = eddies.track[nb_obs_self - 1] + 1
eddies.track[nb_obs_self:] += last_track
eddies.sign_type = self.sign_type
return eddies
def reset(self):
self.observations = zeros(0, dtype=self.dtype)
@property
def obs(self):
"""Return observations."""
return self.observations
def __len__(self):
return len(self.observations)
def __iter__(self):
for obs in self.obs:
yield obs
def iter_on(self, xname, bins=None):
"""
Yield observation group for each bin.
:param str,array xname:
:param array bins: bounds of each bin ,
:return: index or mask, bound low, bound up
.. minigallery:: py_eddy_tracker.EddiesObservations.iter_on
"""
x = self[xname] if isinstance(xname, str) else xname
d = x[1:] - x[:-1]
if bins is None:
bins = arange(x.min(), x.max() + 2)
elif not isinstance(bins, ndarray):
bins = array(bins)
nb_bins = len(bins) - 1
# Not monotonous
if (d < 0).any():
# If bins cover a small part of value
test, translate, x = iter_mode_reduce(x, bins)
# convert value in bins number
i = numba_digitize(x, bins) - 1
# Order by bins
i_sort = i.argsort()
# If in reduced mode we will translate i_sort in full array index
i_sort_ = translate[i_sort] if test else i_sort
# Bound for each bins in sorting view
i0, i1, _ = build_index(i[i_sort])
m = ~(i0 == i1)
i0, i1 = i0[m], i1[m]
for i0_, i1_ in zip(i0, i1):
i_bins = i[i_sort[i0_]]
if i_bins == -1 or i_bins == nb_bins:
continue
yield i_sort_[i0_:i1_], bins[i_bins], bins[i_bins + 1]
else:
i = numba_digitize(x, bins) - 1
i0, i1, _ = build_index(i)
m = ~(i0 == i1)
i0, i1 = i0[m], i1[m]
for i0_, i1_ in zip(i0, i1):
i_bins = i[i0_]
yield slice(i0_, i1_), bins[i_bins], bins[i_bins + 1]
def align_on(self, other, var_name="time", **kwargs):
"""
Align the time indices of two datasets.
.. minigallery:: py_eddy_tracker.EddiesObservations.align_on
"""
iter_self = self.iter_on(var_name, **kwargs)
iter_other = other.iter_on(var_name, **kwargs)
indexs_other, b0_other, b1_other = iter_other.__next__()
for indexs_self, b0_self, b1_self in iter_self:
if b0_self > b0_other:
try:
while b0_other < b0_self:
indexs_other, b0_other, b1_other = iter_other.__next__()
except StopIteration:
break
if b0_self < b0_other:
continue
yield indexs_self, indexs_other, b0_self, b1_self
def insert_observations(self, other, index):
"""Insert other obs in self at the given index."""
if not self.coherence(other):
raise Exception("Observations with no coherence")
insert_size = len(other.obs)
self_size = len(self.obs)
new_size = self_size + insert_size
if self_size == 0:
self.observations = other.obs
return self
elif insert_size == 0:
return self
if index < 0:
index = self_size + index + 1
eddies = self.new_like(self, new_size)
eddies.obs[:index] = self.obs[:index]
eddies.obs[index : index + insert_size] = other.obs
eddies.obs[index + insert_size :] = self.obs[index:]
self.observations = eddies.obs
return self
def append(self, other):
"""Merge."""
return self + other
def __add__(self, other):
return self.insert_observations(other, -1)
def distance(self, other):
"""Use haversine distance for distance matrix between every self and
other eddies."""
return distance_grid(self.lon, self.lat, other.lon, other.lat)
def __copy__(self):
eddies = self.new_like(self, len(self))
for k in self.obs.dtype.names:
eddies[k][:] = self[k][:]
eddies.sign_type = self.sign_type
return eddies
def copy(self):
return self.__copy__()
@classmethod
def new_like(cls, eddies, new_size: int):
return cls(
new_size,
track_extra_variables=eddies.track_extra_variables,
track_array_variables=eddies.track_array_variables,
array_variables=eddies.array_variables,
only_variables=eddies.only_variables,
raw_data=eddies.raw_data,
)
def index(self, index, reverse=False):
"""Return obs from self at the index."""
if reverse:
index = reverse_index(index, len(self))
size = 1
if hasattr(index, "__iter__"):
size = len(index)
elif isinstance(index, slice):
size = index.stop - index.start
eddies = self.new_like(self, size)
eddies.obs[:] = self.obs[index]
eddies.sign_type = self.sign_type
return eddies
@staticmethod
def zarr_dimension(filename):
if isinstance(filename, zarr.storage.MutableMapping):
h = filename
else:
h = zarr.open(filename)
dims = list()
for varname in h:
shape = getattr(h, varname).shape
if len(shape) > len(dims):
dims = shape
return dims
@classmethod
def load_file(cls, filename, **kwargs):
"""
Load the netcdf or the zarr file.
Load only latitude and longitude on the first 300 obs :
.. code-block:: python
kwargs_latlon_300 = dict(
include_vars=["longitude", "latitude",], indexs=dict(obs=slice(0, 300)),
)
small_dataset = TrackEddiesObservations.load_file(
filename, **kwargs_latlon_300
)
For `**kwargs` look at :py:meth:`load_from_zarr` or :py:meth:`load_from_netcdf`
"""
filename_ = (
filename.filename if isinstance(filename, ExFileObject) else filename
)
if isinstance(filename, zarr.storage.MutableMapping):
return cls.load_from_zarr(filename, **kwargs)
if isinstance(filename, (bytes, str)):
end = b".zarr" if isinstance(filename_, bytes) else ".zarr"
zarr_file = filename_.endswith(end)
else:
zarr_file = False
logger.info(f"loading file '{filename}'")
if zarr_file:
return cls.load_from_zarr(filename, **kwargs)
else:
return cls.load_from_netcdf(filename, **kwargs)
@classmethod
def load_from_zarr(
cls,
filename,
raw_data=False,
remove_vars=None,
include_vars=None,
indexs=None,
buffer_size=5000000,
**class_kwargs,
):
"""Load data from zarr.
:param str,store filename: path or store to load data
:param bool raw_data: If true load data without scale_factor and add_offset
:param None,list(str) remove_vars: List of variable name that will be not loaded
:param None,list(str) include_vars: If defined only this variable will be loaded
:param None,dict indexs: Indexes to load only a slice of data
:param int buffer_size: Size of buffer used to load zarr data
:param class_kwargs: argument to set up observations class
:return: Obsevations selected
:return type: class
"""
# FIXME
if isinstance(filename, zarr.storage.MutableMapping):
h_zarr = filename
else:
if not isinstance(filename, str):
filename = filename.astype(str)
h_zarr = zarr.open(filename)
_check_versions(h_zarr.attrs.get("framework_version", None))
var_list = cls.build_var_list(list(h_zarr.keys()), remove_vars, include_vars)
nb_obs = getattr(h_zarr, var_list[0]).shape[0]
track_array_variables = h_zarr.attrs["track_array_variables"]
if indexs is not None and "obs" in indexs:
sl = indexs["obs"]
sl = slice(sl.start, min(sl.stop, nb_obs))
if sl.stop is not None:
nb_obs = sl.stop
if sl.start is not None:
nb_obs -= sl.start
if sl.step is not None:
indexs["obs"] = slice(sl.start, sl.stop)
logger.warning("step of slice won't be use")
logger.debug("%d observations will be load", nb_obs)
kwargs = dict()
kwargs["track_array_variables"] = h_zarr.attrs.get(
"track_array_variables", track_array_variables
)
array_variables = list()
for variable in var_list:
if len(h_zarr[variable].shape) > 1:
var_inv = VAR_DESCR_inv[variable]
array_variables.append(var_inv)
kwargs["array_variables"] = array_variables
track_extra_variables = []
for variable in var_list:
var_inv = VAR_DESCR_inv[variable]
if var_inv not in cls.ELEMENTS and var_inv not in array_variables:
track_extra_variables.append(var_inv)
kwargs["track_extra_variables"] = track_extra_variables
kwargs["raw_data"] = raw_data
kwargs["only_variables"] = (
None if include_vars is None else [VAR_DESCR_inv[i] for i in include_vars]
)
kwargs.update(class_kwargs)
eddies = cls(size=nb_obs, **kwargs)
for i_var, variable in enumerate(var_list):
var_inv = VAR_DESCR_inv[variable]
logger.debug("%s will be loaded (%d/%d)", variable, i_var, len(var_list))
# find unit factor
input_unit = h_zarr[variable].attrs.get("unit", None)
if input_unit is None:
input_unit = h_zarr[variable].attrs.get("units", None)
output_unit = VAR_DESCR[var_inv]["nc_attr"].get("units", None)
factor = cls.compare_units(input_unit, output_unit, variable)
sl_obs = slice(None) if indexs is None else indexs.get("obs", slice(None))
scale_factor = VAR_DESCR[var_inv].get("scale_factor", None)
add_offset = VAR_DESCR[var_inv].get("add_offset", None)
cls.copy_data_to_zarr(
h_zarr[variable],
eddies.obs[var_inv],
sl_obs,
buffer_size,
factor,
raw_data,
scale_factor,
add_offset,
)
eddies.sign_type = int(h_zarr.attrs.get("rotation_type", 0))
if eddies.sign_type == 0:
logger.debug("File come from another algorithm of identification")
eddies.sign_type = -1
return eddies
@staticmethod
def copy_data_to_zarr(
handler_zarr,
handler_eddies,
sl_obs,
buffer_size,
factor,
raw_data,
scale_factor,
add_offset,
):
"""
Copy with buffer for zarr.
Zarr need to get real value, and size could be huge, so we use a buffer to manage memory
:param zarr_dataset handler_zarr:
:param array handler_eddies:
:param slice zarr_dataset sl_obs:
:param int buffer_size:
:param float factor:
:param bool raw_data:
:param None,float scale_factor:
:param None,float add_offset:
"""
i_start, i_stop = sl_obs.start, sl_obs.stop
if i_start is None:
i_start = 0
if i_stop is None:
i_stop = handler_zarr.shape[0]
for i in range(i_start, i_stop, buffer_size):
sl_in = slice(i, min(i + buffer_size, i_stop))
data = handler_zarr[sl_in]
if factor != 1:
data *= factor
if raw_data:
if add_offset is not None:
data -= add_offset
if scale_factor is not None:
data /= scale_factor
sl_out = slice(i - i_start, i - i_start + buffer_size)
handler_eddies[sl_out] = data
@classmethod
def load_from_netcdf(
cls,
filename,
raw_data=False,
remove_vars=None,
include_vars=None,
indexs=None,
**class_kwargs,
):
"""Load data from netcdf.
:param str,ExFileObject filename: path or handler to load data
:param bool raw_data: If true load data without apply scale_factor and add_offset
:param None,list(str) remove_vars: List of variable name which will be not loaded
:param None,list(str) include_vars: If defined only this variable will be loaded
:param None,dict indexs: Indexes to load only a slice of data
:param class_kwargs: argument to set up observations class
:return: Obsevations selected
:return type: class
"""
array_dim = "NbSample"
if isinstance(filename, bytes):
filename = filename.astype(str)
if isinstance(filename, (ExFileObject, BufferedReader, BytesIO)):
filename.seek(0)
args, kwargs = ("in-mem-file",), dict(memory=filename.read())
else:
args, kwargs = (filename,), dict()
with Dataset(*args, **kwargs) as h_nc:
_check_versions(getattr(h_nc, "framework_version", None))
var_list = cls.build_var_list(
list(h_nc.variables.keys()), remove_vars, include_vars
)
obs_dim = cls.obs_dimension(h_nc)
nb_obs = len(h_nc.dimensions[obs_dim])
if indexs is not None and obs_dim in indexs:
sl = indexs[obs_dim]
sl = slice(sl.start, min(sl.stop, nb_obs))
if sl.stop is not None:
nb_obs = sl.stop
if sl.start is not None:
nb_obs -= sl.start
if sl.step is not None:
indexs[obs_dim] = slice(sl.start, sl.stop)
logger.warning("step of slice won't be use")
logger.debug("%d observations will be load", nb_obs)
kwargs = dict()
if array_dim in h_nc.dimensions:
kwargs["track_array_variables"] = len(h_nc.dimensions[array_dim])
kwargs["array_variables"] = list()
for variable in var_list:
if array_dim in h_nc.variables[variable].dimensions:
var_inv = VAR_DESCR_inv[variable]
kwargs["array_variables"].append(var_inv)
array_variables = kwargs.get("array_variables", list())
kwargs["track_extra_variables"] = []
for variable in var_list:
var_inv = VAR_DESCR_inv[variable]
if var_inv not in cls.ELEMENTS and var_inv not in array_variables:
kwargs["track_extra_variables"].append(var_inv)
kwargs["raw_data"] = raw_data
kwargs["only_variables"] = (
None
if include_vars is None
else [VAR_DESCR_inv[i] for i in include_vars]
)
kwargs.update(class_kwargs)
eddies = cls(size=nb_obs, **kwargs)
for variable in var_list:
var_inv = VAR_DESCR_inv[variable]
# Patch
h_nc.variables[variable].set_auto_maskandscale(not raw_data)
logger.debug(
"Up load %s variable%s",
variable,
", with raw mode" if raw_data else "",
)
# find unit factor
factor = 1
if not raw_data:
input_unit = getattr(h_nc.variables[variable], "unit", None)
if input_unit is None:
input_unit = getattr(h_nc.variables[variable], "units", None)
output_unit = VAR_DESCR[var_inv]["nc_attr"].get("units", None)
factor = cls.compare_units(input_unit, output_unit, variable)
if indexs is None:
indexs = dict()
var_sl = [
indexs.get(dim, slice(None))
for dim in h_nc.variables[variable].dimensions
]
if factor != 1:
eddies.obs[var_inv] = h_nc.variables[variable][var_sl] * factor
else:
eddies.obs[var_inv] = h_nc.variables[variable][var_sl]