QuLab 2.0.1__cp310-cp310-macosx_11_0_x86_64.whl

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()

qulab/visualization/__main__.py

@@ -0,0 +1,71 @@
+import pathlib
+import pickle
+
+import click
+import dill
+import matplotlib.pyplot as plt
+import numpy as np
+
+from .qdat import draw as draw_qdat
+
+default_draw_methods = {
+    '.qdat': draw_qdat,
+}
+
+
+def load_data(fname):
+    try:
+        from home.hkxu.tools import get_record_by_id
+        record_id = int(str(fname))
+        return get_record_by_id(record_id).data
+    except:
+        pass
+    with open(fname, 'rb') as f:
+        try:
+            data = pickle.load(f)
+        except:
+            f.seek(0)
+            data = dill.load(f)
+    return data
+
+
+def draw_common(data):
+    try:
+        script = data['meta']['plot_script']
+        assert script.strip()
+        global_namespace = {'plt': plt, 'np': np, 'result': data}
+        exec(script, global_namespace)
+    except:
+        from home.hkxu.tools import plot_record
+        plot_record(data['meta']['id'])
+
+
+def draw_error(data, text="No validate plot script found"):
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.text(0.5, 0.5, text, ha='center', va='center')
+    ax.set_axis_off()
+    return fig
+
+
+@click.command()
+@click.argument('fname', default='')
+def plot(fname):
+    """Plot the data in the file."""
+    try:
+        fname = pathlib.Path(fname)
+        data = load_data(fname)
+        try:
+            draw_common(data)
+        except:
+            default_draw_methods.get(fname.suffix, draw_error)(data)
+    except FileNotFoundError:
+        draw_error(None, text=f"File {fname} not found.")
+    except pickle.UnpicklingError:
+        draw_error(None, text=f"File {fname} is not a pickle file.")
+
+    plt.show()
+
+
+if __name__ == '__main__':
+    plot()

qulab/visualization/_autoplot.py

