JAX: JIT compatible sparse matrix slicing - python

I have a boolean sparse matrix that I represent with row indices and column indices of True values.
import numpy as np
import jax
from jax import numpy as jnp
N = 10000
M = 1000
X = np.random.randint(0, 100, size=(N, M)) == 0 # data setup
rows, cols = np.where(X == True)
rows = jax.device_put(rows)
cols = jax.device_put(cols)
I want to get a column slice of the matrix like X[:, 3], but just from rows indices and column indices.
I managed to do that by using jnp.isin like below, but the problem is that this is not JIT compatible because of the data-dependent shaped array rows[cols == m].
def not_jit_compatible_slice(rows, cols, m):
return jnp.isin(jnp.arange(N), rows[cols == m])
I could make it JIT compatible by using jnp.where in the three-argument form, but this operation is much slower than the previous one.
def jit_compatible_but_slow_slice(rows, cols, m):
return jnp.isin(jnp.arange(N), jnp.where(cols == m, rows, -1))
Is there any fast and JIT compatible solution to acheive the same output?

You can do a bit better than the first answer by using the mode argument of set() to drop out-of-bound indices, eliminating the final slice:
out = jnp.zeros(N, bool).at[jnp.where(cols==3, rows, N)].set(True, mode='drop')

I figured out that the implementation below returns the same output much faster, and it’s JIT compatible.
def slice(rows, cols, m):
res = jnp.zeros(N + 1, dtype=bool)
res = res.at[jnp.where(cols == m, rows, -1)].set(True)
return res[:-1]

Related

Which is efficient (cache/memory) in writting a function handling a numpy array?

I have a question: suppose I have a numpy 1D array Arr, whose length is N. Now I want to manipulate it. The demo is:
import numpy as np
N = 100
Arr = np.zeros(N, dtype=np.float64)
def func_1(Arr):
Arr_out = Arr + 1.0
return Arr_out
def func_2(Arr, N):
# resembling allocating memory
Arr_out = np.zeros(N, dtype=np.float64)
Arr_out = Arr + 1.0
return Arr_out
As you can see, there are two functions to manipulating Arr and returning Arr_out. The difference is that in func_2, I first specify the shape of Arr_out. Is this necessary in terms of efficiency concerning cache/memory? Thanks.

Numpy mask array multiple times and fill nans in 3D array with values from another 3D array

I have the following code:
import numpy as np
def fill(arr1, arr2, arr3, arr4, thresh= 0.5):
out_arr = np.zeros(arr1.shape)
for i in range(0,len(arr1)):
arr1[i] = np.where(np.abs(arr1[i])<=thresh,np.nan,arr1[i])
mask = np.isnan(arr1[i])
arr1[i] = np.nan_to_num(arr1[i])
merged1 = (arr2[i]*mask)+arr1[i]
merged2 = np.where(np.abs(merged1)<=thresh,np.nan,merged1)
mask = np.isnan(merged2)
merged2 = np.nan_to_num(merged2)
merged3 = (arr3[i]*mask)+merged2
merged3 = np.where(np.abs(merged3)<=thresh,np.nan,merged3)
mask = np.isnan(merged3)
merged3 = np.nan_to_num(merged3)
merged4 = (arr4[i]*mask)+merged3
out_arr[i] = merged4
return(out_arr)
arr1 = np.random.rand(10, 10, 10)
arr2 = np.random.rand(10, 10, 10)
arr3 = np.random.rand(10, 10, 10)
arr4 = np.random.rand(10, 10, 10)
arr = fill(arr1, arr2, arr3, arr4, 0.5)
I wonder if there is a more efficient way of doing this maybe with masked arrays? Basically what I am doing is to replace values below the threshold in each layer of the 3D array with the next array, and this over 4 arrays. How would this look like for n arrays?
Thanks!
Your function can be simplified in several ways. In terms of efficiency, the most significant aspect is that you do not need to iterate over the first dimension, you can operate on the whole arrays directly. Besides that, you can refactor the replacement logic to something much simpler, and use a a loop to avoid repeating the same code over and over:
import numpy as np
# Function accepts as many arrays as wanted, with at least one
# (threshold needs to be passed as keyword parameter)
def fill(arr1, *arrs, thresh=0.5):
# Output array
out_arr = arr1.copy()
for arr in arrs:
# Replace values that are still below threshold
mask = np.abs(out_arr) <= thresh
out_arr[mask] = arr[mask]
return out_arr
Since thresh needs to be passed as keyword parameter in this function, you would call it as:
arr = fill(arr1, arr2, arr3, arr4, thresh=0.5)

How to iteratively nest a nested function

