I am looking to translate a bit of code from a NumPy version listed here, to a JAX compatible version. The NumPy code iteratively calculates the value of a matrix, E from the values of other matrices, A, B, D, as well as the value of E from the previous iteration: E_jm1.
Both the NumPy and JAX version work in their listed forms and produce identical results. How can I get the JAX version to work when passing A, B, D as a tuple instead of as a concatenated array? I have a specific use case where a tuple would be more useful.
I found a question asking something similar, but it just confirmed that this should be possible. There are no examples in the documentation or elsewhere that I could find.
Original NumPy version
import numpy as np
import jax
import jax.numpy as jnp
def BAND_J(A, B, D, E_jm1):
'''
output: E(N x N)
input: A(N x N), B(N x N), D(N x N), E_jm1(N x N)
πβ±Ό = -[π + ππβ±Όββ]β»ΒΉ π
'''
B_inv = np.linalg.inv(B + np.dot( A, E_jm1 ))
E = -np.dot(B_inv, D)
return E
key = jax.random.PRNGKey(0)
N = 2
NJ = 4
# initialize matrices with random values
A, B, D = [ jax.random.normal(key, shape=(N,N,NJ)),
jax.random.normal(key, shape=(N,N,NJ)),
jax.random.normal(key, shape=(N,N,NJ)) ]
A_np, B_np, D_np = [np.asarray(A), np.asarray(B), np.asarray(D)]
# initialize E_0
E_0 = jax.random.normal(key+2, shape=(N,N))
E_np = np.empty((N,N,NJ))
E_np[:,:,0] = np.asarray(E_0)
# iteratively calculate E from A, B, D, and πβ±Όββ
for j in range(1,NJ):
E_jm1 = E_np[:,:,j-1]
E_np[:,:,j] = BAND_J(A_np[:,:,j], B_np[:,:,j], D_np[:,:,j], E_jm1)
JAX scan version
def BAND_J(E, ABD):
'''
output: E(N x N)
input: A(N x N), B(N x N), D(N x N), E_jm1(N x N)
'''
A, B, D = ABD
B_inv = jnp.linalg.inv(B + jnp.dot( A, E ))
E = -jnp.dot(B_inv, D)
return E, E # ("carryover", "accumulated")
abd = jnp.asarray([(A[:,:,j], B[:,:,j], D[:,:,j]) for j in range(NJ)])
# abd = tuple([(A[:,:,j], B[:,:,j], D[:,:,j]) for j in range(NJ)]) # this produces error
# ValueError: too many values to unpack (expected 3)
_, E = lax.scan(BAND_J, E_0, abd)
for j in range(1, NJ):
print(np.isclose(E[j-1], E_np[:,:,j]))
The short answer is "you can't". By design, jax.scan can scan over axes of arrays, not entries of arbitrary Python collections.
So if you want to use scan, you'll have to stack your entires into an array.
That said, since your tuple only has three elements, a good alternative would be to skip the scan and simply JIT-compile the for loop approach. JAX tracing will effectively unroll the loop and optimize the flattened sequence of operations. While this can lead to long compile times for large loops, since your application is only 3 iterations it shouldn't be problematic.
Related
Hello I am working on a problem that seems to be out of my league so any tips, pointers to reading materials etc. are really appreciated. That being said here is the problem:
given 3 subsets of numbers a, b, c β {0, ..., n}. In nlog(n) check if there exists numbers n1, n2 in a, b and n3 in c where n1 + n2 = n3.
I am given the hint to convert a and b to polynomial coefficients and to use polynomial multiplication using ftt to multiply the coefficients of a and b.
Now where I am stuck is after getting the result of the polynomial multiplication, what do I do next?
Thank you in advanced.
from numpy.fft import fft, ifft
from numpy import real, imag
def polynomial_multiply(a_coeff_list, b_coeff_list):
# Return the coefficient list of the multiplication
# of the two polynomials
# Returned list must be a list of floating point numbers.
# list from complex to reals by using the
# real function in numpy
len_a = len(a_coeff_list)
len_b = len(b_coeff_list)
for i in range(len_a-1):
b_coeff_list.append(0)
for i in range(len_b-1):
a_coeff_list.append(0)
a_fft = fft(a_coeff_list)
b_fft = fft(b_coeff_list)
c = []
for i in range(len(a_fft)):
c.append(a_fft[i] * b_fft[i])
inverse_c = ifft(c)
return real(inverse_c)
# inputs sets a, b, c
# return True if there exist n1 in a, n2 in B such that n1+n2 in C
# return False otherwise
# number n which signifies the maximum number in a, b, c
def check_sum_exists(a, b, c, n):
a_coeffs = [0]*n
b_coeffs = [0]*n
# convert sets a, b into polynomials as provided in the hint
# a_coeffs and b_coeffs should contain the result
i = 0
for item in a:
a_coeffs[i] = item
i += 1
i = 0
for item in b:
b_coeffs[i] = item
i += 1
# multiply them together
c_coeffs = polynomial_multiply(a_coeffs, b_coeffs)
# now this is where i am lost
# how to determine with c_coeffs?
return False
# return True/False
Thanks to all who helped. I figured it out and hopefully this can help anyone who runs into a similar problem. The issue I had was I incorrectly assigned the coefficients for a_coeffs and b_coeffs.
Here is the solution which passed the tests for those interested.
from numpy.fft import fft, ifft
from numpy import real, imag
def check_sum_exists(a, b, c, n):
a_coeffs = [0] * n
b_coeffs = [0] * n
# convert sets a, b into polynomials as provided in the hint
# a_coeffs and b_coeffs should contain the result
for coeff in a:
a_coeffs[coeff] = 1
for coeff in b:
b_coeffs[coeff] = 1
# multiply them together
c_coeffs = polynomial_multiply(a_coeffs, b_coeffs)
# use the result to solve the problem at hand
for coeff in c:
if c_coeffs[coeff] >= .5:
return True
return False
# return True/False
def polynomial_multiply(a_coeff_list, b_coeff_list):
# Return the coefficient list of the multiplication
# of the two polynomials
# Returned list must be a list of floating point numbers.
# Please convert list from complex to reals by using the
# real function in numpy.
for i in range(len(a_coeff_list) - 1):
b_coeff_list.append(0)
for i in range(len(b_coeff_list) - 1):
a_coeff_list.append(0)
a_fft = fft(a_coeff_list)
b_fft = fft(b_coeff_list)
c = []
for i in range(len(a_fft)):
c.append(a_fft[i] * b_fft[i])
return real(ifft(c))
I have the following for loop that operates over three numpy arrays of the same length:
n = 100
a = np.random.random(n)
b = np.random.random(n)
c = np.random.random(n)
valid = np.empty(n)
for i in range(n):
valid[i] = np.any(a[i] > b[i:] + c[i:].cumsum())
Is there a way to replace this for loop with some vectorized numpy operations?
For example, because I only care if a[i] is larger than any value in b[i:], I can do np.minimum.accumulate(b[::-1])[::-1] which gets the smallest value of b at every index and onwards, and then compare it to a like this:
a > np.minimum.accumulate(b[::-1])[::-1]
but I still would need a way to vectorize the c[i:].cumsum() into a single array calculation.
Your goal is to find the minimum of b[i:] + c[i:].cumsum() for each i. Clearly you can compare that to a directly.
You can write the elements of c[i:].cumsum() as the upper triangle of a matrix. Let's look at a toy case with n = 3:
c = [c1, c2, c3]
s1 = c.cumsum()
s0 = np.r_[0, s1[:-1]]
You can write the elements of the cumulative sum as
c1, c1 + c2, c1 + c2 + c3 s1[0:] s1[0:] - s0[0]
c2, c2 + c3 = s1[1:] - c1 = s1[1:] - s0[1]
c3 s1[2:] - (c1 + c2) s1[2:] - s0[2]
You can use np.triu_indices to construct these sums as a raveled array:
r, c = np.triu_indices(n)
diff = s1[c] - s0[r] + b[c]
Since np.minimum is a ufunc, you can accumulate diff for each run defined by r using minimum.reduceat. The locations are given roughly by np.flatnonzero(np.diff(r)) + 1, but you can generate them faster with np.arange:
m = np.minimum.reduceat(diff, np.r_[0, np.arange(n, 1, -1).cumsum()])
So finally, you have:
valid = a > m
TL;DR
s1 = c.cumsum()
s0 = np.r_[0, s1[:-1]]
r, c = np.triu_indices(n)
valid = a > np.minimum.reduceat(s1[c] - s0[r] + b[c], np.r_[0, np.arange(n, 1, -1).cumsum()])
I assume you want to vectorize it to decrease the running time. Since you are only using pure NumPy operations, you can use numba: see 5 Minutes Guide to Numba
it will look something like this:
import numba
#numba.njit()
def valid_for_single_idx(idx, a, b, c):
return np.any(a[idx] > b[idx:] + c[idx:].cumsum())
valid = np.empty(n)
for i in range(n):
valid[i] = valid_for_single_idx(i, a, b, c)
So far it isn't really vectorization (as the loop still happens), but it translate the the numpy line into llvm, so it happens as fast as probably possible.
Although it's not increasing the speed, but looks a bit nicer, you can use .map:
import numba
from functools import partial
#numba.njit()
def valid_for_single_idx(idx, a, b, c):
return np.any(a[idx] > b[idx:] + c[idx:].cumsum())
valid = map(partial(valid_for_single_idx, a=a, b=b, c=c), range(n))
I have three variables called a, b and c, each of these can assume a different value defined in a range. I'd like to create a function that tests every possible variable value and gives me their best combination for the output 'f'.
a = list(range(1, 10, 2))
b = list(range(5, 8, 1))
c = list(range(1, 3, 1))
def all_combinations (a, b, c):
#something
f = a + (b * a) - (c*(a ^ b))
return BEST a, b, c for my f
it's possible to do it ? what is the best way to do it?
You can use itertools.product() to get all the possible combinations of a, b, and c.
Then calculate your formula for each unique combination of a b c, keep track of the result, and if the result is better than the previous best, save the current values of a b c.
import itertools
def all_combinations (alist, blist, clist):
best_a = 0
best_b = 0
best_c = 0
best_f = 0
for a,b,c in itertools.product(alist, blist, clist):
f = a + (b * a) - (c*(a ^ b))
if f > best_f: # use your own definition of "better"
best_a = a
best_b = b
best_c = c
best_f = f
return best_a, best_b, best_c
First of all, you said I have three variables called a, b and c, each of these can assume a different value defined in a range. Note that the variables in your code are actually equal to three lists of integers, not three integers.
The naive algorithm to test all possible combinations is 3 nested for loops. Here I assume that by "best" you mean "maximum value":
def all_combinations (list1, list2, list3):
best_f, best_a, best_b, best_c = None, None, None, None
for a in list1:
for b in list2:
for c in list3:
f = a + (b * a) - (c*(a ^ b))
# here you have to define what f being "better" than best_f means:
if not f or f > best_f:
best_f = f
best_a = a
best_b = b
best_c = c
return best_a, best_b, best_c
If you're sure those are the only values you want to test, then the following will work. Otherwise you might want to look into scipy.optimize.
from itertools import product
import numpy as np
parameters = list(product(a, b, c))
results = [my_fun(*x) for x in parameters]
print(parameters[np.argmax(results)])
obviously replace np.argmax with np.argmin if you want to minimize the function
I programmed class which looks something like this:
import numpy as np
class blank():
def __init__(self,a,b,c):
self.a=a
self.b=b
self.c=c
n=5
c=a/b*8
if (a>b):
y=c+a*b
else:
y=c-a*b
p = np.empty([1,1])
k = np.empty([1,1])
l = np.empty([1,1])
p[0]=b
k[0]=b*(c-1)
l[0]=p+k
for i in range(1, n, 1):
p=np.append(p,l[i-1])
k=np.append(k,(p[i]*(c+1)))
l=np.append(l,p[i]+k[i])
komp = np.zeros(shape=(n, 1))
for i in range(0, n):
pl_avg = (p[i] + l[i]) / 2
h=pl_avg*3
komp[i]=pl_avg*h/4
self.tot=komp+l
And when I call it like this:
from ex1 import blank
import numpy as np
res=blank(1,2,3)
print(res.tot)
everything works well.
BUT I want to call it like this:
res = blank(np.array([1,2,3]), np.array([3,4,5]), 3)
Is there an easy way to call it for each i element of this two arrays without editing class code?
You won't be able to instantiate a class with NumPy arrays as inputs without changing the class code. #PabloAlvarez and #NagaKiran already provided alternative: iterate with zip over arrays and instantiate class for each pair of elements. While this is pretty simple solution, it defeats the purpose of using NumPy with its efficient vectorized operations.
Here is how I suggest you to rewrite the code:
from typing import Union
import numpy as np
def total(a: Union[float, np.ndarray],
b: Union[float, np.ndarray],
n: int = 5) -> np.array:
"""Calculates what your self.tot was"""
bc = 8 * a
c = bc / b
vectorized_geometric_progression = np.vectorize(geometric_progression,
otypes=[np.ndarray])
l = np.stack(vectorized_geometric_progression(bc, c, n))
l = np.atleast_2d(l)
p = np.insert(l[:, :-1], 0, b, axis=1)
l = np.squeeze(l)
p = np.squeeze(p)
pl_avg = (p + l) / 2
komp = np.array([0.75 * pl_avg ** 2]).T
return komp + l
def geometric_progression(bc, c, n):
"""Calculates array l"""
return bc * np.logspace(start=0,
stop=n - 1,
num=n,
base=c + 2)
And you can call it both for sole numbers and NumPy arrays like that:
>>> print(total(1, 2))
[[2.6750000e+01 6.6750000e+01 3.0675000e+02 1.7467500e+03 1.0386750e+04]
[5.9600000e+02 6.3600000e+02 8.7600000e+02 2.3160000e+03 1.0956000e+04]
[2.1176000e+04 2.1216000e+04 2.1456000e+04 2.2896000e+04 3.1536000e+04]
[7.6205600e+05 7.6209600e+05 7.6233600e+05 7.6377600e+05 7.7241600e+05]
[2.7433736e+07 2.7433776e+07 2.7434016e+07 2.7435456e+07 2.7444096e+07]]
>>> print(total(3, 4))
[[1.71000000e+02 3.39000000e+02 1.68300000e+03 1.24350000e+04 9.84510000e+04]
[8.77200000e+03 8.94000000e+03 1.02840000e+04 2.10360000e+04 1.07052000e+05]
[5.59896000e+05 5.60064000e+05 5.61408000e+05 5.72160000e+05 6.58176000e+05]
[3.58318320e+07 3.58320000e+07 3.58333440e+07 3.58440960e+07 3.59301120e+07]
[2.29323574e+09 2.29323590e+09 2.29323725e+09 2.29324800e+09 2.29333402e+09]]
>>> print(total(np.array([1, 3]), np.array([2, 4])))
[[[2.67500000e+01 6.67500000e+01 3.06750000e+02 1.74675000e+03 1.03867500e+04]
[1.71000000e+02 3.39000000e+02 1.68300000e+03 1.24350000e+04 9.84510000e+04]]
[[5.96000000e+02 6.36000000e+02 8.76000000e+02 2.31600000e+03 1.09560000e+04]
[8.77200000e+03 8.94000000e+03 1.02840000e+04 2.10360000e+04 1.07052000e+05]]
[[2.11760000e+04 2.12160000e+04 2.14560000e+04 2.28960000e+04 3.15360000e+04]
[5.59896000e+05 5.60064000e+05 5.61408000e+05 5.72160000e+05 6.58176000e+05]]
[[7.62056000e+05 7.62096000e+05 7.62336000e+05 7.63776000e+05 7.72416000e+05]
[3.58318320e+07 3.58320000e+07 3.58333440e+07 3.58440960e+07 3.59301120e+07]]
[[2.74337360e+07 2.74337760e+07 2.74340160e+07 2.74354560e+07 2.74440960e+07]
[2.29323574e+09 2.29323590e+09 2.29323725e+09 2.29324800e+09 2.29333402e+09]]]
You can see that results are in compliance.
Explanation:
First of all I'd like to note that your calculation of p, k, and l doesn't have to be in the loop. Moreover, calculating k is unnecessary. If you see carefully, how elements of p and l are calculated, they are just geometric progressions (except the 1st element of p):
p = [b, b*c, b*c*(c+2), b*c*(c+2)**2, b*c*(c+2)**3, b*c*(c+2)**4, ...]
l = [b*c, b*c*(c+2), b*c*(c+2)**2, b*c*(c+2)**3, b*c*(c+2)**4, b*c*(c+2)**5, ...]
So, instead of that loop, you can use np.logspace. Unfortunately, np.logspace doesn't support base parameter as an array, so we have no other choice but to use np.vectorize which is just a loop under the hood...
Calculating of komp though is easily vectorized. You can see it in my example. No need for loops there.
Also, as I already noted in a comment, your class doesn't have to be a class, so I took a liberty of changing it to a function.
Next, note that input parameter c is overwritten, so I got rid of it. Variable y is never used. (Also, you could calculate it just as y = c + a * b * np.sign(a - b))
And finally, I'd like to remark that creating NumPy arrays with np.append is very inefficient (as it was pointed out by #kabanus), so you should always try to create them at once - no loops, no appending.
P.S.: I used np.atleast_2d and np.squeeze in my code and it could be unclear why I did it. They are necessary to avoid if-else clauses where we would check dimensions of array l. You can print intermediate results to see what is really going on there. Nothing difficult.
if it is just calling class with two different list elements, loop can satisfies well
res = [blank(i,j,3) for i,j in zip(np.array([1,2,3]),np.array([3,4,5]))]
You can see list of values for res variable
The only way I can think of iterating lists of arrays is by using a function on the main program for iteration and then do the operations you need to do inside the loop.
This solution works for each element of both arrays (note to use zip function for making the iteration in both lists if they have a small size as listed in this answer here):
for n,x in zip(np.array([1,2,3]),np.array([3,4,5])):
res=blank(n,x,3)
print(res.tot)
Hope it is what you need!
I'm looking to implement in python a simple algorithm which takes as input an array and a sum, and finds a number X where if all elements in the array > X are converted to X, all the elements in the array should add up to the sum.
How do I do this efficiently?
Here is my code:
result = []
for _ in range(int(raw_input())):
input_array = map(int,raw_input().split())
sum_target = raw_input()
for e in input_array:
test_array = input_array
test_array[test_array > e] = e // supposed to replace all elements > e with e, but what's wrong here?
if sum(test_array) == sum_target:
result.append(e)
print result
Using the Numpy library (import numpy), you could replace the line
input_array = map(int,raw_input().split())
with
input_array = numpy.array(raw_input().split()).astype(int)
Then
test_array[test_array > e] = e
just works. Then, you could also do test_array.sum().
(That is, if you want to alter the array in-place, else you could replace
test_array = input_array
with
test_array = np.array(input_array)