QuLab 2.4.0__cp312-cp312-macosx_10_13_universal2.whl

qulab/visualization/qdat.py

@@ -0,0 +1,152 @@
+"""
+This module provides the visualization of the qdat file.
+"""
+import logging
+import math
+import re
+from functools import reduce
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import cm
+
+log = logging.getLogger(__name__)
+
+
+def draw(record_dict, fig=None, transpose=True, remove=True):
+    """Draw a 2D or 3D plot from a record dict.
+
+    Args:
+        record_dict (dict): The record dict.
+        fig (matplotlib.figure.Figure): The figure to draw on.
+        transpose (bool): Whether to transpose the data.
+    """
+    dim = record_dict['info']['dim']
+    zshape = record_dict['info']['zshape']
+    zsize = reduce(lambda a, b: a * b, zshape, 1)
+
+    axis_label = np.asarray(record_dict['info'].get('axis_label',
+                                                    [])).flatten()
+    z_label = np.asarray(record_dict['info'].get('z_label', [])).flatten()
+    axis_unit = np.asarray(record_dict['info'].get('axis_unit', [])).flatten()
+    z_unit = np.asarray(record_dict['info'].get('z_unit', [])).flatten()
+
+    fig = plt.figure() if fig is None else fig
+    if zsize < 4:
+        plot_shape = (1, zsize)
+    else:
+        n = int(np.sqrt(zsize))
+        plot_shape = (math.ceil(zsize / n), n)
+    if dim < 3 and zsize < 100:
+        figsize = plot_shape[1] * 8, plot_shape[0] * 6
+        fig.set_size_inches(*figsize)
+        axes = fig.subplots(*plot_shape)
+        axes = np.array(axes).flatten()
+    else:
+        axes = []
+
+    cb_list = []
+    if dim == 1 and zsize < 101:
+        for i in range(zsize):
+            title = record_dict[
+                'name'] + f' {i}' if zsize > 1 else record_dict['name']
+            try:
+                xlabel = axis_label[0]
+                ylabel = z_label[0] if len(z_label) == 1 else z_label[i]
+            except:
+                xlabel, ylabel = 'X', 'Y'
+            try:
+                x_unit = axis_unit[0]
+                y_unit = z_unit[0] if len(z_unit) == 1 else z_unit[i]
+            except:
+                x_unit, y_unit = 'a.u.', 'a.u.'
+            axes[i].set_title(title)
+            axes[i].set_xlabel(f'{xlabel} ({x_unit})')
+            axes[i].set_ylabel(f'{ylabel} ({y_unit})')
+    elif dim == 2 and zsize < 101:
+        for i in range(zsize):
+            smp = cm.ScalarMappable(norm=None, cmap=None)
+            cb = fig.colorbar(smp, ax=axes[i])  # 默认色谱的colorbar
+            cb_list.append(cb)
+            try:
+                title = record_dict['name'] + \
+                    f': {z_label[i]}' if zsize > 1 else record_dict['name']
+            except:
+                title = record_dict[
+                    'name'] + f' {i}' if zsize > 1 else record_dict['name']
+            try:
+                xlabel, ylabel = axis_label[1], axis_label[0]
+                if transpose:
+                    xlabel, ylabel = ylabel, xlabel
+            except:
+                xlabel, ylabel = 'X', 'Y'
+            try:
+                x_unit, y_unit = axis_unit[1], axis_unit[0]
+                if transpose:
+                    x_unit, y_unit = y_unit, x_unit
+            except:
+                x_unit, y_unit = 'a.u.', 'a.u.'
+            axes[i].set_title(title)
+            axes[i].set_xlabel(f'{xlabel} ({x_unit})')
+            axes[i].set_ylabel(f'{ylabel} ({y_unit})')
+    else:
+        message1 = f'dim {dim} is too large (>2)! ' if dim > 2 else f'dim={dim}; '
+        message2 = f'zsize {zsize} is too large (>101)!' if zsize > 101 else f'zsize={zsize}'
+        log.warning('PASS: ' + message1 + message2)
+    #########################################################################################
+    try:
+        tags = [
+            tag.strip(r'\*')
+            for tag in record_dict['ParaSpace'].get('tags', [])
+            if re.match(r'\*', tag)
+        ]
+        tag_text = ','.join(tags)
+        if tag_text:
+            axes[0].text(-0.1,
+                         1.1,
+                         'TAG: ' + tag_text,
+                         horizontalalignment='left',
+                         verticalalignment='bottom',
+                         transform=axes[0].transAxes)  # fig.transFigure)#
+    except:
+        pass
+    #########################################################################################
+    fig.tight_layout()
+
+    dim = record_dict['info']['dim']
+    zshape = record_dict['info']['zshape']
+    datashape = record_dict['info']['datashape']
+    zsize = reduce(lambda a, b: a * b, zshape, 1)
+    datashape_r = (*datashape[:dim], zsize)
+
+    if dim == 1 and zsize < 101:
+        x, z = record_dict['data']
+        z = z.reshape(datashape_r)
+        z = np.abs(z) if np.any(np.iscomplex(z)) else z.real
+        for i in range(zsize):
+            _ = [a.remove() for a in axes[i].get_lines()] if remove else []
+            axes[i].plot(x, z[:, i], 'C0')
+            axes[i].plot(x, z[:, i], 'C0.')
+    elif dim == 2 and zsize < 101:
+        x, y, z = record_dict['data']
+        x_step, y_step = x[1] - x[0], y[1] - y[0]
+        if transpose:
+            x, y = y, x
+            x_step, y_step = y_step, x_step
+        z = z.reshape(datashape_r)
+        z = np.abs(z) if np.any(np.iscomplex(z)) else z.real
+        for i in range(zsize):
+            _z = z[:, :, i]
+            _ = [a.remove() for a in axes[i].get_images()] if remove else []
+            if transpose:
+                _z = _z.T
+            im = axes[i].imshow(_z,
+                                extent=(y[0] - y_step / 2, y[-1] + y_step / 2,
+                                        x[0] - x_step / 2, x[-1] + x_step / 2),
+                                origin='lower',
+                                aspect='auto')
+            cb_list[i].update_normal(im)  # 更新对应的colorbar
+    else:
+        pass
+
+    return fig, axes

