fix zip(strict=True) for 2D problems

This commit is contained in:
Jan Petykiewicz 2025-01-07 00:05:19 -08:00
parent 47415a0beb
commit 4f2433320d

View File

@ -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')