Search code examples
pythonnumpy

How to group/loop over numpy arrays


I need to instantiate a large number of numpy arrays. The code snippet reads:

def computa(self):
    self.variable1 = ma.zeros([self.nx, self.ny, self.nz], dtype=float)
    self.variable2 = ma.zeros([self.nx, self.ny, self.nz], dtype=float)
    self.variable3 = ma.zeros([self.nx, self.ny, self.nz], dtype=float)

The problem is that this becomes uglier as the number of variables increases. Using a list is not a valid solution because it simply makes copies of the objects but does not instantiate the arrays.

def computa(self):
    self.variables_list = [self.variable1, self.variable2, self.variable3]
    for variables in (self.variables_list):
        variables = ma.zeros([self.nx, self.ny, self.nz], dtype=float)

Is there a pythonic way to group or loop over all the numpy arrays?


Solution

  • Make an array with the an added leading dimension:

    def computa(n):
        return ma.zeros((n, nx, ny, nz), dtype=float)
    

    The nth element of this array has shape (nx, ny, nz) and serves as the index of the variable.