Skip to content
Snippets Groups Projects
Commit 90d7b29f authored by Christoph Groth's avatar Christoph Groth
Browse files

cythonize hamiltonian_submatrix

parent fca192f8
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
kwant/_system.pyx 0 → 100644
+ 289
− 0
View file @ 90d7b29f
cimport cython
import tinyarray as ta
import numpy as np
from scipy import sparse as sp
from itertools import chain
from kwant.graph.core cimport CGraph, gintArraySlice
from kwant.graph.defs cimport gint
from kwant.graph.defs import gint_dtype
msg = 'Hopping from site {0} to site {1} does not match the ' \
'dimensions of onsite Hamiltonians of these sites.'
@cython.boundscheck(False)
def make_sparse(ham, CGraph gr, diag, gint [:] from_sites, n_by_to_site,
gint [:] to_norb, gint [:] to_off,
gint [:] from_norb, gint [:] from_off):
"""For internal use by hamiltonian_submatrix."""
cdef gintArraySlice nbors
cdef gint n_fs, fs, n_ts, ts
cdef gint i, j, num_entries
cdef complex [:, :] h
cdef gint [:, :] rows_cols
cdef complex [:] data
matrix = ta.matrix
# Calculate the data size.
num_entries = 0
for n_fs in xrange(len(from_sites)):
fs = from_sites[n_fs]
if fs in n_by_to_site:
num_entries += from_norb[n_fs] * from_norb[n_fs]
nbors = gr.out_neighbors(fs)
for ts in nbors.data[:nbors.size]:
if ts in n_by_to_site:
n_ts = n_by_to_site[ts]
num_entries += to_norb[n_ts] * from_norb[n_fs]
rows_cols = np.empty((2, num_entries), gint_dtype)
data = np.empty(num_entries, complex)
cdef gint pos = 0
for n_fs in xrange(len(from_sites)):
fs = from_sites[n_fs]
if fs in n_by_to_site:
n_ts = n_by_to_site[fs]
h = diag[n_fs]
if not (h.shape[0] == h.shape[1] == from_norb[n_fs]):
raise ValueError(msg.format(fs, fs))
for i in xrange(h.shape[0]):
for j in xrange(h.shape[1]):
data[pos] = h[i, j]
rows_cols[0, pos] = i + to_off[n_ts]
rows_cols[1, pos] = j + from_off[n_fs]
pos += 1
nbors = gr.out_neighbors(fs)
for ts in nbors.data[:nbors.size]:
if ts not in n_by_to_site:
continue
n_ts = n_by_to_site[ts]
h = matrix(ham(ts, fs), complex)
if h.shape[0] != to_norb[n_ts] or h.shape[1] != from_norb[n_fs]:
raise ValueError(msg.format(fs, ts))
for i in xrange(h.shape[0]):
for j in xrange(h.shape[1]):
data[pos] = h[i, j]
rows_cols[0, pos] = i + to_off[n_ts]
rows_cols[1, pos] = j + from_off[n_fs]
pos += 1
return sp.coo_matrix((data, rows_cols), shape=(to_off[-1], from_off[-1]))
@cython.boundscheck(False)
def make_sparse_full(ham, CGraph gr, diag,
gint [:] to_norb, gint [:] to_off,
gint [:] from_norb, gint [:] from_off):
"""For internal use by hamiltonian_submatrix."""
cdef gintArraySlice nbors
cdef gint n, fs, ts
cdef gint i, j, num_entries
cdef complex [:, :] h
cdef gint [:, :] rows_cols
cdef complex [:] data
matrix = ta.matrix
n = gr.num_nodes
# Calculate the data size.
num_entries = 0
for fs in xrange(n):
num_entries += from_norb[fs] * from_norb[fs]
nbors = gr.out_neighbors(fs)
for ts in nbors.data[:nbors.size]:
num_entries += to_norb[ts] * from_norb[fs]
rows_cols = np.empty((2, num_entries), gint_dtype)
data = np.empty(num_entries, complex)
cdef gint pos = 0
for fs in xrange(n):
h = diag[fs]
if not (h.shape[0] == h.shape[1] == from_norb[fs]):
raise ValueError(msg.format(fs, fs))
for i in xrange(h.shape[0]):
for j in xrange(h.shape[1]):
data[pos] = h[i, j]
rows_cols[0, pos] = i + to_off[fs]
rows_cols[1, pos] = j + from_off[fs]
pos += 1
nbors = gr.out_neighbors(fs)
for ts in nbors.data[:nbors.size]:
h = matrix(ham(ts, fs), complex)
if h.shape[0] != to_norb[ts] or h.shape[1] != from_norb[fs]:
raise ValueError(msg.format(fs, ts))
for i in xrange(h.shape[0]):
for j in xrange(h.shape[1]):
data[pos] = h[i, j]
rows_cols[0, pos] = i + to_off[ts]
rows_cols[1, pos] = j + from_off[fs]
pos += 1
return sp.coo_matrix((data, rows_cols), shape=(to_off[-1], from_off[-1]))
@cython.boundscheck(False)
def make_dense(ham, CGraph gr, diag, gint [:] from_sites, n_by_to_site,
gint [:] to_norb, gint [:] to_off,
gint [:] from_norb, gint [:] from_off):
"""For internal use by hamiltonian_submatrix."""
cdef gintArraySlice nbors
cdef gint n_fs, fs, n_ts, ts
cdef complex [:, :] h_sub_view
cdef complex [:, :] h
matrix = ta.matrix
h_sub = np.zeros((to_off[-1], from_off[-1]), complex)
h_sub_view = h_sub
for n_fs in xrange(len(from_sites)):
fs = from_sites[n_fs]
if fs in n_by_to_site:
n_ts = n_by_to_site[fs]
h = diag[n_fs]
if not (h.shape[0] == h.shape[1] == from_norb[n_fs]):
raise ValueError(msg.format(fs, fs))
h_sub_view[to_off[n_ts] : to_off[n_ts + 1],
from_off[n_fs] : from_off[n_fs + 1]] = h
nbors = gr.out_neighbors(fs)
for ts in nbors.data[:nbors.size]:
if ts not in n_by_to_site:
continue
n_ts = n_by_to_site[ts]
h = matrix(ham(ts, fs), complex)
if h.shape[0] != to_norb[n_ts] or h.shape[1] != from_norb[n_fs]:
raise ValueError(msg.format(fs, ts))
h_sub_view[to_off[n_ts] : to_off[n_ts + 1],
from_off[n_fs] : from_off[n_fs + 1]] = h
return h_sub
@cython.boundscheck(False)
def make_dense_full(ham, CGraph gr, diag,
gint [:] to_norb, gint [:] to_off,
gint [:] from_norb, gint [:] from_off):
"""For internal use by hamiltonian_submatrix."""
cdef gintArraySlice nbors
cdef gint n, fs, ts
cdef complex [:, :] h_sub_view, h
matrix = ta.matrix
n = gr.num_nodes
h_sub = np.zeros((to_off[-1], from_off[-1]), complex)
h_sub_view = h_sub
for fs in xrange(n):
h = diag[fs]
if not (h.shape[0] == h.shape[1] == from_norb[fs]):
raise ValueError(msg.format(fs, fs))
h_sub_view[to_off[fs] : to_off[fs + 1],
from_off[fs] : from_off[fs + 1]] = h
nbors = gr.out_neighbors(fs)
for ts in nbors.data[:nbors.size]:
h = matrix(ham(ts, fs), complex)
if h.shape[0] != to_norb[ts] or h.shape[1] != from_norb[fs]:
raise ValueError(msg.format(fs, ts))
h_sub_view[to_off[ts] : to_off[ts + 1],
from_off[fs] : from_off[fs + 1]] = h
return h_sub
def hamiltonian_submatrix(self, to_sites=None, from_sites=None,
sparse=False):
"""Return a submatrix of the system Hamiltonian.
Parameters
----------
to_sites : sequence of sites or None (default)
from_sites : sequence of sites or None (default)
sparse : bool
Whether to return a sparse or a dense matrix. Defaults to `False`.
Returns
-------
hamiltonian_part : numpy.ndarray or scipy.sparse.coo_matrix
Submatrix of Hamiltonian of the system.
to_norb : array of integers
Numbers of orbitals on each site in to_sites.
from_norb : array of integers
Numbers of orbitals on each site in from_sites.
Notes
-----
The returned submatrix contains all the Hamiltonian matrix elements
from `from_sites` to `to_sites`. The default for `from_sites` and
`to_sites` is `None` which means to use all sites of the system in the
order in which they appear.
"""
cdef gint [:] to_norb, from_norb
cdef gint site, n_site, n
ham = self.hamiltonian
n = self.graph.num_nodes
matrix = ta.matrix
if from_sites is None:
diag = n * [None]
from_norb = np.empty(n, gint_dtype)
for site in xrange(n):
diag[site] = h = matrix(ham(site, site), complex)
from_norb[site] = h.shape[0]
else:
diag = len(from_sites) * [None]
from_norb = np.empty(len(from_sites), gint_dtype)
for n_site, site in enumerate(from_sites):
if site < 0 or site >= n:
raise IndexError('Site number out of range.')
diag[n_site] = h = matrix(ham(site, site), complex)
from_norb[n_site] = h.shape[0]
from_off = np.empty(from_norb.shape[0] + 1, gint_dtype)
from_off[0] = 0
from_off[1 :] = np.cumsum(from_norb)
if to_sites is from_sites:
to_norb = from_norb
to_off = from_off
else:
if to_sites is None:
to_norb = np.empty(n, gint_dtype)
for site in xrange(n):
h = matrix(ham(site, site), complex)
to_norb[site] = h.shape[0]
else:
to_norb = np.empty(len(to_sites), gint_dtype)
for n_site, site in enumerate(to_sites):
if site < 0 or site >= n:
raise IndexError('Site number out of range.')
h = matrix(ham(site, site), complex)
to_norb[n_site] = h.shape[0]
to_off = np.empty(to_norb.shape[0] + 1, gint_dtype)
to_off[0] = 0
to_off[1 :] = np.cumsum(to_norb)
if to_sites is from_sites is None:
func = make_sparse_full if sparse else make_dense_full
mat = func(ham, self.graph, diag, to_norb, to_off, from_norb, from_off)
else:
if to_sites is None:
to_sites = np.arange(n)
n_by_to_site = dict((site, site) for site in to_sites)
else:
n_by_to_site = dict((site, n_site)
for n_site, site in enumerate(to_sites))
if from_sites is None:
from_sites = np.arange(n)
else:
from_sites = np.asarray(from_sites, gint_dtype)
func = make_sparse if sparse else make_dense
mat = func(ham, self.graph, diag, from_sites, n_by_to_site,
to_norb, to_off, from_norb, from_off)
return mat, to_norb, from_norb
kwant/system.py
+ 5
− 149
View file @ 90d7b29f
......@@ -3,11 +3,10 @@
from __future__ import division
__all__ = ['System', 'FiniteSystem', 'InfiniteSystem']
import abc, math
import abc, math, types
import numpy as np
from scipy import sparse as sp
from itertools import chain
from kwant import physics
import _system
class System(object):
"""Abstract general low-level system.
......@@ -46,152 +45,9 @@ class System(object):
"""
pass
def hamiltonian_submatrix(self, to_sites=None, from_sites=None,
sparse=False):
"""Return a submatrix of the system Hamiltonian.
Parameters
----------
to_sites : sequence of sites or None (default)
from_sites : sequence of sites or None (default)
sparse : bool
Whether to return a sparse or a dense matrix. Defaults to `False`.
Returns
-------
hamiltonian_part : numpy.ndarray or scipy.sparse.coo_matrix
Submatrix of Hamiltonian of the system.
to_norb : array of integers
Numbers of orbitals on each site in to_sites.
from_norb : array of integers
Numbers of orbitals on each site in from_sites.
Notes
-----
The returned submatrix contains all the Hamiltonian matrix elements
from `from_sites` to `to_sites`. The default for `from_sites` and
`to_sites` is `None` which means to use all sites of the system in the
order in which they appear.
"""
msg = 'Hopping from site {0} to site {1} does not match the ' \
'dimensions of onsite Hamiltonians of these sites.'
def make_sparse():
# Calculate the data size.
num_entries = 0
for n_i, i in enumerate(from_sites):
for j in chain((i,), gr.out_neighbors(i)):
if j in to_coord:
n_j = to_coord[j]
num_entries += to_norb[n_j] * from_norb[n_i]
ij = np.empty((2, num_entries), dtype=int)
data = np.empty(num_entries, dtype=complex)
pos = 0
for n_i, i in enumerate(from_sites):
for j in chain((i,), gr.out_neighbors(i)):
n_j = to_coord.get(j)
if n_j is None:
continue
h = ham(j, i) if j != i else diag[i]
# The shape check is here to prevent data corruption.
shape = (1, 1) if np.isscalar(h) else h.shape
if shape != (to_norb[n_j], from_norb[n_i]):
raise ValueError(msg.format(i, j))
if np.isscalar(h):
data[pos] = h
ij[0, pos] = n_j
ij[1, pos] = n_i
pos += 1
else:
h = np.ravel(h, order='F')
coord = slice(pos, pos + h.size)
data[coord] = h
jtmp = np.arange(to_norb[n_j]) + to_off[n_j]
itmp = np.arange(from_norb[n_i]) + from_off[n_i]
jtmp, itmp = np.meshgrid(jtmp, itmp)
ij[0, coord] = jtmp.ravel()
ij[1, coord] = itmp.ravel()
pos += h.shape[0]
return sp.coo_matrix((data, ij), shape=result_shape)
def make_dense():
# Shape checks of arrays are performed by numpy upon subblock
# assignment.
h_sub = np.zeros(result_shape, dtype='complex')
for n_i, i in enumerate(from_sites):
for j in chain((i,), gr.out_neighbors(i)):
n_j = to_coord.get(j)
if n_j is None:
continue
h = ham(j, i) if j != i else diag[i]
shape = (1, 1) if np.isscalar(h) else h.shape
if shape != (to_norb[n_j], from_norb[n_i]):
raise ValueError(msg.format(i, j))
h_sub[to_off[n_j] : to_off[n_j + 1],
from_off[n_i] : from_off[n_i + 1]] = h
return h_sub
gr = self.graph
ham = self.hamiltonian
n = self.graph.num_nodes
if not ((from_sites is None or
all(0 <= site < n for site in from_sites)) and
(to_sites is None or
all(0 <= site < n for site in to_sites))):
raise IndexError('Site number out of range.')
# Cache diagonal entries.
isscalar = np.isscalar
if from_sites is None or to_sites is None:
# np.fromiter does not work with dtype=object.
diag = np.array([ham(i, i) for i in xrange(n)], dtype=object)
norb = np.fromiter(
(1 if isscalar(h) else h.shape[0] for h in diag), int, n)
else:
if (max(len(from_sites), len(to_sites)) < n // 4 > 4):
diag = {}
get = diag.get
else:
diag = np.empty(n, dtype=object)
get = diag.__getitem__
# Make from_norb and from_off.
if from_sites is None:
from_sites = xrange(n)
from_norb = norb
else:
from_norb = np.empty(len(from_sites), dtype=int)
for n_i, i in enumerate(from_sites):
h = get(i)
if h is None:
diag[i] = h = ham(i, i)
from_norb[n_i] = 1 if isscalar(h) else h.shape[0]
from_off = np.zeros(from_norb.shape[0] + 1, int)
from_off[1 :] = np.cumsum(from_norb)
# Make to_norb and to_off.
if to_sites is None:
to_sites = xrange(n)
to_norb = norb
else:
to_norb = np.empty(len(to_sites), dtype=int)
for n_i, i in enumerate(to_sites):
h = get(i)
if h is None:
diag[i] = h = ham(i, i)
to_norb[n_i] = 1 if isscalar(h) else h.shape[0]
to_off = np.zeros(to_norb.shape[0] + 1, int)
to_off[1 :] = np.cumsum(to_norb)
# Instead of doing a double loop over from_sites and to_sites it is
# more efficient to check if neighbors of from_sites are in to_sites.
to_coord = dict((i[1], i[0]) for i in enumerate(to_sites))
result_shape = (to_off[-1], from_off[-1])
return make_sparse() if sparse else make_dense(), to_norb, from_norb
# Add a C-implemented function as an unbound method to class System.
System.hamiltonian_submatrix = types.MethodType(
_system.hamiltonian_submatrix, None, System)
class FiniteSystem(System):
......
setup.py
+ 2
− 0
View file @ 90d7b29f
......@@ -66,6 +66,8 @@ else:
# replacing ".pyx" with ".c" if Cython is not to be used.
extensions = [ # (["kwant.graph.scotch", ["kwant/graph/scotch.pyx"]],
# {"libraries" : ["scotch", "scotcherr"]}),
(["kwant._system", ["kwant/_system.pyx"]],
{"include_dirs" : ["kwant/graph"]}),
(["kwant.graph.core", ["kwant/graph/core.pyx"]],
{"depends" : ["kwant/graph/core.pxd", "kwant/graph/defs.h",
"kwant/graph/defs.pxd"]}),
......
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