8000 Merge pull request #11011 from eric-wieser/pad-singleton · numpy/numpy@342609b · GitHub
[go: up one dir, main page]
Skip to content

Commit 342609b

Browse files
charrischarris
authored
Merge pull request #11011 from eric-wieser/pad-singleton
MAINT: Simplify dimension-juggling in np.pad
2 parents 896363f + 4bdcbab commit 342609b

File tree

1 file cha 8000 nged

+34
-109
lines changed

1 file changed

+34
-109
lines changed

‎numpy/lib/arraypad.py

Lines changed: 34 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -156,13 +156,9 @@ def _prepend_edge(arr, pad_amt, axis=-1):
156156
if pad_amt == 0:
157157
return arr
158158

159-
edge_slice = tuple(slice(None) if i != axis else 0
159+
edge_slice = tuple(slice(None) if i != axis else slice(0, 1)
160160
for (i, x) in enumerate(arr.shape))
161-
162-
# Shape to restore singleton dimension after slicing
163-
pad_singleton = tuple(x if i != axis else 1
164-
for (i, x) in enumerate(arr.shape))
165-
edge_arr = arr[edge_slice].reshape(pad_singleton)
161+
edge_arr = arr[edge_slice]
166162
return np.concatenate((edge_arr.repeat(pad_amt, axis=axis), arr),
167163
axis=axis)
168164

@@ -190,13 +186,9 @@ def _append_edge(arr, pad_amt, axis=-1):
190186
if pad_amt == 0:
191187
return arr
192188

193-
edge_slice = tuple(slice(None) if i != axis else arr.shape[axis] - 1
189+
edge_slice = tuple(slice(None) if i != axis else slice(x - 1, x)
194190
for (i, x) in enumerate(arr.shape))
195-
196-
# Shape to restore singleton dimension after slicing
197-
pad_singleton = tuple(x if i != axis else 1
198-
for (i, x) in enumerate(arr.shape))
199-
edge_arr = arr[edge_slice].reshape(pad_singleton)
191+
edge_arr = arr[edge_slice]
200192
return np.concatenate((arr, edge_arr.repeat(pad_amt, axis=axis)),
201193
axis=axis)
202194

@@ -236,15 +228,11 @@ def _prepend_ramp(arr, pad_amt, end, axis=-1):
236228
reverse=True).astype(np.float64)
237229

238230
# Appropriate slicing to extract n-dimensional edge along `axis`
239-
edge_slice = tuple(slice(None) if i != axis else 0
231+
edge_slice = tuple(slice(None) if i != axis else slice(0, 1)
240232
for (i, x) in enumerate(arr.shape))
241233

242-
# Shape to restore singleton dimension after slicing
243-
pad_singleton = tuple(x if i != axis else 1
244-
for (i, x) in enumerate(arr.shape))
245-
246-
# Extract edge, reshape to original rank, and extend along `axis`
247-
edge_pad = arr[edge_slice].reshape(pad_singleton).repeat(pad_amt, axis)
234+
# Extract edge, and extend along `axis`
235+
edge_pad = arr[edge_slice].repeat(pad_amt, axis)
248236

249237
# Linear ramp
250238
slope = (end - edge_pad) / float(pad_amt)
@@ -291,15 +279,11 @@ def _append_ramp(arr, pad_amt, end, axis=-1):
291279
reverse=False).astype(np.float64)
292280

293281
# Slice a chunk from the edge to calculate stats on
294-
edge_slice = tuple(slice(None) if i != axis else -1
282+
edge_slice = tuple(slice(None) if i != axis else slice(x - 1, x)
295283
for (i, x) in enumerate(arr.shape))
296284

297-
# Shape to restore singleton dimension after slicing
298-
pad_singleton = tuple(x if i != axis else 1
299-
for (i, x) in enumerate(arr.shape))
300-
301-
# Extract edge, reshape to original rank, and extend along `axis`
302-
edge_pad = arr[edge_slice].reshape(pad_singleton).repeat(pad_amt, axis)
285+
# Extract edge, and extend along `axis`
286+
edge_pad = arr[edge_ 1E80 slice].repeat(pad_amt, axis)
303287

304288
# Linear ramp
305289
slope = (end - edge_pad) / float(pad_amt)
@@ -351,12 +335,8 @@ def _prepend_max(arr, pad_amt, num, axis=-1):
351335
max_slice = tuple(slice(None) if i != axis else slice(num)
352336
for (i, x) in enumerate(arr.shape))
353337

354-
# Shape to restore singleton dimension after slicing
355-
pad_singleton = tuple(x if i != axis else 1
356-
for (i, x) in enumerate(arr.shape))
357-
358-
# Extract slice, calculate max, reshape to add singleton dimension back
359-
max_chunk = arr[max_slice].max(axis=axis).reshape(pad_singleton)
338+
# Extract slice, calculate max
339+
max_chunk = arr[max_slice].max(axis=axis, keepdims=True)
360340

361341
# Concatenate `arr` with `max_chunk`, extended along `axis` by `pad_amt`
362342
return np.concatenate((max_chunk.repeat(pad_amt, axis=axis), arr),
@@ -407,12 +387,8 @@ def _append_max(arr, pad_amt, num, axis=-1):
407387
else:
408388
max_slice = tuple(slice(None) for x in arr.shape)
409389

410-
# Shape to restore singleton dimension after slicing
411-
pad_singleton = tuple(x if i != axis else 1
412-
for (i, x) in enumerate(arr.shape))
413-
414-
# Extract slice, calculate max, reshape to add singleton dimension back
415-
max_chunk = arr[max_slice].max(axis=axis).reshape(pad_singleton)
390+
# Extract slice, calculate max
391+
max_chunk = arr[max_slice].max(axis=axis, keepdims=True)
416392

417393
# Concatenate `arr` with `max_chunk`, extended along `axis` by `pad_amt`
418394
return np.concatenate((arr, max_chunk.repeat(pad_amt, axis=axis)),
@@ -458,12 +434,8 @@ def _prepend_mean(arr, pad_amt, num, axis=-1):
458434
mean_slice = tuple(slice(None) if i != axis else slice(num)
459435
for (i, x) in enumerate(arr.shape))
460436

461-
# Shape to restore singleton dimension after slicing
462-
pad_singleton = tuple(x if i != axis else 1
463-
for (i, x) in enumerate(arr.shape))
464-
465-
# Extract slice, calculate mean, reshape to add singleton dimension back
466-
mean_chunk = arr[mean_slice].mean(axis).reshape(pad_singleton)
437+
# Extract slice, calculate mean
438+
mean_chunk = arr[mean_slice].mean(axis, keepdims=True)
467439
_round_ifneeded(mean_chunk, arr.dtype)
468440

469441
# Concatenate `arr` with `mean_chunk`, extended along `axis` by `pad_amt`
@@ -515,12 +487,8 @@ def _append_mean(arr, pad_amt, num, axis=-1):
515487
else:
516488
mean_slice = tuple(slice(None) for x in arr.shape)
517489

518-
# Shape to restore singleton dimension after slicing
519-
pad_singleton = tuple(x if i != axis else 1
520-
for (i, x) in enumerate(arr.shape))
521-
522-
# Extract slice, calculate mean, reshape to add singleton dimension back
523-
mean_chunk = arr[mean_slice].mean(axis=axis).reshape(pad_singleton)
490+
# Extract slice, calculate mean
491+
mean_chunk = arr[mean_slice].mean(axis=axis, keepdims=True)
524492
_round_ifneeded(mean_chunk, arr.dtype)
525493

526494
# Concatenate `arr` with `mean_chunk`, extended along `axis` by `pad_amt`
@@ -567,12 +535,8 @@ def _prepend_med(arr, pad_amt, num, axis=-1):
567535
med_slice = tuple(slice(None) if i != axis else slice(num)
568536
for (i, x) in enumerate(arr.shape))
569537

570-
# Shape to restore singleton dimension after slicing
571-
pad_singleton = tuple(x if i != axis else 1
572-
for (i, x) in enumerate(arr.shape))
573-
574-
# Extract slice, calculate median, reshape to add singleton dimension back
575-
med_chunk = np.median(arr[med_slice], axis=axis).reshape(pad_singleton)
538+
# Extract slice, calculate median
539+
med_chunk = np.median(arr[med_slice], axis=axis, keepdims=True)
576540
_round_ifneeded(med_chunk, arr.dtype)
577541

578542
# Concatenate `arr` with `med_chunk`, extended along `axis` by `pad_amt`
@@ -624,12 +588,8 @@ def _append_med(arr, pad_amt, num, axis=-1):
624588
else:
625589
med_slice = tuple(slice(None) for x in arr.shape)
626590

627-
# Shape to restore singleton dimension after slicing
628-
pad_singleton = tuple(x if i != axis else 1
629-
for (i, x) in enumerate(arr.shape))
630-
631-
# Extract slice, calculate median, reshape to add singleton dimension back
632-
med_chunk = np.median(arr[med_slice], axis=axis).reshape(pad_singleton)
591+
# Extract slice, calculate median
592+
med_chunk = np.median(arr[med_slice], axis=axis, keepdims=True)
633593
_round_ifneeded(med_chunk, arr.dtype)
634594

635595
# Concatenate `arr` with `med_chunk`, extended along `axis` by `pad_amt`
@@ -677,12 +637,8 @@ def _prepend_min(arr, pad_amt, num, axis=-1):
677637
min_slice = tuple(slice(None) if i != axis else slice(num)
678638
for (i, x) in enumerate(arr.shape))
679639

680-
# Shape to restore singleton dimension after slicing
681-
pad_singleton = tuple(x if i != axis else 1
682-
for (i, x) in enumerate(arr.shape))
683-
684-
# Extract slice, calculate min, reshape to add singleton dimension back
685-
min_chunk = arr[min_slice].min(axis=axis).reshape(pad_singleton)
640+
# Extract slice, calculate min
641+
min_chunk = arr[min_slice].min(axis=axis, keepdims=True)
686642

687643
# Concatenate `arr` with `min_chunk`, extended along `axis` by `pad_amt`
688644
return np.concatenate((min_chunk.repeat(pad_amt, axis=axis), arr),
@@ -733,12 +689,8 @@ def _append_min(arr, pad_amt, num, axis=-1):
733689
else:
734690
min_slice = tuple(slice(None) for x in arr.shape)
735691

736-
# Shape to restore singleton dimension after slicing
737-
pad_singleton = tuple(x if i != axis else 1
738-
for (i, x) in enumerate(arr.shape))
739-
740-
# Extract slice, calculate min, reshape to add singleton dimension back
741-
min_chunk = arr[min_slice].min(axis=axis).reshape(pad_singleton)
692+
# Extract slice, calculate min
693+
min_chunk = arr[min_slice].min(axis=axis, keepdims=True)
742694

743695
# Concatenate `arr` with `min_chunk`, extended along `axis` by `pad_amt`
744696
return np.concatenate((arr, min_chunk.repeat(pad_amt, axis=axis)),
@@ -790,17 +742,11 @@ def _pad_ref(arr, pad_amt, method, axis=-1):
790742

791743
ref_chunk1 = arr[ref_slice]
792744

793-
# Shape to restore singleton dimension after slicing
794-
pad_singleton = tuple(x if i != axis else 1
795-
for (i, x) in enumerate(arr.shape))
796-
if pad_amt[0] == 1:
797-
ref_chunk1 = ref_chunk1.reshape(pad_singleton)
798-
799745
# Memory/computationally more expensive, only do this if `method='odd'`
800746
if 'odd' in method and pad_amt[0] > 0:
801-
edge_slice1 = tuple(slice(None) if i != axis else 0
747+
edge_slice1 = tuple(slice(None) if i != axis else slice(0, 1)
802748
for (i, x) in enumerate(arr.shape))
803-
edge_chunk = arr[edge_slice1].reshape(pad_singleton)
749+
edge_chunk = arr[edge_slice1]
804750
ref_chunk1 = 2 * edge_chunk - ref_chunk1
805751
del edge_chunk
806752

@@ -816,13 +762,10 @@ def _pad_ref(arr, pad_amt, method, axis=-1):
816762
for (i, x) in enumerate(arr.shape))
817763
ref_chunk2 = arr[ref_slice][rev_idx]
818764

819-
if pad_amt[1] == 1:
820-
ref_chunk2 = ref_chunk2.reshape(pad_singleton)
821-
822765
if 'odd' in method:
823-
edge_slice2 = tuple(slice(None) if i != axis else -1
766+
edge_slice2 = tuple(slice(None) if i != axis else slice(x - 1, x)
824767
for (i, x) in enumerate(arr.shape))
825-
edge_chunk = arr[edge_slice2].reshape(pad_singleton)
768+
edge_chunk = arr[edge_slice2]
826769
ref_chunk2 = 2 * edge_chunk - ref_chunk2
827770
del edge_chunk
828771

@@ -876,17 +819,11 @@ def _pad_sym(arr, pad_amt, method, axis=-1):
876819
for (i, x) in enumerate(arr.shape))
877820
sym_chunk1 = arr[sym_slice][rev_idx]
878821

879-
# Shape to restore singleton dimension after slicing
880-
pad_singleton = tuple(x if i != axis else 1
881-
for (i, x) in enumerate(arr.shape))
882-
if pad_amt[0] == 1:
883-
sym_chunk1 = sym_chunk1.reshape(pad_singleton)
884-
885822
# Memory/computationally more expensive, only do this if `method='odd'`
886823
if 'odd' in method and pad_amt[0] > 0:
887-
edge_slice1 = tuple(slice(None) if i != axis else 0
824+
edge_slice1 = tupl F438 e(slice(None) if i != axis else slice(0, 1)
888825
for (i, x) in enumerate(arr.shape))
889-
edge_chunk = arr[edge_slice1].reshape(pad_singleton)
826+
edge_chunk = arr[edge_slice1]
890827
sym_chunk1 = 2 * edge_chunk - sym_chunk1
891828
del edge_chunk
892829

@@ -900,13 +837,10 @@ def _pad_sym(arr, pad_amt, method, axis=-1):
900837
for (i, x) in enumerate(arr.shape))
901838
sym_chunk2 = arr[sym_slice][rev_idx]
902839

903-
if pad_amt[1] == 1:
904-
sym_chunk2 = sym_chunk2.reshape(pad_singleton)
905-
906840
if 'odd' in method:
907-
edge_slice2 = tuple(slice(None) if i != axis else -1
841+
edge_slice2 = tuple(slice(None) if i != axis else slice(x - 1, x)
908842
for (i, x) in enumerate(arr.shape))
909-
edge_chunk = arr[edge_slice2].reshape(pad_singleton)
843+
edge_chunk = arr[edge_slice2]
910844
sym_chunk2 = 2 * edge_chunk - sym_chunk2
911845
del edge_chunk
912846

@@ -957,12 +891,6 @@ def _pad_wrap(arr, pad_amt, axis=-1):
957891
for (i, x) in enumerate(arr.shape))
958892
wrap_chunk1 = arr[wrap_slice]
959893

960-
# Shape to restore singleton dimension after slicing
961-
pad_singleton = tuple(x if i != axis else 1
962-
for (i, x) in enumerate(arr.shape))
963-
if pad_amt[0] == 1:
964-
wrap_chunk1 = wrap_chunk1.reshape(pad_singleton)
965-
966894
##########################################################################
967895
# Appended region
968896

@@ -971,9 +899,6 @@ def _pad_wrap(arr, pad_amt, axis=-1):
971899
for (i, x) in enumerate(arr.shape))
972900
wrap_chunk2 = arr[wrap_slice]
973901

974-
if pad_amt[1] == 1:
975-
wrap_chunk2 = wrap_chunk2.reshape(pad_singleton)
976-
977902
# Concatenate `arr` with both chunks, extending along `axis`
978903
return np.concatenate((wrap_chunk1, arr, wrap_chunk2), axis=axis)
979904

0 commit comments

Comments
 (0)
0