qulab/visualization/rot3d.py

@@ -0,0 +1,23 @@
+import numpy as np
+from scipy.linalg import expm
+
+
+def _rot(theta, v):
+    x, y, z = np.asarray(v) / np.sqrt(np.sum(np.asarray(v)**2))
+    c, s = np.cos(theta), np.sin(theta)
+
+    return np.array([[
+        c + (1 - c) * x**2, (1 - c) * x * y - s * z, (1 - c) * x * z + s * y
+    ], [(1 - c) * x * y + s * z, c + (1 - c) * y**2, (1 - c) * y * z - s * x],
+                     [(1 - c) * x * z - s * y, (1 - c) * y * z + s * x,
+                      c + (1 - c) * z**2]])
+
+
+def rot_round(x, y, z, v=[0, 0, 1], theta=0, c=[0, 0, 0]):
+    ret = (np.array([x, y, z]).T - np.array(c)) @ _rot(theta, v) + np.array(c)
+    return ret.T
+
+
+def projection(x, y, z, y0=1):
+    d = np.sqrt(x**2 + y**2 + z**2)
+    return x * y0 / y, z * y0 / y, d

qulab/visualization/widgets.py

@@ -0,0 +1,86 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+class DataPicker():
+
+    def __init__(self, ax=None):
+        self.points_and_text = {}
+        self.points = None
+        self.hline = None
+        self.vline = None
+        self.text = None
+        if ax is None:
+            ax = plt.gca()
+        self.ax = ax
+        self.ax.figure.canvas.mpl_connect('button_press_event', self.on_click)
+        # self.ax.figure.canvas.mpl_connect('motion_notify_event', self.on_move)
+        self.ax.figure.canvas.mpl_connect('key_press_event', self.on_key_press)
+        self.mode = 'pick'
+
+    def on_key_press(self, event):
+        if event.key == 'a':
+            if self.mode != 'pick':
+                self.mode = 'pick'
+            else:
+                self.mode = 'default'
+
+    def on_move(self, event):
+        if event.inaxes is self.ax:
+            # self.hline = self.ax.axhline(y=np.nan, color='r', lw=1)
+            # self.vline = self.ax.axvline(x=np.nan, color='r', lw=1)
+            # self.text = self.ax.text(0, 0, '', verticalalignment='center')
+            self.hline.set_ydata(event.ydata)
+            self.vline.set_xdata(event.xdata)
+            self.text.set_position((event.xdata, event.ydata))
+            self.text.set_text(f'({event.xdata:.2f}, {event.ydata:.2f})')
+            self.ax.draw()
+
+    def on_click(self, event):
+        if self.mode != 'pick':
+            return
+        # 鼠标左键的button值为1
+        if event.button == 1 and event.inaxes is self.ax:
+            point = (event.xdata, event.ydata)
+            text = self.ax.text(point[0],
+                                point[1],
+                                f'({point[0]:.2f}, {point[1]:.2f})',
+                                verticalalignment='center')
+            self.points_and_text[point] = text
+            x, y = self.get_xy()
+            if self.points is None:
+                self.points, = self.ax.plot(x, y, 'ro')
+            else:
+                self.points.set_data(x, y)
+            self.ax.draw()
+
+        elif event.button == 3 and event.inaxes is self.ax:
+            for point, text in list(self.points_and_text.items()):
+                point_xdisplay, point_ydisplay = self.ax.transData.transform_point(
+                    point)
+
+                distance = np.sqrt((point_xdisplay - event.x)**2 +
+                                   (point_ydisplay - event.y)**2)
+                if distance < 10:
+                    text.remove()
+                    self.points_and_text.pop(point)
+                    if self.points_and_text:
+                        x, y = self.get_xy()
+                        self.points.set_data(x, y)
+                    else:
+                        self.points.remove()
+                        self.points = None
+                    self.ax.draw()
+                    break
+
+    def get_xy(self):
+        if self.points_and_text:
+            data = np.asarray(list(self.points_and_text.keys()))
+            x, y = data[:, 0], data[:, 1]
+
+            index = np.argsort(x)
+            x = np.asarray(x)[index]
+            y = np.asarray(y)[index]
+            return x, y
+        else:
+            return np.array([]), np.array([])