CODEWITHSUNDEEP
pythonnumpy
Michael
Michael

Reputation: 486

Python method for array with variable dimensions

This is a simple question, but I couldn't find a 'best practice', and wondered if there was something simpler than lots of if statements. Say I have a method which takes in a variable 'data', which is of shape (N,M), where N can vary. Specifically sometimes data is a 1D array of shape (M,), other times N=100 and data is shape (100,M) for example.

Below is skeleton code of what the method does, for when N>1. How can I adopt this for a general case, when N>1 or N can equal 1 (or preferably when data.shape=(M,), not just (1,M)) ? I can put lots of if statements, but I was hoping for a cleaner solution.

#start with variables data.shape=(N,M), vol.shape=(M,), jstarts and jends .shape=(4,)
N=3
#N=1 #uncomment to test
M=20
jstarts = np.array([0,5,12,15])
jends = np.array([3,10,14,18])
data = np.arange(0,N,M).reshape(N,M)
data_new = np.empty((N,M))

for i in range(0,N):
    for j in range(0,jstarts.size):
        jstart = jstarts[j]
        jend = jends[j]
        tmp = np.sum(data[i,jstart:jend]*vol[jstart:jend])/np.sum(vol[jstart:jend])
        data_new[i,jstart:jend] = tmp

*NOTE: jstart and jend depend on j, but don't depend on i

Upvotes: 1

Views: 625

Answers (1)

mgilson
mgilson

Reputation: 309929

It's been a while since I've spent a ton of time with numpy, but IIRC, you should be able to drop at least the outer loop by using Ellipsis:

for j in range(10):
    jstart = jstarts[j]
    jend = jends[j]
    tmp = np.sum(data[...,jstart:jend]*vol[jstart:jend], axis=1)/np.sum(vol[jstart:jend])
    data_new[..., jstart:jend] = tmp

Upvotes: 2

Related Questions