fix zip(strict=True) for 2D problems
This commit is contained in:
parent
47415a0beb
commit
4f2433320d
@ -34,7 +34,7 @@ def shift_circ(
|
||||
if axis not in range(len(shape)):
|
||||
raise Exception(f'Invalid direction: {axis}, shape is {shape}')
|
||||
|
||||
shifts = [abs(shift_distance) if a == axis else 0 for a in range(3)]
|
||||
shifts = [abs(shift_distance) if a == axis else 0 for a in range(len(shape))]
|
||||
shifted_diags = [(numpy.arange(n) + s) % n for n, s in zip(shape, shifts, strict=True)]
|
||||
ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
|
||||
|
||||
@ -82,7 +82,7 @@ def shift_with_mirror(
|
||||
v = numpy.where(v < 0, - 1 - v, v)
|
||||
return v
|
||||
|
||||
shifts = [shift_distance if a == axis else 0 for a in range(3)]
|
||||
shifts = [shift_distance if a == axis else 0 for a in range(len(shape))]
|
||||
shifted_diags = [mirrored_range(n, s) for n, s in zip(shape, shifts, strict=True)]
|
||||
ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user