diff --git a/micromagneticdata/abstract_drive.py b/micromagneticdata/abstract_drive.py index d4e999e..81f7136 100644 --- a/micromagneticdata/abstract_drive.py +++ b/micromagneticdata/abstract_drive.py @@ -177,6 +177,7 @@ def n(self): def _step_files(self): """List of filenames of individual snapshots.""" + @abc.abstractmethod def __getitem__(self, item): """Magnetisation field of an individual step. @@ -200,10 +201,6 @@ def __getitem__(self, item): Field(...) """ - field = df.Field.fromfile(filename=self._step_files[item]) - with contextlib.suppress(FileNotFoundError): - field.mesh.load_subregions(self._m0_path) - return self._apply_callbacks(field) def __iter__(self): """Iterator. diff --git a/micromagneticdata/combined_drive.py b/micromagneticdata/combined_drive.py index 086f14b..ecc04dd 100644 --- a/micromagneticdata/combined_drive.py +++ b/micromagneticdata/combined_drive.py @@ -1,4 +1,5 @@ import functools +import numbers import operator import ubermagutil as uu @@ -140,6 +141,16 @@ def table(self): def _step_files(self): return sum((drive._step_files for drive in self.drives), start=[]) + def __getitem__(self, item): + if not isinstance(item, numbers.Integral): + raise TypeError(f"{type(item)=} is not supported") + if item >= self.n or item < -self.n: + raise IndexError("List index out of range.") + for drive in self.drives: + if item < drive.n: + return self._apply_callbacks(drive[item]) + item -= drive.n + def __lshift__(self, other): if isinstance(other, md.Drive): return self.__class__(*self.drives, other) diff --git a/micromagneticdata/drive.py b/micromagneticdata/drive.py index ff7ab48..7e5d563 100644 --- a/micromagneticdata/drive.py +++ b/micromagneticdata/drive.py @@ -1,4 +1,5 @@ import abc +import contextlib import copy import json import numbers @@ -212,7 +213,10 @@ def __getitem__(self, item): """ if isinstance(item, numbers.Integral): - return super().__getitem__(item) + field = df.Field.fromfile(filename=self._step_files[item]) + with contextlib.suppress(FileNotFoundError): + field.mesh.load_subregions(self._m0_path) + return self._apply_callbacks(field) elif isinstance(item, slice): step_files = self._step_files[item] table = copy.copy(self.table)