lab-cosmo/librascal

Add (Num)Pythonic indexing and slicing for AtomsLists

Opened this issue · 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.