def test_impl(): n = 20 arr = np.ones((n, n)) val = 0 for idx in pndindex(arr.shape): val += idx[0] * idx[1] return val
def do_scale(X, maxv, nthr): with numba.objmode(t0='f8'): t0 = time.time() #nthr = numba.get_num_threads() s = np.zeros((nthr, X.shape[1])) ss = np.zeros((nthr, X.shape[1])) mean = np.zeros(X.shape[1]) std = np.zeros(X.shape[1]) n = X.shape[0] for i in numba.prange(nthr): for r in range(i, n, nthr): for c in range(X.shape[1]): v = X[r, c] s[i, c] += v ss[i, c] += v * v for c in numba.prange(X.shape[1]): s0 = s[:, c].sum() mean[c] = s0 / n std[c] = np.sqrt((ss[:, c].sum() - s0 * s0 / n) / (n - 1)) with numba.objmode(): print("finshed getting means, stddev", time.time() - t0) for (r, c) in numba.pndindex(X.shape): v = (X[r, c] - mean[c]) / std[c] X[r, c] = maxv if v > maxv else v with numba.objmode(): print("finshed scaling values", time.time() - t0)
def normal_cone_angles(normals: np.ndarray, threshold: float, min_norm): assert normals.ndim == 4 size = normals.shape[0] assert normals.shape == (size, size, size, 3) assert size > 2 result = np.zeros((size - 2, size - 2, size - 2), dtype=np.bool8) for i in numba.pndindex((size - 2, size - 2, size - 2)): # Collect normals for position i current = np.empty((26, 3), dtype=np.float32) # 26 possible neighbors ci: numba.uint32 = 0 for o in np.ndindex((3, 3, 3)): if o != (1, 1, 1): x, y, z = i[0] + o[0], i[1] + o[1], i[2] + o[2] value = normals[x, y, z] norm = np.linalg.norm(value) if norm > min_norm: # Only add if norm is valid current[ci] = value / norm ci += 1 if ci > 2: valid = current[:ci] # Check angle between all valid normals result[i[0], i[1], i[2]] = np.any(np.arccos(valid @ valid.T) > threshold) return result
def set_array_normals(target: np.ndarray, mask: np.ndarray, values: np.ndarray): assert target.shape == mask.shape and mask.shape == values.shape for i in numba.pndindex(target.shape): x, y, z = i if mask[x, y, z]: target[x, y, z] = values[x, y, z]
def _collapse_dz(deltas: np.ndarray): col = np.zeros(deltas.shape, dtype=deltas.dtype) for xy in numba.pndindex((deltas.shape[0], deltas.shape[1])): x, y = xy z0 = 0 value0 = np.nan d0 = 1 if value0 > 0 else -1 for z in range(deltas.shape[2]): current = deltas[x, y, z] if current and value0 == current: col[x, y, z0] += d0 continue z0 = z value0 = current d0 = 1 if value0 > 0 else -1 if current: col[x, y, z0] = d0 return col
def _collapse_dx(deltas: np.ndarray): col = np.zeros(deltas.shape, dtype=deltas.dtype) for yz in numba.pndindex((deltas.shape[1], deltas.shape[2])): y, z = yz x0 = 0 value0 = np.nan d0 = 1 if value0 > 0 else -1 for x in range(deltas.shape[0]): current = deltas[x, y, z] if current and value0 == current: col[x0, y, z] += d0 continue x0 = x value0 = current d0 = 1 if value0 > 0 else -1 if current: col[x0, y, z] = d0 return col
def _sum(X): s = 0 for i in numba.pndindex(X.shape): s += X[i] return s