python - numpy array of objects

Question

I'm trying to implement a simulation for a lattice model (lattice boltzmann) in Python. Each site of the lattice has a number of properties, and interact with neighboring sites according to certain rules. I figured that it might be clever to make a class with all the properties and make a grid of instances of that class. (As I'm inexperienced with Python, this might not be a good idea at all, so feel free to comment on my approach.)

Here is a toy example of what I'm doing

class site:
    def __init__(self,a,...):
        self.a = a
        .... other properties ...
    def set_a(self, new_a):
        self.a = new_a

Now I want to deal with a 2D/3D lattice (grid) of such sites so I tried to do the following (here is a 2D 3x3 grid as an example, but in simulation I would need the order of >1000x1000X1000)

lattice = np.empty( (3,3), dtype=object)
lattice[:,:] = site(3)

Now, the problem is that each lattice point refer to the same instance, for example

lattice[0,0].set_a(5)

will also set the value of lattice[0,2].a to 5. This behavior is unwanted. To avoid the problem i can loop over each grid point and assign the objects element by element, like

for i in range(3):
    for j in range(3):
        lattice[i,j] = site(a)

But is there a better way (not involving the loops) to assign objects to a multidimensional array?

Thanks

Answer 1

You can vectorize the class's __init__ function:

import numpy as np

class Site:
    def __init__(self, a):
        self.a = a
    def set_a(self, new_a):
        self.a = new_a

vSite = np.vectorize(Site)

init_arry = np.arange(9).reshape((3,3))

lattice = np.empty((3,3), dtype=object)
lattice[:,:] = vSite(init_arry)

This may look cleaner but has no performance advantage over your looping solution. The list comprehension answers create an intermediate python list which would cause a performance hit.