How to vmap over specific funciton in jax? - python

I have this function which works for single vector:
def vec_to_board(vector, player, dim, reverse=False):
player_board = np.zeros(dim * dim)
player_pos = np.argwhere(vector == player)
if not reverse:
player_board[mapping[player_pos.T]] = 1
else:
player_board[reverse_mapping[player_pos.T]] = 1
return np.reshape(player_board, [dim, dim])
However, I want it to work for a batch of vectors.
What I have tried so far:
states = jnp.array([[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2], [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2]])
batch_size = 1
b_states = vmap(vec_to_board)((states, 1, 4), batch_size)
This doesn't work. However, if I understand correctly vmap should be able to handle this transformation for batches?

There are a couple issues you'll run into when trying to vmap this function:
This function is defined in terms of numpy arrays, not jax arrays. How do I know? JAX arrays are immutable, so things like arr[idx] = 1 will raise errors. You need to replace these with equivalent JAX operations (see JAX Sharp Bits: in-place updates) and ensure your function works with JAX array operations rather than numpy array operations.
Your function makes used of dynamically-shaped arrays; e.g. player_pos, has a shape dependent on the number of nonzero entries in vector == player. You'll have to rewrite your function in terms of statically-shaped arrays. There is some discussion of this in the jnp.argwhere docstring; for example, if you know a priori how many True entries you expect in the array, you can specify the size to make this work.
Good luck!

Related

Python convert nested array to cython c array

