WIP: Code update

This MR has a dual goal:

• Updating to current holoviews to avoid a bunch of deprecation errors
• Updating to Kwant 1.4:

• Eliminating args=[SimpleNamespace(...)] in favor of params={...}.
• Using the Builder.fill/Builder.substituted pattern when defining systems.

For now we do not intend to use discretizer yet.

code/functions.py

@@ -7,104 +7,143 @@ import kwant
 import numpy as np
 import holoviews as hv
 
-if tuple(int(i) for i in np.__version__.split('.')[:3]) <= (1, 8, 0):
-    raise RuntimeError("numpy >= (1, 8, 0) is required")
-
-__all__ = ['spectrum', 'hamiltonian_array', 'h_k', 'pauli']
-
-pauli = SimpleNamespace(s0=np.array([[1., 0.], [0., 1.]]),
-                        sx=np.array([[0., 1.], [1., 0.]]),
-                        sy=np.array([[0., -1j], [1j, 0.]]),
-                        sz=np.array([[1., 0.], [0., -1.]]))
-
-pauli.s0s0 = np.kron(pauli.s0, pauli.s0)
-pauli.s0sx = np.kron(pauli.s0, pauli.sx)
-pauli.s0sy = np.kron(pauli.s0, pauli.sy)
-pauli.s0sz = np.kron(pauli.s0, pauli.sz)
-pauli.sxs0 = np.kron(pauli.sx, pauli.s0)
-pauli.sxsx = np.kron(pauli.sx, pauli.sx)
-pauli.sxsy = np.kron(pauli.sx, pauli.sy)
-pauli.sxsz = np.kron(pauli.sx, pauli.sz)
-pauli.sys0 = np.kron(pauli.sy, pauli.s0)
-pauli.sysx = np.kron(pauli.sy, pauli.sx)
-pauli.sysy = np.kron(pauli.sy, pauli.sy)
-pauli.sysz = np.kron(pauli.sy, pauli.sz)
-pauli.szs0 = np.kron(pauli.sz, pauli.s0)
-pauli.szsx = np.kron(pauli.sz, pauli.sx)
-pauli.szsy = np.kron(pauli.sz, pauli.sy)
-pauli.szsz = np.kron(pauli.sz, pauli.sz)
-
-
-def spectrum(syst, p=None, k_x=None, k_y=None, k_z=None, title=None, xdim=None,
-             ydim=None, zdim=None, xticks=None, yticks=None, zticks=None,
-             xlims=None, ylims=None, zlims=None, num_bands=None, return_energies=False):
-    """Function that plots system spectrum for varying parameters or momenta.
+__all__ = ['spectrum', 'hamiltonian_array', 'pauli']
+
+pauli = dict(s0=np.array([[1., 0.], [0., 1.]]),
+             sx=np.array([[0., 1.], [1., 0.]]),
+             sy=np.array([[0., -1j], [1j, 0.]]),
+             sz=np.array([[1., 0.], [0., -1.]]))
+
+for first, second in itertools.product(pauli.items(), pauli.items()):
+    pauli[first[0] + second[0]] = np.kron(first[1], second[1])
+
+pauli = SimpleNamespace(**pauli)
+
+
+def hamiltonian_array(syst, params):
+    """Evaluate the Hamiltonian of a system over a grid of parameters.
+
+    Parameters
+    ----------
+    syst : kwant.System or a callable
+    params : dict
+        A container of Hamiltonian parameters. The parameters that are
+        sequences with length larger than 1 are used to loop over.
+
+    Returns
+    -------
+    hamiltonians : numpy.ndarray
+        An array with the Hamiltonians. The first n-2 dimensions correspond to
+        the expanded variables.
+
+    Examples
+    --------
+    >>> hamiltonian_array(syst, dict(t=1, mu=np.linspace(-2, 2),
+    ...                              k_x=np.linspace(-np.pi, np.pi)))
+
+    """
+    if isinstance(syst, kwant.system.System):
+        def ham(**kwargs):
+            return syst.hamiltonian_submatrix(params=kwargs, sparse=False)
+    else:
+        ham = syst
+
+    ham = np.vectorize(
+        ham,
+        signature = ','.join(['()'] * len(params)) + '->(n,n)'
+    )
+
+    return ham(**params)
+
+
+def spectrum(
+    syst, params, title=None,
+    num_bands=None, return_energies=False
+):
+    """Plot system spectrum for varying parameters or momenta.
 
     Parameters:
     -----------
