ialdev-maths 0.1.0__py3-none-any.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.
iad/maths/hist.py ADDED
@@ -0,0 +1,1458 @@
1
+ from copy import deepcopy
2
+ from typing import Literal, Callable, Any, Iterable, overload, Collection, get_args, Mapping
3
+
4
+ import numba as nb
5
+ import numpy as np
6
+ import pandas as pd
7
+ from numpy.typing import NDArray
8
+
9
+ from iad.core import as_list
10
+ from iad.core.datatools import select_from
11
+
12
+ __all__ = ['Sampler', 'log_compress', 'bins_edges', 'StatGather', 'StatGather2D']
13
+
14
+ _cache_nb = True
15
+
16
+ Number = float | int | np.number
17
+ _UNDEF = object()
18
+
19
+ _NormT = Literal['total', 'range', 'nonan', None, False, True]
20
+ _Axis = Literal[0, 1]
21
+
22
+ HistT = np.uint32 # counters type
23
+
24
+ SUM_METRICS = Literal['s1', 's2', 'finite', 'total']
25
+ BASIC_METRICS = Literal['mean', 'rmse', 'std']
26
+ METRICS = Literal[SUM_METRICS, BASIC_METRICS]
27
+
28
+ _sum_metrics = get_args(SUM_METRICS)
29
+ _basic_metrics = get_args(BASIC_METRICS)
30
+
31
+
32
+ def pack_bit_arrays_64(arrays):
33
+ """
34
+ Pack list of binary arrays into single array of same shape,
35
+ with every initial arrays represented by its bit-plane.
36
+
37
+ Supports up to 64 arrays, raise ValueError if given more than that.
38
+
39
+ Resulting array dtype is u[8,16,32, 64] depending on the number of inputs.
40
+
41
+ :param arrays:
42
+ :return: bit-packed array
43
+ """
44
+ from iad.core.binary import align_type_bits
45
+ n = len(arrays)
46
+ shapes = [a.shape for a in arrays]
47
+ shape = shapes[0]
48
+ if not all(s == shape for s in shapes[1:]):
49
+ raise ValueError(f"All arrays must be of same shape: {shapes}")
50
+
51
+ bits = align_type_bits(n)
52
+ dtype = np.dtype(f'uint{bits}')
53
+
54
+ out = np.zeros(shape, dtype=dtype)
55
+ _b = np.zeros(1, dtype=dtype)
56
+ arrays = tuple(np.ascontiguousarray(a) for a in arrays)
57
+ _pack_bit_array(arrays, out, _b)
58
+ return out
59
+
60
+
61
+ @nb.njit
62
+ def _pack_bit_array(arrays, out, _b):
63
+ out = out.flat
64
+ for shift, array in enumerate(arrays):
65
+ for i, a in enumerate(array.flat):
66
+ _b[0] = a << shift # force casting into out type
67
+ out[i] += _b[0]
68
+
69
+
70
+ def init_sum_stats(*, num_regions: int, ndim: Literal[1, 2], dtype=np.float64) -> dict[SUM_METRICS, NDArray]:
71
+ """
72
+ Prepare accumulating arrays data structure for different configurations.
73
+
74
+ If ``num_regions`` > 0 than creates 2D arrays to have separate dimension for regions.
75
+
76
+ **Notice!** that size of regions dimension is ``num_regions`` + 1 for separate 'full' statistics.
77
+
78
+ Return dict {metric: zeros[``num_regions + 1`` × ``ndim``] | zeros[``ndim``]}
79
+
80
+ :param num_regions: number of selection masks collecting separate statistics
81
+ :param ndim: number of variables to collect statistics for
82
+ :param dtype: type of values accumulator (use default)
83
+ :return:
84
+ """
85
+ shape = (
86
+ (1 + num_regions, ndim) if num_regions else
87
+ (ndim,) if ndim > 1 else ()
88
+ ) # ToDo: masks_num > 1 and ndim == 1?
89
+
90
+ return {'s1': np.zeros(shape, dtype=dtype),
91
+ 's2': np.zeros(shape, dtype=dtype),
92
+ 'finite': np.zeros(shape, dtype=np.int64),
93
+ 'total': np.zeros(shape, dtype=np.int64)}
94
+
95
+
96
+ def _basic_stats(s1, s2, finite, **_) -> dict[BASIC_METRICS, NDArray | Number]:
97
+ """Calculate `mean, msre, std` from the accumulated sums.
98
+ If nothing is accumulated, set them to nan.
99
+ :param s1: sum of finite values
100
+ :param s2: sum of their squares
101
+ :param finite: their number
102
+ :return: dict(mean=..., std=...)
103
+ """
104
+ if not np.asarray(finite).any():
105
+ return select_from({}, _basic_metrics, default=np.nan)
106
+ with np.errstate(all='ignore'):
107
+ m1, m2 = s1 / finite, s2 / finite
108
+ return {'mean': m1,
109
+ 'rmse': np.sqrt(m2),
110
+ 'std': np.sqrt(m2 - m1 * m1)} # ToDo: Add test of matching keys to _basic_measures
111
+
112
+
113
+ @nb.jit(nopython=True, cache=_cache_nb)
114
+ def _equal_bins_stats(a, low, bin_size, hist: np.ndarray,
115
+ log_scale: float = 0, sums=False):
116
+ """
117
+ Low level optimized function to collect various typical statistical measurement over array:
118
+ 1. Histogram with bins *equidistant* in either the `data` or `log1p(data)` space.
119
+ - *The histogram array must be pre-allocated and is filled by the function.*
120
+ 2. Basic statistical measurements over array of numbers:
121
+ - Counts distribution between equal-sized bins distribution
122
+ - Calculates `sum(x)`, `sum(x**2)`, `count(x)` for *finite* `x`.
123
+
124
+ Special 3 bin are created for:
125
+ values outside the given range and nan, with indices:
126
+ * 0 for x <= low (below)
127
+ * -2 for x > high (inf)
128
+ * -1 for nans (nans)
129
+
130
+ **Note**: for case of `data >= 0` to use the first (below) bin for `x > 0`, set `low=bin_size`.
131
+ ::
132
+ low high inf nan <- right edges[]
133
+ _________|_________|____...____|_________|_________|_________|
134
+ low_bin=0 high_bin inf_bin nan_bin <- bins indices
135
+
136
+ :param a: array - could be float
137
+ :param low: right edge of the lowest bin (must be finite)
138
+ :param bin_size: size of every bin
139
+ :param hist: hist array to be updated inside
140
+ :param sums: if `True` also calculate for finite values their number, sum and sum of squares
141
+ :param log_scale: flag to pass the data through log before distributing between the bins
142
+ :return (s1, s2, finite, total)
143
+ """
144
+ bins = len(hist)
145
+ high_bin, inf_bin, nan_bin = range(bins - 3, bins) # high_bin == number of bins in [low, high]
146
+
147
+ finite: int = 0 # count finite values
148
+ s1 = s2 = 0. # accumulate sum and sum of squares over finite values
149
+ for x in a.flat:
150
+ if np.isnan(x):
151
+ bin_id = nan_bin
152
+ elif np.isinf(x):
153
+ bin_id = inf_bin if x > 0 else 0 # below
154
+ else:
155
+ if sums: # statistics is calculated on uncompressed data
156
+ s1 += x # requires attention in parallel mode to avoid race condition!
157
+ s2 += x * x
158
+ finite += 1
159
+
160
+ if log_scale != 0:
161
+ x = _log_compress(x, log_scale)
162
+
163
+ # here x is finite, but still can fall into [< low] or [> high] bins.
164
+ bin_id = int((x - low) / bin_size) + 1 # +1 since 0 for < low
165
+ if bin_id > high_bin:
166
+ bin_id = high_bin
167
+ elif bin_id < 0:
168
+ bin_id = 0
169
+
170
+ hist[bin_id] += 1
171
+
172
+ return s1, s2, finite, a.size
173
+ # return dict(
174
+ # [('s1', s1), ('s2', s2), ('finite', finite), ('total', a.size)]) # this form is eeasy for nb :-
175
+
176
+
177
+ @nb.njit(inline='always', cache=_cache_nb)
178
+ def _find_bin(x, low, bin_size, nan_bin, log_scale):
179
+ """find bin id ∈ `[0, bins-1]` for `x` ∈ `[low-bins_size, high]` with 3 special bins:
180
+ ::
181
+ 0: x < low
182
+ nan_bin - 1: x > high
183
+ nan_bin: for x == nan
184
+
185
+ :param x: value *before compression*
186
+ :param low: right edge of first bin inside the `x` range (bin id = 1)
187
+ :param bin_size: size of the bin *after compression*
188
+ :param nan_bin: id of the bin for `nan`s ( = bins - 1)
189
+ :param log_scale: values compression factor, 0 for no compression
190
+ """
191
+ if np.isnan(x): # first deal with nan to avoid nan arithmetics
192
+ bin_id = nan_bin
193
+ else:
194
+ if log_scale != 0:
195
+ x = _log_compress(x, log_scale)
196
+ bin_id = int((x - low) / bin_size) + 1 # +1 since 0 for < low
197
+
198
+ if bin_id < 0: # apply range clipping on the calculated bin_id
199
+ bin_id = 0
200
+ elif bin_id >= nan_bin:
201
+ bin_id = nan_bin - 1
202
+ return bin_id
203
+
204
+
205
+ def _equal_bins2d_stats(a1, a2, *,
206
+ low: tuple[Number, Number],
207
+ bin_size: tuple[Number, Number],
208
+ masks: Collection[np.ndarray] | None = None,
209
+ hist: np.ndarray = None,
210
+ log_scale: tuple[Number, Number] = (0, 0),
211
+ stats: dict[SUM_METRICS, NDArray[Number]] | None = None,
212
+ sums=True) -> dict[SUM_METRICS, NDArray]:
213
+ """
214
+ Wraps low level compiled versions to:
215
+ 1. support both with and without ``masks`` cases
216
+ 2. bring input data into the required shape
217
+ 3. create ``stats`` data structure if not passed in
218
+
219
+ :param a1:
220
+ :param a2:
221
+ :param low:
222
+ :param bin_size:
223
+ :param masks:
224
+ :param hist:
225
+ :param log_scale:
226
+ :param stats:
227
+ :param sums:
228
+ :return: stats
229
+ """
230
+
231
+ assert all(len(_) == 2 for _ in (low, bin_size, log_scale))
232
+ a1, a2 = map(np.asarray, [a1, a2])
233
+ assert a1.shape == a2.shape
234
+
235
+ num_masks = 0 if masks is None else len(masks)
236
+ if not stats:
237
+ stats = init_sum_stats(num_regions=num_masks, ndim=2)
238
+
239
+ a1, a2 = (a.reshape(-1) for a in [a1, a2]) # _equal.. requires 1D for a and
240
+ if num_masks:
241
+ assert hist.ndim == 3 and hist.shape[0] == num_masks + 1
242
+ masks = pack_bit_arrays_64(masks).reshape(-1)
243
+ _nb_equal_bins2m_stats(a1, a2, masks, low=low, bin_size=bin_size, hist=hist,
244
+ **stats, log_scale=log_scale, sums=sums)
245
+ else:
246
+ assert hist.ndim == 2
247
+ a1, a2 = (a.reshape(-1) for a in [a1, a2]) # _equal.. requires 1D for a and
248
+ _nb_equal_bins2_stats(a1, a2, low=low, bin_size=bin_size, hist=hist,
249
+ **stats, log_scale=log_scale, sums=sums)
250
+ return stats
251
+
252
+
253
+ @nb.jit(nopython=True, cache=_cache_nb)
254
+ def _nb_equal_bins2m_stats(a1: np.ndarray, a2: np.ndarray, masks: np.ndarray,
255
+ low: tuple, bin_size: tuple, hist: np.ndarray,
256
+ s1: np.ndarray, s2: np.ndarray, finite: np.ndarray,
257
+ total: np.ndarray, log_scale: tuple = (0, 0), sums=True):
258
+ """
259
+ Compiled low level version with special requirements for some of its inputs.
260
+
261
+ Also, there are some additional conventions:
262
+ - results are provided by *incrementing* input arguments: ``hist``, ``s1``, ``s2``, ``finite``
263
+ - ``hist`` array is a stack of 2d arrays (for every mask + 1 for *full*)
264
+ - accumulating arrays (``s1``, ``s2``, ``finite``, ``total``) also have additional row for full
265
+ - number of bins is deduced from the shape of ``hist``
266
+ - to calculate *without* ``masks`` ensure `masks.shape[0] == 0`
267
+ - ``low`` tuple[2] represents right edges of first range bin for 2 dimensions (a1, a2)
268
+ - ``bin_size`` tuple[2] also represents 2 dimensions
269
+
270
+ :param a1: ndarray[`size`] 1D
271
+ :param a2: ndarray[`size`] 1D
272
+ :param masks: ndarray[masks_num × `size`] or ndarray[`masks_num = 0` × any]
273
+ :param low: right edges of the first inside the values range (per dimension)
274
+ :param bin_size: sine of bin > 0, per dimension
275
+ :param hist: ndarray[1 + `masks_num` × bins1 × bins2]
276
+ :param s1: ndarray[1 + `masks_num` × 2] sum of finite values per mask per dimension
277
+ :param s2: ndarray[1 + `masks_num` × 2] sum of squares of finite values per mask per dimension
278
+ :param finite: ndarray[1 + `masks_num` × 2] number of the summed finite values per mask per dimension
279
+ :param log_scale: compression scale factor per dimension
280
+ :param sums: calculate stats by updating arguments: ``s1``, ``s2``, ``finite``
281
+ :return:
282
+ """
283
+ assert a1.shape == a2.shape and a1.ndim == 1 and hist.ndim == 3
284
+ # num_masks = masks.shape[0] # always add first (0) "full" mask
285
+ # assert masks.size == num_masks * a1.size
286
+ num_hist, bins = hist.shape[0], hist.shape[1:]
287
+ # assert num_hist == 1 + num_masks
288
+
289
+ one = HistT(1)
290
+
291
+ for i in range(a1.size): # over data elements (pixels)
292
+ v1, v2, m = a1[i], a2[i], masks[i]
293
+ b1 = _find_bin(v1, low[0], bin_size[0], bins[0] - 1, log_scale[0])
294
+ b2 = _find_bin(v2, low[1], bin_size[1], bins[1] - 1, log_scale[1])
295
+
296
+ for ii in range(num_hist): # over 1(full) + len(masks) regions
297
+ if not ii or ((m >> (ii - 1)) & 1): # first (ii == 0) is the virtual "full" region
298
+ hist[ii, b1, b2] += one
299
+ for d in range(2):
300
+ v = v2 if d else v1
301
+ if sums: # ToDo: move above for d
302
+ if np.isfinite(v): # statistics is calculated on uncompressed data
303
+ s1[ii, d] += v
304
+ s2[ii, d] += v * v
305
+ finite[ii, d] += one
306
+ total[ii, d] += one
307
+
308
+
309
+ @nb.jit(nopython=True, cache=_cache_nb)
310
+ def _nb_equal_bins2_stats(a1: np.ndarray, a2: np.ndarray,
311
+ low: tuple, bin_size: tuple, hist: np.ndarray,
312
+ s1: np.ndarray, s2: np.ndarray, finite: np.ndarray,
313
+ total: np.ndarray, log_scale: tuple = (0, 0), sums=True):
314
+ """
315
+ Compiled low level version with special requirements for some of its inputs.
316
+
317
+ Also, there are some additional conventions:
318
+ - results are provided by *incrementing* input arguments: ``hist``, ``s1``, ``s2``, ``finite``
319
+ - hist array is a stack of 2d arrays per every masks + 1 for *full*
320
+ - number of bins is deduced from the shape of ``hist``
321
+ - to calculate *without* ``masks`` ensure `masks.shape[0] == 0`
322
+ - ``low`` tuple[2] represents right edges of first range bin for 2 dimensions (a1, a2)
323
+ - ``bin_size`` tuple[2] also represents 2 dimensions
324
+
325
+ :param a1: ndarray[`size`] 1D
326
+ :param a2: ndarray[`size`] 1D
327
+ :param low: right edges of the first inside the values range (per dimension)
328
+ :param bin_size: sine of bin > 0, per dimension
329
+ :param hist: ndarray[1 + `masks_num` × bins1 × bins2]
330
+ :param s1: ndarray[1 + `masks_num` × 2] sum of finite values per mask per dimension
331
+ :param s2: ndarray[1 + `masks_num` × 2] sum of squares of finite values per mask per dimension
332
+ :param finite: ndarray[1 + `masks_num` × 2] number of the summed finite values per mask per dimension
333
+ :param total: total number of data elements processed per mask per dimension
334
+ :param log_scale: compression scale factor per dimension
335
+ :param sums: calculate stats by updating arguments: ``s1``, ``s2``, ``finite``
336
+ :return:
337
+ """
338
+ bins = hist.shape
339
+ for v1, v2 in zip(a1, a2):
340
+ b1 = _find_bin(v1, low[0], bin_size[0], bins[0] - 1, log_scale[0])
341
+ b2 = _find_bin(v2, low[1], bin_size[1], bins[1] - 1, log_scale[1])
342
+ hist[b1, b2] += 1
343
+ for d in range(2):
344
+ v = v2 if d else v1
345
+ if sums:
346
+ if np.isfinite(v): # statistics is calculated on uncompressed data
347
+ s1[d] += v
348
+ s2[d] += v * v
349
+ finite[d] += 1
350
+ total[d] += 1
351
+
352
+
353
+ #
354
+ # @nb.jit(nopython=True)
355
+ # def _equal_bins_stats(a, low, bin_size, hist: np.ndarray,
356
+ # log_scale: float = 0, sums=False):
357
+ # """
358
+ # Low level optimized function to collect various typical statistical measurement over array:
359
+ # 1. Histogram with bins *equidistant* in either the `data` or `log1p(data)` space.
360
+ # - *The histogram array must be pre-allocated and is filled by the function.*
361
+ # 2. Basic statistical measurements over array of numbers:
362
+ # - Counts distribution between equal-sized bins distribution
363
+ # - Calculates `sum(x)`, `sum(x**2)`, `count(x)` for *finite* `x`.
364
+ # *Stats are returned explicitly, as dict, even if empty array is given*.
365
+ #
366
+ # Special 3 bin are created for:
367
+ # values outside the given range and nan, with indices:
368
+ # * 0 for x <= low (below)
369
+ # * -2 for x > high (inf)
370
+ # * -1 for nans (nans)
371
+ #
372
+ # **Note**: for case of `data >= 0` to use the first (below) bin for `x > 0`, set `low=bin_size`.
373
+ # ::
374
+ # low high inf nan <- right edges[]
375
+ # _________|_________|____...____|_________|_________|_________|
376
+ # low_bin=0 high_bin inf_bin nan_bin <- bins indices
377
+ #
378
+ # :param a: array - could be float
379
+ # :param low: right edge of the lowest bin (must be finite)
380
+ # :param bin_size: size of every bin
381
+ # :param sums: if `True` also calculate for finite values their number, sum and sum of squares
382
+ # :param log_scale: flag to pass the data through log before distributing between the bins
383
+ # :return dict.keys() = (s1, s2, finite, total)
384
+ # """
385
+ # bins = len(hist)
386
+ # high_bin, inf_bin, nan_bin = range(bins - 3, bins) # high_bin == number of bins in [low, high]
387
+ #
388
+ # finite: int = 0 # count finite values
389
+ # s1 = s2 = 0. # accumulate sum and sum of squares over finite values
390
+ # for x in a.flat:
391
+ # if np.isnan(x):
392
+ # bin_id = nan_bin
393
+ # elif np.isinf(x):
394
+ # bin_id = inf_bin if x > 0 else 0 # below
395
+ # else:
396
+ # if sums: # statistics is calculated on uncompressed data
397
+ # s1 += x # requires attention in parallel mode to avoid race condition!
398
+ # s2 += x * x
399
+ # finite += 1
400
+ #
401
+ # if log_scale != 0:
402
+ # x = _log_compress(x, log_scale)
403
+ #
404
+ # # here x is finite, but still can fall into [< low] or [> high] bins.
405
+ # bin_id = int((x - low) / bin_size) + 1 # +1 since 0 for < low
406
+ # if bin_id > high_bin:
407
+ # bin_id = high_bin
408
+ # elif bin_id < 0:
409
+ # bin_id = 0
410
+ #
411
+ # hist[bin_id] += 1
412
+ #
413
+ # return dict([('s1', s1), ('s2', s2),
414
+ # ('finite', finite), ('total', a.size)]) # easy to compile
415
+ #
416
+
417
+
418
+ @nb.njit(cache=_cache_nb)
419
+ def index_count(indices, num):
420
+ """Count histogram of integer numbers representing indices ranging from 0 to num-1."""
421
+ counts = np.zeros(num, dtype=np.int32)
422
+
423
+ for i in indices.flat:
424
+ counts[i] += 1
425
+ return counts
426
+
427
+
428
+ @nb.njit(cache=_cache_nb)
429
+ def _log_compress(v, scale) -> NDArray[np.float64] | np.float64:
430
+ """Apply logarithmic scale compression or de-compression to array or number.
431
+ Signed data is processed as well according to:
432
+ ::
433
+ compressed(x) = sign(x) * log(1 + abs(x / scale)) (decompression is inverse of that)
434
+
435
+ :param v: array or number to compress decompress
436
+ :param scale: direction and scale of compression:
437
+ `+` (*positive*): compress,
438
+ `-` (*negative*): decompress,
439
+ `0`: not supported for scale == 0
440
+
441
+ :return transformed array
442
+ """
443
+ scale = nb.float64(scale)
444
+
445
+ direct = scale > 0
446
+ scale = np.fabs(scale)
447
+ if direct:
448
+ out = np.log1p(np.fabs(v / scale))
449
+ else:
450
+ out = np.expm1(np.fabs(v)) * scale
451
+ return np.copysign(out, v)
452
+
453
+
454
+ def log_compress(v: NDArray[Number] | Number, scale: float = 1) -> NDArray[np.float64]:
455
+ """Apply logarithmic scale compression or de-compression to array or number.
456
+ Signed data is processed as well according to:
457
+ ::
458
+ compressed(x) = sign(x) * log(1 + abs(x / scale)) (decompression is inverse of that)
459
+
460
+ :param v: array or number to compress decompress
461
+ :param scale: direction and scale of compression:
462
+ `+` (*positive*): compress,
463
+ `-` (*negative*): decompress,
464
+ `0`: disable and return the input
465
+ :return transformed array
466
+ """
467
+ if not (np.isscalar(v) or isinstance(v, np.ndarray)):
468
+ v = np.fromiter(v, np.float32)
469
+ if not scale:
470
+ return v
471
+ return _log_compress(v, scale)
472
+
473
+
474
+ def bins_edges(low, high, bins: int, below=True) -> NDArray[Number]:
475
+ """
476
+ Creates array of bins *right* edges for histogram, given:
477
+ - range limits (`low, `high`)
478
+ - and number of bins to cover this range.
479
+ That is the *bin size* is always `(high-low)/bins`
480
+
481
+ Resulting `edges` array includes two *special* elements at the end: `inf`, `nan`,
482
+ to reserve bins for correspondingly values > high, and 'nan' (requires `high < inf` !)
483
+
484
+ Additional bin for `values < low` is controlled by `below` argument:
485
+ - ``below`` is `True`, first bin has its right `edges[0] = low`
486
+ ::
487
+
488
+ low high inf nan <- right edge value
489
+ <..._|______|_ ... _|______|_...> | ~~~~ |
490
+ 0 1 bins bins+1 bins+2 <- bin/edge id
491
+
492
+ - ``below`` is `False`, first bin collects [low, low+step], `edges[0] = low + bin`
493
+ ::
494
+
495
+ low+step high inf nan <- right edge value
496
+ |______|_ ... _|______|_...> | ~~~~ |
497
+ 0 bins-1 bins bins+1 <- bin/edge id
498
+
499
+ That is, the size of the resulting `edges` array: `bins + 2 + int(below)`.
500
+ In all the cases, it *always* represents *right* edges of the bins.
501
+
502
+ :param low: **left** edge of the binning range
503
+ :param high: *right* edge of the binning range (must obey `high > low`)
504
+ :param bins: number of *whole* bins between `low` and `high` values
505
+ :param below: expect and count data below the ``low`` value
506
+ :return: edges
507
+ """
508
+ x_range = high - low
509
+ bin_size = x_range / bins
510
+
511
+ assert np.isfinite(x_range) and x_range > 0
512
+ assert bins > 0
513
+
514
+ edges = np.arange(
515
+ start=low + (0 if below else bin_size), # create a bin for x < low
516
+ stop=high + (0.01 + 2) * bin_size, # +0.1 to ensure high is included
517
+ step=bin_size
518
+ )
519
+ edges[-2:] = np.inf, np.nan # special bins
520
+ return edges
521
+
522
+
523
+ # class SamplerModel(YamlModel, extra='forbid'):
524
+ # low: float
525
+ # high: float
526
+ # bins: Optional[int]
527
+ # step: Optional[float]
528
+ # log_scale: float = 0
529
+ # below: bool = True
530
+ # zero: bool | None
531
+ # relax: Literal['low', 'high', 'step'] = 'relax'
532
+ # name: str = ''
533
+ #
534
+ # @root_validator(pre=True)
535
+ # def _params(cls, values):
536
+ # from iad.core.datatools import select_from
537
+ # values['high'] = Sampler._val_range(**select_from(values, ['high', 'low', 'bins', 'step']))
538
+ # return values
539
+ #
540
+
541
+ class Sampler:
542
+ @classmethod
543
+ def from_range(cls, low: Any, high: Any, step_or_bins: float | int, name: str = ''):
544
+ """Helper constructor useful when constructing from tuples:
545
+ - (`low`, `high`, `step`: ``float``), or
546
+ - (`low`, `high`, `bins`: ``int``)
547
+
548
+ Distinguished by the `step` type!
549
+ """
550
+ return cls(low, high, name=name, **{
551
+ 'step' if isinstance(step_or_bins, float) else 'bins': step_or_bins})
552
+
553
+ def __init__(self, low, high=None, *, below=True, step=None, bins=None, zero: bool | None = None,
554
+ log_scale: Number | bool = False, name='', relax: Literal['low', 'high', 'step'] = 'high'):
555
+ """
556
+
557
+ Create *equidistant* sampling of given interval for binning, meshes and similar needs.
558
+
559
+ Allows defining the sampling using either:
560
+ - min, max and bins
561
+ - min, max and step
562
+
563
+ In the second case provided parameters may be inconsistent with even sampling.
564
+ In this case number bins is rounded up and one of the *min, max, step* is relaxed to fit.
565
+
566
+ ----------
567
+
568
+ **Sampling of 0**
569
+
570
+ Argument ``zero`` requires `0` value to be placed:
571
+ - *on* a sample if `True`
572
+ - in the *middle* between two successive samples if `False`
573
+ - `None` for no requirements
574
+
575
+ It adjusts bins number and is applicable ONLY for symmetric ranges: `-low=high`.
576
+
577
+ **Range Log Compression**
578
+
579
+ ``Sampler`` supports *log compression* of the sampled range to increase inter-sampling distance
580
+ as values grow.
581
+
582
+ That is achieved by applying ``log_compress`` function on both:
583
+ - the range parameters ``low``, ``high``, ``step``, and
584
+ - the incoming data being processed
585
+
586
+ In this case sampling remains *equidistant* in the compressed (not the actual values) space.
587
+ That is reflected in the `edges` attribute, which will produce exponentially uncompressed samples.
588
+
589
+ *Notice* specific meaning of ``step`` argument in this case, such that `log_compress(step)`
590
+ corresponds to the equal samples distance in the compressed space.
591
+
592
+ :param low: low edge of the sampling interval
593
+ :param high: high edge of the sampling interval
594
+ :param step: sampling step
595
+ :param bins: number of bins (either that or steps must be given)
596
+ :param name: Optional name of the sampler
597
+ :param below: True to allow values < `low`
598
+ :param zero: ``True``: sample at **0**, ``False``: **0** at half `step`, ``None`` - don't care
599
+ :param relax: when step is given - which one of `min`, `max`, `step` to fit.
600
+ """
601
+ # --------- validations --------------
602
+ high = self._val_range(low, high, bins, step)
603
+
604
+ # ------------ initialization -------------
605
+ self.low, self.high, self.step = low, high, step
606
+ self.log_scale = float(log_scale)
607
+ self.below = below # collect below the range low limit
608
+ self.bins = bins
609
+ self.name = str(name)
610
+ self._adjust_sampling(zero, relax)
611
+
612
+ @staticmethod
613
+ def _val_range(low, high, bins, step):
614
+ if high is None:
615
+ high = low + step * bins
616
+ else:
617
+ assert high > low
618
+ assert (bins and not step) or (step and not bins), "Both bins and step are provided"
619
+ assert not step or step < high - low
620
+ assert not bins or int(bins) == bins and bins > 0
621
+
622
+ return high
623
+
624
+ def low_high_step(self, compress=False):
625
+ """Return tuple `(low, high, step)` in the original or compressed space."""
626
+ params = (self.low, self.high, self.step)
627
+ if compress and self.log_scale:
628
+ return tuple(log_compress(_, self.log_scale) for _ in params)
629
+ return params
630
+
631
+ def _adjust_sampling(self, zero, relax):
632
+ """Adjust sampling to the constraints in the COMPRESSED space!"""
633
+ low, high, step = self.low_high_step(compress=True)
634
+
635
+ def bins_zero_correction(_bins):
636
+ # make even or odd depending on the goal
637
+ if zero is None or np.isclose(low, 0): # we always allow [0, high]
638
+ return 0
639
+ elif zero not in (True, False):
640
+ raise ValueError(f"Invalid value {zero=}")
641
+
642
+ if not np.isclose(low, -high):
643
+ raise ValueError(f"Sampling at 0 can't may be adjusted with range {[low, high]}")
644
+
645
+ fix_even = _bins % 2 # in symmetric case if 1/2 sample ensure
646
+ return fix_even if zero else 1 - fix_even
647
+
648
+ if self.bins: # high, low, bins -> step
649
+ self.bins += bins_zero_correction(self.bins)
650
+ step = (high - low) / self.bins
651
+ else:
652
+ bins = (high - low) / step
653
+ self.bins = int(bins + .5) + bins_zero_correction(int(bins + .5))
654
+ if not np.isclose(self.bins, bins):
655
+ match relax:
656
+ case 'low':
657
+ low = high - self.bins * step
658
+ case 'high':
659
+ high = low + self.bins * step
660
+ case 'step':
661
+ step = (high - low) / self.bins
662
+ case False | None | 'none':
663
+ raise ValueError(f"Can't meet sampling constraints with {relax=}!")
664
+ case _:
665
+ raise ValueError(f"Unsupported 'relax' value {relax}")
666
+ # return to the uncompressed space
667
+ self.low, self.high, self.step = log_compress([low, high, step], -self.log_scale)
668
+
669
+ def __eq__(self, other):
670
+ return all(np.isclose(getattr(self, attr), getattr(other, attr))
671
+ for attr in ['low', 'high', 'step', 'bins'])
672
+
673
+ def bins_edges(self, compress=True) -> NDArray[np.number]:
674
+ """Create right edges for binning.
675
+
676
+ Includes two elements at the end:
677
+ [-2] = `inf` for bin counting `values > high`
678
+ [-1] = `nan` for bin counting `nan`s
679
+
680
+ **Note**, `compress` has effect only if `Sampler` is log-compressed.
681
+
682
+ By default, (`True`) return equidistant
683
+
684
+ :param compress: relevant if ``Sampler`` has `log compression` activated.
685
+ returned `edges` the *compressed* space or the original values space.
686
+ """
687
+ low, high, _ = self.low_high_step(compress=True) # bins_edges operates in equidistant space!
688
+ edges = bins_edges(low=low, high=high, bins=self.bins, below=self.below)
689
+ if self.log_scale and compress is False: # if compression is active
690
+ return log_compress(edges, -self.log_scale)
691
+ return edges
692
+
693
+ def near_index_below(self, values, compressed: bool = False, fp_index=False):
694
+ """Return indices from the nearest edges below the given values, and the offsets from them.
695
+
696
+ If offsets are returned in same compression state as the input values
697
+
698
+ :param values: values to search the nearest index for
699
+ :param compressed: indicate if values has been already compressed (relevant if log_scale)
700
+ :param fp_index: return only index with fp precision
701
+ :return: indices (int), offsets (right from edges), index with fp precision
702
+ """
703
+ if not compressed:
704
+ values = log_compress(values, self.log_scale)
705
+ fp_idx = (values - self.low) / self.step
706
+ if fp_index:
707
+ return fp_idx
708
+
709
+ idx = np.int32(fp_idx)
710
+ offsets = fp_idx - idx
711
+ if compressed:
712
+ offsets = log_compress(offsets, -self.log_scale)
713
+ return np.int32(fp_idx), offsets, fp_idx
714
+
715
+ @property
716
+ def bins_centers(self):
717
+ """
718
+ Return position of the centers of the bins in original data space.
719
+
720
+ :param below: if `below` is `False`, center of the first bin > `low`, otherwise < `low`.
721
+ """
722
+ centers = self.bins_edges(compress=True)[:-2] - self.step / 2
723
+ if self.below:
724
+ centers = centers[1:]
725
+ return log_compress(centers, -self.log_scale)
726
+
727
+ @property
728
+ def steps(self):
729
+ prepend = {} if self.below else {'prepend': self.low}
730
+ return np.diff(self.bins_edges(compress=False)[:-2], **prepend)
731
+
732
+ def __repr__(self):
733
+ name = f"<{self.name}>" if self.name else ''
734
+ tos = lambda _: f"{_:.3g}"
735
+ items = self.low_high_step(compress=True)
736
+ left = '⍇' if self.below else '['
737
+ rng = f"{left}{':'.join(map(tos, items))}]"
738
+
739
+ if self.log_scale:
740
+ org = self.bins_edges(compress=False)[:-2]
741
+ if len(org) < 8:
742
+ seq = ','.join(map('{:.3g}'.format, org))
743
+ else:
744
+ seq = (', '.join(map('{:.3g}'.format, org[:3]))
745
+ + ', ..., ' +
746
+ ', '.join(map('{:.3g}'.format, org[-3:])))
747
+ rng = f"{rng} → [{seq}]"
748
+
749
+ return f"{self.__class__.__name__}{name}[{self.bins}] {rng}"
750
+
751
+
752
+ def _decode_metrics_arg(metrics):
753
+ match metrics:
754
+ case True:
755
+ return _sum_metrics + _basic_metrics
756
+ case False:
757
+ return ()
758
+ case None:
759
+ return _basic_metrics
760
+ case [_] | (_):
761
+ if unknown := set(metrics).difference(_basic_metrics + _sum_metrics):
762
+ raise NameError(f"Requested metrics are {unknown = }")
763
+ return as_list(metrics, collect=tuple)
764
+
765
+
766
+ class StatGather:
767
+ """
768
+ Basic Statistics Calculator.
769
+
770
+ Operates in accumulative and stateless modes.
771
+
772
+ Supports 3 main types of statistics:
773
+ 1. histogram, with bins defined by the `sampler` argument
774
+ 2. basic statistical metrics:
775
+ - `mean`, `std`, `rmse`, counters over all and finite values
776
+ 3. counters of values under specific named *levels*
777
+
778
+ Levels counters are
779
+ """
780
+
781
+ # FixMe: @Ilya align names with 2D
782
+ # FixMe: @Ilya reuse with 2D
783
+ # FixMe: stats=None does not work
784
+ def __init__(self, sampler: Sampler,
785
+ arrays: Iterable[np.ndarray] | np.ndarray = (), *,
786
+ name='', stats: bool | list[METRICS] = True,
787
+ levels: Mapping[str, Number] | None = None):
788
+ """
789
+ Initialize data gatherer for histogram and statistical metrics:
790
+ ::
791
+
792
+ SUM_METRICS = Literal['s1', 's2', 'finite', 'total']
793
+ BASIC_METRICS = Literal['mean', 'rmse', 'std']
794
+
795
+ Allows extension of the basic statistical metrics with accumulative hist sampled at given levels.
796
+
797
+ :param sampler: defines binning for the histogram (see ``Sampler`` for details)
798
+ :param arrays: optionally one or `Iterable` over multiple arrays to process.
799
+ :param stats: if ``False`` skip gathering and calculating the basic statistical `measures`
800
+ :param levels: add to the `stats` accumulations by levels
801
+ :param name: useful for logistics of multiple instances of gatherers.
802
+ """
803
+ self._sampler = sampler
804
+
805
+ self._equid_edges = sampler.bins_edges()
806
+ self._calc_stats = _decode_metrics_arg(stats)
807
+ self._appends = None
808
+
809
+ self.name = name
810
+ self.levels = levels
811
+ self.measures = {}
812
+ self.hist = None
813
+
814
+ self.reset() # initializes all three
815
+
816
+ if isinstance(arrays, np.ndarray):
817
+ arrays = [arrays]
818
+ for a in map(np.asarray, arrays):
819
+ if a.ndim and a.size:
820
+ self.process(a, update=True)
821
+
822
+ def __repr__(self):
823
+ from iad.core.strings import dict_str
824
+ sm = self._sampler
825
+ compress, step = ('', f"{sm.step:.3g}") if not sm.log_scale else (
826
+ "🗜", "{:.3g}→{:.3g}".format(*sm.steps[[0, -1]])
827
+ )
828
+ rng = f"{sm.low:.3g}:{sm.high:.3g}:{step}"
829
+ measures = dict_str(select_from(self.measures, _basic_metrics), prec=3, sep='|', to=':')
830
+ below = "<-|" if sm.below else "["
831
+
832
+ name = f"<{self.name}>" if self.name else ''
833
+ return (f"{self.__class__.__name__}{name} "
834
+ f"{below}{rng}]{compress} "
835
+ f"∑{self._appends}{{{measures}}}")
836
+
837
+ @property
838
+ def edges(self):
839
+ """Edges in the original values space (decompressed)"""
840
+ return self._sampler.bins_edges(compress=False)
841
+
842
+ def _empty_hist(self, hist=None):
843
+ """Prepare empty histogram array, by creating new or zeroing existing"""
844
+ if hist is None:
845
+ hist = np.zeros_like(self._equid_edges, dtype=HistT)
846
+ else:
847
+ hist[:] = 0
848
+ return hist
849
+
850
+ def _calc_hist_stats(self, a: np.ndarray | None, hist: np.ndarray, prev_stats: dict):
851
+ """Update aggregated measurements and calculate integrating statistics.
852
+
853
+ The core of the statistical computations of the class.
854
+
855
+ Relays on external flow providing it with proper data context: `hist` and `prev_stats`.
856
+
857
+ If `a` is None or empty - valid structure of stats is returned with values `0` or `nan`
858
+
859
+ :param a: data array to measure
860
+ :param hist: histogram array to accumulate into (must be pre-allocated)
861
+ :param prev_stats: previous stats to update
862
+ :return: updated `hist`, `stats`
863
+ """
864
+ stats = (
865
+ select_from({}, _sum_metrics, default=0)
866
+ if a is None or not a.size else
867
+ dict(zip(_sum_metrics,
868
+ _equal_bins_stats(a, low=self._equid_edges[0],
869
+ bin_size=self._sampler.step,
870
+ log_scale=self._sampler.log_scale,
871
+ sums=bool(self._calc_stats), hist=hist)
872
+ )) # stats contain ONLY aggregable metrics!
873
+ )
874
+ if prev_stats: # aggregate what may be aggregated
875
+ stats = {k: prev_stats[k] + v for k, v in stats.items()}
876
+
877
+ # append or overwrite fields which must be calculated anew
878
+ stats |= _basic_stats(**stats) # averaging over accumulated measures
879
+
880
+ if self.levels:
881
+ interp = self.interp_hist(self.levels.values(), cum=True, norm=True)
882
+ stats |= dict(zip(self.levels, interp))
883
+
884
+ if self._calc_stats is not True: # select only requested stats
885
+ stats = select_from(stats, self._calc_stats, strict=False)
886
+ return stats
887
+
888
+ def reset(self):
889
+ """Resets accumulated statistics"""
890
+ # reflects stages in _calc_hist_stats
891
+ self.hist = self._empty_hist(self.hist)
892
+ self.measures = self._calc_hist_stats(None, self.hist, prev_stats={})
893
+ self._appends = 0
894
+
895
+ def process(self, a: np.ndarray, *, restart=False):
896
+ """Process data array and update the histogram and integral stats.
897
+ :param a: array to process
898
+ :param restart: if True `reset` the stats
899
+ :return:
900
+ """
901
+ restart and self.reset()
902
+ self.measures = self._calc_hist_stats(a, self.hist, prev_stats=self.measures)
903
+ self._appends += 1
904
+ return self.hist.copy(), self.measures
905
+
906
+ # ToDo: @Ilya stat_table, hist_table
907
+
908
+ def __getattr__(self, item):
909
+ if (ms := self.measures) and (val := ms.get(item, _UNDEF)) is not _UNDEF:
910
+ return val
911
+ raise KeyError(f"{self.__class__.__name__}.measures has no field '{item}'")
912
+
913
+ def norm_hist(self, total=False):
914
+ """Return normalized histogram instead of integer counters"""
915
+ norm = self.measures.get('total' if total else 'finite', None)
916
+ return self.hist / norm if norm else None
917
+
918
+ def cum_hist(self, norm=True, total=False):
919
+ """Return cumulative histogram optionally `normalized` over `total` or `finite` counts"""
920
+ hist = self.norm_hist(total=total) if norm else self.hist
921
+ return None if hist is None else np.cumsum(hist)
922
+
923
+ def interp_hist(self, values: Collection, cum=False, norm=True, total=False):
924
+ """Interpolate computed histogram in given values.
925
+
926
+ If requested normalization is not available (no total stats) return array of `nan`.
927
+
928
+ :param values: Iterable of location to interpolate at.
929
+ :param cum: if `True` - interpolate the cumulative histogram
930
+ :param norm: normalize the histogram values or use the original counters
931
+ :param total: use total number of values to normalize or only the finite ones.
932
+ :return: selected form of histogram interpolated for the provided values.
933
+ """
934
+ if norm:
935
+ if not self.measures.get(total and 'total' or 'finite', 0):
936
+ return np.full(len(values), np.nan)
937
+ h = self.cum_hist(norm=norm, total=total) if cum else self.norm_hist(total=total)
938
+ else:
939
+ h = self.cum_hist(norm=False) if cum else self.hist
940
+
941
+ values = log_compress(values, self._sampler.log_scale)
942
+ return np.interp(values, self._equid_edges, h, left=0)
943
+
944
+
945
+ # def _tuple(itr: Iterable, func: Callable | str, *args, **kwargs):
946
+ # """Create tuple by applying given function to the given iterable.
947
+ #
948
+ # If ``func`` is str, apply the obj method of this name.
949
+ # :param itr: iterable over objects (with method 'func' if func is str)
950
+ # :param func: function or method name to apply
951
+ # :param args: additional func arguments
952
+ # :param kwargs: additional func keyword arguments
953
+ # """
954
+ # gen = (
955
+ # (getattr(_, func)(*args, **kwargs) for _ in itr)
956
+ # if func is str else
957
+ # (func(_, *args, **kwargs) for _ in itr)
958
+ # )
959
+ # return tuple(gen)
960
+
961
+
962
+ class StatGather2D:
963
+ """
964
+ Basic Statistics Calculator.
965
+
966
+ Operates in accumulative and stateless modes.
967
+
968
+ Supports 3 main types of statistics:
969
+ 1. histogram, with bins defined by the `sampler` argument
970
+ 2. basic statistical metrics:
971
+ - `mean`, `std`, `rmse`, counters over all and finite values
972
+ 3. counters of values under specific named *levels*
973
+
974
+ Levels counters are
975
+ """
976
+
977
+ def __init__(self, *samplers: Sampler,
978
+ regions: Iterable[str] | None = None,
979
+ name='', stats: bool | Collection[METRICS] | Literal['all'] = None,
980
+ levels: tuple[Mapping[str, Number] | None] = ()):
981
+ """
982
+ Initialize statistics gatherer.
983
+
984
+ ``stats`` argument defines which statistical metrics to collect:
985
+ - ``True`` - all possible
986
+ - ``None`` - default, except the internal counters
987
+ - ``False`` - none
988
+ - `collection` of specific *metrics* selected from the supported:
989
+ {'s1', 's2', 'finite', 'total', 'mean', 'rmse', 'std'}
990
+
991
+ Allows extension of the basic statistical metrics with accumulative hist sampled at given levels.
992
+
993
+ :param sampler1, sampler2: binning for every histogram axis (see ``Sampler`` for details)
994
+ :param stats: if ``False`` skip gathering and calculating the basic statistical `measures`
995
+ :param levels: add to the `stats` accumulations by levels, for every axis
996
+ :param name: useful for logistics of multiple instances of gatherers.
997
+ """
998
+ if len(samplers) > 2:
999
+ raise NotImplemented(f"Provided {len(samplers)} samplers, supported only 2")
1000
+
1001
+ self.samplers = samplers
1002
+ self._equid_edges = tuple(_.bins_edges() for _ in samplers)
1003
+
1004
+ self._stats = {} # actual collection of statistics
1005
+ self._metrics = _decode_metrics_arg(stats)
1006
+ self.regions = as_list(regions, collect=tuple)
1007
+ self.levels: tuple[dict[str, Collection]] = as_list(levels, collect=tuple)
1008
+ assert len(self.levels) in (0, 2)
1009
+
1010
+ self.hist = None
1011
+ self.name = name
1012
+ self._appends = None
1013
+
1014
+ self.reset() # initializes all three
1015
+
1016
+ def get_stats(self, *metrics: str, copy=True) -> dict[METRICS, NDArray[Number]]:
1017
+ """
1018
+ Return selected metrics from accumulated stats.
1019
+
1020
+ To override metrics provided in constructor pass:
1021
+ - desired list as arguments
1022
+ - or *one* `True` for all.
1023
+ :param metrics:
1024
+ :param copy: if `True` return a copy of the data, to avoid contamination of the internal state
1025
+ :return:
1026
+ """
1027
+ if not metrics:
1028
+ metrics = self._metrics
1029
+ elif metrics[0] in (True, None):
1030
+ assert len(metrics) == 1
1031
+ metrics = metrics[0]
1032
+ if metrics is False: return None
1033
+ metrics = _decode_metrics_arg(metrics)
1034
+ else:
1035
+ metrics = _decode_metrics_arg(metrics)
1036
+
1037
+ stats = select_from(self._stats, metrics)
1038
+ if copy:
1039
+ stats = deepcopy(stats)
1040
+ return stats
1041
+
1042
+ @property
1043
+ def stats(self):
1044
+ return self.get_stats()
1045
+
1046
+ def __repr__(self):
1047
+ from iad.core.strings import dict_str
1048
+ samplers = ' × '.join(
1049
+ f"{sm.name}[{sm.bins}]{'🗜' if sm.log_scale else ''}"
1050
+ for sm in self.samplers
1051
+ )
1052
+ measures = ''
1053
+ if self._metrics:
1054
+ measures = select_from(self._stats, _basic_metrics)
1055
+ measures = f"{{{dict_str(measures, prec=3, sep='|', to=':')}}}"
1056
+
1057
+ name = f"{self.name}: " if self.name else ''
1058
+ return (f"{self.__class__.__name__} <{name}{samplers}> "
1059
+ f"∑{self._appends}{measures}")
1060
+
1061
+ @property
1062
+ def edges(self):
1063
+ """Edges in the original values space (decompressed)"""
1064
+ return tuple(s.bins_edges(compress=False) for s in self.samplers)
1065
+
1066
+ def _empty_hist(self, hist=None):
1067
+ """Prepare empty histogram array, by creating new or zeroing provided"""
1068
+
1069
+ shape = tuple(_.size for _ in self._equid_edges)
1070
+ if self.regions:
1071
+ shape = (1 + len(self.regions), *shape)
1072
+
1073
+ if hist is None:
1074
+ hist = np.zeros(shape, dtype=HistT)
1075
+ else:
1076
+ assert hist.shape == shape
1077
+ hist[:] = 0
1078
+ return hist
1079
+
1080
+ def _calc_hist_stats(self, a12: tuple | None, masks=None, *,
1081
+ hist: NDArray, prev_stats: dict):
1082
+ """Update aggregated measurements and calculate integrating statistics.
1083
+
1084
+ The core of the statistical computations of the class.
1085
+
1086
+ Relays on external flow providing it with proper data context: `hist` and `prev_stats`.
1087
+
1088
+ If `a` is None or empty - valid structure of stats is returned with values `0` or `nan`
1089
+
1090
+ :param a12: 2 data arrays to measure
1091
+ :param hist: histogram array to accumulate into (must be pre-allocated)
1092
+ :param prev_stats: previous stats to update
1093
+ :return: updated `hist`, `stats`
1094
+ """
1095
+ # noinspection PyTypeChecker
1096
+ stats = (
1097
+ init_sum_stats(num_regions=0, ndim=2) if a12 is None or not a12[0].size else
1098
+ _equal_bins2d_stats(
1099
+ *a12,
1100
+ masks=masks,
1101
+ low=tuple(_[0] for _ in self._equid_edges),
1102
+ bin_size=tuple(_.step for _ in self.samplers),
1103
+ log_scale=tuple(_.log_scale for _ in self.samplers),
1104
+ sums=self._metrics,
1105
+ hist=hist) # stats contain ONLY aggregable metrics!
1106
+ ) # only summations here
1107
+
1108
+ if prev_stats: # aggregate with previous
1109
+ stats = {k: prev_stats[k] + v for k, v in stats.items()}
1110
+ stats |= _basic_stats(**stats) # add averaging statistics to the summations
1111
+
1112
+ for levels in self.levels: # for each variable
1113
+ interp = self.interp_hist(levels.values(), cum=True, norm=True)
1114
+ stats |= dict(zip(levels, interp))
1115
+
1116
+ return stats
1117
+
1118
+ def reset(self):
1119
+ """Resets accumulated statistics"""
1120
+ # reflects stages in _calc_hist_stats
1121
+ self.hist = self._empty_hist(self.hist)
1122
+ self._stats = self._calc_hist_stats(None, hist=self.hist, prev_stats={})
1123
+ self._appends = 0
1124
+
1125
+ def process(self, a1: Iterable[Number], a2: Iterable[Number],
1126
+ masks: Iterable[Iterable[Number]] = None, *, reset: bool = False):
1127
+ """Process 2 data arrays to calculate the histogram and integral stats.
1128
+
1129
+ Masks can be provided in multiple ways:
1130
+ - dict `{region_name: region_mask}` with subset of regions defined in constructor,
1131
+ the rest are assumed zero-masks
1132
+ - iterable over masks in the same order and number as defined regions
1133
+ - array with stack of masks with first dimension sweeping the regions
1134
+
1135
+ All the masks must be objects convertable by ``asarray`` into array of
1136
+
1137
+ :param a1: first data collection
1138
+ :param a2: second data collection
1139
+ :param masks: optional masks matching regions defined in constructor,
1140
+ :param reset: reset accumulated statistics
1141
+ """
1142
+ reset and self.reset()
1143
+ a1, a2 = map(np.asarray, (a1, a2))
1144
+ assert a1.size == a2.size
1145
+ num_regions = len(self.regions)
1146
+ assert bool(num_regions) + bool(masks is None) == 1, "Musks are required if regions are defined"
1147
+
1148
+ if masks is not None:
1149
+ if isinstance(masks, dict): # convert into array of masks
1150
+ if unknown := set(masks).difference(self.regions):
1151
+ raise KeyError(f"Unknown masks names {unknown}")
1152
+ if len(masks) != num_regions: # missing masks set as zeros
1153
+ zero_mask = np.zeros_like(a1, dtype=bool)
1154
+ masks = (masks.get(k, zero_mask) for k in self.regions)
1155
+
1156
+ masks = tuple(map(np.ascontiguousarray, masks))
1157
+
1158
+ if len(shapes := set(m.shape for m in masks)) != 1:
1159
+ raise f"masks have different {shapes = }!"
1160
+ shape = shapes.pop()
1161
+
1162
+ if a1.shape != shape:
1163
+ raise ValueError(f"Masks {shape=} differs from data array's {a1.shape}!")
1164
+ if (num_masks := len(masks)) != num_regions:
1165
+ raise ValueError(f"Masks {num_masks=} differs from {num_regions}")
1166
+
1167
+ self._stats = self._calc_hist_stats((a1, a2), masks, hist=self.hist, prev_stats=self._stats)
1168
+ self._appends += 1
1169
+ return self.hist, self._stats
1170
+
1171
+ def __getattr__(self, item):
1172
+ if (ms := self._stats) and (val := ms.get(item, _UNDEF)) is not _UNDEF:
1173
+ return val
1174
+ raise KeyError(f"{self.__class__.__name__}.measures has no field '{item}'")
1175
+
1176
+ def norm_hist(self, axis: _Axis = None, over: _NormT = True):
1177
+ """
1178
+ Return 2D histogram normalized over specified domain, and
1179
+ optionally collapsed into specified axis by summing along the other.
1180
+ ::
1181
+ - axis is None: [num_regions × bins1 × bins2]
1182
+ - axis in 0, 1: [range × axis_bins]
1183
+
1184
+ Supported domains to normalize `over`:
1185
+ ::
1186
+ 'total' (or True), 'range', 'nonan', None (or False)
1187
+
1188
+ So that
1189
+ - `(over=None, axis=None)` returns the original histogram.
1190
+ - `(over=None, axis=1)` returns not normalized collapsed into axis 1.
1191
+
1192
+ :param axis: Optional axis along which histogram is calculated
1193
+ :param over: domain kind over which to normalize
1194
+ """
1195
+ match over:
1196
+ case 'total' | True:
1197
+ slc = np.s_[:]
1198
+ case 'range':
1199
+ slc = np.s_[self.samplers[0].below:-2, self.samplers[1].below:-2]
1200
+ if self.hist.ndim == 3:
1201
+ slc = np.s_[:, *slc]
1202
+ case 'nonan':
1203
+ slc = np.s_[self.samplers[0].below:-1, self.samplers[1].below:-1]
1204
+ if self.hist.ndim == 3:
1205
+ slc = np.s_[:, *slc]
1206
+ case None | False:
1207
+ pass
1208
+ case _:
1209
+ raise ValueError(f"Invalid norm method {over}")
1210
+
1211
+ if over:
1212
+ hist = self.hist[slc]
1213
+ hist = hist / hist.sum()
1214
+ if axis is not None:
1215
+ hist = hist.sum(1 - axis) # sum over 1 if axis 0, over 0 if axis 1
1216
+ return hist
1217
+
1218
+ def cum_hist(self, axis: _Axis, norm: _NormT = True):
1219
+ """Return cumulative histogram optionally `normalized` over `total` or `finite` counts.
1220
+
1221
+ Supported ``norm`` values:
1222
+ ::
1223
+ 'total' (or True), 'range', 'nonan', None (or False)
1224
+
1225
+ :param axis: select along which axis calculate teh histogram (collapse into)
1226
+ :param norm: kind of normalization to use (See `norm_hist`)
1227
+
1228
+ """
1229
+ hist = self.norm_hist(axis=axis, over=norm)
1230
+ return np.cumsum(hist)
1231
+
1232
+ def interp_hist(self, axis: _Axis, values: Collection, *,
1233
+ cum=False, norm: _NormT = True):
1234
+ """Interpolate computed histogram in given values.
1235
+
1236
+ :param axis: axis to collapse histogram 2D → 1D into before interpolation.
1237
+ :param values: Iterable of location to interpolate at.
1238
+ :param cum: if `True` - interpolate the cumulative histogram
1239
+ :param norm: normalize the histogram values or use the original counters
1240
+ :return: selected form of histogram interpolated for the provided values.
1241
+ """
1242
+ assert axis in (0, 1), "Must specify valid axis to collapse the histogram into"
1243
+
1244
+ hist = self.cum_hist(axis=axis, norm=norm) if cum else self.norm_hist(axis=axis, over=norm)
1245
+
1246
+ values = log_compress(values, self._sampler.log_scale)
1247
+ from scipy.interpolate import interp1d
1248
+ return interp1d(self._equid_edges[axis], hist, assume_sorted=True)(values)
1249
+
1250
+ def hist_table(self):
1251
+ if not self._appends:
1252
+ return None
1253
+
1254
+ axis = dict(zip(['index', 'columns'], (
1255
+ pd.Index(edg, name=sm.name) for edg, sm in zip(self.edges, self.samplers))))
1256
+
1257
+ frame = lambda h2d: pd.DataFrame(h2d, **axis)
1258
+
1259
+ if self.regions:
1260
+ return pd.concat(map(frame, self.hist), keys=['full', *self.regions],names=['region', 'erg'])
1261
+ else:
1262
+ return frame(self.hist)
1263
+
1264
+ def stats_table(self, metrics: bool | METRICS | None = ()):
1265
+ if not (self._appends and self._metrics):
1266
+ return None
1267
+
1268
+ if self.regions:
1269
+ #FixMe: Why are we adding 'full'? It is not requested by user.
1270
+ frame = lambda s: pd.DataFrame(
1271
+ s, index=pd.Index(['full', *self.regions], name='region'),
1272
+ columns=pd.Index([s.name for s in self.samplers], name='var'))
1273
+
1274
+ stats = self.get_stats(*as_list(metrics), copy=False)
1275
+ df = pd.concat(map(frame, stats.values()), keys=stats, names=['metric'])
1276
+ return df.stack('var').unstack(['region', 'var'])
1277
+
1278
+ def plot(self, region=None, title: str = None, ticks: int | tuple[int, int] = 5,
1279
+ cmap='terrain', norm: Literal['asinh', 'log', 'logit', 'symlog'] = None, **kws):
1280
+ """
1281
+ Plot image of 2D histogram
1282
+ :param title: to put on the top of the plot
1283
+ :param ticks: number of ticks of both or each of the axes
1284
+ :param norm: normalization scaling for visualization of z-axis (colors compression)
1285
+ :param cmap: color map to use
1286
+ :param kws: `imshow` anf `figure` keywords
1287
+ :return: fig, axes[2d]
1288
+ """
1289
+ all_regions = ['full', *self.regions]
1290
+ regions = as_list(region) or all_regions
1291
+ regions = [all_regions[r] if isinstance(r, int) else r for r in regions] # indices into names
1292
+ if region is not None and (inv := set(regions).issubset(all_regions)):
1293
+ raise ValueError(f"Unknown regions {inv}")
1294
+ hists = self.hist
1295
+ if hists.ndim == 2:
1296
+ hists = hists[None, :, :]
1297
+ named_hists = {r: hists[all_regions.index(r)] for r in regions}
1298
+
1299
+ from iad.vis.insight import hist_grid
1300
+ axis_names = [s.name or f'v{i}' for i, s in enumerate(self.samplers)]
1301
+ return hist_grid(named_hists, dict(zip(axis_names, self._equid_edges)),
1302
+ title=title, cmap=cmap, norm=norm, **kws)
1303
+
1304
+
1305
+ @overload
1306
+ def equal_bins_stats(a, sampler: Sampler, stats=False, edges=True,
1307
+ hist: np.ndarray | None = None): ...
1308
+
1309
+
1310
+ @overload
1311
+ def equal_bins_stats(a, low: float, high, *, bins: int, below=True,
1312
+ log_scale: float | bool = False, stats=False,
1313
+ edges=True, hist: np.ndarray | None = None): ...
1314
+
1315
+
1316
+ @overload
1317
+ def equal_bins_stats(a, low: float, high, *, bin_size: float, below=True,
1318
+ log_scale: float | bool = False, stats=False,
1319
+ edges=True, hist: np.ndarray | None = None): ...
1320
+
1321
+
1322
+ # ToDo: remove this function after integration of Gatherers
1323
+ def equal_bins_stats(a, low: float | Sampler, high=None, *, below=True,
1324
+ bins: int = None, bin_size=None,
1325
+ log_scale: float | bool = False, stats=False,
1326
+ edges=True, hist: np.ndarray | None = None):
1327
+ """
1328
+ Calculates equal bins histogram and basic statistics over given array.
1329
+
1330
+ Wraps in simplified interface usage of more flexible machinery of
1331
+ - more flexible ``Sampler`` class
1332
+ - fast low level ``_equal_bins_stats`` function
1333
+
1334
+ :param a: nd array to measure
1335
+ :param low: right edge of the lower bin (if ``below`` is True, otherwise its left edge).
1336
+ :param high: right edge of the high bin
1337
+ :param below: controls if ``low`` denotes *left* or *right* edge of the first bin
1338
+ :param bins: number of bins in the histogram
1339
+ :param bin_size: an alternative to the `bins` argument
1340
+ :param log_scale: compress data range by log-scaling bins intervals
1341
+ :param stats: request collect additionally some basic statistics
1342
+ :param edges: return edges if True
1343
+ :return: hist, [stats], [edges] - return if corresponding arguments are True
1344
+ """
1345
+ if isinstance(low, Sampler):
1346
+ sampler = low
1347
+ else:
1348
+ assert not (bins and bin_size)
1349
+ sampler = Sampler(low, high, log_scale=log_scale,
1350
+ **({'bins': bins} if bins else {'step': bin_size}))
1351
+
1352
+ edges = sampler.bins_edges(compress=True) # first equidistant for hist params
1353
+
1354
+ if hist is None:
1355
+ hist = np.zeros(len(edges), HistT)
1356
+ else:
1357
+ if hist.ndim != 1 and hist.size != edges.size:
1358
+ raise ValueError("Invalid size of hist array")
1359
+
1360
+ sums = stats # stats request translates to request for sums
1361
+ stats = dict(zip(_sum_metrics,
1362
+ _equal_bins_stats(
1363
+ a, low=edges[0], bin_size=np.diff(edges[:2]).item(),
1364
+ hist=hist, sums=sums, log_scale=log_scale)
1365
+ ))
1366
+ if sums:
1367
+ stats |= _basic_stats(**stats)
1368
+
1369
+ if edges:
1370
+ edges = sampler.bins_edges(compress=False) # edges in the original data space
1371
+ return hist, edges, stats if sums else hist, edges
1372
+ else:
1373
+ return hist, stats if sums else hist
1374
+
1375
+
1376
+ _Range = tuple[float, float, float] # (start, stop, step)
1377
+
1378
+
1379
+ def hist2d(a1: np.ndarray, a2: np.ndarray, range_1: _Range, range_2: _Range):
1380
+ """
1381
+ :param a1: array 1
1382
+ :param a2: array 2
1383
+ :param range_1: (start, stop, step)
1384
+ :param range_2: (start, stop, step)
1385
+ """
1386
+ bins = [int((r[1] - r[0]) / r[2]) for r in (range_1, range_2)]
1387
+ ranges = [r[:2] for r in (range_1, range_2)]
1388
+ assert all(_ > 1 for _ in bins)
1389
+ return np.histogram2d(a1.ravel(), a2.ravel(), bins=bins, range=ranges)
1390
+
1391
+
1392
+ class Hist2D:
1393
+
1394
+ def __repr__(self):
1395
+ return f"Hist2D{self.samplers}"
1396
+
1397
+ def clear(self):
1398
+ """Clear histogram counts"""
1399
+ self.counts, _, _ = np.histogram2d([], [], bins=self.bins)
1400
+
1401
+ def __init__(self, bins_1: _Range, bins_2: _Range,
1402
+ a1: np.ndarray = None, a2: np.ndarray = None):
1403
+ """
1404
+ `bins_*` can be represented as tuple:
1405
+ - `(min, max, step: float | bins: int, [name: str])`
1406
+ - or ``Sampler`` instance.
1407
+
1408
+ `arrays` are optional at the initialization, could be added later using ``add`` method.
1409
+ """
1410
+
1411
+ self.samplers = [Sampler.from_range(*r[:3], name=str(r[3:]) or f'a{i}')
1412
+ for i, r in enumerate([bins_1, bins_2])]
1413
+ self.bins = tuple(sampler.bins_edges() for sampler in self.samplers)
1414
+
1415
+ self.counts = None
1416
+ self.clear()
1417
+
1418
+ if a1 is None:
1419
+ assert a2 is None
1420
+ else:
1421
+ self.add(a1, a2)
1422
+
1423
+ def add(self, a1: np.ndarray, a2: np.ndarray):
1424
+ """Add to the histogram accumulators pair of arrays (must be of same size)"""
1425
+ arrays = list(np.asarray(a).ravel() for a in [a1, a2])
1426
+ counts, _, _ = np.histogram2d(*arrays, bins=self.bins)
1427
+ self.counts += counts
1428
+
1429
+ def plot(self, title: str = None, ticks: int | tuple[int, int] = 20,
1430
+ transform: Callable | None = None, **imgrid_kws):
1431
+ """
1432
+ Plot image of 2D histogram
1433
+ :param title: to put on the top of the plot
1434
+ :param ticks: number of ticks of both or each of the axes
1435
+ :param transform: optional function to apply on the hist2d counts before forming an image
1436
+ :param imgrid_kws: relevant arguments to imgrid used internally to plot the image
1437
+ :return:
1438
+ """
1439
+ from iad.vis.insight import imgrid
1440
+
1441
+ title = title or 'Hist2D'
1442
+ if transform and (name := transform.__name__) != (lambda _: _).__name__:
1443
+ title += f' ( {name}(counts) )'
1444
+ counts = transform(self.counts) if transform else self.counts
1445
+ imgrid_kws = dict(cmap='terrain') | imgrid_kws
1446
+ (ax,) = imgrid(counts, out='axs', titles=[title], **imgrid_kws)
1447
+ ax.invert_yaxis()
1448
+
1449
+ ticks = [ticks] * 2 if isinstance(ticks, int) else ticks
1450
+ labels = [s.name for s in self.samplers]
1451
+ for axis, label, edges, num_ticks in zip(['y', 'x'], labels, self.bins, ticks):
1452
+ num_ticks = min((num_edges := len(edges)), num_ticks)
1453
+ tick_ixs = np.linspace(0, num_edges - 1, num_ticks, endpoint=True, dtype=int)
1454
+ tick_edges = [f"{edges[i]:.2f}" for i in tick_ixs]
1455
+ getattr(ax, f'set_{axis}ticks')(tick_ixs, tick_edges)
1456
+ getattr(ax, f'set_{axis}label')(label)
1457
+ ax.figure.set_tight_layout(True)
1458
+ return ax