diff --git a/kwant/builder.py b/kwant/builder.py
index b6712a2ed260febef7251846e3cda4e2f7988bc3..24e2c9d6d5d1868ba46d296e4fbe074cce5a6d57 100644
--- a/kwant/builder.py
+++ b/kwant/builder.py
@@ -95,7 +95,10 @@ class SiteFamily(metaclass=abc.ABCMeta):
representation and a name. The canonical representation will be returned as
the objects representation and must uniquely identify the site family
instance. The name is a string used to distinguish otherwise identical site
- families. It may be empty.
+ families. It may be empty. ``norbs`` defines the number of orbitals
+ on sites associated with this site family; it may be `None`, in which case
+ the number of orbitals is not specified.
+
All site families must define the method `normalize_tag` which brings a tag
to the standard format for this site family.
@@ -104,22 +107,37 @@ class SiteFamily(metaclass=abc.ABCMeta):
method `pos(tag)`, which returns a vector with real-space coordinates of the
site belonging to this family with a given tag.
+ If the ``norbs`` of a site family are provided, and sites of this family
+ are used to populate a `~kwant.builder.Builder`, then the associated
+ Hamiltonian values must have the correct shape. That is, if a site family
+ has ``norbs = 2``, then any on-site terms for sites belonging to this
+ family should be 2x2 matrices. Similarly, any hoppings to/from sites
+ belonging to this family must have a matrix structure where there are two
+ rows/columns. This condition applies equally to Hamiltonian values that
+ are given by functions. If this condition is not satisfied, an error will
+ be raised.
"""
- def __init__(self, canonical_repr, name):
+ def __init__(self, canonical_repr, name, norbs):
self.canonical_repr = canonical_repr
self.hash = hash(canonical_repr)
self.name = name
+ if norbs is not None:
+ if int(norbs) != norbs or norbs <= 0:
+ raise ValueError('The norbs parameter must be an integer > 0.')
+ norbs = int(norbs)
+ self.norbs = norbs
def __repr__(self):
return self.canonical_repr
def __str__(self):
if self.name:
- msg = '<{0} site family {1}>'
+ msg = '<{0} site family {1}{2}>'
else:
- msg = '<unnamed {0} site family>'
- return msg.format(self.__class__.__name__, self.name)
+ msg = '<unnamed {0} site family{2}>'
+ orbs = ' with {0} orbitals'.format(self.norbs) if self.norbs else ''
+ return msg.format(self.__class__.__name__, self.name, orbs)
def __hash__(self):
return self.hash
@@ -171,9 +189,10 @@ class SimpleSiteFamily(SiteFamily):
`SimpleSiteFamily` when `kwant.lattice.Monatomic` would also work.
"""
- def __init__(self, name=None):
- canonical_repr = '{0}({1})'.format(self.__class__, repr(name))
- super().__init__(canonical_repr, name)
+ def __init__(self, name=None, norbs=None):
+ canonical_repr = '{0}({1}, {2})'.format(self.__class__, repr(name),
+ repr(norbs))
+ super().__init__(canonical_repr, name, norbs)
def normalize_tag(self, tag):
tag = tuple(tag)
diff --git a/kwant/lattice.py b/kwant/lattice.py
index 22fc0320d863c14c16edf27efcbf0542cdba6a01..c9de68b47900276074f1c7c313169cfce5d21ac4 100644
--- a/kwant/lattice.py
+++ b/kwant/lattice.py
@@ -18,7 +18,7 @@ from .linalg import lll
from ._common import ensure_isinstance
-def general(prim_vecs, basis=None, name=''):
+def general(prim_vecs, basis=None, name='', norbs=None):
"""
Create a Bravais lattice of any dimensionality, with any number of sites.
@@ -32,6 +32,9 @@ def general(prim_vecs, basis=None, name=''):
Name of the lattice, or sequence of names of all of the sublattices.
If the name of the lattice is given, the names of sublattices (if any)
are obtained by appending their number to the name of the lattice.
+ norbs : int or sequence of ints, optional
+ The number of orbitals per site on the lattice, or a sequence
+ of the number of orbitals of sites on each of the sublattices.
Returns
-------
@@ -44,9 +47,9 @@ def general(prim_vecs, basis=None, name=''):
lattices.
"""
if basis is None:
- return Monatomic(prim_vecs, name=name)
+ return Monatomic(prim_vecs, name=name, norbs=norbs)
else:
- return Polyatomic(prim_vecs, basis, name=name)
+ return Polyatomic(prim_vecs, basis, name=name, norbs=norbs)
class Polyatomic:
@@ -62,18 +65,21 @@ class Polyatomic:
The primitive vectors of the Bravais lattice
basis : 2d array-like of floats
The coordinates of the basis sites inside the unit cell.
- name : string or sequence of strings
+ name : string or sequence of strings, optional
The name of the lattice, or a sequence of the names of all the
sublattices. If the name of the lattice is given, the names of
sublattices are obtained by appending their number to the name of the
lattice.
+ norbs : int or sequence of ints, optional
+ The number of orbitals per site on the lattice, or a sequence
+ of the number of orbitals of sites on each of the sublattices.
Raises
------
ValueError
If dimensionalities do not match.
"""
- def __init__(self, prim_vecs, basis, name=''):
+ def __init__(self, prim_vecs, basis, name='', norbs=None):
prim_vecs = ta.array(prim_vecs, float)
if prim_vecs.ndim != 2:
raise ValueError('`prim_vecs` must be a 2d array-like object.')
@@ -82,6 +88,7 @@ class Polyatomic:
name = ''
if isinstance(name, str):
name = [name + str(i) for i in range(len(basis))]
+
if prim_vecs.shape[0] > dim:
raise ValueError('Number of primitive vectors exceeds '
'the space dimensionality.')
@@ -91,8 +98,17 @@ class Polyatomic:
if basis.shape[1] != dim:
raise ValueError('Basis dimensionality does not match '
'the space dimensionality.')
- self.sublattices = [Monatomic(prim_vecs, offset, sname)
- for offset, sname in zip(basis, name)]
+
+ try:
+ norbs = list(norbs)
+ if len(norbs) != len(basis):
+ raise ValueError('Length of `norbs` is not the same as '
+ 'the number of basis vectors')
+ except TypeError:
+ norbs = [norbs] * len(basis)
+
+ self.sublattices = [Monatomic(prim_vecs, offset, sname, norb)
+ for offset, sname, norb in zip(basis, name, norbs)]
# Sequence of primitive vectors of the lattice.
self._prim_vecs = prim_vecs
# Precalculation of auxiliary arrays for real space calculations.
@@ -405,7 +421,7 @@ class Monatomic(builder.SiteFamily, Polyatomic):
Displacement of the lattice origin from the real space coordinates origin
"""
- def __init__(self, prim_vecs, offset=None, name=''):
+ def __init__(self, prim_vecs, offset=None, name='', norbs=None):
prim_vecs = ta.array(prim_vecs, float)
if prim_vecs.ndim != 2:
raise ValueError('`prim_vecs` must be a 2d array-like object.')
@@ -423,11 +439,12 @@ class Monatomic(builder.SiteFamily, Polyatomic):
raise ValueError('Dimensionality of offset does not match '
'that of the space.')
- msg = '{0}({1}, {2}, {3})'
+ msg = '{0}({1}, {2}, {3}, {4})'
cl = self.__module__ + '.' + self.__class__.__name__
canonical_repr = msg.format(cl, short_array_repr(prim_vecs),
- short_array_repr(offset), repr(name))
- super().__init__(canonical_repr, name)
+ short_array_repr(offset),
+ repr(name), repr(norbs))
+ super().__init__(canonical_repr, name, norbs)
self.sublattices = [self]
self._prim_vecs = prim_vecs
@@ -442,12 +459,14 @@ class Monatomic(builder.SiteFamily, Polyatomic):
self.lattice_dim = len(prim_vecs)
if name != '':
- msg = "<Monatomic lattice {0}>"
- self.cached_str = msg.format(name)
+ msg = "<Monatomic lattice {0}{1}>"
+ orbs = ' with {0} orbitals'.format(self.norbs) if self.norbs else ''
+ self.cached_str = msg.format(name, orbs)
else:
- msg = "<unnamed Monatomic lattice, vectors {0}, origin [{1}]>"
+ msg = "<unnamed Monatomic lattice, vectors {0}, origin [{1}]{2}>"
+ orbs = ', with {0} orbitals'.format(norbs) if norbs else ''
self.cached_str = msg.format(short_array_str(self._prim_vecs),
- short_array_str(self.offset))
+ short_array_str(self.offset), orbs)
def __str__(self):
return self.cached_str
@@ -684,33 +703,35 @@ class TranslationalSymmetry(builder.Symmetry):
################ Library of lattices
-def chain(a=1, name=''):
+def chain(a=1, name='', norbs=None):
"""Make a one-dimensional lattice."""
- return Monatomic(((a,),), name=name)
+ return Monatomic(((a,),), name=name, norbs=norbs)
-def square(a=1, name=''):
+def square(a=1, name='', norbs=None):
"""Make a square lattice."""
- return Monatomic(((a, 0), (0, a)), name=name)
+ return Monatomic(((a, 0), (0, a)), name=name, norbs=norbs)
tri = ta.array(((1, 0), (0.5, 0.5 * sqrt(3))))
-def triangular(a=1, name=''):
+def triangular(a=1, name='', norbs=None):
"""Make a triangular lattice."""
- return Monatomic(a * tri, name=name)
+ return Monatomic(a * tri, name=name, norbs=norbs)
-def honeycomb(a=1, name=''):
+def honeycomb(a=1, name='', norbs=None):
"""Make a honeycomb lattice."""
- lat = Polyatomic(a * tri, ((0, 0), (0, a / sqrt(3))), name=name)
+ lat = Polyatomic(a * tri, ((0, 0), (0, a / sqrt(3))),
+ name=name, norbs=norbs)
lat.a, lat.b = lat.sublattices
return lat
-def kagome(a=1, name=''):
+def kagome(a=1, name='', norbs=None):
"""Make a kagome lattice."""
- lat = Polyatomic(a * tri, ((0, 0),) + tuple(0.5 * a * tri), name=name)
+ lat = Polyatomic(a * tri, ((0, 0),) + tuple(0.5 * a * tri),
+ name=name, norbs=norbs)
lat.a, lat.b, lat.c = lat.sublattices
return lat
diff --git a/kwant/tests/test_lattice.py b/kwant/tests/test_lattice.py
index baa108edd2293ffefb8209ddfc9878da060848f6..92a36ca21996c4fa1d1a92b413eae0374b3080c4 100644
--- a/kwant/tests/test_lattice.py
+++ b/kwant/tests/test_lattice.py
@@ -197,3 +197,36 @@ def test_monatomic_lattice():
lat2 = lattice.general(np.identity(2))
lat3 = lattice.square(name='no')
assert len(set([lat, lat2, lat3, lat(0, 0), lat2(0, 0), lat3(0, 0)])) == 4
+
+
+def test_norbs():
+ id_mat = np.identity(2)
+ # Monatomic lattices
+ assert_equal(lattice.general(id_mat).norbs, None)
+ assert_equal(lattice.general(id_mat, norbs=2).norbs, 2)
+ # Polyatomic lattices
+ lat = lattice.general(id_mat, basis=id_mat, norbs=None)
+ for l in lat.sublattices:
+ assert_equal(l.norbs, None)
+ lat = lattice.general(id_mat, basis=id_mat, norbs=2)
+ for l in lat.sublattices:
+ assert_equal(l.norbs, 2)
+ lat = lattice.general(id_mat, basis=id_mat, norbs=[1, 2])
+ for l, n in zip(lat.sublattices, [1, 2]):
+ assert_equal(l.norbs, n)
+ # should raise ValueError for # of norbs different to length of `basis`
+ assert_raises(ValueError, lattice.general, id_mat, id_mat, norbs=[])
+ assert_raises(ValueError, lattice.general, id_mat, id_mat, norbs=[1, 2, 3])
+ # TypeError if Monatomic lattice
+ assert_raises(TypeError, lattice.general, id_mat, norbs=[])
+ # should raise ValueError if norbs not an integer
+ assert_raises(ValueError, lattice.general, id_mat, norbs=1.5)
+ assert_raises(ValueError, lattice.general, id_mat, id_mat, norbs=1.5)
+ assert_raises(ValueError, lattice.general, id_mat, id_mat, norbs=[1.5, 1.5])
+ # test that lattices with different norbs are compared `not equal`
+ lat = lattice.general(id_mat, basis=id_mat, norbs=None)
+ lat1 = lattice.general(id_mat, basis=id_mat, norbs=1)
+ lat2 = lattice.general(id_mat, basis=id_mat, norbs=2)
+ assert_not_equal(lat, lat1)
+ assert_not_equal(lat, lat2)
+ assert_not_equal(lat1, lat2)