librascal icon indicating copy to clipboard operation
librascal copied to clipboard

Add (Num)Pythonic indexing and slicing for AtomsLists

Open max-veit opened this issue 3 years ago • 0 comments

Right now, it seems the only way to get subsets of AtomsLists (a set of structures along with their computed representations, returned e.g. by all ) is to call the get_subset() function, which takes only a list of integers giving the frame indices to select. Slicing does not seem to be possible, despite the naming suggesting this should work like a regular Python list:

...
>>> soap = SphericalInvariants(**hypers)
>>> managers = soap.transform(frames)
>>> managers[:50]

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-26-ec3072fe1540> in <module>
----> 1 managers[:50]

/local/scratch/mveit/codes/librascal/build/rascal/neighbourlist/structure_manager.py in __getitem__(self, key)
     92 
     93     def __getitem__(self, key):
---> 94         return self.managers[key]
     95 
     96     def get_subset(self, selected_ids):

TypeError: __getitem__(): incompatible function arguments. The following argument types are supported:
    1. (self: rascal.lib._rascal.neighbour_list.ManagerCollection_Strict_CenterContribution_NeighbourList_Centers, arg0: int) -> rascal.lib._rascal.neighbour_list.Adaptor.Strict_CenterContribution_NeighbourList_Centers

Invoked with: <rascal.lib._rascal.neighbour_list.ManagerCollection_Strict_CenterContribution_NeighbourList_Centers object at 0x7fb7699b84f0>, slice(None, 50, None)

I would suggest we modify AtomsList.__getitem__() to take Python slice objects, and also Boolean arrays (in the NumPy slicing style) while we're at it, and transform them into the list of integers that the underlying bindings expect.

max-veit avatar Apr 23 '21 16:04 max-veit