FazBrowse GitHub Viewer
|
Trending
|
URL:
|
Home
Tools:
[Download Repo ZIP]
[View Raw Code]
[Original HTTPS Page]
xarray-spatial/xrspatial/resample.py at main · xarray-contrib/xarray-spatial · GitHub
Uh oh!
There was an error while loading.
Please reload this page
.
xarray-contrib
/
xarray-spatial
Public
Notifications
You must be signed in to change notification settings
Fork
91
Star
969
Code
Issues
73
Pull requests
8
Discussions
Actions
Projects
Security and quality
0
Insights
Additional navigation options
Code
Issues
Pull requests
Discussions
Actions
Projects
Security and quality
Insights
Expand file tree
Breadcrumbs
xarray-spatial
/
xrspatial
/
resample.py
Copy path
More file actions
More file actions
Latest commit
History
History
History
1669 lines (1448 loc) · 66.9 KB
Breadcrumbs
xarray-spatial
/
xrspatial
/
resample.py
Copy path
File metadata and controls
1669 lines (1448 loc) · 66.9 KB
Raw
Copy raw file
Download raw file
Open symbols panel
Edit and raw actions
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
"""Raster resampling -- resolution change without reprojection.
Provides :func:`resample` for changing raster cell size using
interpolation or block-aggregation methods.
"""
from
__future__
import
annotations
from
functools
import
partial
import
numpy
as
np
import
xarray
as
xr
from
scipy
.
ndimage
import
map_coordinates
as
_scipy_map_coords
from
scipy
.
ndimage
import
spline_filter
as
_scipy_spline_filter
try
:
import
dask
.
array
as
da
except
ImportError
:
da
=
None
try
:
import
cupy
except
ImportError
:
cupy
=
None
from
xrspatial
.
dataset_support
import
supports_dataset
from
xrspatial
.
utils
import
(
ArrayTypeFunctionMapping
,
_dask_task_name_kwargs
,
_validate_raster
,
calc_res
,
ngjit
,
)
# -- Constants ---------------------------------------------------------------
INTERP_METHODS
=
{
'nearest'
:
0
,
'bilinear'
:
1
,
'cubic'
:
3
}
AGGREGATE_METHODS
=
{
'average'
,
'min'
,
'max'
,
'median'
,
'mode'
}
ALL_METHODS
=
set
(
INTERP_METHODS
)
|
AGGREGATE_METHODS
# Overlap depth (input pixels) each interpolation kernel needs from
# neighbouring chunks when processing dask arrays. Cubic requires extra
# depth because the B-spline prefilter is a global IIR filter whose
# boundary transient decays as ~0.268^n. Depth 16 puts the residual at
# ~7e-10, comfortably below float32 epsilon so chunk-seam parity rounds
# to zero in the float32 output. The dask drivers clamp this per axis
# down to ``axis_total - 1`` when the array is too small to absorb the
# full depth; see ``_run_dask_numpy`` for the rationale.
_INTERP_DEPTH
=
{
'nearest'
:
1
,
'bilinear'
:
1
,
'cubic'
:
16
}
# Approximate working-set size per output cell for the eager backends:
# one float64 working buffer (8 B) plus a float64 output cell (8 B) in
# the worst case. scipy.ndimage.map_coordinates also allocates a
# temporary of the same size during higher-order spline evaluation; the
# 0.5 * available bound below leaves room for that.
_BYTES_PER_OUTPUT_CELL
=
16
# -- Working / output dtype selection ----------------------------------------
def
_working_dtype
(
input_dtype
):
"""Pick the working float dtype for resampling.
float64 inputs stay in float64 to preserve precision; everything else
(smaller floats, integers, bool) uses float32.
"""
dt
=
np
.
dtype
(
input_dtype
)
if
dt
.
kind
==
'f'
and
dt
.
itemsize
>=
8
:
return
np
.
float64
return
np
.
float32
def
_output_dtype
(
input_dtype
):
"""Pick the output dtype for resampling.
Float inputs keep their dtype. Integer / bool inputs return float32
because NaN-sentinel resampling needs a float type.
"""
dt
=
np
.
dtype
(
input_dtype
)
if
dt
.
kind
==
'f'
:
return
dt
.
type
return
np
.
float32
def
_maybe_astype
(
arr
,
dtype
):
"""astype copy that no-ops when already at the requested dtype."""
return
arr
if
arr
.
dtype
==
np
.
dtype
(
dtype
)
else
arr
.
astype
(
dtype
)
# -- Memory guard ------------------------------------------------------------
def
_available_memory_bytes
():
"""Best-effort estimate of available host memory in bytes."""
# Try /proc/meminfo (Linux)
try
:
with
open
(
'/proc/meminfo'
,
'r'
)
as
f
:
for
line
in
f
:
if
line
.
startswith
(
'MemAvailable:'
):
return
int
(
line
.
split
()[
1
])
*
1024
except
(
OSError
,
ValueError
,
IndexError
):
pass
# Try psutil
try
:
import
psutil
return
psutil
.
virtual_memory
().
available
except
(
ImportError
,
AttributeError
):
pass
# Fallback: 2 GB
return
2
*
1024
**
3
def
_available_gpu_memory_bytes
():
"""Best-effort estimate of free GPU memory in bytes.
Returns 0 when CuPy / CUDA is unavailable or the query fails -- callers
treat that as a sentinel meaning "no GPU info, skip the guard".
"""
try
:
import
cupy
as
_cp
free
,
_total
=
_cp
.
cuda
.
runtime
.
memGetInfo
()
return
int
(
free
)
except
Exception
:
return
0
def
_check_resample_memory
(
out_h
,
out_w
):
"""Raise MemoryError if the eager output buffer would exceed RAM.
The numpy and cupy-eager backends allocate a single (out_h, out_w)
float64 working buffer plus a float32 output before any actual work.
A user passing a huge ``scale_factor`` (or a tiny ``target_resolution``)
would otherwise OOM the process before this function returns.
"""
required
=
int
(
out_h
)
*
int
(
out_w
)
*
_BYTES_PER_OUTPUT_CELL
available
=
_available_memory_bytes
()
if
required
>
0.5
*
available
:
raise
MemoryError
(
f"resample output of
{
out_h
}
x
{
out_w
}
would need "
f"~
{
required
/
1e9
:.1f
}
GB of working memory but only "
f"~
{
available
/
1e9
:.1f
}
GB is available. "
f"Use a smaller scale_factor / larger target_resolution, "
f"or pass a dask-backed DataArray for out-of-core processing."
)
def
_check_resample_gpu_memory
(
out_h
,
out_w
):
"""Raise MemoryError if the cupy-eager output buffer would exceed VRAM.
Skips the check (returns silently) when free GPU memory cannot be
queried -- the kernel will fail later at the cupy.empty boundary
anyway.
"""
available
=
_available_gpu_memory_bytes
()
if
available
<=
0
:
return
required
=
int
(
out_h
)
*
int
(
out_w
)
*
_BYTES_PER_OUTPUT_CELL
if
required
>
0.5
*
available
:
raise
MemoryError
(
f"resample output of
{
out_h
}
x
{
out_w
}
would need "
f"~
{
required
/
1e9
:.1f
}
GB of GPU working memory but only "
f"~
{
available
/
1e9
:.1f
}
GB is free on the active device. "
f"Use a smaller scale_factor / larger target_resolution, "
f"or pass a dask+cupy DataArray for out-of-core processing."
)
# -- Input-validation helpers ------------------------------------------------
def
_validate_resample_scalar_or_pair
(
value
,
param_name
):
"""Validate a scalar-or-2-tuple resolution / scale parameter.
Accepts either a real scalar or a length-2 tuple/list of scalars.
Each component must be finite (not NaN, not inf) and strictly
positive. Raises ``ValueError`` with a message naming the parameter
and the offending value.
"""
is_pair
=
isinstance
(
value
, (
tuple
,
list
))
if
is_pair
:
if
len
(
value
)
!=
2
:
raise
ValueError
(
f"
{
param_name
}
must have length 2, got length
{
len
(
value
)
}
"
)
components
=
value
else
:
components
=
(
value
,)
for
i
,
comp
in
enumerate
(
components
):
# Suffix points at the bad slot when the input was a pair, so
# `(0.0, 1.0)` reports "got 0.0 at index 0 of (0.0, 1.0)"
# instead of dumping the whole tuple.
where
=
f"
{
comp
!r
}
at index
{
i
}
of
{
value
!r
}
"
if
is_pair
else
f"
{
value
!r
}
"
try
:
f
=
float
(
comp
)
except
(
TypeError
,
ValueError
):
raise
ValueError
(
f"
{
param_name
}
must be a finite positive number "
f"(or length-2 sequence of them), got
{
where
}
"
)
from
None
if
not
np
.
isfinite
(
f
):
raise
ValueError
(
f"
{
param_name
}
must be finite and > 0, got
{
where
}
"
)
if
f
<=
0
:
raise
ValueError
(
f"
{
param_name
}
must be > 0, got
{
where
}
"
)
def
_validate_monotonic_regular_coords
(
agg
):
"""Reject inputs whose spatial coords are not regular and monotonic.
``resample`` assumes a regular, monotonic grid: ``calc_res`` derives
the input resolution from the full coordinate extent while the output
coordinates are rebuilt from first/last neighbour spacing. On an
irregular or non-monotonic grid those two views of "resolution"
disagree and the function silently produces inconsistent output
geometry (wrong width, coords spilling past the input range). Fail
fast here instead.
Only 1-D coords that actually exist on the spatial dims are checked;
an input without spatial coords is left to the existing code paths.
For 3-D inputs ``resample`` recurses per band, so this runs once per
band on identical coords -- a cheap, harmless repeat.
"""
for
dim
in
agg
.
dims
[
-
2
:]:
if
dim
not
in
agg
.
coords
:
continue
vals
=
np
.
asarray
(
agg
[
dim
].
values
,
dtype
=
np
.
float64
)
if
vals
.
ndim
!=
1
or
vals
.
size
<
2
:
continue
diffs
=
np
.
diff
(
vals
)
if
not
(
np
.
all
(
diffs
>
0
)
or
np
.
all
(
diffs
<
0
)):
raise
ValueError
(
f"resample(): `agg` coordinate
{
dim
!r
}
must be strictly "
f"monotonic (consistently increasing or decreasing); "
f"resample only supports regular monotonic rasters"
)
# Allow floating-point jitter but reject genuinely uneven spacing
# (e.g. [0, 1, 4]). Compare every step to the mean step. The
# tolerance scales with the step size via ``rtol`` so it tracks
# the coordinate magnitude.
step
=
diffs
.
mean
()
if
not
np
.
allclose
(
diffs
,
step
,
rtol
=
1e-5
,
atol
=
0.0
):
raise
ValueError
(
f"resample(): `agg` coordinate
{
dim
!r
}
must be evenly "
f"spaced; resample only supports regular monotonic "
f"rasters, not irregular grids"
)
# -- Output-geometry helpers -------------------------------------------------
def
_output_shape
(
in_h
,
in_w
,
scale_y
,
scale_x
):
return
max
(
1
,
round
(
in_h
*
scale_y
)),
max
(
1
,
round
(
in_w
*
scale_x
))
def
_output_chunks
(
in_chunks
,
scale
):
"""Compute per-chunk output sizes via cumulative rounding.
Guarantees ``sum(result) == round(sum(in_chunks) * scale)``.
"""
cum
=
np
.
cumsum
([
0
]
+
list
(
in_chunks
))
out_cum
=
np
.
round
(
cum
*
scale
).
astype
(
int
)
return
tuple
(
int
(
max
(
1
,
out_cum
[
i
+
1
]
-
out_cum
[
i
]))
for
i
in
range
(
len
(
in_chunks
)))
# -- Block-centered coordinate mapping ---------------------------------------
def
_block_centered_coords
(
n_in
,
n_out
):
"""Return input coordinates for each output pixel using block-centered mapping.
Maps output pixel ``o`` to input pixel ``(o + 0.5) * (n_in / n_out) - 0.5``.
This places each output pixel at the center of its spatial footprint,
matching the convention used by ``_new_coords`` for output coordinate
metadata.
"""
o
=
np
.
arange
(
n_out
,
dtype
=
np
.
float64
)
return
(
o
+
0.5
)
*
(
n_in
/
n_out
)
-
0.5
# -- Spline prefilter helpers -----------------------------------------------
#
# scipy.ndimage.map_coordinates(prefilter=True) silently does three things:
# (1) edge-pad the input by 12 pixels for mode='nearest' / 'grid-constant'
# so the IIR transient stabilises before reaching real data,
# (2) call spline_filter on that padded array, and
# (3) shift the sample coordinates by the same offset. The padding step
# is private (``_prepad_for_spline_filter``) and is needed for the
# explicit-prefilter path to match the implicit one bit-for-bit.
#
# We replicate it here so callers can prefilter once per array (e.g. the
# NaN-aware filled / weights pair) and pass ``prefilter=False`` to
# map_coordinates without changing the boundary semantics. Doing the
# prefilter explicitly also makes the per-block dask path deterministic --
# the same spline coefficients are computed in eager and chunked modes
# (modulo the IIR transient that the depth=10 overlap already absorbs).
_SPLINE_PREPAD_NEAREST
=
12
def
_prepad_and_filter_np
(
arr
,
order
):
"""Edge-pad and spline-filter *arr* for an explicit ``mode='nearest'``
prefilter pass. Returns ``(filtered, npad)``; the caller adds *npad*
to its sample coordinates.
"""
npad
=
_SPLINE_PREPAD_NEAREST
padded
=
np
.
pad
(
arr
,
npad
,
mode
=
'edge'
)
filtered
=
_scipy_spline_filter
(
padded
,
order
=
order
,
mode
=
'nearest'
)
return
filtered
,
npad
def
_prepad_and_filter_cupy
(
arr
,
order
,
spline_filter_fn
):
"""CuPy variant of :func:`_prepad_and_filter_np`."""
npad
=
_SPLINE_PREPAD_NEAREST
padded
=
cupy
.
pad
(
arr
,
npad
,
mode
=
'edge'
)
filtered
=
spline_filter_fn
(
padded
,
order
=
order
,
mode
=
'nearest'
)
return
filtered
,
npad
# -- NaN-aware interpolation (NumPy) ----------------------------------------
def
_nan_aware_interp_np
(
data
,
out_h
,
out_w
,
order
):
"""Interpolate *data* to *(out_h, out_w)* with NaN-aware weighting.
Uses ``scipy.ndimage.map_coordinates`` with block-centered coordinate
mapping so that sample positions match the output coordinate metadata.
For *order* 0 (nearest-neighbour) NaN propagates naturally.
For higher orders the zero-fill / weight-mask trick is used so that
NaN pixels do not corrupt their neighbours.
"""
iy
=
_block_centered_coords
(
data
.
shape
[
0
],
out_h
)
ix
=
_block_centered_coords
(
data
.
shape
[
1
],
out_w
)
yy
,
xx
=
np
.
meshgrid
(
iy
,
ix
,
indexing
=
'ij'
)
coords
=
np
.
array
([
yy
.
ravel
(),
xx
.
ravel
()])
if
order
==
0
:
result
=
_scipy_map_coords
(
data
,
coords
,
order
=
0
,
mode
=
'nearest'
)
return
result
.
reshape
(
out_h
,
out_w
)
# For order >= 2 run the spline prefilter explicitly so the IIR boundary
# transient is computed once per array instead of implicitly inside each
# map_coordinates call. Bilinear (order == 1) prefilter is a no-op.
use_explicit
=
order
>=
2
mask
=
np
.
isnan
(
data
)
if
not
mask
.
any
():
if
use_explicit
:
src
,
npad
=
_prepad_and_filter_np
(
data
,
order
)
result
=
_scipy_map_coords
(
src
,
coords
+
npad
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
result
=
_scipy_map_coords
(
data
,
coords
,
order
=
order
,
mode
=
'nearest'
)
return
result
.
reshape
(
out_h
,
out_w
)
filled
=
np
.
where
(
mask
,
0.0
,
data
)
weights
=
(
~
mask
).
astype
(
data
.
dtype
)
if
use_explicit
:
filled
,
npad
=
_prepad_and_filter_np
(
filled
,
order
)
weights
,
_
=
_prepad_and_filter_np
(
weights
,
order
)
sample_coords
=
coords
+
npad
z_data
=
_scipy_map_coords
(
filled
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
z_wt
=
_scipy_map_coords
(
weights
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
z_data
=
_scipy_map_coords
(
filled
,
coords
,
order
=
order
,
mode
=
'nearest'
)
z_wt
=
_scipy_map_coords
(
weights
,
coords
,
order
=
order
,
mode
=
'nearest'
)
# Gate on majority weight: an output pixel is valid only when more
# than half of the resampling kernel weight came from valid input
# pixels. This rejects pixels lit only by cubic-kernel sidelobes
# leaking small positive weight from a single neighbour.
result
=
np
.
where
(
z_wt
>
0.5
,
z_data
/
np
.
maximum
(
z_wt
,
1e-10
),
np
.
nan
)
return
result
.
reshape
(
out_h
,
out_w
)
# -- NaN-aware interpolation (CuPy) -----------------------------------------
def
_nan_aware_interp_cupy
(
data
,
out_h
,
out_w
,
order
):
"""CuPy variant of :func:`_nan_aware_interp_np`."""
from
cupyx
.
scipy
.
ndimage
import
map_coordinates
as
_cupy_map_coords
from
cupyx
.
scipy
.
ndimage
import
spline_filter
as
_cupy_spline_filter
iy
=
cupy
.
asarray
(
_block_centered_coords
(
data
.
shape
[
0
],
out_h
))
ix
=
cupy
.
asarray
(
_block_centered_coords
(
data
.
shape
[
1
],
out_w
))
yy
,
xx
=
cupy
.
meshgrid
(
iy
,
ix
,
indexing
=
'ij'
)
coords
=
cupy
.
array
([
yy
.
ravel
(),
xx
.
ravel
()])
if
order
==
0
:
result
=
_cupy_map_coords
(
data
,
coords
,
order
=
0
,
mode
=
'nearest'
)
return
result
.
reshape
(
out_h
,
out_w
)
use_explicit
=
order
>=
2
mask
=
cupy
.
isnan
(
data
)
if
not
mask
.
any
():
if
use_explicit
:
src
,
npad
=
_prepad_and_filter_cupy
(
data
,
order
,
_cupy_spline_filter
)
result
=
_cupy_map_coords
(
src
,
coords
+
npad
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
result
=
_cupy_map_coords
(
data
,
coords
,
order
=
order
,
mode
=
'nearest'
)
return
result
.
reshape
(
out_h
,
out_w
)
filled
=
cupy
.
where
(
mask
,
0.0
,
data
)
weights
=
(
~
mask
).
astype
(
data
.
dtype
)
if
use_explicit
:
filled
,
npad
=
_prepad_and_filter_cupy
(
filled
,
order
,
_cupy_spline_filter
)
weights
,
_
=
_prepad_and_filter_cupy
(
weights
,
order
,
_cupy_spline_filter
)
sample_coords
=
coords
+
npad
z_data
=
_cupy_map_coords
(
filled
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
z_wt
=
_cupy_map_coords
(
weights
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
z_data
=
_cupy_map_coords
(
filled
,
coords
,
order
=
order
,
mode
=
'nearest'
)
z_wt
=
_cupy_map_coords
(
weights
,
coords
,
order
=
order
,
mode
=
'nearest'
)
# Majority-weight gate (see _nan_aware_interp_np for rationale).
result
=
cupy
.
where
(
z_wt
>
0.5
,
z_data
/
cupy
.
maximum
(
z_wt
,
1e-10
),
cupy
.
nan
)
return
result
.
reshape
(
out_h
,
out_w
)
# -- Block-aggregation kernels (NumPy, numba) --------------------------------
@
ngjit
def
_agg_mean
(
data
,
out_h
,
out_w
):
h
,
w
=
data
.
shape
out
=
np
.
empty
((
out_h
,
out_w
),
dtype
=
np
.
float64
)
for
oy
in
range
(
out_h
):
y0
=
int
(
oy
*
h
/
out_h
)
y1
=
max
(
y0
+
1
,
int
((
oy
+
1
)
*
h
/
out_h
))
for
ox
in
range
(
out_w
):
x0
=
int
(
ox
*
w
/
out_w
)
x1
=
max
(
x0
+
1
,
int
((
ox
+
1
)
*
w
/
out_w
))
total
=
0.0
count
=
0
for
y
in
range
(
y0
,
y1
):
for
x
in
range
(
x0
,
x1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
):
total
+=
v
count
+=
1
out
[
oy
,
ox
]
=
total
/
count
if
count
>
0
else
np
.
nan
return
out
@
ngjit
def
_agg_min
(
data
,
out_h
,
out_w
):
h
,
w
=
data
.
shape
out
=
np
.
empty
((
out_h
,
out_w
),
dtype
=
np
.
float64
)
for
oy
in
range
(
out_h
):
y0
=
int
(
oy
*
h
/
out_h
)
y1
=
max
(
y0
+
1
,
int
((
oy
+
1
)
*
h
/
out_h
))
for
ox
in
range
(
out_w
):
x0
=
int
(
ox
*
w
/
out_w
)
x1
=
max
(
x0
+
1
,
int
((
ox
+
1
)
*
w
/
out_w
))
best
=
np
.
inf
found
=
False
for
y
in
range
(
y0
,
y1
):
for
x
in
range
(
x0
,
x1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
)
and
v
<
best
:
best
=
v
found
=
True
out
[
oy
,
ox
]
=
best
if
found
else
np
.
nan
return
out
@
ngjit
def
_agg_max
(
data
,
out_h
,
out_w
):
h
,
w
=
data
.
shape
out
=
np
.
empty
((
out_h
,
out_w
),
dtype
=
np
.
float64
)
for
oy
in
range
(
out_h
):
y0
=
int
(
oy
*
h
/
out_h
)
y1
=
max
(
y0
+
1
,
int
((
oy
+
1
)
*
h
/
out_h
))
for
ox
in
range
(
out_w
):
x0
=
int
(
ox
*
w
/
out_w
)
x1
=
max
(
x0
+
1
,
int
((
ox
+
1
)
*
w
/
out_w
))
best
=
-
np
.
inf
found
=
False
for
y
in
range
(
y0
,
y1
):
for
x
in
range
(
x0
,
x1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
)
and
v
>
best
:
best
=
v
found
=
True
out
[
oy
,
ox
]
=
best
if
found
else
np
.
nan
return
out
@
ngjit
def
_agg_median
(
data
,
out_h
,
out_w
):
h
,
w
=
data
.
shape
out
=
np
.
empty
((
out_h
,
out_w
),
dtype
=
np
.
float64
)
for
oy
in
range
(
out_h
):
y0
=
int
(
oy
*
h
/
out_h
)
y1
=
max
(
y0
+
1
,
int
((
oy
+
1
)
*
h
/
out_h
))
for
ox
in
range
(
out_w
):
x0
=
int
(
ox
*
w
/
out_w
)
x1
=
max
(
x0
+
1
,
int
((
ox
+
1
)
*
w
/
out_w
))
buf
=
np
.
empty
((
y1
-
y0
)
*
(
x1
-
x0
),
dtype
=
np
.
float64
)
n
=
0
for
y
in
range
(
y0
,
y1
):
for
x
in
range
(
x0
,
x1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
):
buf
[
n
]
=
v
n
+=
1
if
n
==
0
:
out
[
oy
,
ox
]
=
np
.
nan
else
:
s
=
np
.
sort
(
buf
[:
n
])
if
n
%
2
==
1
:
out
[
oy
,
ox
]
=
s
[
n
//
2
]
else
:
out
[
oy
,
ox
]
=
(
s
[
n
//
2
-
1
]
+
s
[
n
//
2
])
/
2.0
return
out
@
ngjit
def
_agg_mode
(
data
,
out_h
,
out_w
):
h
,
w
=
data
.
shape
out
=
np
.
empty
((
out_h
,
out_w
),
dtype
=
np
.
float64
)
for
oy
in
range
(
out_h
):
y0
=
int
(
oy
*
h
/
out_h
)
y1
=
max
(
y0
+
1
,
int
((
oy
+
1
)
*
h
/
out_h
))
for
ox
in
range
(
out_w
):
x0
=
int
(
ox
*
w
/
out_w
)
x1
=
max
(
x0
+
1
,
int
((
ox
+
1
)
*
w
/
out_w
))
buf
=
np
.
empty
((
y1
-
y0
)
*
(
x1
-
x0
),
dtype
=
np
.
float64
)
n
=
0
for
y
in
range
(
y0
,
y1
):
for
x
in
range
(
x0
,
x1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
):
buf
[
n
]
=
v
n
+=
1
if
n
==
0
:
out
[
oy
,
ox
]
=
np
.
nan
continue
s
=
np
.
sort
(
buf
[:
n
])
best_val
=
s
[
0
]
best_cnt
=
1
cur_val
=
s
[
0
]
cur_cnt
=
1
for
i
in
range
(
1
,
n
):
if
s
[
i
]
==
cur_val
:
cur_cnt
+=
1
else
:
if
cur_cnt
>
best_cnt
:
best_cnt
=
cur_cnt
best_val
=
cur_val
cur_val
=
s
[
i
]
cur_cnt
=
1
if
cur_cnt
>
best_cnt
:
best_val
=
cur_val
out
[
oy
,
ox
]
=
best_val
return
out
_AGG_FUNCS
=
{
'average'
:
_agg_mean
,
'min'
:
_agg_min
,
'max'
:
_agg_max
,
'median'
:
_agg_median
,
'mode'
:
_agg_mode
,
}
# -- Block-aggregation kernels for dask chunks -------------------------------
#
# These mirror the eager `_agg_mean / _agg_min / ...` family but compute
# per-pixel windows from the *global* input/output geometry and a chunk
# offset, rather than from the local block shape. The whole chunk runs
# inside a single jitted call, instead of one numba dispatch per output
# pixel as the previous `func(sub, 1, 1)[0, 0]` loop did.
#
# Window bounds for output pixel `go` (a *global* output index):
# gy0 = int(go * global_in_h / global_out_h) - in_y0
# gy1 = max(gy0 + 1,
# int((go + 1) * global_in_h / global_out_h) - in_y0)
# where `in_y0` is the global input index of the chunk's first row
# (negative if `_add_overlap` extended the chunk past the input edge).
@
ngjit
def
_agg_block_mean_nb
(
data
,
target_h
,
target_w
,
go_y0
,
go_x0
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
in_y0
,
in_x0
):
out
=
np
.
empty
((
target_h
,
target_w
),
dtype
=
np
.
float64
)
for
lo_y
in
range
(
target_h
):
go_y
=
go_y0
+
lo_y
gy0
=
int
(
go_y
*
global_in_h
/
global_out_h
)
-
in_y0
gy1
=
int
((
go_y
+
1
)
*
global_in_h
/
global_out_h
)
-
in_y0
if
gy1
<
gy0
+
1
:
gy1
=
gy0
+
1
for
lo_x
in
range
(
target_w
):
go_x
=
go_x0
+
lo_x
gx0
=
int
(
go_x
*
global_in_w
/
global_out_w
)
-
in_x0
gx1
=
int
((
go_x
+
1
)
*
global_in_w
/
global_out_w
)
-
in_x0
if
gx1
<
gx0
+
1
:
gx1
=
gx0
+
1
total
=
0.0
count
=
0
for
y
in
range
(
gy0
,
gy1
):
for
x
in
range
(
gx0
,
gx1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
):
total
+=
v
count
+=
1
out
[
lo_y
,
lo_x
]
=
total
/
count
if
count
>
0
else
np
.
nan
return
out
@
ngjit
def
_agg_block_min_nb
(
data
,
target_h
,
target_w
,
go_y0
,
go_x0
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
in_y0
,
in_x0
):
out
=
np
.
empty
((
target_h
,
target_w
),
dtype
=
np
.
float64
)
for
lo_y
in
range
(
target_h
):
go_y
=
go_y0
+
lo_y
gy0
=
int
(
go_y
*
global_in_h
/
global_out_h
)
-
in_y0
gy1
=
int
((
go_y
+
1
)
*
global_in_h
/
global_out_h
)
-
in_y0
if
gy1
<
gy0
+
1
:
gy1
=
gy0
+
1
for
lo_x
in
range
(
target_w
):
go_x
=
go_x0
+
lo_x
gx0
=
int
(
go_x
*
global_in_w
/
global_out_w
)
-
in_x0
gx1
=
int
((
go_x
+
1
)
*
global_in_w
/
global_out_w
)
-
in_x0
if
gx1
<
gx0
+
1
:
gx1
=
gx0
+
1
best
=
np
.
inf
found
=
False
for
y
in
range
(
gy0
,
gy1
):
for
x
in
range
(
gx0
,
gx1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
)
and
v
<
best
:
best
=
v
found
=
True
out
[
lo_y
,
lo_x
]
=
best
if
found
else
np
.
nan
return
out
@
ngjit
def
_agg_block_max_nb
(
data
,
target_h
,
target_w
,
go_y0
,
go_x0
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
in_y0
,
in_x0
):
out
=
np
.
empty
((
target_h
,
target_w
),
dtype
=
np
.
float64
)
for
lo_y
in
range
(
target_h
):
go_y
=
go_y0
+
lo_y
gy0
=
int
(
go_y
*
global_in_h
/
global_out_h
)
-
in_y0
gy1
=
int
((
go_y
+
1
)
*
global_in_h
/
global_out_h
)
-
in_y0
if
gy1
<
gy0
+
1
:
gy1
=
gy0
+
1
for
lo_x
in
range
(
target_w
):
go_x
=
go_x0
+
lo_x
gx0
=
int
(
go_x
*
global_in_w
/
global_out_w
)
-
in_x0
gx1
=
int
((
go_x
+
1
)
*
global_in_w
/
global_out_w
)
-
in_x0
if
gx1
<
gx0
+
1
:
gx1
=
gx0
+
1
best
=
-
np
.
inf
found
=
False
for
y
in
range
(
gy0
,
gy1
):
for
x
in
range
(
gx0
,
gx1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
)
and
v
>
best
:
best
=
v
found
=
True
out
[
lo_y
,
lo_x
]
=
best
if
found
else
np
.
nan
return
out
@
ngjit
def
_agg_block_median_nb
(
data
,
target_h
,
target_w
,
go_y0
,
go_x0
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
in_y0
,
in_x0
):
out
=
np
.
empty
((
target_h
,
target_w
),
dtype
=
np
.
float64
)
for
lo_y
in
range
(
target_h
):
go_y
=
go_y0
+
lo_y
gy0
=
int
(
go_y
*
global_in_h
/
global_out_h
)
-
in_y0
gy1
=
int
((
go_y
+
1
)
*
global_in_h
/
global_out_h
)
-
in_y0
if
gy1
<
gy0
+
1
:
gy1
=
gy0
+
1
for
lo_x
in
range
(
target_w
):
go_x
=
go_x0
+
lo_x
gx0
=
int
(
go_x
*
global_in_w
/
global_out_w
)
-
in_x0
gx1
=
int
((
go_x
+
1
)
*
global_in_w
/
global_out_w
)
-
in_x0
if
gx1
<
gx0
+
1
:
gx1
=
gx0
+
1
buf
=
np
.
empty
((
gy1
-
gy0
)
*
(
gx1
-
gx0
),
dtype
=
np
.
float64
)
n
=
0
for
y
in
range
(
gy0
,
gy1
):
for
x
in
range
(
gx0
,
gx1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
):
buf
[
n
]
=
v
n
+=
1
if
n
==
0
:
out
[
lo_y
,
lo_x
]
=
np
.
nan
else
:
s
=
np
.
sort
(
buf
[:
n
])
if
n
%
2
==
1
:
out
[
lo_y
,
lo_x
]
=
s
[
n
//
2
]
else
:
out
[
lo_y
,
lo_x
]
=
(
s
[
n
//
2
-
1
]
+
s
[
n
//
2
])
/
2.0
return
out
@
ngjit
def
_agg_block_mode_nb
(
data
,
target_h
,
target_w
,
go_y0
,
go_x0
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
in_y0
,
in_x0
):
out
=
np
.
empty
((
target_h
,
target_w
),
dtype
=
np
.
float64
)
for
lo_y
in
range
(
target_h
):
go_y
=
go_y0
+
lo_y
gy0
=
int
(
go_y
*
global_in_h
/
global_out_h
)
-
in_y0
gy1
=
int
((
go_y
+
1
)
*
global_in_h
/
global_out_h
)
-
in_y0
if
gy1
<
gy0
+
1
:
gy1
=
gy0
+
1
for
lo_x
in
range
(
target_w
):
go_x
=
go_x0
+
lo_x
gx0
=
int
(
go_x
*
global_in_w
/
global_out_w
)
-
in_x0
gx1
=
int
((
go_x
+
1
)
*
global_in_w
/
global_out_w
)
-
in_x0
if
gx1
<
gx0
+
1
:
gx1
=
gx0
+
1
buf
=
np
.
empty
((
gy1
-
gy0
)
*
(
gx1
-
gx0
),
dtype
=
np
.
float64
)
n
=
0
for
y
in
range
(
gy0
,
gy1
):
for
x
in
range
(
gx0
,
gx1
):
v
=
data
[
y
,
x
]
if
not
np
.
isnan
(
v
):
buf
[
n
]
=
v
n
+=
1
if
n
==
0
:
out
[
lo_y
,
lo_x
]
=
np
.
nan
continue
s
=
np
.
sort
(
buf
[:
n
])
best_val
=
s
[
0
]
best_cnt
=
1
cur_val
=
s
[
0
]
cur_cnt
=
1
for
i
in
range
(
1
,
n
):
if
s
[
i
]
==
cur_val
:
cur_cnt
+=
1
else
:
if
cur_cnt
>
best_cnt
:
best_cnt
=
cur_cnt
best_val
=
cur_val
cur_val
=
s
[
i
]
cur_cnt
=
1
if
cur_cnt
>
best_cnt
:
best_val
=
cur_val
out
[
lo_y
,
lo_x
]
=
best_val
return
out
_AGG_BLOCK_FUNCS
=
{
'average'
:
_agg_block_mean_nb
,
'min'
:
_agg_block_min_nb
,
'max'
:
_agg_block_max_nb
,
'median'
:
_agg_block_median_nb
,
'mode'
:
_agg_block_mode_nb
,
}
# -- Dask block helpers ------------------------------------------------------
#
# Interpolation uses map_coordinates with *global* coordinate mapping so
# that results are identical regardless of chunk layout. Each block
# receives the cumulative chunk boundaries and computes which global
# output pixels it is responsible for, maps them back to global input
# coordinates, then converts to local (within-block) coordinates.
def
_interp_block_np
(
block
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
cum_in_y
,
cum_in_x
,
cum_out_y
,
cum_out_x
,
depth_y
,
depth_x
,
order
,
work_dtype
,
out_dtype
,
block_info
=
None
):
"""Interpolate one (possibly overlapped) numpy block."""
yi
,
xi
=
block_info
[
0
][
'chunk-location'
]
target_h
=
int
(
cum_out_y
[
yi
+
1
]
-
cum_out_y
[
yi
])
target_w
=
int
(
cum_out_x
[
xi
+
1
]
-
cum_out_x
[
xi
])
block
=
_maybe_astype
(
block
,
work_dtype
)
# Global output pixel indices for this chunk
oy
=
np
.
arange
(
cum_out_y
[
yi
],
cum_out_y
[
yi
+
1
],
dtype
=
np
.
float64
)
ox
=
np
.
arange
(
cum_out_x
[
xi
],
cum_out_x
[
xi
+
1
],
dtype
=
np
.
float64
)
# Map to global input coordinates using block-centered formula
iy
=
(
oy
+
0.5
)
*
(
global_in_h
/
global_out_h
)
-
0.5
ix
=
(
ox
+
0.5
)
*
(
global_in_w
/
global_out_w
)
-
0.5
# Convert to local block coordinates (overlap shifts the origin)
iy_local
=
iy
-
(
cum_in_y
[
yi
]
-
depth_y
)
ix_local
=
ix
-
(
cum_in_x
[
xi
]
-
depth_x
)
yy
,
xx
=
np
.
meshgrid
(
iy_local
,
ix_local
,
indexing
=
'ij'
)
coords
=
np
.
array
([
yy
.
ravel
(),
xx
.
ravel
()])
# NaN-aware interpolation. For order >= 2 we run the spline prefilter
# explicitly per array (block / filled / weights) so the IIR boundary
# transient is identical between eager and chunked paths.
use_explicit
=
order
>=
2
mask
=
np
.
isnan
(
block
)
if
order
==
0
or
not
mask
.
any
():
if
use_explicit
:
src
,
npad
=
_prepad_and_filter_np
(
block
,
order
)
result
=
_scipy_map_coords
(
src
,
coords
+
npad
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
result
=
_scipy_map_coords
(
block
,
coords
,
order
=
order
,
mode
=
'nearest'
)
else
:
filled
=
np
.
where
(
mask
,
0.0
,
block
)
weights
=
(
~
mask
).
astype
(
block
.
dtype
)
if
use_explicit
:
filled
,
npad
=
_prepad_and_filter_np
(
filled
,
order
)
weights
,
_
=
_prepad_and_filter_np
(
weights
,
order
)
sample_coords
=
coords
+
npad
z_data
=
_scipy_map_coords
(
filled
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
z_wt
=
_scipy_map_coords
(
weights
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
z_data
=
_scipy_map_coords
(
filled
,
coords
,
order
=
order
,
mode
=
'nearest'
)
z_wt
=
_scipy_map_coords
(
weights
,
coords
,
order
=
order
,
mode
=
'nearest'
)
# Majority-weight gate (see _nan_aware_interp_np for rationale).
result
=
np
.
where
(
z_wt
>
0.5
,
z_data
/
np
.
maximum
(
z_wt
,
1e-10
),
np
.
nan
)
return
_maybe_astype
(
result
.
reshape
(
target_h
,
target_w
),
out_dtype
)
def
_interp_block_cupy
(
block
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
cum_in_y
,
cum_in_x
,
cum_out_y
,
cum_out_x
,
depth_y
,
depth_x
,
order
,
work_dtype
,
out_dtype
,
block_info
=
None
):
"""CuPy variant of :func:`_interp_block_np`."""
from
cupyx
.
scipy
.
ndimage
import
map_coordinates
as
_cupy_map_coords
from
cupyx
.
scipy
.
ndimage
import
spline_filter
as
_cupy_spline_filter
yi
,
xi
=
block_info
[
0
][
'chunk-location'
]
target_h
=
int
(
cum_out_y
[
yi
+
1
]
-
cum_out_y
[
yi
])
target_w
=
int
(
cum_out_x
[
xi
+
1
]
-
cum_out_x
[
xi
])
if
block
.
dtype
!=
cupy
.
dtype
(
work_dtype
):
block
=
block
.
astype
(
work_dtype
)
oy
=
cupy
.
arange
(
int
(
cum_out_y
[
yi
]),
int
(
cum_out_y
[
yi
+
1
]),
dtype
=
cupy
.
float64
)
ox
=
cupy
.
arange
(
int
(
cum_out_x
[
xi
]),
int
(
cum_out_x
[
xi
+
1
]),
dtype
=
cupy
.
float64
)
# Map to global input coordinates using block-centered formula
iy
=
(
oy
+
0.5
)
*
(
global_in_h
/
global_out_h
)
-
0.5
ix
=
(
ox
+
0.5
)
*
(
global_in_w
/
global_out_w
)
-
0.5
iy_local
=
iy
-
float
(
cum_in_y
[
yi
]
-
depth_y
)
ix_local
=
ix
-
float
(
cum_in_x
[
xi
]
-
depth_x
)
yy
,
xx
=
cupy
.
meshgrid
(
iy_local
,
ix_local
,
indexing
=
'ij'
)
coords
=
cupy
.
array
([
yy
.
ravel
(),
xx
.
ravel
()])
use_explicit
=
order
>=
2
mask
=
cupy
.
isnan
(
block
)
if
order
==
0
or
not
mask
.
any
():
if
use_explicit
:
src
,
npad
=
_prepad_and_filter_cupy
(
block
,
order
,
_cupy_spline_filter
)
result
=
_cupy_map_coords
(
src
,
coords
+
npad
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
result
=
_cupy_map_coords
(
block
,
coords
,
order
=
order
,
mode
=
'nearest'
)
else
:
filled
=
cupy
.
where
(
mask
,
0.0
,
block
)
weights
=
(
~
mask
).
astype
(
block
.
dtype
)
if
use_explicit
:
filled
,
npad
=
_prepad_and_filter_cupy
(
filled
,
order
,
_cupy_spline_filter
)
weights
,
_
=
_prepad_and_filter_cupy
(
weights
,
order
,
_cupy_spline_filter
)
sample_coords
=
coords
+
npad
z_data
=
_cupy_map_coords
(
filled
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
z_wt
=
_cupy_map_coords
(
weights
,
sample_coords
,
order
=
order
,
mode
=
'nearest'
,
prefilter
=
False
)
else
:
z_data
=
_cupy_map_coords
(
filled
,
coords
,
order
=
order
,
mode
=
'nearest'
)
z_wt
=
_cupy_map_coords
(
weights
,
coords
,
order
=
order
,
mode
=
'nearest'
)
# Majority-weight gate (see _nan_aware_interp_np for rationale).
result
=
cupy
.
where
(
z_wt
>
0.5
,
z_data
/
cupy
.
maximum
(
z_wt
,
1e-10
),
cupy
.
nan
)
result
=
result
.
reshape
(
target_h
,
target_w
)
if
result
.
dtype
!=
cupy
.
dtype
(
out_dtype
):
result
=
result
.
astype
(
out_dtype
)
return
result
def
_agg_block_np
(
block
,
method
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
cum_in_y
,
cum_in_x
,
cum_out_y
,
cum_out_x
,
depth_y
,
depth_x
,
out_dtype
,
block_info
=
None
):
"""Block-aggregate one (possibly overlapped) numpy chunk.
Runs the entire chunk inside one numba dispatch via the
`_agg_block_*_nb` kernels. Earlier versions called a 1x1 jitted
aggregate per output pixel, which scaled badly for large rasters.
"""
yi
,
xi
=
block_info
[
0
][
'chunk-location'
]
target_h
=
int
(
cum_out_y
[
yi
+
1
]
-
cum_out_y
[
yi
])
target_w
=
int
(
cum_out_x
[
xi
+
1
]
-
cum_out_x
[
xi
])
# _AGG_FUNCS kernels are @ngjit-compiled with hard-coded float64
# working buffers; cast accordingly so numba dispatch matches.
block
=
_maybe_astype
(
block
,
np
.
float64
)
# The overlapped block starts depth pixels before the original chunk
in_y0
=
int
(
cum_in_y
[
yi
])
-
depth_y
in_x0
=
int
(
cum_in_x
[
xi
])
-
depth_x
go_y0
=
int
(
cum_out_y
[
yi
])
go_x0
=
int
(
cum_out_x
[
xi
])
kernel
=
_AGG_BLOCK_FUNCS
[
method
]
out
=
kernel
(
block
,
target_h
,
target_w
,
go_y0
,
go_x0
,
int
(
global_in_h
),
int
(
global_in_w
),
int
(
global_out_h
),
int
(
global_out_w
),
in_y0
,
in_x0
)
return
_maybe_astype
(
out
,
out_dtype
)
def
_agg_block_cupy
(
block
,
method
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
cum_in_y
,
cum_in_x
,
cum_out_y
,
cum_out_x
,
depth_y
,
depth_x
,
out_dtype
,
block_info
=
None
):
"""Block-aggregate one cupy chunk (falls back to CPU)."""
cpu
=
cupy
.
asnumpy
(
block
)
result
=
_agg_block_np
(
cpu
,
method
,
global_in_h
,
global_in_w
,
global_out_h
,
global_out_w
,
cum_in_y
,
cum_in_x
,
cum_out_y
,
cum_out_x
,
depth_y
,
depth_x
,
out_dtype
,
block_info
=
block_info
,
)
return
cupy
.
asarray
(
result
)
# -- Per-backend runners -----------------------------------------------------
def
_run_numpy
(
data
,
scale_y
,
scale_x
,
method
):
work_dt
=
_working_dtype
(
data
.
dtype
)
out_dt
=
_output_dtype
(
data
.
dtype
)
data
=
_maybe_astype
(
data
,
work_dt
)
out_h
,
out_w
=
_output_shape
(
*
data
.
shape
,
scale_y
,
scale_x
)
if
method
in
INTERP_METHODS
:
result
=
_nan_aware_interp_np
(
data
,
out_h
,
out_w
,
INTERP_METHODS
[
method
])
return
_maybe_astype
(
result
,
out_dt
)
result
=
_AGG_FUNCS
[
method
](
data
,
out_h
,
out_w
)
return
_maybe_astype
(
result
,
out_dt
)
def
_run_cupy
(
data
,
scale_y
,
scale_x
,
method
):
work_dt
=
_working_dtype
(
data
.
dtype
)
out_dt
=
_output_dtype
(
data
.
dtype
)
data
=
data
if
data
.
dtype
==
cupy
.
dtype
(
work_dt
)
else
data
.
astype
(
work_dt
)
out_h
,
out_w
=
_output_shape
(
*
data
.
shape
,
scale_y
,
scale_x
)
if
method
in
INTERP_METHODS
:
View remainder of file in raw view
Back
|
FazBrowse Home
|
New Git URL