add spectra function to work on a sequence of leads, restrict spectrum to only a single lead

kwant_spectrum.py

@@ -12,7 +12,7 @@ from scipy.optimize import linear_sum_assignment
 
 from kwant.physics import Bands
 
-__all__ = ['spectrum', 'intersect_intervals', 'BandSketching']
+__all__ = ['spectrum', 'spectra', 'intersect_intervals', 'BandSketching']
 
 
 def _scale_estimate(onsite_hamiltonian, hopping_elements):
@@ -134,13 +134,6 @@ def _is_not_empty(a):
         return False
 
 
-def _make_iterable(variable):
-    """If a variable is present, make it iterable"""
-    if not isinstance(variable, collections.abc.Iterable):
-        return [variable]
-    return variable
-
-
 def intersect_intervals(interval_a, interval_b, tol=1E-16):
     """Return the intersecting intervals between two lists of intervals."""
     def large_enough(a, b):
@@ -585,35 +578,6 @@ def spectrum(syst, args=(), *, params=None, kmin=-np.pi, kmax=np.pi,
     subsequent interpolation of the spectrum for consistency.
     """
 
-    def lead_spectrum(lead_syst):
-
-        # match the bands continously
-        def array_function(func):
-            """Return energies and first derivative (velocities) as array"""
-            @wraps(func)
-            def wrapper(*args, **kwargs):
-                return np.array(func(*args, **kwargs, derivative_order=1))
-            return wrapper
-
-        bands = Bands(lead_syst, args, params=params)
-        tol_eff = tol * np.sum(np.abs(_scale_estimate(bands.ham, bands.hop)))
-        x, y, dy, ordering = match(array_function(bands), kmin, kmax, tol_eff)
-
-        # order bands according to their energy at momentum `orderpoint`
-        band_order = np.argsort(y[np.abs(x-orderpoint).argmin()])
-        y = y[:, band_order]
-        dy = dy[:, band_order]
-        ordering = ordering[:, band_order]
-
-        # provide a function that calculates the open modes for a given energy
-        mode_function = partial(lead_syst.modes, args=args, params=params)
-
-        # return a cubic interpolation of the spectrum, provide access to
-        # ordering
-        band_sketching = BandSketching(x, y, dy, mode_function, tol)
-        band_sketching.ordering = ordering
-        return band_sketching
-
     # type and input checks
     assert _is_type(kmin, 'real_number')
     assert _is_type(kmax, 'real_number')
@@ -622,11 +586,56 @@ def spectrum(syst, args=(), *, params=None, kmin=-np.pi, kmax=np.pi,
     assert kmin <= kmax, 'bounds swapped'
     assert tol > 0
 
-    syst = _make_iterable(syst)
-    spec = [lead_spectrum(lead_syst) for lead_syst in syst]
-    if len(spec) == 1:
-        return spec[0]
-    return spec
+    # match the bands continously
+    def array_function(func):
+        """Return energies and first derivative (velocities) as array"""
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return np.array(func(*args, **kwargs, derivative_order=1))
+        return wrapper
+
+    bands = Bands(syst, args, params=params)
+    tol_eff = tol * np.sum(np.abs(_scale_estimate(bands.ham, bands.hop)))
+    x, y, dy, ordering = match(array_function(bands), kmin, kmax, tol_eff)
+
+    # order bands according to their energy at momentum `orderpoint`
+    band_order = np.argsort(y[np.abs(x-orderpoint).argmin()])
+    y = y[:, band_order]
+    dy = dy[:, band_order]
+    ordering = ordering[:, band_order]
+
+    # provide a function that calculates the open modes for a given energy
+    mode_function = partial(syst.modes, args=args, params=params)
+
+    # return a cubic interpolation of the spectrum, provide access to
+    # ordering
+    band_sketching = BandSketching(x, y, dy, mode_function, tol)
+    band_sketching.ordering = ordering
+    return band_sketching
+
+
+def spectra(syst, *args, **kwargs):
+    """Interpolates a sequence of dispersion functions and provide methods to
+    simplify curve sketching and analyzation the periodic spectra.
+
+    Parameters
+    ----------
+    syst : :class:`kwant.system.InfiniteSystem` or sequence thereof
+        The low level infinite systems.
+
+    Returns
+    -------
+    specs : list of :class:`~kwant_spectrum.BandSketching`
+
+    Notes
+    -----
+    The function is similar to :class:`~kwant_spectrum.spectrum` but works as well
+    for a sequence of leads. It always returns a list of spectra.
+    See :class:`~kwant_spectrum.spectrum` for optional function arguments.
+    """
+    if not isinstance(syst, collections.abc.Iterable):
+        syst = [syst]
+    return [spectrum(sys_, *args, **kwargs) for sys_ in syst]
 
 
 class BandSketching: