diff --git a/doc/source/whatsnew/0.3.rst b/doc/source/whatsnew/0.3.rst
index eaaf78ce93fbb80452ceff0c7dcc00bd85a1c85e..594270133bfd9f5b95f1ac7df2ad579871ceaafe 100644
--- a/doc/source/whatsnew/0.3.rst
+++ b/doc/source/whatsnew/0.3.rst
@@ -133,3 +133,8 @@ responsibility.
The new class `~kwant.builder.ModesLead` allows to attach leads that have a
custom way of calculating their modes (e.g. ideal leads) directly to a
`~kwant.builder.Builder`.
+
+Modes or self-energies can now be precomputed before passing the system to a
+solver, using the method `~kwant.system.FiniteSystem.precalculate`. This may
+save time, when the linear system has to be solved many times with the same
+lead parameters.
diff --git a/kwant/system.py b/kwant/system.py
index d434b3cca953b8eee7168098edff767ebd1bb550..41b84f78456e6ab346a0583baf48e4404af29ffc 100644
--- a/kwant/system.py
+++ b/kwant/system.py
@@ -13,6 +13,7 @@ __all__ = ['System', 'FiniteSystem', 'InfiniteSystem']
import abc
import types
+from copy import copy
import numpy as np
from . import physics, _system
@@ -90,6 +91,57 @@ class FiniteSystem(System):
"""
__metaclass__ = abc.ABCMeta
+ def precalculate(self, energy, args=(), leads=None, check=True,
+ calculate_selfenergy=True):
+ """
+ Precalculate modes or self-energies in the leads.
+
+ Construct a copy of the system, with the lead modes precalculated,
+ which may significantly speed up calculations where only the system
+ is changing.
+
+ Parameters
+ ----------
+ energy : float
+ Energy at which the modes or self-energies have to be
+ evaluated.
+ args : sequence
+ Additional parameters required for calculating the Hamiltionians
+ leads : list of integers or None
+ Numbers of the leads to be precalculated. If `None`, all are
+ precalculated.
+ check : bool
+ Whether the resulting system should raise an error if solving
+ is attempted at parameter or energy values other than used
+ during precalculation.
+ calculate_selfenergy : bool
+ Whether to calculate self-energy if modes are available.
+
+ Returns
+ -------
+ sys : FiniteSystem
+ A copy of the original system with some leads precalculated.
+ """
+ result = copy(self)
+ if leads is None:
+ leads = range(len(self.leads))
+ new_leads = []
+ res_energy, res_args = (energy, args) if check else (None, None)
+ for nr, lead in enumerate(self.leads):
+ if nr not in leads:
+ new_leads.append(lead)
+ continue
+ modes, selfenergy = None, None
+ try:
+ modes = lead.modes(energy, args)
+ if calculate_selfenergy:
+ selfenergy = physics.selfenergy_from_modes(modes)
+ except AttributeError:
+ selfenergy = lead.selfenergy(energy, args)
+ new_leads.append(PrecalculatedLead(modes, selfenergy, res_energy,
+ res_args))
+ result.leads = new_leads
+ return result
class InfiniteSystem(System):
"""
@@ -175,3 +227,60 @@ class InfiniteSystem(System):
# Subtract energy from the diagonal.
ham.flat[::ham.shape[0] + 1] -= energy
return physics.selfenergy(ham, self.inter_slice_hopping(args=args))
+
+
+class PrecalculatedLead(object):
+ def __init__(self, modes=None, selfenergy=None, energy=None, args=None):
+ """A general lead defined by its self energy.
+
+ Parameters
+ ----------
+ modes : kwant.physics.ModesTuple
+ Modes of the lead.
+ selfenergy : numpy array
+ Lead self-energy.
+ energy : float
+ Energy at which modes or self-energy were calculated.
+ args : sequence
+ Extra parameters used in the calculation.
+
+ Raises
+ ------
+ ValueError
+ If modes or self-energy are requested, at the energy or argument
+ values different from the ones specified during initialization.
+ If the initial values were `None`, no check is performed.
+
+ Notes
+ -----
+ At least one of `modes` and `selfenergy` must be provided.
+ """
+ if modes is None and selfenergy is None:
+ raise ValueError("No precalculated values provided.")
+ self._modes = modes
+ self._selfenergy = selfenergy
+ self._energy = energy
+ self._args = args
+
+ def _check_params(self, energy, args=()):
+ cond1 = self._energy is not None and self._energy != energy
+ cond2 = self._args is not None and tuple(self._args) != tuple(args)
+ if cond1 or cond2:
+ raise ValueError("The lead modes or self-energy are requested "
+ "at parameter values different from the ones "
+ "at which they were evaluated.")
+
+ def modes(self, energy, args=()):
+ self._check_params(energy, args)
+ if self._modes is not None:
+ return self._modes
+ else:
+ raise ValueError("No precalculated modes were provided.")
+
+ def selfenergy(self, energy, args=()):
+ self._check_params(energy, args)
+ return self._selfenergy
+ if self._selfenergy is not None:
+ return self._selfenergy
+ else:
+ raise ValueError("No precalculated self-energy was provided.")
diff --git a/kwant/tests/test_system.py b/kwant/tests/test_system.py
index 72d1caf7921307c3b9c657bb655166a4ad5d19c5..736c2e5576af367a6f80b702a9977a1492c65fc6 100644
--- a/kwant/tests/test_system.py
+++ b/kwant/tests/test_system.py
@@ -56,6 +56,19 @@ def test_hamiltonian_submatrix():
mat_sp = sys2.hamiltonian_submatrix(sparse=True).todense()
np.testing.assert_array_equal(mat_sp, mat_dense)
+ # Test precalculation of modes.
+ np.random.seed(5)
+ lead = kwant.Builder(kwant.TranslationalSymmetry((-1,)))
+ lead[gr(0)] = np.zeros((2, 2))
+ lead[gr(0), gr(1)] = np.random.randn(2, 2)
+ sys.attach_lead(lead)
+ sys2 = sys.finalized()
+ smatrix = kwant.solve(sys2, .1).data
+ sys3 = sys2.precalculate(.1)
+ smatrix2 = kwant.solve(sys3, .1).data
+ np.testing.assert_almost_equal(smatrix, smatrix2)
+ assert_raises(ValueError, kwant.solve, sys3, 0.2)
+
# Test for shape errors.
sys[gr(0), gr(2)] = np.array([[1, 2]])
sys2 = sys.finalized()
@@ -86,3 +99,4 @@ def test_hamiltonian_submatrix():
mat = mat[perm, :]
mat = mat[:, perm]
np.testing.assert_array_equal(mat, mat_should_be)
+