I want to convert array of array which have default value of zeros in each array. The reason to do this is that I want to convert a high computing algorithm from python to cython to speed up computation more. The code for array is like this for python:
self.v = [[0 for i in range(self.D)] for j in range(self.NP)] #velocity
self.Sol = [[0 for i in range(self.D)] for j in range(self.NP)]
self.D and self.NP could be any integer values. The format of the sample data in python is like this, for self.D=4 and self.NP=3:
[[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]

Optimizing execution time for mapping array to value with dictionary and numpy

I am trying to implement a simple mapping to a set of values from an array created with numpy of 2-D.
For each row in the array I need to choose the correct value corresponding with the set of values and add it to a array.
For example:
[0, 1, 0, 0] -> 3
...
[1, 0, 1, 0] -> 2
But, my first implementation made me wonder if I'm doing something really wrong or not efficient at all because of the size of my dataset, so I did this workaround without using for loops and optimize speed execution using dictionary lookup.
import numpy as np
# function to perform the search and return the index accordingly (it is supposed to be fast because of data structure)
def get_val(n):
map_list = {0: [0, 1, 0], 1: [0, 1, 0], 2: [1, 0, 0], 3: [0, 0, 1]}
map_vals = list(map_list.values())
index = map_vals.index(list(n))
return(index)
# set of arbitrary arrays
li = np.array([[0, 1, 0], [0, 0, 1]])
# here is the performance improvement attempt with the help of the function above
arr = [get_val(n) for n in li]
print(arr)
I'm not completely sure if this is the correct way to do it for getting the needed value for a set like this. If there is a better way, please let me know.
Otherwise, I refer to my main question:
what is the best way possible to optimize the code?
Thanks so much for your help.
You can try use matrix multiplication (dot product):
a=np.array([[0, 0, 0],[0, 1, 0], [1, 0, 0], [0, 0, 1]]) # dict values
c=np.array([0,1,2,3]) # dict keys
li = np.array([[0, 1, 0], [0, 0, 1]])
b=np.linalg.pinv(a)#c # decoding table
result=li#b
print(result)

Crop empty arrays (padding) from a volume

What I want to do is crop a volume to remove all irrelevant data. For example, say I have a 100x100x100 volume filled with zeros, except for a 50x50x50 volume within that is filled with ones.
How do I obtain the cropped 50x50x50 volume from the original ?
Here's the naive method I came up with.
import numpy as np
import tensorflow as tf
test=np.zeros((100,100,100)) # create an empty 100x100x100 volume
rand=np.random.rand(66,25,34) # create a 66x25x34 filled volume
test[10:76, 20:45, 30:64] = rand # partially fill the empty volume
# initialize the cropping coordinates
minx=miny=minz=0
maxx=maxy=maxz=0
maxx,maxy,maxz=np.subtract(test.shape,1)
# compute the optimal cropping coordinates
dimensions=test.shape
while(tf.reduce_max(test[minx,:,:]) == 0): # check for empty slices along the x axis
minx+=1
while(tf.reduce_max(test[:,miny,:]) == 0): # check for empty slices along the y axis
miny+=1
while(tf.reduce_max(test[:,:,minz]) == 0): # check for empty slices along the z axis
minz+=1
while(tf.reduce_max(test[maxx,:,:]) == 0):
maxx-=1
while(tf.reduce_max(test[:,maxy,:]) == 0):
maxy-=1
while(tf.reduce_max(test[:,:,maxz]) == 0):
maxz-=1
maxx,maxy,maxz=np.add((maxx,maxy,maxz),1)
crop = test[minx:maxx,miny:maxy,minz:maxz]
print(minx,miny,minz,maxx,maxy,maxz)
print(rand.shape)
print(crop.shape)
This prints:
10 20 30 76 45 64
(66, 25, 34)
(66, 25, 34)
, which is correct. However, it takes too long and is probably suboptimal. I'm looking for better ways to achieve the same thing.
NB:
The subvolume wouldn't necessarily be a cuboid, it could be any shape.
I want to keep gaps within the subvolume, only remove what's "outside" the shape to be cropped.
(Edit)
Oops, I hadn't seen the comment about keeping the so-called "gaps" between elements! This should be the one, finally.
def get_nonzero_sub(arr):
arr_slices = tuple(np.s_[curr_arr.min():curr_arr.max() + 1] for curr_arr in arr.nonzero())
return arr[arr_slices]
While you wait for a sensible response (I would guess this is a builtin function in an image processing library somewhere), here's a way
y, x = np.where(np.any(test, 0))
z, _ = np.where(np.any(test, 1))
test[min(z):max(z)+1, min(y):max(y)+1, min(x):max(x)+1]
I think leaving tf out of this should up your performance.
Explanation (based on 2D array)
test = np.array([
[0, 0, 0, 0, 0, ],
[0, 0, 1, 2, 0, ],
[0, 0, 3, 0, 0, ],
[0, 0, 0, 0, 0, ],
[0, 0, 0, 0, 0, ],
])
We want to crop it to get
[[1, 2]
[3, 0]]
np.any(..., 0) this will 'iterate' over axis 0 and return True if any of the elements in the slice are truthy. I show the result of this in the comments here:
np.array([
[0, 0, 0, 0, 0, ], # False
[0, 0, 1, 2, 0, ], # True
[0, 0, 3, 0, 0, ], # True
[0, 0, 0, 0, 0, ], # False
[0, 0, 0, 0, 0, ], # False
])
i.e. it returns np.array([False, True, True, False, False])
np.any(..., 1) does the same as step 2 but over axis 1 instead of axis zero i.e.
np.array([
[0, 0, 0, 0, 0, ],
[0, 0, 1, 2, 0, ],
[0, 0, 3, 0, 0, ],
[0, 0, 0, 0, 0, ],
[0, 0, 0, 0, 0, ],
# False False True True False
])
Note that in the case of a 3D array, these steps return 2D arrays
(x,) = np.where(...) this returns the index values of the truthy values in an array. So np.where([False, True, True, False, False]) returns (array([1, 2]),). Note that this is a tuple so in the 2D case we would need to call (x,) = ... so x is just the array array([1, 2]). The syntax is nicer in the 2D case as we can use tuple-unpacking i.e x, y = ...
Note that in the 3D case, np.where can give us the value for 2 axes at a time. I chose to do x-y in one go and then z-? in the second go. The ? is either x or y, I can't be bothered to work out which and since we don't need it I throw it away in a variable named _ which by convention is a reasonable place to store junk output you don't actually want. Note I need to do z, _ = as I want the tuple-unpacking and not just z = otherwise z become the tuple with both arrays.
Well, this step is pretty much the same as what you did at the end of your answer so I assume you understand it. Simple slicing in each dimension from the first element with a value in that dimension to the last. You need the + 1 because slicing in python are not inclusive of the index after the :.
Hopefully that's clear?

How to plot eigenvalues representing symbolic functions in Python?

I need to calculate the eigenvalues of an 8x8-matrix and plot each of the eigenvalues for a symbolic variable occuring in the matrix. For the matrix I'm using I get 8 different eigenvalues where each is representing a function in "W", which is my symbolic variable.
Using python I tried calculating the eigenvalues with Scipy and Sympy which worked kind of, but the results are stored in a weird way (at least for me as a newbie not understanding much of programming so far) and I didn't find a way to extract just one eigenvalue in order to plot it.
import numpy as np
import sympy as sp
W = sp.Symbol('W')
w0=1/780
wl=1/1064
# This is my 8x8-matrix
A= sp.Matrix([[w0+3*wl, 2*W, 0, 0, 0, np.sqrt(3)*W, 0, 0],
[2*W, 4*wl, 0, 0, 0, 0, 0, 0],
[0, 0, 2*wl+w0, np.sqrt(3)*W, 0, 0, 0, np.sqrt(2)*W],
[0, 0, np.sqrt(3)*W, 3*wl, 0, 0, 0, 0],
[0, 0, 0, 0, wl+w0, np.sqrt(2)*W, 0, 0],
[np.sqrt(3)*W, 0, 0, 0, np.sqrt(2)*W, 2*wl, 0, 0],
[0, 0, 0, 0, 0, 0, w0, W],
[0, 0, np.sqrt(2)*W, 0, 0, 0, W, wl]])
# Calculating eigenvalues
eva = A.eigenvals()
evaRR = np.array(list(eva.keys()))
eva1p = evaRR[0] # <- this is my try to refer to the first eigenvalue
In the end I hope to get a plot over "W" where the interesting range is [-0.002 0.002]. For the ones interested it's about atomic physics and W refers to the rabi frequency and I'm looking at so called dressed states.
You're not doing anything incorrectly -- I think you're just caught up since your eigenvalues look so jambled and complicated.
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
W = sp.Symbol('W')
w0=1/780
wl=1/1064
# This is my 8x8-matrix
A= sp.Matrix([[w0+3*wl, 2*W, 0, 0, 0, np.sqrt(3)*W, 0, 0],
[2*W, 4*wl, 0, 0, 0, 0, 0, 0],
[0, 0, 2*wl+w0, np.sqrt(3)*W, 0, 0, 0, np.sqrt(2)*W],
[0, 0, np.sqrt(3)*W, 3*wl, 0, 0, 0, 0],
[0, 0, 0, 0, wl+w0, np.sqrt(2)*W, 0, 0],
[np.sqrt(3)*W, 0, 0, 0, np.sqrt(2)*W, 2*wl, 0, 0],
[0, 0, 0, 0, 0, 0, w0, W],
[0, 0, np.sqrt(2)*W, 0, 0, 0, W, wl]])
# Calculating eigenvalues
eva = A.eigenvals()
evaRR = np.array(list(eva.keys()))
# The above is copied from your question
# We have to answer what exactly the eigenvalue is in this case
print(type(evaRR[0])) # >>> Piecewise
# Okay, so it's a piecewise function (link to documentation below).
# In the documentation we see that we can use the .subs method to evaluate
# the piecewise function by substituting a symbol for a value. For instance,
print(evaRR[0].subs(W, 0)) # Will substitute 0 for W
# This prints out something really nasty with tons of fractions..
# We can evaluate this mess with sympy's numerical evaluation method, N
print(sp.N(evaRR[0].subs(W, 0)))
# >>> 0.00222190090611143 - 6.49672880062804e-34*I
# That's looking more like it! Notice the e-34 exponent on the imaginary part...
# I think it's safe to assume we can just trim that off.
# This is done by setting the chop keyword to True when using N:
print(sp.N(evaRR[0].subs(W, 0), chop=True)) # >>> 0.00222190090611143
# Now let's try to plot each of the eigenvalues over your specified range
fig, ax = plt.subplots(3, 3) # 3x3 grid of plots (for our 8 e.vals)
ax = ax.flatten() # This is so we can index the axes easier
plot_range = np.linspace(-0.002, 0.002, 10) # Range from -0.002 to 0.002 with 10 steps
for n in range(8):
current_eigenval = evaRR[n]
# There may be a way to vectorize this computation, but I'm not familiar enough with sympy.
evaluated_array = np.zeros(np.size(plot_range))
# This will be our Y-axis (or W-value). It is set to be the same shape as
# plot_range and is initally filled with all zeros.
for i in range(np.size(plot_range)):
evaluated_array[i] = sp.N(current_eigenval.subs(W, plot_range[i]),
chop=True)
# The above line is evaluating your eigenvalue at a specific point,
# approximating it numerically, and then chopping off the imaginary.
ax[n].plot(plot_range, evaluated_array, "c-")
ax[n].set_title("Eigenvalue #{}".format(n))
ax[n].grid()
plt.tight_layout()
plt.show()
And as promised, the Piecewise documentation.

How do I perform dimensionality reduction on two independent XOR gates?

Take the probability distribution of a XOR gate in which every configuration is equally probable (configurations are given by outcomes_sub; the probability mass function by pmf_xor_sub):
import numpy as np
import itertools as it
outcomes_sub = [list(item) for item in list(it.product([0,1], repeat=3))]
pmf_xor_sub = np.array([1/4, 0, 0, 1/4, 0, 1/4, 1/4, 0])
Now take the probability distribution corresponding to two uncorrelated such XORs:
outcomes = [outcome1 + outcome2 for (outcome1, outcome2)
in it.product(outcomes_sub, outcomes_sub)]
pmf_xor = [pmf1 * pmf2 for (pmf1, pmf2)
in it.product(pmf_xor_sub, pmf_xor_sub)]
And create some data based on it:
indices = np.random.choice(len(outcomes), 10000, p=pmf_xor)
data_xor = np.array([outcomes[index] for index in indices])
data_xor looks like this:
array([[1, 1, 0, 0, 0, 0],
[1, 0, 1, 0, 0, 0],
[0, 1, 1, 1, 1, 0],
...,
[0, 1, 1, 1, 1, 0],
[1, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0]])
I.e., two independent XORs back to back. What's the right way to perform dimensionality reduction on it? PCA won't work (because the dependence is non-linear, right?):
from sklearn import decomposition
pca_xor = decomposition.PCA()
pca_xor.fit(data_xor)
Now, pca_xor.explained_variance_ratio_ gives:
array([ 0.17145045, 0.17018817, 0.16758773, 0.16575979, 0.16410862,
0.16090524], dtype=float32)
No two components stand out. I understand that a non-linear method such as kernel PCA should work here, but I am struggling to find pointers to ways of applying it to my problem.
To give a bit more context: what I am actually after is ways to bring out the structure in data_xor: two big XOR blobs, each of which is composed of some finer-grained stuff. If I am going about it all wrong, feel free to point that out too.

Categories