convert operators example to jupyter-sphinx

doc/source/tutorial/operators.rst

@@ -13,8 +13,43 @@ texture.
 
 .. seealso::
     The complete source code of this example can be found in
-    :download:`magnetic_texture.py </code/download/magnetic_texture.py>`
-
+    :jupyter-download:script:`magnetic_texture`
+
+.. jupyter-kernel::
+    :id: magnetic_texture
+
+.. jupyter-execute::
+    :hide-code:
+
+    # Tutorial 2.7. Spin textures
+    # ===========================
+    #
+    # Physics background
+    # ------------------
+    #  - Spin textures
+    #  - Skyrmions
+    #
+    # Kwant features highlighted
+    # --------------------------
+    #  - operators
+    #  - plotting vector fields
+
+    from math import sin, cos, tanh, pi
+    import itertools
+    import numpy as np
+    import tinyarray as ta
+    import matplotlib.pyplot as plt
+
+    import kwant
+
+    sigma_0 = ta.array([[1, 0], [0, 1]])
+    sigma_x = ta.array([[0, 1], [1, 0]])
+    sigma_y = ta.array([[0, -1j], [1j, 0]])
+    sigma_z = ta.array([[1, 0], [0, -1]])
+
+    # vector of Pauli matrices σ_αiβ where greek
+    # letters denote spinor indices
+    sigma = np.rollaxis(np.array([sigma_x, sigma_y, sigma_z]), 1)
 
 Introduction
 ------------
@@ -47,28 +82,119 @@ of site :math:`i`, and :math:`r_i = \sqrt{x_i^2 + y_i^2}`.
 To define this model in Kwant we start as usual by defining functions that
 depend on the model parameters:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_model
-    :end-before: #HIDDEN_END_model
+.. jupyter-execute::
+
+    def field_direction(pos, r0, delta):
+        x, y = pos
+        r = np.linalg.norm(pos)
+        r_tilde = (r - r0) / delta
+        theta = (tanh(r_tilde) - 1) * (pi / 2)
+
+        if r == 0:
+            m_i = [0, 0, -1]
+        else:
+            m_i = [
+                (x / r) * sin(theta),
+                (y / r) * sin(theta),
+                cos(theta),
+            ]
+
+        return np.array(m_i)
+
+
+    def scattering_onsite(site, r0, delta, J):
+        m_i = field_direction(site.pos, r0, delta)
+        return J * np.dot(m_i, sigma)
+
+
+    def lead_onsite(site, J):
+        return J * sigma_z
 
 and define our system as a square shape on a square lattice with two orbitals
 per site, with leads attached on the left and right:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_syst
-    :end-before: #HIDDEN_END_syst
+.. jupyter-execute::
+
+    lat = kwant.lattice.square(norbs=2)
+
+    def make_system(L=80):
+
+        syst = kwant.Builder()
+
+        def square(pos):
+            return all(-L/2 < p < L/2 for p in pos)
+
+        syst[lat.shape(square, (0, 0))] = scattering_onsite
+        syst[lat.neighbors()] = -sigma_0
+
+        lead = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)),
+                             conservation_law=-sigma_z)
+
+        lead[lat.shape(square, (0, 0))] = lead_onsite
+        lead[lat.neighbors()] = -sigma_0
+
+        syst.attach_lead(lead)
+        syst.attach_lead(lead.reversed())
+
+        return syst
 
 Below is a plot of a projection of :math:`\mathbf{m}_i` onto the x-y plane
 inside the scattering region. The z component is shown by the color scale:
 
-.. image:: /code/figure/mag_field_direction.*
+.. jupyter-execute::
+    :hide-code:
+
+    def plot_vector_field(syst, params):
+        xmin, ymin = min(s.tag for s in syst.sites)
+        xmax, ymax = max(s.tag for s in syst.sites)
+        x, y = np.meshgrid(np.arange(xmin, xmax+1), np.arange(ymin, ymax+1))
+
+        m_i = [field_direction(p, **params) for p in zip(x.flat, y.flat)]
+        m_i = np.reshape(m_i, x.shape + (3,))
+        m_i = np.rollaxis(m_i, 2, 0)
+
+        fig, ax = plt.subplots(1, 1)
+        im = ax.quiver(x, y, *m_i, pivot='mid', scale=75)
+        fig.colorbar(im)
+        plt.show()
+
+
+    def plot_densities(syst, densities):
+        fig, axes = plt.subplots(1, len(densities), figsize=(13, 10))
+        for ax, (title, rho) in zip(axes, densities):
+            kwant.plotter.map(syst, rho, ax=ax, a=4)
+            ax.set_title(title)
+        plt.show()
+
+
+    def plot_currents(syst, currents):
+        fig, axes = plt.subplots(1, len(currents), figsize=(13, 10))
+        if not hasattr(axes, '__len__'):
+            axes = (axes,)
+        for ax, (title, current) in zip(axes, currents):
+            kwant.plotter.current(syst, current, ax=ax, colorbar=False,
+                                  fig_size=(13, 10))
+            ax.set_title(title)
+        plt.show()
+
+.. jupyter-execute::
+    :hide-code:
+
+    syst = make_system().finalized()
+
+.. jupyter-execute::
+    :hide-code:
+
+    plot_vector_field(syst, dict(r0=20, delta=10))
 
 We will now be interested in analyzing the form of the scattering states
 that originate from the left lead:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_wavefunction