I have an array arr_multi_dim which is multi-dimensional. Every time when I increase a parameter n, there will be more entries created in the array results and the array will get larger.
With each increase in n, I need to perform the function np.concatenate() on the array arr_multi_dim, in such a way that there will be more np.concatenate() function nested every time n increases.
For eg.,
when n=2:
arr_multi_dim = np.concatenate(np.concatenate(arr_multi_dim, axis=1), axis=1)
when n=3:
arr_multi_dim = np.concatenate(np.concatenate(
np.concatenate(np.concatenate(arr_multi_dim, axis=1), axis=1), axis=1), axis=1)
when n=4:
arr_multi_dim = np.concatenate(np.concatenate(
np.concatenate(np.concatenate(
np.concatenate(np.concatenate(arr_multi_dim, axis=1), axis=1), axis=1), axis=1), axis=1), axis=1)
etc.
where at each increment of n, a pair of np.concatenate() (ie. two) gets added into the function.
How do I write a function, loops (or something similar), so that when I specify any values for n, the appropriate np.concatenate() function will be used?
Many thanks in advance.
Edit:
This is the full code that I have written which uses the above np.concatenate() function.
from itertools import product
from joblib import Parallel, delayed
from functools import reduce
from operator import mul
import numpy as np
lst = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
arr = np.array(lst)
n = 2
def test1(arr, n):
flat = np.ravel(arr).tolist()
gen = (list(a) for a in product(flat, repeat=n))
results = Parallel(n_jobs=-1)(delayed(reduce)(mul, x) for (x) in gen)
nrows = arr.shape[0]
ncols = arr.shape[1]
arr_multi_dim = np.array(results).reshape((nrows, ncols)*n)
arr_final = np.concatenate(np.concatenate(arr_multi_dim, axis=1), axis=1) # need to generalise this
return arr_final
The above code only works for n=2. I am trying to generalize the np.concatenate part of the code so that it would work for any n as mentioned above.
If i understood you correctly its pretty simple:
arr_multi_dim = results
for i in range(n):
if i < 2:
arr_multi_dim = np.concatenate(arr_multi_dim , axis=1)
else:
arr_multi_dim = np.concatenate(np.concatenate(arr_multi_dim , axis=1), axis=1)
becase the first two iteration only add a single layer while the rest add two layers

Initialize a numpy sparse matrix efficiently

I have an array with m rows and arrays as values, which indicate the index of columns and are bounded to a large number n.
E.g:
Y = [[1,34,203,2032],...,[2984]]
Now I want an efficient way to initialize a sparse numpy matrix X with dimensions m,n and values corresponding to Y (X[i,j] = 1, if j is in Y[i], = 0 otherwise).
Your data are already close to csr format, so I suggest using that:
import numpy as np
from scipy import sparse
from itertools import chain
# create an example
m, n = 20, 10
X = np.random.random((m, n)) < 0.1
Y = [list(np.where(y)[0]) for y in X]
# construct the sparse matrix
indptr = np.fromiter(chain((0,), map(len, Y)), int, len(Y) + 1).cumsum()
indices = np.fromiter(chain.from_iterable(Y), int, indptr[-1])
data = np.ones_like(indices)
S = sparse.csr_matrix((data, indices, indptr), (m, n))
# or
S = sparse.csr_matrix((data, indices, indptr))
# check
assert np.all(S==X)

Sort numpy array in-place using order from a second array

Let two ndarrays: A of shape (n, *m), and B of shape (n, ). Is there a way to sort A in-place using the order that would sort B?
Sorting A with B is easy using np.argsort, but this is not done in-place:
A = A[np.argsort(B)]
Comments:
A and B have different dtypes, and A can have more than two dimensions. Hence they can’t be stacked to use ndarray.sort().
A takes up a lot of space, which is why it needs to be sorted in-place. Any solution requiring twice the space occupied by A would therefore defeat this purpose.
The title of this question “Re-arranging numpy array in place” may sound related, but the question itself is not very clear, and the answers do not match my question.
Here is a solution that works by following cycles in the index array. It can optionally be compiled using pythran giving a significant speedup if rows are small (80x for 10 elements) and a small speedup if rows are large (30% for 1000 elements).
To keep it pythran compatible I had to simplify it a bit, so it only accepts 2D arrays and it only sorts along axis 0.
Code:
import numpy as np
#pythran export take_inplace(float[:, :] or int[:, :], int[:])
def take_inplace(a, idx):
n, m = a.shape
been_there = np.zeros(n, bool)
keep = np.empty(m, a.dtype)
for i in range(n):
if been_there[i]:
continue
keep[:] = a[i]
been_there[i] = True
j = i
k = idx[i]
while not been_there[k]:
a[j] = a[k]
been_there[k] = True
j = k
k = idx[k]
a[j] = keep
Sample run using compiled version. As indicated above compilation is only required for small rows, for larger rows pure python should be fast enough.
>>> from timeit import timeit
>>> import numpy as np
>>> import take_inplace
>>>
>>> a = np.random.random((1000, 10))
>>> idx = a[:, 4].argsort()
>>>
>>> take_inplace.take_inplace(a, idx)
>>>
# correct
>>> np.all(np.arange(1000) == a[:, 4].argsort())
True
>>>
# speed
>>> timeit(lambda: take_inplace.take_inplace(a, idx), number=1000)
0.011950935004279017
>>>
# for comparison
>>> timeit(lambda: a[idx], number=1000)
0.02985276997787878
If you can set A beforehand as a structured array whose datatype is composed of a subarray of shape (m, ) and a scalar of the same type (e.g., np.int32), then you can sort it in-place with respect to B. For example:
import numpy as np
B = np.array([3, 1, 2])
A = np.array([[10, 11], [20, 21], [30, 31]])
(n, m) = A.shape
dt = np.dtype([('a', np.int32, (m, )), ('b', int)])
A2 = np.array([(a, b) for a, b in zip(A, B)], dtype=dt)
A2.sort(order='b')
print A2

Categories

Resources