What is JaxNumpy-compatible equivalent to this Python function? - python

How do I implement the below in a JAX-compatable way (e.g., using jax.numpy)?
def actions(state: tuple[int, ...]) -> list[tuple[int, ...]]:
l = []
iterables = [range(1, i+1) for i in state]
ns = list(range(len(iterables)))
for i, iterable in enumerate(iterables):
for value in iterable:
action = tuple(value if n == i else 0 for n in ns)
l.append(action)
return l
>>> state = (3, 1, 2)
>>> actions(state)
[(1, 0, 0), (2, 0, 0), (3, 0, 0), (0, 1, 0), (0, 0, 1), (0, 0, 2)]

Jax, like numpy, cannot efficiently operate on Python container types like lists and tuples, so there's not really any JAX-compatible way to create a function with the exact signature you specify above.
But if you're alright with the return value being a two-dimensional array, you could do something like this, based on jnp.vstack:
from typing import Tuple
import jax.numpy as jnp
from jax import jit, partial
#partial(jit, static_argnums=0)
def actions(state: Tuple[int, ...]) -> jnp.ndarray:
return jnp.vstack([
jnp.zeros((val, len(state)), int).at[:, i].set(jnp.arange(1, val + 1))
for i, val in enumerate(state)])
>>> state = (3, 1, 2)
>>> actions(state)
DeviceArray([[1, 0, 0],
[2, 0, 0],
[3, 0, 0],
[0, 1, 0],
[0, 0, 1],
[0, 0, 2]], dtype=int32)
Note that because the size of the output array depends on the content of state, state must be a static quantity, so a tuple is a good option for the input.

Related

In a numpy array (a list of tuples), the processing speed is slow by extending () many times. I want to make that part faster

There is a numpy array that can be formed by combining an array of tuples in a for loop like "res" in this code. (Variable names and contents are simplified from the actual code.)
If you take a closer look at this, a for loop is executed for the length of arr_2, and the array extends () is executed.It turns out that the processing speed becomes extremely heavy when arr_2 becomes long.
Wouldn't it be possible to process at high speed by making array creation well?
# -*- coding: utf-8 -*-
import numpy as np
arr_1 = np.array([[0, 0, 1], [0, 0.5, -1], [-1, 0, -1], [0, -0.5, -1], [1, 0, -1]])
arr_2 = np.array([[0, 1, 2], [0, 1, 2]])
all_arr = []
for p in arr_2:
all_arr = [
(arr_1[0], p), (arr_1[1], p), (arr_1[2], p),
(arr_1[0], p), (arr_1[1], p), (arr_1[4], p),
(arr_1[0], p), (arr_1[2], p), (arr_1[3], p),
(arr_1[0], p), (arr_1[3], p), (arr_1[4], p),
(arr_1[1], p), (arr_1[2], p), (arr_1[4], p),
(arr_1[2], p), (arr_1[3], p), (arr_1[4], p)]
all_arr.extend(all_arr)
vtype = [('type_a', np.float32, 3), ('type_b', np.float32, 3)]
res = np.array(all_arr, dtype=vtype)
print(res)
I couldn't figure out why you used this indexing for arr_1 so I just copied it
import numpy as np
arr_1 = np.array([[0, 0, 1], [0, 0.5, -1], [-1, 0, -1], [0, -0.5, -1], [1, 0, -1]])
arr_2 = np.array([[0, 1, 2], [0, 1, 2]])
weird_idx = np.array([0,1,2,0,1,4,0,2,3,0,3,4,1,2,4,2,3,4])
weird_arr1 = arr_1[weird_idx]
all_arr = [(wiered_arr1[i],arr_2[j]) for j in range(len(arr_2)) for i in range(len(wiered_arr1)) ]
vtype = [('type_a', np.float32, 3), ('type_b', np.float32, 3)]
res = np.array(all_arr, dtype=vtype)
you can also repeat the arrays
arr1_rep = np.tile(weird_arr1.T,2).T
arr2_rep = np.repeat(arr_2,weird_arr1.shape[0],0)
res = np.empty(arr1_rep.shape[0],dtype=vtype)
res['type_a']=arr1_rep
res['type_b']=arr2_rep
Often with structured arrays it is faster to assign by field instead of the list of tuples approach:
In [388]: idx = [0,1,2,0,1,4,0,2,3,0,3,4,1,2,4,2,3,4]
In [400]: res1 = np.zeros(36, dtype=vtype)
In [401]: res1['type_a'][:18] = arr_1[idx]
In [402]: res1['type_a'][18:] = arr_1[idx]
In [403]: res1['type_b'][:18] = arr_2[0]
In [404]: res1['type_b'][18:] = arr_2[1]
In [405]: np.allclose(res['type_a'], res1['type_a'])
Out[405]: True
In [406]: np.allclose(res['type_b'], res1['type_b'])
Out[406]: True