-    syst : kwant.Builder object
-        The un-finalized (in)finite system.
-    p : SimpleNamespace object
-        A container used to store Hamiltonian parameters. The parameters that
-        are sequences are used as plot axes.
-    k_x, k_y, k_z : floats or sequences of floats
-        Real space momenta at which the Hamiltonian has to be evaluated.
-        If the system dimensionality is low, extra momenta are ignored.
+    syst : kwant.System or callable
+        kwant System or a function returning the Hamiltonian
+    params : dict
+        dictionary of parameters. The keys are either strings
+        or instances of `hv.Dimension` for better formatting.
+        Iterable values are expanded over and used as axes.
+        If the system expects extra momenta arguments, these are added.
     title : function or str
-        Function that takes p as argument and generates a string. If a string
-        it's used as static title.
-    xdim, ydim, zdim : holoviews.Dimension or string
-        The labels of the axes. Default to best guess, and extra ones
-        are ignored.
-    xticks, yticks zticks : list
-        Lists of axes xticks, extra ones are ignored.
-    xlims, ylims, zlims : tuple
-        Upper and lower plot limit of the axes. If None the upper and lower
-        limits of the ticks are used. Extra ones are ignored.
+        Function that takes params as argument and returns the title.
+        If a string it's used as a static title.
     num_bands : int
-        Number of bands that should be plotted, only works for 2D plots. If
-        None all bands are plotted.
+        Number of bands near the middle that should be plotted.
     return_energies : bool
-        If True the function only returns the energies in an array.
+        If True the function only returns the energies in an array
+        and do not produce a plot.
 
     Returns:
     --------
