QuLab 2.4.0__cp312-cp312-macosx_10_13_universal2.whl

qulab/sys/rpc/zmq_socket.py

@@ -0,0 +1,220 @@
+from typing import Optional
+
+import zmq
+import zmq.asyncio
+import zmq.auth
+from watchdog.events import FileSystemEventHandler
+from watchdog.observers import Observer
+from zmq.auth.asyncio import AsyncioAuthenticator
+from zmq.auth.thread import ThreadAuthenticator
+
+
+class ReloadCertificatesHandler(FileSystemEventHandler):
+
+    def __init__(self, zmq: 'ZMQContextManager'):
+        self.zmq = zmq
+
+    def on_modified(self, event):
+        self.zmq.reload_certificates()
+
+
+class ZMQContextManager:
+    """
+    A context manager for managing ZeroMQ sockets with asynchronous support.
+    It handles the creation, connection, binding, and security configuration of
+    the sockets and ensures proper resource cleanup.
+
+    The context manager can be used as a synchronous context manager or an
+    asynchronous context manager. When used as an asynchronous context manager,
+    the socket is created with an asyncio context and can be used with the
+    `await` keyword.
+
+    The security settings for the socket can be configured using the secret key
+    and public key parameters. If the secret key is provided, the socket is
+    configured to use ZeroMQ curve encryption. The public key of the server can
+    also be provided for client sockets to connect to a server with curve
+    encryption. You can also provide the paths to the secret key and public key
+    files, and the public key of the server to load the keys from files. The
+    keys can be reloaded automatically when the files are modified by setting
+    the `public_keys_location` parameter.
+
+    To generate a secret key and public key pair, you can use the following
+    commands:
+    
+        ```bash
+        # Generate a secret key and public key pair
+        python -c "import zmq.auth; zmq.auth.create_certificates('.', 'filename')"
+        ```
+
+    Attributes:
+        socket_type: zmq.SocketType
+            The type of the socket to create (e.g., zmq.REP).
+        bind: str, optional
+            The address to bind the socket to.
+        connect: str, optional
+            The address to connect the socket to.
+        secret_key_file: str, optional
+            The path to the secret key file for ZeroMQ curve encryption.
+        server_public_key_file: str, optional
+            The path to the public key file for the server in ZeroMQ curve
+            encryption.
+        public_keys_location: str, optional
+            The location to store the public keys for ZeroMQ curve encryption.
+        secret_key: bytes, optional
+            The secret key for ZeroMQ curve encryption.
+        public_key: bytes, optional
+            The public key for ZeroMQ curve encryption.
+        server_public_key: bytes, optional
+            The public key for the server in ZeroMQ curve encryption.
+
+    Methods:
+        _create_socket: zmq.Socket
+            Creates and configures a ZeroMQ socket.
+        _close_socket: None
+            Closes the ZeroMQ socket and the context, and stops the authenticator
+            if it was started.
+
+    Examples:
+        Create a REP socket and bind it to an address:
+
+        >>> async with ZMQContextManager(zmq.REP, bind='tcp://*:5555') as socket:
+        ...     while True:
+        ...         message = await socket.recv()
+        ...         await socket.send(message)
+
+        Create a REQ socket and connect it to an address:
+
+        >>> with ZMQContextManager(zmq.REQ, connect='tcp://localhost:5555') as socket:
+        ...     socket.send(b'Hello')
+        ...     message = socket.recv()
+    """
+
+    def __init__(self,
+                 socket_type: zmq.SocketType,
+                 bind: Optional[str] = None,
+                 connect: Optional[str] = None,
+                 secret_key_file: Optional[str] = None,
+                 server_public_key_file: Optional[str] = None,
+                 public_keys_location: Optional[str] = None,
+                 secret_key: Optional[bytes] = None,
+                 public_key: Optional[bytes] = None,
+                 server_public_key: Optional[bytes] = None,
+                 socket: Optional[zmq.Socket] = None):
+        self.socket_type = socket_type
+        if bind is None and connect is None:
+            raise ValueError("Either 'bind' or 'connect' must be specified.")
+        if bind is not None and connect is not None:
+            raise ValueError("Both 'bind' and 'connect' cannot be specified.")
+        self.bind = bind
+        self.connect = connect
+        self.secret_key = secret_key
+        self.public_key = public_key
+        self.server_public_key = server_public_key
+
+        if secret_key_file:
+            self.public_key, self.secret_key = zmq.auth.load_certificate(
+                secret_key_file)
+
+        if (self.secret_key is not None and self.public_key is None
+                or self.secret_key is None and self.public_key is not None):
+            raise ValueError(
+                "Both secret key and public key must be specified.")
+
+        if server_public_key_file:
+            self.server_public_key = zmq.auth.load_certificate(
+                server_public_key_file)[0]
+
+        self.public_keys_location = public_keys_location
+
+        self.observer = None
+        self.auth = None
+        self.context = None
+        self.socket = None
+        self._external_socket = None
+        try:
+            if not socket.closed:
+                self._external_socket = socket
+        except:
+            pass
+
+    def _create_socket(self, asyncio=False) -> zmq.Socket:
+        """
+        Creates and configures a ZeroMQ socket. Sets up security if required,
+        and binds or connects the socket according to the specified settings.
+
+        Returns:
+            zmq.Socket: The configured ZeroMQ socket.
+        """
+        if self._external_socket:
+            return self._external_socket
+        if asyncio:
+            self.context = zmq.asyncio.Context()
+        else:
+            self.context = zmq.Context.instance()
+
+        self.socket = self.context.socket(self.socket_type)
+        self.auth = None
+
+        if self.bind and self.secret_key:
+            if asyncio:
+                self.auth = AsyncioAuthenticator(self.context)
+            else:
+                self.auth = ThreadAuthenticator(self.context)
+            self.auth.start()
+            self.reload_certificates()
+            self.auto_reload_certificates()
+            self.socket.curve_server = True  # must come before bind
+
+        if self.secret_key:
+            self.socket.curve_secretkey = self.secret_key
+            self.socket.curve_publickey = self.public_key
+
+        if self.bind:
+            self.socket.bind(self.bind)
+        if self.connect:
+            if self.server_public_key:
+                self.socket.curve_serverkey = self.server_public_key
+            self.socket.connect(self.connect)
+        return self.socket
+
+    def reload_certificates(self):
+        if self.public_keys_location and self.auth:
+            self.auth.configure_curve(domain='*',
+                                      location=self.public_keys_location)
+
+    def auto_reload_certificates(self):
+        self.observer = Observer()
+        self.observer.schedule(ReloadCertificatesHandler(self),
+                               self.public_keys_location,
+                               recursive=False)
+        self.observer.start()
+
+    def _close_socket(self) -> None:
+        """
+        Closes the ZeroMQ socket and the context, and stops the authenticator
+        if it was started.
+        """
+        if self._external_socket:
+            return
+        if self.observer:
+            self.observer.stop()
+            self.observer.join()
+        self.socket.close()
+        if self.auth:
+            self.auth.stop()
+        self.context.term()
+
+    def __enter__(self) -> zmq.Socket:
+        return self._create_socket(asyncio=False)
+
+    def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception],
+                 exc_tb: Optional[type]) -> None:
+        self._close_socket()
+
+    async def __aenter__(self) -> zmq.Socket:
+        return self._create_socket(asyncio=True)
+
+    async def __aexit__(self, exc_type: Optional[type],
+                        exc_val: Optional[Exception],
+                        exc_tb: Optional[type]) -> None:
+        self._close_socket()

qulab/version.py

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

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