-    :end-before: #HIDDEN_END_wavefunction
+.. jupyter-execute::
+
+    params = dict(r0=20, delta=10, J=1)
+    wf = kwant.wave_function(syst, energy=-1, params=params)
+    psi = wf(0)[0]
 
 Local densities
 ---------------
@@ -82,34 +208,45 @@ When there are multiple degrees of freedom per site, however, one has to be
 more careful. In the present case with two (spin) degrees of freedom per site
 one could calculate the per-site density like:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_ldos
-    :end-before: #HIDDEN_END_ldos
+.. jupyter-execute::
+
+    # even (odd) indices correspond to spin up (down)
+    up, down = psi[::2], psi[1::2]
+    density = np.abs(up)**2 + np.abs(down)**2
 
 With more than one degree of freedom per site we have more freedom as to what
 local quantities we can meaningfully compute. For example, we may wish to
 calculate the local z-projected spin density. We could calculate
 this in the following way:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_lsdz
-    :end-before: #HIDDEN_END_lsdz
+.. jupyter-execute::
+
+    # spin down components have a minus sign
+    spin_z = np.abs(up)**2 - np.abs(down)**2
 
 If we wanted instead to calculate the local y-projected spin density, we would
 need to use an even more complicated expression:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_lsdy
-    :end-before: #HIDDEN_END_lsdy
+.. jupyter-execute::
+
+    # spin down components have a minus sign
+    spin_y = 1j * (down.conjugate() * up - up.conjugate() * down)
 
 The `kwant.operator` module aims to alleviate somewhat this tedious
 book-keeping by providing a simple interface for defining operators that act on
 wavefunctions. To calculate the above quantities we would use the
 `~kwant.operator.Density` operator like so:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_lden
-    :end-before: #HIDDEN_END_lden
+.. jupyter-execute::
+
+    rho = kwant.operator.Density(syst)
+    rho_sz = kwant.operator.Density(syst, sigma_z)
+    rho_sy = kwant.operator.Density(syst, sigma_y)
+
+    # calculate the expectation values of the operators with 'psi'
+    density = rho(psi)
+    spin_z = rho_sz(psi)
+    spin_y = rho_sy(psi)
 
 `~kwant.operator.Density` takes a `~kwant.system.System` as its first parameter
 as well as (optionally) a square matrix that defines the quantity that you wish
@@ -126,8 +263,14 @@ Below we can see colorplots of the above-calculated quantities. The array that
 is returned by evaluating a `~kwant.operator.Density` can be used directly with
 `kwant.plotter.density`:
 
-.. image:: /code/figure/spin_densities.*
+.. jupyter-execute::
+    :hide-code:
 
+    plot_densities(syst, [
+        ('$σ_0$', density),
+        ('$σ_z$', spin_z),
+        ('$σ_y$', spin_y),
+    ])
 
 .. specialnote:: Technical Details
 
@@ -180,14 +323,27 @@ where :math:`\mathbf{H}_{ab}` is the hopping matrix from site :math:`b` to site
 :math:`a`.  For example, to calculate the local current and
 spin current:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_current
-    :end-before: #HIDDEN_END_current
+.. jupyter-execute::
+
+    J_0 = kwant.operator.Current(syst)
+    J_z = kwant.operator.Current(syst, sigma_z)
+    J_y = kwant.operator.Current(syst, sigma_y)
+
+    # calculate the expectation values of the operators with 'psi'
+    current = J_0(psi)
+    spin_z_current = J_z(psi)
+    spin_y_current = J_y(psi)
 
 Evaluating a `~kwant.operator.Current` operator on a wavefunction returns a
 1D array of values that can be directly used with `kwant.plotter.current`:
 
-.. image:: /code/figure/spin_currents.*
+.. jupyter-execute::
+
+    plot_currents(syst, [
+        ('$J_{σ_0}$', current),
+        ('$J_{σ_z}$', spin_z_current),
+        ('$J_{σ_y}$', spin_y_current),
+    ])
 
 .. note::
 
@@ -262,9 +418,16 @@ scattering region.
 Doing this is as simple as passing a *function* when instantiating
 the `~kwant.operator.Current`, instead of a constant matrix:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_following
