QuLab 2.0.0__tar.gz → 2.0.1__tar.gz

{QuLab-2.0.0 → QuLab-2.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: QuLab
-Version: 2.0.0
+Version: 2.0.1
 Summary: contral instruments and manage data
 Author-email: feihoo87 <feihoo87@gmail.com>
 Maintainer-email: feihoo87 <feihoo87@gmail.com>

{QuLab-2.0.0 → QuLab-2.0.1}/QuLab.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: QuLab
-Version: 2.0.0
+Version: 2.0.1
 Summary: contral instruments and manage data
 Author-email: feihoo87 <feihoo87@gmail.com>
 Maintainer-email: feihoo87 <feihoo87@gmail.com>

{QuLab-2.0.0 → QuLab-2.0.1}/QuLab.egg-info/SOURCES.txt

@@ -6,19 +6,22 @@ setup.py
 QuLab.egg-info/PKG-INFO
 QuLab.egg-info/SOURCES.txt
 QuLab.egg-info/dependency_links.txt
+QuLab.egg-info/entry_points.txt
 QuLab.egg-info/requires.txt
 QuLab.egg-info/top_level.txt
 qulab/__init__.py
+qulab/__main__.py
 qulab/version.py
 qulab/monitor/__init__.py
-qulab/monitor/multiploter/__init__.py
-qulab/monitor/multiploter/config.py
-qulab/monitor/multiploter/dataset.py
-qulab/monitor/multiploter/event_queue.py
-qulab/monitor/multiploter/main.py
-qulab/monitor/multiploter/ploter.py
-qulab/monitor/multiploter/qt_compat.py
-qulab/monitor/multiploter/toolbar.py
+qulab/monitor/__main__.py
+qulab/monitor/config.py
+qulab/monitor/dataset.py
+qulab/monitor/event_queue.py
+qulab/monitor/mainwindow.py
+qulab/monitor/monitor.py
+qulab/monitor/ploter.py
+qulab/monitor/qt_compat.py
+qulab/monitor/toolbar.py
 qulab/scan/__init__.py
 qulab/scan/base.py
 qulab/scan/dataset.py
@@ -74,5 +77,12 @@ qulab/sys/rpc/server.py
 qulab/sys/rpc/socket.py
 qulab/sys/rpc/utils.py
 qulab/sys/rpc/worker.py
+qulab/visualization/__init__.py
+qulab/visualization/__main__.py
+qulab/visualization/_autoplot.py
+qulab/visualization/plot_layout.py
+qulab/visualization/plot_seq.py
+qulab/visualization/qdat.py
+qulab/visualization/widgets.py
 src/qulab.h
 tests/test_scan_iter.py

QuLab-2.0.1/QuLab.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+qulab = qulab.__main__:main

{QuLab-2.0.0 → QuLab-2.0.1}/pyproject.toml

@@ -41,6 +41,9 @@ dependencies = [
 ]
 dynamic = ["version"]
 
+[project.scripts]
+"qulab" = "qulab.__main__:main"
+
 [project.urls]
 Homepage = "https://github.com/feihoo87/QuLab"
 "Bug Reports" = "https://github.com/feihoo87/QuLab/issues"

QuLab-2.0.1/qulab/__main__.py

@@ -0,0 +1,24 @@
+import click
+
+from .monitor.__main__ import main as monitor
+from .sys.net.cli import dht
+from .visualization.__main__ import plot
+
+
+@click.group()
+def main():
+    pass
+
+
+@main.command()
+def hello():
+    """Print hello world."""
+    click.echo('hello, world')
+
+
+main.add_command(monitor)
+main.add_command(plot)
+main.add_command(dht)
+
+if __name__ == '__main__':
+    main()

QuLab-2.0.1/qulab/monitor/__init__.py

@@ -0,0 +1 @@
+from .monitor import Monitor, get_monitor

QuLab-2.0.1/qulab/monitor/__main__.py

@@ -0,0 +1,8 @@
+import click
+
+from .monitor import Monitor
+
+
+@click.command(name='monitor')
+def main():
+    pass

QuLab-2.0.0/qulab/monitor/__init__.py → QuLab-2.0.1/qulab/monitor/monitor.py

@@ -11,8 +11,8 @@ def main(queue: mp.Queue,
          ncols: int = 4,
          minimum_height: int = 400,
          colors: list[tuple[int, int, int]] = []):
-    from .multiploter import MainWindow
-    from .multiploter.qt_compat import QtWidgets
+    from .mainwindow import MainWindow
+    from .qt_compat import QtWidgets
 
     app = QtWidgets.QApplication(sys.argv)
     main = MainWindow(queue, ncols, minimum_height, colors)

{QuLab-2.0.0 → QuLab-2.0.1}/qulab/scan/base.py

@@ -1,4 +1,3 @@
-
 import inspect
 import logging
 import warnings
@@ -7,7 +6,6 @@ from concurrent.futures import Executor, Future
 from dataclasses import dataclass, field
 from graphlib import TopologicalSorter
 from itertools import chain, count
-
 from queue import Empty, Queue
 from typing import Any, Callable, Iterable, Sequence, Type
 
@@ -323,20 +321,27 @@ def _args_generator(loops: list,
             kw = _generate_kwds(keys, current_iters, kwds, i, limit)
         except StopIteration:
             break
+        if vars:
+            vars2 = [
+                *vars[:-1],
+                tuple([*vars[-1], *local_vars]),
+                tuple(kw.keys())
+            ]
+        else:
+            vars2 = [tuple([*local_vars, *kw.keys()])]
         yield Begin(level=level,
                     pos=pos + (i, ),
                     kwds=kwds | kw,
-                    vars=[*vars, tuple([*local_vars, *kw.keys()])],
+                    vars=vars2,
                     _pipes=pipes,
                     _trackers=trackers)
-        yield from _args_generator(
-            loops, kwds | kw, level + 1, pos + (i, ),
-            [*vars, tuple([*local_vars, *kw.keys()])], filter, functions,
-            trackers, pipes, order)
+        yield from _args_generator(loops, kwds | kw, level + 1, pos + (i, ),
+                                   vars2, filter, functions, trackers, pipes,
+                                   order)
         yield End(level=level,
                   pos=pos + (i, ),
                   kwds=kwds | kw,
-                  vars=[*vars, tuple([*local_vars, *kw.keys()])],
+                  vars=vars2,
                   _pipes=pipes,
                   _trackers=trackers)
         _feedback(current_iters)
@@ -541,4 +546,3 @@ def scan_iters(loops: dict[str | tuple[str, ...],
         step.unchanged = i
         yield step
         last_step = step
-

{QuLab-2.0.0 → QuLab-2.0.1}/qulab/scan/expression.py

@@ -8,6 +8,7 @@ from pyparsing import (CaselessKeyword, Combine, Forward, Group, Keyword,
                        alphanums, alphas, delimitedList, nums, oneOf, opAssoc,
                        pyparsing_common, restOfLine, srange, stringEnd,
                        stringStart)
+from scipy import special
 
 LPAREN, RPAREN, LBRACK, RBRACK, LBRACE, RBRACE, DOT, TILDE, BANG, PLUS, MINUS = map(
     Suppress, "()[]{}.~!+-")
@@ -60,15 +61,45 @@ class Env():
         self.consts = {}
         self.variables = {}
         self.refs = {}
+        self.functions = {
+            'sin': np.sin,
+            'cos': np.cos,
+            'tan': np.tan,
+            'pi': np.pi,
+            'e': np.e,
+            'log': np.log,
+            'log2': np.log2,
+            'log10': np.log10,
+            'exp': np.exp,
+            'sqrt': np.sqrt,
+            'abs': np.abs,
+            'sinh': np.sinh,
+            'cosh': np.cosh,
+            'tanh': np.tanh,
+            'arcsin': np.arcsin,
+            'arccos': np.arccos,
+            'arctan': np.arctan,
+            'arctan2': np.arctan2,
+            'arcsinh': np.arcsinh,
+            'arccosh': np.arccosh,
+            'arctanh': np.arctanh,
+            'sinc': np.sinc,
+            'sign': np.sign,
+            'heaviside': np.heaviside,
+            'erf': special.erf,
+            'erfc': special.erfc,
+        }
 
     def __contains__(self, key):
-        return key in self.consts or key in self.variables or key in self.refs
+        return key in self.consts or key in self.variables or key in self.functions or key in self.refs
 
     def __getitem__(self, key):
         if key in self.consts:
             return self.consts[key]
         if key in self.variables:
             return self.variables[key]
+        if key in self.functions:
+            return self.functions[key]
         if key in self.refs:
             return self[self.refs[key]]
         raise KeyError(f"Key {key} not found")
@@ -110,6 +141,9 @@ class Env():
         return key in self.consts
 
 
+_default_env = Env()
+
+
 class Expression():
 
     def __init__(self):
@@ -127,33 +161,85 @@ class Expression():
         raise NotImplementedError
 
     def __add__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return self
         return BinaryExpression(self, other, operator.add)
 
     def __radd__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return self
         return BinaryExpression(other, self, operator.add)
 
     def __sub__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return self
         return BinaryExpression(self, other, operator.sub)
 
     def __rsub__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return -self
         return BinaryExpression(other, self, operator.sub)
 
     def __mul__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return 0
+        if isinstance(other, ConstType) and other == 1:
+            return self
+        if isinstance(other, ConstType) and other == -1:
+            return -self
         return BinaryExpression(self, other, operator.mul)
 
     def __rmul__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return 0
+        if isinstance(other, ConstType) and other == 1:
+            return self
+        if isinstance(other, ConstType) and other == -1:
+            return -self
         return BinaryExpression(other, self, operator.mul)
 
     def __truediv__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 1:
+            return self
+        if isinstance(other, ConstType) and other == -1:
+            return -self
         return BinaryExpression(self, other, operator.truediv)
 
     def __rtruediv__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return 0
         return BinaryExpression(other, self, operator.truediv)
 
     def __pow__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return 1
+        if isinstance(other, ConstType) and other == 1:
+            return self
         return BinaryExpression(self, other, operator.pow)
 
     def __rpow__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
+        if isinstance(other, ConstType) and other == 0:
+            return 0
         return BinaryExpression(other, self, operator.pow)
 
     def __neg__(self):
@@ -163,32 +249,52 @@ class Expression():
         return UnaryExpression(self, operator.pos)
 
     def __eq__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return BinaryExpression(self, other, operator.eq)
 
     def __ne__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return BinaryExpression(self, other, operator.ne)
 
     def __lt__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return BinaryExpression(self, other, operator.lt)
 
     def __le__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return BinaryExpression(self, other, operator.le)
 
     def __gt__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return BinaryExpression(self, other, operator.gt)
 
     def __ge__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return BinaryExpression(self, other, operator.ge)
 
     def __getitem__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return ObjectMethod(self, '__getitem__', other)
 
     def __getattr__(self, other):
+        if isinstance(other, Expression):
+            other = other.eval(_default_env)
         return ObjectMethod(self, '__getattr__', other)
 
     def __call__(self, *args):
+        args = [
+            o.eval(_default_env) if isinstance(o, Expression) else o
+            for o in args
+        ]
         return ObjectMethod(self, '__call__', *args)
-    
+
     def __round__(self, n=None):
         return self
 
@@ -243,7 +349,7 @@ class UnaryExpression(Expression):
             return self.op(self.a.d(x))
         else:
             return 0
-        
+
     def __repr__(self) -> str:
         return f"{self.op.__name__}({self.a!r})"
 
@@ -300,7 +406,7 @@ class BinaryExpression(Expression):
                 return 0
         else:
             return 0
-        
+
     def __repr__(self) -> str:
         return f"({self.a!r} {self.op.__name__} {self.b!r})"
 
@@ -333,6 +439,12 @@ class ObjectMethod(Expression):
             return ObjectMethod(obj, self.method, *args)
         else:
             return getattr(obj, self.method)(*args)
+        
+    def __repr__(self):
+        if self.method == '__call__':
+            return f"{self.obj!r}({', '.join(map(repr, self.args))})"
+        else:
+            return f"{self.obj!r}.{self.method}({', '.join(map(repr, self.args))})"
 
 
 class Symbol(Expression):
@@ -355,6 +467,6 @@ class Symbol(Expression):
             return 1
         else:
             return 0
-    
+
     def __repr__(self) -> str:
         return self.name

{QuLab-2.0.0 → QuLab-2.0.1}/qulab/scan/scanner.py

@@ -1,4 +1,5 @@
 import ast
+import itertools
 
 import numpy as np
 
@@ -73,6 +74,19 @@ class Optimizer():
                                                        **self.kwds)
 
 
+class Action():
+
+    __slots__ = ('name', 'args', 'kwds')
+
+    def __init__(self, name, *args, **kwds):
+        self.name = name
+        self.args = args
+        self.kwds = kwds
+
+    def __repr__(self):
+        return f"Action({self.name!r}, {self.args}, {self.kwds})"
+
+
 class Scan():
 
     def __new__(cls, *args, mixin=None, **kwds):
@@ -99,6 +113,8 @@ class Scan():
         self.filter = None
         self.scan_info = {'loops': {}}
         self._tmp = {}
+        self.actions = {}
+        self.iteration = 0
 
     @property
     def name(self):
@@ -114,12 +130,12 @@ class Scan():
         return default
 
     def set(self, key, value):
-        print(f"Set {key} to {value}")
+        self.add_action(Action('set', key, value))
 
     def _mapping(self, key, value):
         tmpkey = f"__tmp_{self._mapping_i}__"
         self._mapping_i += 1
-        self[tmpkey] = value
+        self.__setitem(tmpkey, value)
         self.mapping[key] = tmpkey
 
     def __setitem__(self, key, value):
@@ -219,17 +235,27 @@ class Scan():
         self.assemble()
         for step in self.scan():
             for k, v in self.mapping.items():
-                self.set(k, step.kwds[v])
+                if v in set(
+                        itertools.chain.from_iterable(
+                            step.vars[step.unchanged:])
+                ) or step.iteration == 0:
+                    self.set(k, step.kwds[v])
             self.process(step)
 
     def process(self, step):
-        print(step.kwds)
+        for k, v in step.kwds.items():
+            if not k.startswith('__tmp_') and not k.startswith('__'):
+                self.add_action(Action('write', k, v))
 
     def scan(self):
         for step in scan_iters(**self.scan_info):
             for k, v in self.mapping.items():
                 step.kwds[k] = step.kwds[v]
             yield step
+            self.iteration += 1
+
+    def add_action(self, action: Action):
+        self.actions.setdefault(self.iteration, []).append(action)
 
     def run(self, dry_run=False):
         pass

QuLab-2.0.1/qulab/version.py

@@ -0,0 +1 @@
+__version__ = "2.0.1"

QuLab-2.0.1/qulab/visualization/__init__.py

@@ -0,0 +1,188 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ._autoplot import autoplot
+
+
+def plotLine(c0, c1, ax, **kwargs):
+    t = np.linspace(0, 1, 11)
+    c = (c1 - c0) * t + c0
+    ax.plot(c.real, c.imag, **kwargs)
+
+
+def plotCircle(c0, r, ax, **kwargs):
+    t = np.linspace(0, 1, 1001) * 2 * np.pi
+    s = c0 + r * np.exp(1j * t)
+    ax.plot(s.real, s.imag, **kwargs)
+
+
+def plotEllipse(c0, a, b, phi, ax, **kwargs):
+    t = np.linspace(0, 1, 1001) * 2 * np.pi
+    c = np.exp(1j * t)
+    s = c0 + (c.real * a + 1j * c.imag * b) * np.exp(1j * phi)
+    ax.plot(s.real, s.imag, **kwargs)
+
+
+def plotDistribution(s0,
+                     s1,
+                     fig=None,
+                     axes=None,
+                     info=None,
+                     hotThresh=10000,
+                     logy=False):
+    from waveforms.math.fit import get_threshold_info, mult_gaussian_pdf
+
+    if info is None:
+        info = get_threshold_info(s0, s1)
+    else:
+        info = get_threshold_info(s0, s1, info['threshold'], info['phi'])
+    thr, phi = info['threshold'], info['phi']
+    # visibility, p0, p1 = info['visibility']
+    # print(
+    #     f"thr={thr:.6f}, phi={phi:.6f}, visibility={visibility:.3f}, {p0}, {1-p1}"
+    # )
+
+    if axes is not None:
+        ax1, ax2 = axes
+    else:
+        if fig is None:
+            fig = plt.figure()
+        ax1 = fig.add_subplot(121)
+        ax2 = fig.add_subplot(122)
+
+    if (len(s0) + len(s1)) < hotThresh:
+        ax1.plot(np.real(s0), np.imag(s0), '.', alpha=0.2)
+        ax1.plot(np.real(s1), np.imag(s1), '.', alpha=0.2)
+    else:
+        _, *bins = np.histogram2d(np.real(np.hstack([s0, s1])),
+                                  np.imag(np.hstack([s0, s1])),
+                                  bins=50)
+
+        H0, *_ = np.histogram2d(np.real(s0),
+                                np.imag(s0),
+                                bins=bins,
+                                density=True)
+        H1, *_ = np.histogram2d(np.real(s1),
+                                np.imag(s1),
+                                bins=bins,
+                                density=True)
+        vlim = max(np.max(np.abs(H0)), np.max(np.abs(H1)))
+
+        ax1.imshow(H1.T - H0.T,
+                   alpha=(np.fmax(H0.T, H1.T) / vlim).clip(0, 1),
+                   interpolation='nearest',
+                   origin='lower',
+                   cmap='coolwarm',
+                   vmin=-vlim,
+                   vmax=vlim,
+                   extent=(bins[0][0], bins[0][-1], bins[1][0], bins[1][-1]))
+
+    ax1.axis('equal')
+    ax1.set_xticks([])
+    ax1.set_yticks([])
+    for s in ax1.spines.values():
+        s.set_visible(False)
+
+    # c0, c1 = info['center']
+    # a0, b0, a1, b1 = info['std']
+    params = info['params']
+    r0, i0, r1, i1 = params[0][0], params[1][0], params[0][1], params[1][1]
+    a0, b0, a1, b1 = params[0][2], params[1][2], params[0][3], params[1][3]
+    c0 = (r0 + 1j * i0) * np.exp(1j * phi)
+    c1 = (r1 + 1j * i1) * np.exp(1j * phi)
+    phi0 = phi + params[0][6]
+    phi1 = phi + params[1][6]
+    plotEllipse(c0, 2 * a0, 2 * b0, phi0, ax1)
+    plotEllipse(c1, 2 * a1, 2 * b1, phi1, ax1)
+
+    im0, im1 = info['idle']
+    lim = min(im0.min(), im1.min()), max(im0.max(), im1.max())
+    t = (np.linspace(lim[0], lim[1], 3) + 1j * thr) * np.exp(-1j * phi)
+    ax1.plot(t.imag, t.real, 'k--')
+
+    ax1.plot(np.real(c0), np.imag(c0), 'o', color='C3')
+    ax1.plot(np.real(c1), np.imag(c1), 'o', color='C4')
+
+    re0, re1 = info['signal']
+    x, a, b, c = info['cdf']
+
+    xrange = (min(re0.min(), re1.min()), max(re0.max(), re1.max()))
+
+    n0, bins0, *_ = ax2.hist(re0, bins=80, range=xrange, alpha=0.5)
+    n1, bins1, *_ = ax2.hist(re1, bins=80, range=xrange, alpha=0.5)
+
+    x_range = np.linspace(x.min(), x.max(), 1001)
+    *_, cov0, cov1 = info['std']
+    ax2.plot(
+        x_range,
+        np.sum(n0) * (bins0[1] - bins0[0]) *
+        mult_gaussian_pdf(x_range, [r0, r1], [
+            np.sqrt(cov0[0, 0]), np.sqrt(cov1[0, 0])
+        ], [params[0][4], 1 - params[0][4]]))
+    ax2.plot(
+        x_range,
+        np.sum(n1) * (bins1[1] - bins1[0]) *
+        mult_gaussian_pdf(x_range, [r0, r1], [
+            np.sqrt(cov0[0, 0]), np.sqrt(cov1[0, 0])
+        ], [params[0][5], 1 - params[0][5]]))
+    ax2.set_ylabel('Count')
+    ax2.set_xlabel('Projection Axes')
+    if logy:
+        ax2.set_yscale('log')
+        ax2.set_ylim(0.1, max(np.sum(n0), np.sum(n1)))
+
+    ax3 = ax2.twinx()
+    ax3.plot(x, a, '--', lw=1, color='C0')
+    ax3.plot(x, b, '--', lw=1, color='C1')
+    ax3.plot(x, c, 'k--', alpha=0.5, lw=1)
+    ax3.set_ylim(0, 1.1)
+    ax3.vlines(thr, 0, 1.1, 'k', alpha=0.5)
+    ax3.set_ylabel('Integral Probability')
+
+    return info
+
+
+ALLXYSeq = [('I', 'I'), ('X', 'X'), ('Y', 'Y'), ('X', 'Y'), ('Y', 'X'),
+            ('X/2', 'I'), ('Y/2', 'I'), ('X/2', 'Y/2'), ('Y/2', 'X/2'),
+            ('X/2', 'Y'), ('Y/2', 'X'), ('X', 'Y/2'), ('Y', 'X/2'),
+            ('X/2', 'X'), ('X', 'X/2'), ('Y/2', 'Y'), ('Y', 'Y/2'), ('X', 'I'),
+            ('Y', 'I'), ('X/2', 'X/2'), ('Y/2', 'Y/2')]
+
+
+def plotALLXY(data, ax=None):
+    assert len(data) % len(ALLXYSeq) == 0
+
+    if ax is None:
+        ax = plt.gca()
+
+    ax.plot(np.array(data), 'o-')
+    repeat = len(data) // len(ALLXYSeq)
+    ax.set_xticks(np.arange(len(ALLXYSeq)) * repeat + 0.5 * (repeat - 1))
+    ax.set_xticklabels([','.join(seq) for seq in ALLXYSeq], rotation=60)
+    ax.grid(which='major')
+
+
+def plot_mat(rho, title='$\\chi$', cmap='coolwarm'):
+    lim = np.abs(rho).max()
+    N = rho.shape[0]
+
+    fig = plt.figure(figsize=(6, 4))
+    fig.suptitle(title)
+
+    ax1 = plt.subplot(121)
+    cax1 = ax1.imshow(rho.real, vmin=-lim, vmax=lim, cmap=cmap)
+    ax1.set_title('Re')
+    ax1.set_xticks(np.arange(N))
+    ax1.set_yticks(np.arange(N))
+
+    ax2 = plt.subplot(122)
+    cax2 = ax2.imshow(rho.imag, vmin=-lim, vmax=lim, cmap=cmap)
+    ax2.set_title('Im')
+    ax2.set_xticks(np.arange(N))
+    ax2.set_yticks(np.arange(N))
+
+    plt.subplots_adjust(bottom=0.2, right=0.9, top=0.95)
+
+    cbar_ax = fig.add_axes([0.15, 0.15, 0.7, 0.05])
+    cb = fig.colorbar(cax1, cax=cbar_ax, orientation='horizontal')
+    plt.show()