Building an upper triangular matrix recursively

Ive been breaking my head over trying to come up with a recursive way to build the following matrix in python. It is quite a challenge without pointers. Could anyone maybe help me out?
The recursion is the following:
T0 = 1,
Tn+1 = [[Tn, Tn],
[ 0, Tn]]
I have tried many iterations of some recursive function, but I cannot wrap my head around it.
def T(n, arr):
n=int(n)
if n == 0:
return 1
else:
c = 2**(n-1)
Tn = np.zeros((c,c))
Tn[np.triu_indices(n=c)] = self.T(n=n-1, arr=arr)
return Tn
arr = np.zeros((8,8))
T(arr=arr, n=3)
It's not hard to do this, but you need to be careful about the meaning of the zero in the recursion. This isn't really precise for larger values of n:
Tn+1 = [[Tn, Tn],
[ 0, Tn]]
Because that zero can represent a block of zeros for example on the second iteration you have this:
[1, 1, 1, 1],
[0, 1, 0, 1],
[0, 0, 1, 1],
[0, 0, 0, 1]
Those four zeros in the bottom-left are all represented by the one zero in the formula. The block of zeros needs to be the same shape as the blocks around it.
After that it's a matter of making Numpy put thing in the right order and shape for you. numpy.block is really handy for this and makes it pretty simple:
import numpy as np
def makegasket(n):
if n == 0:
return np.array([1], dtype=int)
else:
node = makegasket(n-1)
return np.block([[node, node], [np.zeros(node.shape, dtype=int), node]])
makegasket(3)
Result:
array([[1, 1, 1, 1, 1, 1, 1, 1],
[0, 1, 0, 1, 0, 1, 0, 1],
[0, 0, 1, 1, 0, 0, 1, 1],
[0, 0, 0, 1, 0, 0, 0, 1],
[0, 0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 0, 0, 1, 0, 1],
[0, 0, 0, 0, 0, 0, 1, 1],
[0, 0, 0, 0, 0, 0, 0, 1]])
If you use larger n you might enjoy matplotlib.pyplot.imshow for display:
from matplotlib.pyplot import imshow
# ....
imshow(makegasket(7))
You don't really need a recursive function to implement this recursion. The idea is to start with the UR corner and build outward. You can even start with the UL corner to avoid some of the book-keeping and flip the matrix along either axis, but this won't be as efficient in the long run.
def build_matrix(n):
size = 2**n
# Depending on the application, even dtype=np.bool might work
matrix = np.zeros((size, size), dtype=np.int)
# This is t[0]
matrix[0, -1] = 1
for i in range(n):
k = 2**i
matrix[:k, -2 * k:-k] = matrix[k:2 * k, -k:] = matrix[:k, -k:]
return matrix
Just for fun, here is a plot of timing results for this implementation vs #Mark Meyer's answer. It shows the slight timing advantage (also memory) of using a looping approach in this case:
Both algorithms run out of memory around n=15 on my machine, which is not too surprising.

Allocating array (list) algorithm permutations in python

