Bug when deleting particles with grid attributes (like a 2d abundance array)

[elkogerville]

Describe the bug
In Krome, the particles have an abundance attribute which is a 2D array (N,M) for N particles and M species.
If you then delete a particle, amuse will flatten the abundance array first before deleting, which causes the output abundance array to be incorrect. This will happen to any vector attribute in amuse due to how attributes are deleted.

For example:

>>> p = cloud = new_molecular_cloud(target_number_of_particles=100, convert_nbody=converter)
>>> chem = Krome()
>>> chem.particles.add_particles(p)
>>> print(chem.abundances.shape)
(100, 9) # 9 species per particle
>>> chem.particles.remove_particle(p[0]) # remove particle zero
>>> print(chem.abundances.shape)
(899,) # wrong shape, should be (99, 9) !!!!

When code.particles.remove_particles is called, amuse will eventually call the remove_particles_from_store method of the InMemoryAttributeStorage class in src/amuse/datamodel/memory_storage.py (posted here for convenience):

def remove_particles_from_store(self, indices):
    for attribute, attribute_values in list(
        self.mapping_from_attribute_to_quantities.items()
    ):
        attribute_values.remove_indices(indices) 
    self.particle_keys = numpy.delete(self.particle_keys, indices)
    self.reindex()

    self.__version__ = self.__version__ + 1
    self.index_array = numpy.arange(len(self))

The issue is the line attribute_values.remove_indices(indices), which will then call the remove_indices method of either InMemoryUnitlessAttribute or InMemoryVectorQuantityAttribute depending on whether attribute_values has units or not.

In either case, the method looks more or less like this:

def remove_indices(self, indices):
    self.quantity._number = numpy.delete(self.quantity.number, indices)

The bug is that numpy.delete has 3 args, arr, obj, and axis=None. When axis=None, the array arr is flattened first before indices is applied. This is incorrect for 2D arrays, and typically you would want axis=0.

Expected behavior
A simple fix could be

    def remove_indices(self, indices):
        if isinstance(self.values, numpy.ndarray) and self.values.ndim == 2:
            self.values = numpy.delete(self.values, indices, axis=0)
        else:
            self.values = numpy.delete(self.values, indices)

This checks if the array is 2D, and if so, uses axis=0 in numpy.delete. This seems to fix the bug for my purposes, but it would be good to check if this has any unwanted side effects (for example, axis=0 may not always be the correct axis).

The logic for deleting particles is quite complex and spans many hundreds of lines of code across multiple files; thus, I am not sure that this is the only 'fix' required to solve this bug.