Commit d1644397 authored by Christoph Groth's avatar Christoph Groth
Browse files

make site arrays proper sequences

parent 5fd117c1
...@@ -19,7 +19,7 @@ import abc ...@@ -19,7 +19,7 @@ import abc
import warnings import warnings
import operator import operator
from copy import copy from copy import copy
from collections import namedtuple import collections
from functools import total_ordering, lru_cache from functools import total_ordering, lru_cache
import numpy as np import numpy as np
import tinyarray as ta import tinyarray as ta
...@@ -96,7 +96,7 @@ class Site(tuple): ...@@ -96,7 +96,7 @@ class Site(tuple):
return self.family.pos(self.tag) return self.family.pos(self.tag)
class SiteArray: class SiteArray(collections.abc.Sequence):
"""An array of sites, members of a `SiteFamily`. """An array of sites, members of a `SiteFamily`.
Parameters Parameters
...@@ -137,6 +137,12 @@ class SiteArray: ...@@ -137,6 +137,12 @@ class SiteArray:
def __len__(self): def __len__(self):
return len(self.tags) return len(self.tags)
def __getitem__(self, key):
if isinstance(key, slice):
return SiteArray(self.family, self.tags[key])
else:
return Site(self.family, self.tags[key])
def __eq__(self, other): def __eq__(self, other):
if not isinstance(other, SiteArray): if not isinstance(other, SiteArray):
raise NotImplementedError() raise NotImplementedError()
...@@ -527,7 +533,7 @@ class System(metaclass=abc.ABCMeta): ...@@ -527,7 +533,7 @@ class System(metaclass=abc.ABCMeta):
hamiltonian_submatrix = _system.hamiltonian_submatrix hamiltonian_submatrix = _system.hamiltonian_submatrix
Term = namedtuple( Term = collections.namedtuple(
"Term", "Term",
["subgraph", "symmetry_element", "hermitian", "parameters"], ["subgraph", "symmetry_element", "hermitian", "parameters"],
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment