Add (Num)Pythonic indexing and slicing for AtomsLists
Opened this issue · 0 comments
max-veit commented
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.