@@ -0,0 +1,457 @@
+import math
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import LogNorm
+from matplotlib.ticker import EngFormatter, LogFormatterSciNotation
+from scipy.interpolate import griddata
+
+
+def good_for_logscale(x, threshold=4):
+    if np.any(x <= 0):
+        return False
+    mid = (np.nanmin(x) + np.nanmax(x)) / 2
+    a = np.count_nonzero(np.nan_to_num(x <= mid, nan=0))
+    b = np.count_nonzero(np.nan_to_num(x >= mid, nan=0))
+    if a / b > threshold:
+        return True
+    return False
+
+
+def equal_logspace(x):
+    logx = np.logspace(np.log10(x[0]), np.log10(x[-1]), len(x))
+    return np.allclose(x, logx)
+
+
+def equal_linspace(x):
+    linearx = np.linspace(x[0], x[-1], len(x))
+    return np.allclose(x, linearx)
+
+
+def as_1d_data(x, y, z):
+    x = np.asarray(x)
+    y = np.asarray(y)
+    z = np.asarray(z)
+    if z.ndim == 1:
+        return x, y, z
+
+    if z.ndim == 2:
+        x, y = np.meshgrid(x, y)
+        return x.ravel(), y.ravel(), z.ravel()
+
+    raise ValueError("z must be 1D or 2D")
+
+
+def griddata_logx_logy(x, y, z, shape=(401, 401)):
+    x, y, z = as_1d_data(x, y, z)
+    xspace = np.logspace(np.log10(x.min()), np.log10(x.max()), shape[0])
+    yspace = np.logspace(np.log10(y.min()), np.log10(y.max()), shape[1])
+    xgrid, ygrid = np.meshgrid(xspace, yspace)
+    zgrid = griddata((x, y), z, (xgrid, ygrid), method='nearest')
+    return xspace, yspace, zgrid
+
+
+def griddata_logx_linear_y(x, y, z, shape=(401, 401)):
+    x, y, z = as_1d_data(x, y, z)
+    xspace = np.logspace(np.log10(x.min()), np.log10(x.max()), shape[0])
+    yspace = np.linspace(y.min(), y.max(), shape[1])
+    xgrid, ygrid = np.meshgrid(xspace, yspace)
+    zgrid = griddata((x, y), z, (xgrid, ygrid), method='nearest')
+    return xspace, yspace, zgrid
+
+
+def griddata_linear_x_logy(x, y, z, shape=(401, 401)):
+    x, y, z = as_1d_data(x, y, z)
+    xspace = np.linspace(x.min(), x.max(), shape[0])
+    yspace = np.logspace(np.log10(y.min()), np.log10(y.max()), shape[1])
+    xgrid, ygrid = np.meshgrid(xspace, yspace)
+    zgrid = griddata((x, y), z, (xgrid, ygrid), method='nearest')
+    return xspace, yspace, zgrid
+
+
+def griddata_linear_x_linear_y(x, y, z, shape=(401, 401)):
+    x, y, z = as_1d_data(x, y, z)
+    xspace = np.linspace(x.min(), x.max(), shape[0])
+    yspace = np.linspace(y.min(), y.max(), shape[1])
+    xgrid, ygrid = np.meshgrid(xspace, yspace)
+    zgrid = griddata((x, y), z, (xgrid, ygrid), method='nearest')
+    return xspace, yspace, zgrid
+
+
+def _get_log_ticks(x):
+    log10x = np.log10(x)
+
+    major_ticks = np.array(
+        range(math.floor(log10x[0]) - 1,
+              math.ceil(log10x[-1]) + 1))
+    minor_ticks = np.hstack([
+        np.log10(np.linspace(2, 10, 9, endpoint=False)) + x
+        for x in major_ticks
+    ])
+
+    major_ticks = major_ticks[(major_ticks >= log10x[0]) *
+                              (major_ticks <= log10x[-1])]
+    minor_ticks = minor_ticks[(minor_ticks >= log10x[0]) *
+                              (minor_ticks <= log10x[-1])]
+
+    return log10x, major_ticks, minor_ticks
+
+
+class MyLogFormatter(EngFormatter):
+
+    def format_ticks(self, values):
+        if self.unit is None or self.unit == '':
+            fmt = LogFormatterSciNotation()
+            return [f"${fmt.format_data(10.0**x)}$" for x in values]
+        else:
+            return super().format_ticks(values)
+
+    def format_eng(self, x):
+        if self.unit is None or self.unit == '':
+            self.unit = ''
+            return f"{10.0**x:g}"
+        else:
+            return super().format_eng(10.0**x)
+
+
+def imshow_logx(x, y, z, x_unit=None, ax=None, **kwargs):
+    if ax is None:
+        ax = plt.gca()
+
+    log10x, major_ticks, minor_ticks = _get_log_ticks(x)
+
+    dlogx, dy = log10x[1] - log10x[0], y[1] - y[0]
+    extent = (log10x[0] - dlogx / 2, log10x[-1] + dlogx / 2, y[0] - dy / 2,
+              y[-1] + dy / 2)
+
+    img = ax.imshow(z, extent=extent, **kwargs)
+
+    ax.set_xticks(major_ticks, minor=False)
+    ax.xaxis.set_major_formatter(MyLogFormatter(x_unit))
+    ax.set_xticks(minor_ticks, minor=True)
+
+    return img
+
+
+def imshow_logy(x, y, z, y_unit=None, ax=None, **kwargs):
+    if ax is None:
+        ax = plt.gca()
+
+    log10y, major_ticks, minor_ticks = _get_log_ticks(y)
+
+    dlogy, dx = log10y[1] - log10y[0], x[1] - x[0]
+    extent = (x[0] - dx / 2, x[-1] + dx / 2, log10y[0] - dlogy / 2,
+              log10y[-1] + dlogy / 2)
+
+    img = ax.imshow(z, extent=extent, **kwargs)
+
+    ax.set_yticks(major_ticks, minor=False)
+    ax.yaxis.set_major_formatter(MyLogFormatter(y_unit))
+    ax.set_yticks(minor_ticks, minor=True)
+
+    return img
+
+
+def imshow_loglog(x, y, z, x_unit=None, y_unit=None, ax=None, **kwargs):
+    if ax is None:
+        ax = plt.gca()
+
+    log10x, x_major_ticks, x_minor_ticks = _get_log_ticks(x)
+    log10y, y_major_ticks, y_minor_ticks = _get_log_ticks(y)
+
+    dlogx, dlogy = log10x[1] - log10x[0], log10y[1] - log10y[0]
+    extent = (log10x[0] - dlogx / 2, log10x[-1] + dlogx / 2,
+              log10y[0] - dlogy / 2, log10y[-1] + dlogy / 2)
+
+    img = ax.imshow(z, extent=extent, **kwargs)
+
+    ax.set_xticks(x_major_ticks, minor=False)
+    ax.xaxis.set_major_formatter(MyLogFormatter(x_unit))
+    ax.set_xticks(x_minor_ticks, minor=True)
+
+    ax.set_yticks(y_major_ticks, minor=False)
+    ax.yaxis.set_major_formatter(MyLogFormatter(y_unit))
+    ax.set_yticks(y_minor_ticks, minor=True)
+
+    return img
+
+
+def plot_lines(x,
+               y,
+               z,
+               xlabel,
+               ylabel,
+               zlabel,
+               x_unit,
+               y_unit,
+               z_unit,
+               ax,
+               xscale='linear',
+               yscale='linear',
+               zscale='linear',
+               index=None,
+               **kwds):
+    z = np.asarray(z)
+    if len(y) > len(x):
+        x, y = y, x
+        xlabel, ylabel = ylabel, xlabel
+        xscale, yscale = yscale, xscale
+        z = z.T
+    if index is not None:
+        y = y[index]
+        z = z[index, :]
+
+    for i, l in enumerate(y):
+        if y_unit:
+            label = f"{ylabel}={l:.3} [{y_unit}]"
+        else:
+            if isinstance(l, float):
+                label = f"{ylabel}={l:.3}"
+            else:
+                label = f"{ylabel}={l}"
+        ax.plot(x, z[i, :], label=label, **kwds)
+    ax.legend()
+    xlabel = f"{xlabel} [{x_unit}]" if x_unit else xlabel
+    zlabel = f"{zlabel} [{z_unit}]" if z_unit else zlabel
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(zlabel)
+    ax.set_xscale(xscale)
+    ax.set_yscale(zscale)
+
+
+def plot_img(x,
+             y,
+             z,
+             xlabel,
+             ylabel,
+             zlabel,
+             x_unit,
+             y_unit,
+             z_unit,
+             fig,
+             ax,
+             xscale='linear',
+             yscale='linear',
+             zscale='linear',
+             resolution=None,
+             **kwds):
+    kwds.setdefault('origin', 'lower')
+    kwds.setdefault('aspect', 'auto')
+    kwds.setdefault('interpolation', 'nearest')
+
+    if zscale == 'log':
+        vmim = kwds.get('vmin', np.min(z))
+        vmax = kwds.get('vmax', np.max(z))
+        kwds.setdefault('norm', LogNorm(vmax=vmax, vmin=vmim))
+    zlabel = f"{zlabel} [{z_unit}]" if z_unit else zlabel
+
+    band_area = False
+    if x.ndim == 1 and y.ndim == 2 and y.shape[1] == x.shape[0]:
+        x = np.asarray([x] * y.shape[0])
+        band_area = True
+    elif x.ndim == 2 and y.ndim == 1 and x.shape[0] == y.shape[0]:
+        y = np.asarray([y] * x.shape[1]).T
+        band_area = True
+    if band_area:
+        kwds.pop('origin', None)
+        kwds.pop('aspect', None)
+        kwds.pop('interpolation', None)
+        pc = ax.pcolormesh(x, y, z, **kwds)
+        xlabel = f"{xlabel} [{x_unit}]" if x_unit else xlabel
+        ylabel = f"{ylabel} [{y_unit}]" if y_unit else ylabel
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel(ylabel)
+        cb = fig.colorbar(pc, ax=ax)
+        ax.set_xscale(xscale)
+        ax.set_yscale(yscale)
+        cb.set_label(zlabel)
+        return
+
+    if resolution is None:
+        resolution = (401, 401)
+    elif isinstance(resolution, int):
+        resolution = (resolution, resolution)
+
+    if (z.ndim == 1 or (xscale == 'linear' and not equal_linspace(x))
+            or (yscale == 'linear' and not equal_linspace(y))
+            or (xscale == 'log' and not equal_logspace(x))
+            or (yscale == 'log' and not equal_logspace(y))):
+        griddata = {
+            ('log', 'log'): griddata_logx_logy,
+            ('log', 'linear'): griddata_logx_linear_y,
+            ('linear', 'log'): griddata_linear_x_logy,
+            ('linear', 'linear'): griddata_linear_x_linear_y,
+        }[(xscale, yscale)]
+        x, y, z = griddata(x, y, z, resolution)
+
+    if (xscale, yscale) == ('linear', 'linear'):
+        dx, dy = x[1] - x[0], y[1] - y[0]
+        extent = (x[0] - dx / 2, x[-1] + dx / 2, y[0] - dy / 2, y[-1] + dy / 2)
+        kwds.setdefault('extent', extent)
+        img = ax.imshow(np.asarray(z), **kwds)
+        xlabel = f"{xlabel} [{x_unit}]" if x_unit else xlabel
+        ylabel = f"{ylabel} [{y_unit}]" if y_unit else ylabel
+    elif (xscale, yscale) == ('log', 'linear'):
+        ylabel = f"{ylabel} [{y_unit}]" if y_unit else ylabel
+        img = imshow_logx(x, y, z, x_unit, ax, **kwds)
+    elif (xscale, yscale) == ('linear', 'log'):
+        xlabel = f"{xlabel} [{x_unit}]" if x_unit else xlabel
+        img = imshow_logy(x, y, z, y_unit, ax, **kwds)
+    elif (xscale, yscale) == ('log', 'log'):
+        img = imshow_loglog(x, y, z, x_unit, y_unit, ax, **kwds)
+    else:
+        pass
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    cb = fig.colorbar(img, ax=ax)
+    cb.set_label(zlabel)
+
+
+def plot_scatter(x,
+                 y,
+                 z,
+                 xlabel,
+                 ylabel,
+                 zlabel,
+                 x_unit,
+                 y_unit,
+                 z_unit,
+                 fig,
+                 ax,
+                 xscale='linear',
+                 yscale='linear',
+                 zscale='linear',
+                 **kwds):
+    if np.any(np.iscomplex(z)):
+        s = np.abs(z)
+        c = np.angle(z)
+    else:
+        s = np.abs(z)
+        c = z.real
+    ax.scatter(x, y, s=s, c=c, **kwds)
+    xlabel = f"{xlabel} [{x_unit}]" if x_unit else xlabel
+    ylabel = f"{ylabel} [{y_unit}]" if y_unit else ylabel
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    ax.set_xscale(xscale)
+    ax.set_yscale(yscale)
+
+
+def autoplot(x,
+             y,
+             z,
+             xlabel='x',
+             ylabel='y',
+             zlabel='z',
+             x_unit='',
+             y_unit='',
+             z_unit='',
+             fig=None,
+             ax=None,
+             index=None,
+             xscale='auto',
+             yscale='auto',
+             zscale='auto',
+             max_lines=3,
+             scatter_lim=1000,
+             resolution=None,
+             **kwds):
+    """
+    Plot a 2D array as a line plot or an image.
+
+    Parameters:
+        x (array): x values
+        y (array): y values
+        z (array): z values
+        xlabel (str): x label
+        ylabel (str): y label
+        zlabel (str): z label
+        x_unit (str): x unit
+        y_unit (str): y unit
+        z_unit (str): z unit
+        fig (Figure): figure to plot on
+        ax (Axes): axes to plot on
+        index (int): index of the line to plot
+        xscale (str): x scale 'auto', 'linear' or 'log'
+        yscale (str): y scale 'auto', 'linear' or 'log'
+        zscale (str): z scale 'auto', 'linear' or 'log'
+        max_lines (int): maximum number of lines to plot
+        **kwds: keyword arguments passed to plot_img or plot_lines
+    """
+    if ax is not None:
+        fig = ax.figure
+    if fig is None:
+        fig = plt.gcf()
+    if ax is None:
+        ax = fig.add_subplot(111)
+
+    x = np.asarray(x)
+    y = np.asarray(y)
+    z = np.asarray(z)
+
+    if xscale == 'auto':
+        if good_for_logscale(x):
+            xscale = 'log'
+        else:
+            xscale = 'linear'
+    if yscale == 'auto':
+        if good_for_logscale(y):
+            yscale = 'log'
+        else:
+            yscale = 'linear'
+    if zscale == 'auto':
+        if good_for_logscale(z):
+            zscale = 'log'
+        else:
+            zscale = 'linear'
+
+    if x.shape == y.shape == z.shape and z.size < scatter_lim:
+        plot_scatter(x,
+                     y,
+                     z,
+                     xlabel,
+                     ylabel,
+                     zlabel,
+                     x_unit,
+                     y_unit,
+                     z_unit,
+                     fig,
+                     ax,
+                     xscale=xscale,
+                     yscale=yscale,
+                     zscale=zscale,
+                     **kwds)
+    elif z.ndim == 2 and (len(y) <= max_lines or len(x) <= max_lines
+                          or index is not None):
+        plot_lines(x,
+                   y,
+                   z,
+                   xlabel,
+                   ylabel,
+                   zlabel,
+                   x_unit=x_unit,
+                   y_unit=y_unit,
+                   z_unit=z_unit,
+                   xscale=xscale,
+                   yscale=yscale,
+                   zscale=zscale,
+                   ax=ax,
+                   index=index,
+                   **kwds)
+    else:
+        plot_img(x,
+                 y,
+                 z,
+                 xlabel,
+                 ylabel,
+                 zlabel,
+                 x_unit=x_unit,
+                 y_unit=y_unit,
+                 z_unit=z_unit,
+                 xscale=xscale,
+                 yscale=yscale,
+                 zscale=zscale,
+                 fig=fig,
+                 ax=ax,
+                 resolution=resolution,
+                 **kwds)