-    plot : holoviews.Path object
+    plot : an appropriate holoviews object
         Plot of varying parameter vs. spectrum.
     """
+    energy_dim = hv.Dimension(
+        name='energy',
+        label='$E$'
+    )
+
+    if params is None:
+        params = {}
+
+    # Add momenta to parameters
+    if not callable(syst):
+        momenta = {
+            hv.Dimension(
+                name=k, label=f'${k}$', range=(-np.pi, np.pi)
+            ): np.linspace(-np.pi, np.pi, 101)
+            for k in 'k_x k_y k_z'.split()
+            if k in syst.parameters and not k in params
+        }
+    else:
+        momenta = {}
+        
     pi_ticks = [(-np.pi, r'$-\pi$'), (0, '$0$'), (np.pi, r'$\pi$')]
-    if p is None:
-        p = SimpleNamespace()
-    dimensionality = syst.symmetry.num_directions
-    k = [k_x, k_y, k_z]
-    k = [(np.linspace(-np.pi, np.pi, 101) if i is None else i)
-         for i in k]
-    k = [(i if j < dimensionality else 0) for (j, i) in enumerate(k)]
-    k_x, k_y, k_z = k
-
-    hamiltonians, variables = hamiltonian_array(syst, p, k_x, k_y, k_z, True)
-    # Don't waste effort calculating eigenvalues if we aren't going to plot
-    # anything.
-    if len(variables) in (1, 2):
-        energies = np.linalg.eigvalsh(hamiltonians)
-
-    if len(variables) == 0:
-        raise ValueError("A 0D plot requested")
+
+    params = {**momenta, **params}
+
+    # Obtain the data
+    hamiltonians = hamiltonian_array(
+        syst, params={str(dim): value for dim, value in params.items()}
+    )
+
+    energies = np.linalg.eigvalsh(hamiltonians)
+
+    if num_bands is not None:
+        mid = energies.shape[-1] // 2
+        to_take = (
+            (slice(None),) * (len(energies.shape) - 1)
+            + (slice(mid - num_bands//2, mid + num_bands//2),)
+        )
+        energies = energies[to_take]
 
     if return_energies:
         return energies
 
-    elif len(variables) == 1:
-        # 1D plot.
-        if xdim is None:
-            if variables[0][0] in 'k_x k_y k_z'.split():
-                xdim = r'${}$'.format(variables[0][0])
-            else:
-                xdim = variables[0][0]
-        if ydim is None:
-            ydim = r'$E$'
+    # Separate the parameters into constant and varied
+    constants = {
+        k: v for k, v in params.items()
+        if not hasattr(v, "__len__")
+    }
+    variables = {
+        k: v for k, v in params.items()
+        if hasattr(v, "__len__")        
+    }
 
-        plot = hv.Path((variables[0][1], energies), kdims=[xdim, ydim])
+    # Actual plotting
+    if not variables:
+        raise ValueError("No variable parameters")
+
+    elif len(variables) == 1:
+        xdim, values = variables.popitem()
+        plot = hv.Path((values, energies), kdims=[xdim, energy_dim])
 
         ticks = {}
         if isinstance(xticks, collections.Iterable):
@@ -119,35 +158,32 @@ def spectrum(syst, p=None, k_x=None, k_y=None, k_z=None, title=None, xdim=None,
         elif isinstance(yticks, int):
             ticks['yticks'] = yticks
 
-        xlims = slice(*xlims) if xlims is not None else slice(None)
-        ylims = slice(*ylims) if ylims is not None else slice(None)
-
         if callable(title):
-            plot = plot.relabel(title(p))
+            plot = plot.relabel(title(**params))
         elif isinstance(title, str):
             plot = plot.relabel(title)
 
-        return plot[xlims, ylims](plot={'Path': ticks})
+        return plot[xlims, ylims].opts(plot={'Path': ticks})
 
     elif len(variables) == 2:
         # 2D plot.
         style = {}
-        if xticks is None and variables[0][0] in 'k_x k_y k_z'.split():
+        if xticks is None and variables[0][0] in momenta_keys:
             style['xticks'] = pi_ticks
         elif xticks is not None:
             style['xticks'] = list(xticks)
-        if yticks is None and variables[1][0] in 'k_x k_y k_z'.split():
+        if yticks is None and variables[1][0] in momenta_keys:
             style['yticks'] = pi_ticks
         elif yticks is not None:
             style['yticks'] = list(yticks)
 
         if xdim is None:
-            if variables[0][0] in 'k_x k_y k_z'.split():
+            if variables[0][0] in momenta_keys:
                 xdim = r'${}$'.format(variables[0][0])
             else:
                 xdim = variables[0][0]
         if ydim is None:
-            if variables[1][0] in 'k_x k_y k_z'.split():
+            if variables[1][0] in momenta_keys:
                 ydim = r'${}$'.format(variables[1][0])
             else:
                 ydim = variables[1][0]
@@ -157,148 +193,24 @@ def spectrum(syst, p=None, k_x=None, k_y=None, k_z=None, title=None, xdim=None,
         if zticks is not None:
             style['zticks'] = zticks
 
-        if xlims is None:
-            xlims = np.round([min(variables[0][1]), max(variables[0][1])], 2)
-        if ylims is None:
-            ylims = np.round([min(variables[1][1]), max(variables[1][1])], 2)
-        if zlims is None:
-            zlims = (None, None)
-
-        kwargs = {'extents': (xlims[0], ylims[0], zlims[0],
-                              xlims[1], ylims[1], zlims[1]),
-                  'kdims': [xdim, ydim],
-                  'vdims': [zdim]}
-
-        if num_bands is None:
-            plot = hv.Overlay([hv.Surface(energies[:, :, i], **kwargs)(plot=style) 
-                               for i in range(energies.shape[-1])])
-        else:
-            mid = energies.shape[-1] // 2
-            num_bands //= 2
-            plot = hv.Overlay([hv.Surface(energies[:, :, i], **kwargs)(plot=style) 
-                               for i in range(mid - num_bands, mid + num_bands)])
+        kwargs = {'kdims': [xdim, ydim], 'vdims': [zdim]}
 
-        if callable(title):
-            plot = plot.relabel(title(p))
-        elif isinstance(title, str):
-            plot = plot.relabel(title)
+        plot = hv.Overlay(
+            [
+                hv.Surface(energies[:, :, i], **kwargs)(plot=style)
+                for i in range(energies.shape[-1])
+            ]
+        )
 
-        return plot(plot={'Overlay': {'fig_size': 200}})
 
+        return plot.opts(plot={'Overlay': {'fig_size': 200}})
+    
     else:
         raise ValueError("Cannot make 4D plots yet.")
-
-
-
-def h_k(syst, p, momentum):
-    """Function that returns the Hamiltonian of a kwant 1D system as a momentum.
-    """
-    return hamiltonian_array(syst, p, momentum)[0]
-
-
-def hamiltonian_array(syst, p=None, k_x=0, k_y=0, k_z=0, return_grid=False):
-    """Evaluate the Hamiltonian of a system over a grid of parameters.
-
-    Parameters:
-    -----------
-    syst : kwant.Builder object
-        The un-finalized kwant system whose Hamiltonian is calculated.
-    p : SimpleNamespace object
-        A container of Hamiltonian parameters. The parameters that are
-        sequences are used to loop over.
-    k_x, k_y, k_z : floats or sequences of floats
-        Momenta at which the Hamiltonian has to be evaluated.  If the system
-        only has 1 translation symmetry, only `k_x` is used, and interpreted as
-        lattice momentum. Otherwise the momenta are in reciprocal space.
-    return_grid : bool
-        Whether to also return the names of the variables used for expansion,
-        and their values.
-
-    Returns:
-    --------
-    hamiltonians : numpy.ndarray
-        An array with the Hamiltonians. The first n-2 dimensions correspond to
-        the expanded variables.
-    parameters : list of tuples
-        Names and ranges of values that were used in evaluation of the
-        Hamiltonians.
-
-    Examples:
-    ---------
-    >>> hamiltonian_array(syst, SimpleNamespace(t=1, mu=np.linspace(-2, 2)),
-    ...                   k_x=np.linspace(-np.pi, np.pi))
-    >>> hamiltonian_array(sys_2d, p, np.linspace(-np.pi, np.pi),
-    ...                   np.linspace(-np.pi, np.pi))
-
-    """
-    if p is None:
-        p = SimpleNamespace()
-    try:
-        space_dimensionality = syst.symmetry.periods.shape[-1]
-    except AttributeError:
-        space_dimensionality = 0
-    dimensionality = syst.symmetry.num_directions
-    pars = copy(p)
-    if dimensionality == 0:
-        syst = syst.finalized()
-        def momentum_to_lattice(k):
-            return []
-    else:
-        if len(syst.symmetry.periods) == 1:
-            def momentum_to_lattice(k):
-                if any(k[dimensionality:]):
-                    raise ValueError("Dispersion is 1D, but more momenta are provided.")
-                return [k[0]]
-        else:
-            B = np.array(syst.symmetry.periods).T
-            A = B @ np.linalg.inv(B.T @ B)
-
-            def momentum_to_lattice(k):
-                k, residuals = np.linalg.lstsq(A, k[:space_dimensionality])[:2]
-                if np.any(abs(residuals) > 1e-7):
-                    raise RuntimeError("Requested momentum doesn't correspond"
-                                       " to any lattice momentum.")
-                return list(k)
-        syst = kwant.wraparound.wraparound(syst).finalized()
-
-    changing = dict()
-    for key, value in pars.__dict__.items():
-        if isinstance(value, collections.Iterable):
-            changing[key] = value
-
-    for key, value in [('k_x', k_x), ('k_y', k_y), ('k_z', k_z)]:
-        if key in changing:
-            raise RuntimeError('One of the system parameters is {}, '
-                               'which is reserved for momentum. '
-                               'Please rename it.'.format(key))
-        if isinstance(value, collections.Iterable):
-            changing[key] = value
-
-    if not changing:
-        hamiltonians = syst.hamiltonian_submatrix([pars] +
-                                                  momentum_to_lattice([k_x, k_y, k_z]),
-                                                  sparse=False)[None, ...]
-        if return_grid:
-            return hamiltonians, []
-        else:
-            return hamiltonians
-
-    def hamiltonian(**values):
-        pars.__dict__.update(values)
-        k = [values.get('k_x', k_x), values.get('k_y', k_y),
-             values.get('k_z', k_z)]
-        k = momentum_to_lattice(k)
-        return syst.hamiltonian_submatrix(args=([pars] + k), sparse=False)
-
-    names, values = zip(*sorted(changing.items()))
-    hamiltonians = [hamiltonian(**dict(zip(names, value)))
-                    for value in itertools.product(*values)]
-    size = list(hamiltonians[0].shape)
-
-    hamiltonians = np.array(hamiltonians).reshape([len(value)
-                                                   for value in values] + size)
-
-    if return_grid:
-        return hamiltonians, list(zip(names, values))
-    else:
-        return hamiltonians
+        
+    if title is not None:
+        plot = plot.relabel(
+            title(**params) if callable(title) else title
+        )
+        
+    return plot