This is my code, working with dim=3, but I would like it to work for any dimensionality without having to manually edit code.
I would like to be able to vary the dimensionality between 3 and 20 eventually without manually having to ad for-loops. I was looking at itertools, but don't know how to select the correct values from the tuples created by itertools.product()
to square and add up for my if statement.
arrayshape = (width * 2 + 1,) * dim
funcspace = np.zeros(shape=arrayshape, dtype='b')
x1 = list(range(-int(width), int(width + 1)))
x2 = x1
x3 = x1
for i in range(len(x1)):
for j in range(len(x2)):
for k in range(len(x3)):
if round(np.sqrt(x1[i] ** 2 + x2[j] ** 2 + x3[k] ** 2)) in ranges:
funcspace[i][j][k] = 1
You can use product
on enumerate
of your vectors, which will yield the value and the index:
for ((i,v1),(j,v2),(k,v3)) in itertools.product(enumerate(x1),enumerate(x2),enumerate(x3)):
if round(np.sqrt(v1**2+v2**2+v3**2)) in ranges:
funcspace[i][j][k]=1
as a bonus, you get rid of the unpythonic range(len())
construct.
I've cooked a more general case when you have a vector of vectors. It's a little harder to read because unpacking isn't done in the for
loop.
The square sum is done using sum
on the 1 indexes (the values), and if the condition matches, we loop until we find the "deeper" list to set the value to 1
.
for t in itertools.product(*(enumerate(x) for x in x_list)):
# compute the squared sum of values
sqsum = sum(v[1]**2 for v in t)
if round(sqsum) in ranges:
# traverse the dimensions except the last one
deeper_list = funcspace
for i in range(len(t)-1):
deeper_list = deeper_list[t[i][0]]
# set the flag using the last dimension list
deeper_list[t[-1][0]] = 1
as noted in comments, since x1
seems to be repeated you can replace the first statement by:
for t in itertools.product(enumerate(x1), repeat=dim):
Another comments states that since funcspace
is a numpy ndarray
, we can simplify the "set to 1" loop by passing the list of indexes:
funcspace[[x[0] for x in t]] = 1
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.