-    :end-before: #HIDDEN_END_following
+.. jupyter-execute::
+
+    def following_m_i(site, r0, delta):
+        m_i = field_direction(site.pos, r0, delta)
+        return np.dot(m_i, sigma)
+
+    J_m = kwant.operator.Current(syst, following_m_i)
+
+    # evaluate the operator
+    m_current = J_m(psi, params=dict(r0=25, delta=10))
 
 The function must take a `~kwant.builder.Site` as its first parameter,
 and may optionally take other parameters (i.e. it must have the same
@@ -274,7 +437,7 @@ matrix that defines the operator we wish to calculate.
 .. note::
 
     In the above example we had to pass the extra parameters needed by the
-    ``following_operator`` function via the ``param`` keyword argument.  In
+    ``following_operator`` function via the ``params`` keyword argument.  In
     general you must pass all the parameters needed by the Hamiltonian via
     ``params`` (as you would when calling `~kwant.solvers.default.smatrix` or
     `~kwant.solvers.default.wave_function`).  In the previous examples,
@@ -286,7 +449,13 @@ Using this we can see that the spin current is essentially oriented along
 the direction of :math:`m_i` in the present regime where the onsite term
 in the Hamiltonian is dominant:
 
-.. image:: /code/figure/spin_current_comparison.*
+.. jupyter-execute::
+    :hide-code:
+
+    plot_currents(syst, [
+        (r'$J_{\mathbf{m}_i}$', m_current),
+        ('$J_{σ_z}$', spin_z_current),
+    ])
 
 .. note:: Although this example used exclusively `~kwant.operator.Current`,
           you can do the same with `~kwant.operator.Density`.
@@ -308,11 +477,16 @@ Density of states in a circle
 To calculate the density of states inside a circle of radius
 20 we can simply do:
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_density_cut
-    :end-before: #HIDDEN_END_density_cut
+.. jupyter-execute::
+
+    def circle(site):
+        return np.linalg.norm(site.pos) < 20
 
-.. literalinclude:: /code/figure/circle_dos.txt
+    rho_circle = kwant.operator.Density(syst, where=circle, sum=True)
+
+    all_states = np.vstack((wf(0), wf(1)))
+    dos_in_circle = sum(rho_circle(p) for p in all_states) / (2 * pi)
+    print('density of states in circle:', dos_in_circle)
 
 note that we also provide ``sum=True``, which means that evaluating the
 operator on a wavefunction will produce a single scalar. This is semantically
@@ -325,11 +499,22 @@ Current flowing through a cut
 Below we calculate the probability current and z-projected spin current near
 the interfaces with the left and right leads.
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_current_cut
-    :end-before: #HIDDEN_END_current_cut
+.. jupyter-execute::
+
+    def left_cut(site_to, site_from):
+        return site_from.pos[0] <= -39 and site_to.pos[0] > -39
+
+    def right_cut(site_to, site_from):
+        return site_from.pos[0] < 39 and site_to.pos[0] >= 39
 
-.. literalinclude:: /code/figure/current_cut.txt
+    J_left = kwant.operator.Current(syst, where=left_cut, sum=True)
+    J_right = kwant.operator.Current(syst, where=right_cut, sum=True)
+
+    Jz_left = kwant.operator.Current(syst, sigma_z, where=left_cut, sum=True)
+    Jz_right = kwant.operator.Current(syst, sigma_z, where=right_cut, sum=True)
+
+    print('J_left:', J_left(psi), ' J_right:', J_right(psi))
+    print('Jz_left:', Jz_left(psi), ' Jz_right:', Jz_right(psi))
 
 We see that the probability current is conserved across the scattering region,
 but the z-projected spin current is not due to the fact that the Hamiltonian
@@ -360,8 +545,23 @@ This can be achieved with `~kwant.operator.Current.bind`:
              *different* set of parameters. This will almost certainly give
              incorrect results.
 
-.. literalinclude:: /code/include/magnetic_texture.py
-    :start-after: #HIDDEN_BEGIN_bind
-    :end-before: #HIDDEN_END_bind
+.. jupyter-execute::
+
+    J_m = kwant.operator.Current(syst, following_m_i)
+    J_z = kwant.operator.Current(syst, sigma_z)
+
+    J_m_bound = J_m.bind(params=dict(r0=25, delta=10, J=1))
+    J_z_bound = J_z.bind(params=dict(r0=25, delta=10, J=1))
+
+    # Sum current local from all scattering states on the left at energy=-1
+    wf_left = wf(0)
+    J_m_left = sum(J_m_bound(p) for p in wf_left)
+    J_z_left = sum(J_z_bound(p) for p in wf_left)
+
+.. jupyter-execute::
+    :hide-code:
 
-.. image:: /code/figure/bound_current.*
+    plot_currents(syst, [
+        (r'$J_{\mathbf{m}_i}$ (from left)', J_m_left),
+        (r'$J_{σ_z}$ (from left)', J_z_left),
+    ])