diff --git a/kwant/system.py b/kwant/system.py index 545739d926a700fdcf9b19cc0cd9244812d1d969..2b9a7bd47589a7c0a00b5c8fa91f5ca9a69bd0f9 100644 --- a/kwant/system.py +++ b/kwant/system.py @@ -83,7 +83,7 @@ class FiniteSystem(System): __metaclass__ = abc.ABCMeta def precalculate(self, energy=0, args=(), leads=None, - calculate_selfenergy=True): + what='modes'): """ Precalculate modes or self-energies in the leads. @@ -101,10 +101,9 @@ class FiniteSystem(System): leads : list of integers or None Numbers of the leads to be precalculated. If `None`, all are precalculated. - calculate_selfenergy : bool - Whether to calculate self-energy if modes are available. Defaults - to `False`. Disabling this saves a typically negligible amount of - time and memory. + what : 'modes', 'selfenergy', 'all' + The quantitity to precompute. 'all' will compute both + modes and self-energies. Defaults to 'modes'. Returns ------- @@ -117,6 +116,11 @@ class FiniteSystem(System): they might give wrong results if used to solve the system with different parameter values. Use this function with caution. """ + + if what not in ('modes', 'selfenergy', 'all'): + raise ValueError("Invalid value of argument 'what': " + "{0}".format(what)) + result = copy(self) if leads is None: leads = range(len(self.leads)) @@ -126,12 +130,13 @@ class FiniteSystem(System): new_leads.append(lead) continue modes, selfenergy = None, None - try: + if what in ('modes', 'all'): modes = lead.modes(energy, args) - if calculate_selfenergy: - selfenergy = modes[1].selfenergy() - except AttributeError: - selfenergy = lead.selfenergy(energy, args) + if what in ('selfenergy', 'all'): + if modes: + selfenergy = modes[1].selfenergy() + else: + selfenergy = lead.selfenergy(energy, args) new_leads.append(PrecalculatedLead(modes, selfenergy)) result.leads = new_leads return result @@ -247,10 +252,14 @@ class PrecalculatedLead(object): if self._modes is not None: return self._modes else: - raise ValueError("No precalculated modes were provided.") + raise ValueError("No precalculated modes were provided. " + "Consider using precalculate() with " + "what='modes' or what='all'") def selfenergy(self, energy=0, args=()): if self._selfenergy is not None: return self._selfenergy else: - raise ValueError("No precalculated self-energy was provided.") + raise ValueError("No precalculated selfenergy was provided. " + "Consider using precalculate() with " + "what='selfenergy' or what='all'") diff --git a/kwant/tests/test_comprehensive.py b/kwant/tests/test_comprehensive.py index 8a6f39b9ff6723032d3712a715deef0fe87b1326..918001d1b9472f10d7e10be08ce92d34c9b73aa8 100644 --- a/kwant/tests/test_comprehensive.py +++ b/kwant/tests/test_comprehensive.py @@ -50,7 +50,7 @@ def test_qhe(W=16, L=8): ((5.2, 5.5), 3, 1e-1)]: for r_phi in r_phis: args = (1.0 / r_phi, "") - pc = sys.precalculate(1.0, args) + pc = sys.precalculate(1.0, args, what='all') for result in [kwant.smatrix(pc, 1, args), kwant.solvers.default.greens_function(pc, 1, args)]: assert abs(T_nominal - result.transmission(1, 0)) < max_err diff --git a/kwant/tests/test_system.py b/kwant/tests/test_system.py index 2f3d6bb5f5d9bc78d8473007123d6200a68f8fff..346a2dd4691165c8e08c85e8fd76cf96843010aa 100644 --- a/kwant/tests/test_system.py +++ b/kwant/tests/test_system.py @@ -64,7 +64,7 @@ def test_hamiltonian_submatrix(): sys.attach_lead(lead) sys2 = sys.finalized() smatrix = kwant.smatrix(sys2, .1).data - sys3 = sys2.precalculate(.1, calculate_selfenergy=False) + sys3 = sys2.precalculate(.1, what='modes') smatrix2 = kwant.smatrix(sys3, .1).data np.testing.assert_almost_equal(smatrix, smatrix2) assert_raises(ValueError, kwant.solvers.default.greens_function, sys3, 0.2)