|
import math |
|
import numpy as np |
|
|
|
|
|
|
|
|
|
def bitreversal_po2(n): |
|
""" |
|
From S4 codebase. |
|
|
|
:param n: |
|
:return: |
|
""" |
|
m = int(math.log(n)/math.log(2)) |
|
perm = np.arange(n).reshape(n,1) |
|
for i in range(m): |
|
n1 = perm.shape[0]//2 |
|
perm = np.hstack((perm[:n1],perm[n1:])) |
|
return perm.squeeze(0) |
|
|
|
|
|
def bitreversal_permutation(n): |
|
""" |
|
From S4 codebase. |
|
|
|
:param n: |
|
:return: |
|
""" |
|
m = int(math.ceil(math.log(n)/math.log(2))) |
|
N = 1 << m |
|
perm = bitreversal_po2(N) |
|
return np.extract(perm < n, perm) |
|
|
|
|
|
def transpose_permutation(h, w): |
|
""" |
|
From S4 codebase. |
|
|
|
:param h: |
|
:param w: |
|
:return: |
|
""" |
|
indices = np.arange(h*w) |
|
indices = indices.reshape((h, w)) |
|
indices = indices.T |
|
indices = indices.reshape(h*w) |
|
return indices |
|
|
|
|
|
def snake_permutation(h, w): |
|
""" |
|
From S4 codebase. |
|
|
|
:param h: |
|
:param w: |
|
:return: |
|
""" |
|
indices = np.arange(h*w) |
|
indices = indices.reshape((h, w)) |
|
indices[1::2, :] = indices[1::2, ::-1] |
|
indices = indices.reshape(h*w) |
|
return indices |
|
|
|
|
|
def hilbert_permutation(n): |
|
""" |
|
From S4 codebase. |
|
|
|
From S4 codebase. |
|
:param n: |
|
:return: |
|
""" |
|
|
|
m = int(math.log2(n)) |
|
assert n == 2**m |
|
inds = decode(list(range(n*n)), 2, m) |
|
ind_x, ind_y = inds.T |
|
indices = np.arange(n*n).reshape((n, n)) |
|
indices = indices[ind_x, ind_y] |
|
return(indices) |
|
|
|
|
|
def decode(hilberts, num_dims, num_bits): |
|
''' |
|
From S4 codebase. |
|
|
|
Hilbert curve utilities taken from https://github.com/PrincetonLIPS/numpy-hilbert-curve |
|
|
|
Decode an array of Hilbert integers into locations in a hypercube. |
|
This is a vectorized-ish version of the Hilbert curve implementation by John |
|
Skilling as described in: |
|
Skilling, J. (2004, April). Programming the Hilbert curve. In AIP Conference |
|
Proceedings (Vol. 707, No. 1, pp. 381-387). American Institute of Physics. |
|
Params: |
|
------- |
|
hilberts - An ndarray of Hilbert integers. Must be an integer dtype and |
|
cannot have fewer bits than num_dims * num_bits. |
|
num_dims - The dimensionality of the hypercube. Integer. |
|
num_bits - The number of bits for each dimension. Integer. |
|
Returns: |
|
-------- |
|
The output is an ndarray of unsigned integers with the same shape as hilberts |
|
but with an additional dimension of size num_dims. |
|
''' |
|
|
|
if num_dims*num_bits > 64: |
|
raise ValueError( |
|
''' |
|
num_dims=%d and num_bits=%d for %d bits total, which can't be encoded |
|
into a uint64. Are you sure you need that many points on your Hilbert |
|
curve? |
|
''' % (num_dims, num_bits) |
|
) |
|
|
|
|
|
hilberts = np.atleast_1d(hilberts) |
|
|
|
|
|
orig_shape = hilberts.shape |
|
|
|
|
|
|
|
hh_uint8 = np.reshape(hilberts.ravel().astype('>u8').view(np.uint8), (-1, 8)) |
|
|
|
|
|
|
|
hh_bits = np.unpackbits(hh_uint8, axis=1)[:, -num_dims*num_bits:] |
|
|
|
|
|
gray = binary2gray(hh_bits) |
|
|
|
|
|
|
|
gray = np.swapaxes( |
|
np.reshape(gray, (-1, num_bits, num_dims)), |
|
axis1=1, axis2=2, |
|
) |
|
|
|
|
|
for bit in range(num_bits-1, -1, -1): |
|
|
|
|
|
for dim in range(num_dims-1, -1, -1): |
|
|
|
|
|
mask = gray[:, dim, bit] |
|
|
|
|
|
gray[:,0,bit+1:] = np.logical_xor(gray[:, 0, bit+1:], mask[:, np.newaxis]) |
|
|
|
|
|
to_flip = np.logical_and( |
|
np.logical_not(mask[:, np.newaxis]), |
|
np.logical_xor(gray[:, 0, bit+1:], gray[:, dim, bit+1:]) |
|
) |
|
gray[:,dim,bit+1:] = np.logical_xor(gray[:, dim, bit+1:], to_flip) |
|
gray[:,0,bit+1:] = np.logical_xor(gray[:, 0, bit+1:], to_flip) |
|
|
|
|
|
extra_dims = 64 - num_bits |
|
padded = np.pad(gray, ((0, 0), (0, 0), (extra_dims, 0)), |
|
mode='constant', constant_values=0) |
|
|
|
|
|
locs_chopped = np.reshape(padded[:, :, ::-1], (-1, num_dims, 8, 8)) |
|
|
|
|
|
locs_uint8 = np.squeeze(np.packbits(locs_chopped, bitorder='little', axis=3)) |
|
|
|
|
|
flat_locs = locs_uint8.view(np.uint64) |
|
|
|
|
|
return np.reshape(flat_locs, (*orig_shape, num_dims)) |
|
|
|
|
|
def right_shift(binary, k=1, axis=-1): |
|
''' Right shift an array of binary values. |
|
Parameters: |
|
----------- |
|
binary: An ndarray of binary values. |
|
k: The number of bits to shift. Default 1. |
|
axis: The axis along which to shift. Default -1. |
|
Returns: |
|
-------- |
|
Returns an ndarray with zero prepended and the ends truncated, along |
|
whatever axis was specified. |
|
''' |
|
|
|
|
|
if binary.shape[axis] <= k: |
|
return np.zeros_like(binary) |
|
|
|
|
|
padding = [(0, 0)] * len(binary.shape) |
|
padding[axis] = (k, 0) |
|
|
|
|
|
slicing = [slice(None)] * len(binary.shape) |
|
slicing[axis] = slice(None, -k) |
|
|
|
shifted = np.pad(binary[tuple(slicing)], padding, |
|
mode='constant', constant_values=0) |
|
|
|
return shifted |
|
|
|
|
|
def binary2gray(binary, axis=-1): |
|
''' Convert an array of binary values into Gray codes. |
|
This uses the classic X ^ (X >> 1) trick to compute the Gray code. |
|
Parameters: |
|
----------- |
|
binary: An ndarray of binary values. |
|
axis: The axis along which to compute the gray code. Default=-1. |
|
Returns: |
|
-------- |
|
Returns an ndarray of Gray codes. |
|
''' |
|
shifted = right_shift(binary, axis=axis) |
|
|
|
|
|
gray = np.logical_xor(binary, shifted) |
|
|
|
return gray |
|
|