So i'm having trouble to solve the following problem:
given an array size, lets say for the ease of the question size =20
it is filled with zeros as follows
arr = [0]*20 ==> [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
we have couple of constant sample sizes, such as 4,3,2
SampleA=4, SampleB=3, SampleC= 2
i need to have the permutations/variations of how to allocate the list.
i can put each sample in different place/index
for example, sampleA= 4 i can put it in indexes of 0:3, or 1:4... 15:19..
(as you can see, there are quite a lot of possibilities)
the thing gets complicated once it get more crowded, for example:
3+2+3 +4
[0, x, x, x, 0, 0, x x, 0, x, x, x, 0, 0, 0, 0, x, x,x, x]
what i basically need, is to find all the possibilities to allocate samples,
i get a dictionary:
key = sample size of indexes, and the
value=many times it repeats.
for the upper example: {3:2,2:1,4:1}
and i would like the function to return a list of indexes !=0
for this example:
[0, x, x, x, 0, 0, x x, 0, x, x, x, 0, 0, 0, 0, x, x,x, x]
the function will return:
list_ind = [0,5,6,9,13,14,15,16]
so i asked a colleague to help, and we came into a solution:
i put an example of:
4:2, 3:1, 2:1
or verbally:
two times 4, one time 3, one time 2
the code below:
*if someone can optimize, will be great
size_of_wagon = 20
dct = {4:2,3:1,2:1}
l = [item for sublist in [[k] * v for (k, v) in dct.items()] for item in sublist]
def gen_func(lstOfSamples, length, shift):
try:
# print(f'lstOfSamples={lstOfSamples}')
sample = lstOfSamples[0] # Take first sample
for i in range(length - sample):
for x in gen_func(lstOfSamples[1:], length - (sample + i), shift + sample + i):
yield [(shift + i, sample)] + x
except:
yield []
g = list(gen_func(l, size_of_wagon, 0))
for i in g:
print(i)
print(len(g))

How to fold/accumulate a numpy matrix product (dot)?

With using python library numpy, it is possible to use the function cumprod to evaluate cumulative products, e.g.
a = np.array([1,2,3,4,2])
np.cumprod(a)
gives
array([ 1, 2, 6, 24, 48])
It is indeed possible to apply this function only along one axis.
I would like to do the same with matrices (represented as numpy arrays), e.g. if I have
S0 = np.array([[1, 0], [0, 1]])
Sx = np.array([[0, 1], [1, 0]])
Sy = np.array([[0, -1j], [1j, 0]])
Sz = np.array([[1, 0], [0, -1]])
and
b = np.array([S0, Sx, Sy, Sz])
then I would like to have a cumprod-like function which gives
np.array([S0, S0.dot(Sx), S0.dot(Sx).dot(Sy), S0.dot(Sx).dot(Sy).dot(Sz)])
(This is a simple example, in reality I have potentially large matrices evaluated over n-dimensional meshgrids, so I seek for the most simple and efficient way to evaluate this thing.)
In e.g. Mathematica I would use
FoldList[Dot, IdentityMatrix[2], {S0, Sx, Sy, Sz}]
so I searched for a fold function, and all I found is an accumulate method on numpy.ufuncs. To be honest, I know that I am probably doomed because an attempt at
np.core.umath_tests.matrix_multiply.accumulate(np.array([pauli_0, pauli_x, pauli_y, pauli_z]))
as mentioned in a numpy mailing list yields the error
Reduction not defined on ufunc with signature
Do you have an idea how to (efficiently) do this kind of calculation ?
Thanks in advance.
As food for thought, here are 3 ways of evaluating the 3 sequential dot products:
With the normal Python reduce (which could also be written as a loop)
In [118]: reduce(np.dot,[S0,Sx,Sy,Sz])
array([[ 0.+1.j, 0.+0.j],
[ 0.+0.j, 0.+1.j]])
The einsum equivalent
In [119]: np.einsum('ij,jk,kl,lm',S0,Sx,Sy,Sz)
The einsum index expression looks like a sequence of operations, but it is actually evaluated as a 5d product with summation on 3 axes. In the C code this is done with an nditer and strides, but the effect is as follows:
In [120]: np.sum(S0[:,:,None,None,None] * Sx[None,:,:,None,None] *
Sy[None,None,:,:,None] * Sz[None,None,None,:,:],(1,2,3))
In [127]: np.prod([S0[:,:,None,None,None], Sx[None,:,:,None,None],
Sy[None,None,:,:,None], Sz[None,None,None,:,:]]).sum((1,2,3))
A while back while creating a patch from np.einsum I translated that C code to Python, and also wrote a Cython sum-of-products function(s). This code is on github at
https://github.com/hpaulj/numpy-einsum
einsum_py.py is the Python einsum, with some useful debugging output
sop.pyx is the Cython code, which is compiled to sop.so.
Here's how it could be used for part of your problem. I'm skipping the Sy array since my sop is not coded for complex numbers (but that could be changed).
import numpy as np
import sop
import einsum_py
S0 = np.array([[1., 0], [0, 1]])
Sx = np.array([[0., 1], [1, 0]])
Sz = np.array([[1., 0], [0, -1]])
print np.einsum('ij,jk,kl', S0, Sx, Sz)
# [[ 0. -1.] [ 1. 0.]]
# same thing, but with parsing information
einsum_py.myeinsum('ij,jk,kl', S0, Sx, Sz, debug=True)
"""
{'max_label': 108, 'min_label': 105, 'nop': 3,
'shapes': [(2, 2), (2, 2), (2, 2)],
'strides': [(16, 8), (16, 8), (16, 8)],
'ndim_broadcast': 0, 'ndims': [2, 2, 2], 'num_labels': 4,
....
op_axes [[0, -1, 1, -1], [-1, -1, 0, 1], [-1, 1, -1, 0], [0, 1, -1, -1]]
"""
# take op_axes (for np.nditer) from this debug output
op_axes = [[0, -1, 1, -1], [-1, -1, 0, 1], [-1, 1, -1, 0], [0, 1, -1, -1]]
w = sop.sum_product_cy3([S0,Sx,Sz], op_axes)
print w
As written sum_product_cy3 cannot take an arbitrary number of ops. Plus the iteration space increases with each op and index. But I can imagine calling it repeatedly, either at the Cython level, or from Python. I think it has potential for being faster than repeat(dot...) for lots of small arrays.
A condensed version of the Cython code is:
def sum_product_cy3(ops, op_axes, order='K'):
#(arr, axis=None, out=None):
cdef np.ndarray[double] x, y, z, w
cdef int size, nop
nop = len(ops)
ops.append(None)
flags = ['reduce_ok','buffered', 'external_loop'...]
op_flags = [['readonly']]*nop + [['allocate','readwrite']]
it = np.nditer(ops, flags, op_flags, op_axes=op_axes, order=order)
it.operands[nop][...] = 0
it.reset()
for x, y, z, w in it:
for i in range(x.shape[0]):
w[i] = w[i] + x[i] * y[i] * z[i]
return it.operands[nop]

A 3-D grid of regularly spaced points

I want to create a list containing the 3-D coords of a grid of regularly spaced points, each as a 3-element tuple. I'm looking for advice on the most efficient way to do this.
In C++ for instance, I simply loop over three nested loops, one for each coordinate. In Matlab, I would probably use the meshgrid function (which would do it in one command). I've read about meshgrid and mgrid in Python, and I've also read that using numpy's broadcasting rules is more efficient. It seems to me that using the zip function in combination with the numpy broadcast rules might be the most efficient way, but zip doesn't seem to be overloaded in numpy.
Use ndindex:
import numpy as np
ind=np.ndindex(3,3,2)
for i in ind:
print(i)
# (0, 0, 0)
# (0, 0, 1)
# (0, 1, 0)
# (0, 1, 1)
# (0, 2, 0)
# (0, 2, 1)
# (1, 0, 0)
# (1, 0, 1)
# (1, 1, 0)
# (1, 1, 1)
# (1, 2, 0)
# (1, 2, 1)
# (2, 0, 0)
# (2, 0, 1)
# (2, 1, 0)
# (2, 1, 1)
# (2, 2, 0)
# (2, 2, 1)
Instead of meshgrid and mgrid, you can use ogrid, which is a "sparse" version of mgrid. That is, only the dimension along which the values change are filled in. The others are simply broadcast. This uses much less memory for large grids than the non-sparse alternatives.
For example:
>>> import numpy as np
>>> x, y = np.ogrid[-1:2, -2:3]
>>> x
array([[-1],
[ 0],
[ 1]])
>>> y
array([[-2, -1, 0, 1, 2]])
>>> x**2 + y**2
array([[5, 2, 1, 2, 5],
[4, 1, 0, 1, 4],
[5, 2, 1, 2, 5]])
I would say go with meshgrid or mgrid, in particular if you need non-integer coordinates. I'm surprised that Numpy's broadcasting rules would be more efficient, as meshgrid was designed especially for the problem that you want to solve.
for multi-d (greater than 2) meshgrids, use numpy.lib.index_tricks.nd_grid like so:
import numpy
grid = numpy.lib.index_tricks.nd_grid()
g1 = grid[:3,:3,:3]
g2 = grid[0:1:0.5, 0:1, 0:2]
g3 = grid[0:1:3j, 0:1:2j, 0:2:2j]
where g1 has x values of [0,1,2]
and g2 has x values of [0,.5],
and g3 has x values of [0.0,0.5,1.0] (the 3j defining the step count instead of the step increment. see the documentation for more details.
Here's an efficient option similar to your C++ solution, which I've used for exactly the same purpose:
import numpy, itertools, collections
def grid(xmin, xmax, xstep, ymin, ymax, ystep, zmin, zmax, zstep):
"return nested tuples of grid-sampled coordinates that include maxima"
return collections.deque( itertools.product(
numpy.arange(xmin, xmax+xstep, xstep).tolist(),
numpy.arange(ymin, ymax+ystep, ystep).tolist(),
numpy.arange(zmin, zmax+zstep, zstep).tolist() ) )
Performance is best (in my tests) when using a.tolist(), as shown above, but you can use a.flat instead and drop the deque() to get an iterator that will sip memory. Of course, you can also use a plain old tuple() or list() instead of deque() for a slight performance penalty (again, in my tests).

Categories

Resources