QuLab 2.0.9__tar.gz → 2.1.0__tar.gz

{qulab-2.0.9 → qulab-2.1.0}/PKG-INFO

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

{qulab-2.0.9 → qulab-2.1.0}/QuLab.egg-info/PKG-INFO

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

{qulab-2.0.9 → qulab-2.1.0}/qulab/scan/recorder.py

@@ -69,14 +69,14 @@ class BufferList():
         self._list = []
         self.lu = ()
         self.rd = ()
-        self.inner_shape = None
+        self.inner_shape = ()
         self.file = file
         self._slice = slice
         self._lock = Lock()
-        self._database = None
+        self._data_id = None
 
     def __repr__(self):
-        return f"<BufferList: lu={self.lu}, rd={self.rd}, slice={self._slice}>"
+        return f"<BufferList: shape={self.shape}, lu={self.lu}, rd={self.rd}, slice={self._slice}>"
 
     def __getstate__(self):
         self.flush()
@@ -99,11 +99,12 @@ class BufferList():
         self._list = []
         self._slice = None
         self._lock = Lock()
-        self._database = None
+        self._data_id = None
 
     @property
     def shape(self):
-        return tuple([i - j for i, j in zip(self.rd, self.lu)])
+        return tuple([i - j
+                      for i, j in zip(self.rd, self.lu)]) + self.inner_shape
 
     def flush(self):
         if not self._list:
@@ -143,11 +144,27 @@ class BufferList():
                             break
 
     def iter(self):
-        for pos, value in itertools.chain(self._iter_file(), self._list):
-            if not self._slice:
-                yield pos, value
-            elif all([index_in_slice(s, i) for s, i in zip(self._slice, pos)]):
-                yield pos, value[self._slice[len(pos):]]
+        if self._data_id is None:
+            for pos, value in itertools.chain(self._iter_file(), self._list):
+                if not self._slice:
+                    yield pos, value
+                elif all(
+                    [index_in_slice(s, i) for s, i in zip(self._slice, pos)]):
+                    if self.inner_shape:
+                        yield pos, value[self._slice[len(pos):]]
+                    else:
+                        yield pos, value
+        else:
+            server, record_id, key = self._data_id
+            with ZMQContextManager(zmq.DEALER, connect=server) as socket:
+                socket.send_pyobj({
+                    'method': 'bufferlist_slice',
+                    'record_id': record_id,
+                    'key': key,
+                    'slice': self._slice
+                })
+                ret = socket.recv_pyobj()
+                yield from ret
 
     def value(self):
         d = []
@@ -170,47 +187,111 @@ class BufferList():
 
     def array(self):
         pos, data = self.items()
-        pos = np.asarray(pos) - np.asarray(self.lu)
+        if self._slice:
+            pos = np.asarray(pos)
+            lu = tuple(np.min(pos, axis=0))
+            rd = tuple(np.max(pos, axis=0) + 1)
+            pos = np.asarray(pos) - np.asarray(lu)
+            shape = []
+            for k, (s, i, j) in enumerate(zip(self._slice, rd, lu)):
+                if s.step is not None:
+                    pos[:, k] = pos[:, k] / s.step
+                    shape.append(round(np.ceil((i - j) / s.step)))
+                else:
+                    shape.append(i - j)
+            shape = tuple(shape)
+        else:
+            shape = tuple([i - j for i, j in zip(self.rd, self.lu)])
+            pos = np.asarray(pos) - np.asarray(self.lu)
         data = np.asarray(data)
         inner_shape = data.shape[1:]
-        x = np.full(self.shape + inner_shape, np.nan, dtype=data[0].dtype)
+        x = np.full(shape + inner_shape, np.nan, dtype=data[0].dtype)
         x.__setitem__(tuple(pos.T), data)
         return x
 
     def _full_slice(self, slice_tuple: slice
                     | tuple[slice | int | EllipsisType, ...]):
+        ndim = len(self.lu)
+        if self.inner_shape:
+            ndim += len(self.inner_shape)
+
         if isinstance(slice_tuple, slice):
-            slice_tuple = (slice_tuple, ) + (slice(0, sys.maxsize,
-                                                   1), ) * (len(self.lu) - 1)
+            slice_tuple = (
+                slice_tuple, ) + (slice(0, sys.maxsize, 1), ) * (ndim - 1)
         if slice_tuple is Ellipsis:
-            slice_tuple = (slice(0, sys.maxsize, 1), ) * len(self.lu)
+            slice_tuple = (slice(0, sys.maxsize, 1), ) * ndim
         else:
-            head, tail = [], []
+            head, tail = (), ()
             for i, s in enumerate(slice_tuple):
                 if s is Ellipsis:
                     head = slice_tuple[:i]
                     tail = slice_tuple[i + 1:]
                     break
-            slice_tuple = head + (slice(0, sys.maxsize, 1), ) * (
-                len(self.lu) - len(head) - len(tail)) + tail
+            else:
+                head = slice_tuple
+                tail = ()
+            slice_tuple = head + (slice(
+                0, sys.maxsize, 1), ) * (ndim - len(head) - len(tail)) + tail
         slice_list = []
-        for s in slice_tuple:
+        contract = []
+        reversed = []
+        for i, s in enumerate(slice_tuple):
             if isinstance(s, int):
-                slice_list.append(s)
+                if s >= 0:
+                    slice_list.append(slice(s, s + 1, 1))
+                elif i < len(self.lu):
+                    s = self.rd[i] + s
+                    slice_list.append(slice(s, s + 1, 1))
+                else:
+                    slice_list.append(slice(s, s - 1, -1))
+                contract.append(i)
             else:
                 start, stop, step = s.start, s.stop, s.step
+                if step is None:
+                    step = 1
+                if step < 0 and i < len(self.lu):
+                    step = -step
+                    reversed.append(i)
+                    if start is None and stop is None:
+                        start, stop = 0, sys.maxsize
+                    elif start is None:
+                        start, stop = self.lu[i], sys.maxsize
+                    elif stop is None:
+                        start, stop = 0, start + self.lu[i]
+                    else:
+                        start, stop = stop + self.lu[i] + 1, start + self.lu[
+                            i] + 1
+
                 if start is None:
                     start = 0
+                elif start < 0 and i < len(self.lu):
+                    start = self.rd[i] + start
                 if step is None:
                     step = 1
                 if stop is None:
                     stop = sys.maxsize
+                elif stop < 0 and i < len(self.lu):
+                    stop = self.rd[i] + stop
+
                 slice_list.append(slice(start, stop, step))
-        return tuple(slice_list)
+        return tuple(slice_list), contract, reversed
 
     def __getitem__(self, slice_tuple: slice | EllipsisType
                     | tuple[slice | int | EllipsisType, ...]):
-        return super().__getitem__(self._full_slice(slice_tuple))
+        self._slice, contract, reversed = self._full_slice(slice_tuple)
+        ret = self.array()
+        slices = []
+        for i, s in enumerate(self._slice):
+            if i in contract:
+                slices.append(0)
+            elif isinstance(s, slice):
+                if i in reversed:
+                    slices.append(slice(None, None, -1))
+                else:
+                    slices.append(slice(None, None, 1))
+        ret = ret.__getitem__(tuple(slices))
+        self._slice = None
+        return ret
 
 
 class Record():
@@ -310,9 +391,9 @@ class Record():
         self.flush()
 
     def __getitem__(self, key):
-        return self.get(key)
+        return self.get(key, buffer_to_array=True)
 
-    def get(self, key, default=_notgiven, buffer_to_array=True, slice=None):
+    def get(self, key, default=_notgiven, buffer_to_array=False, slice=None):
         if self.is_remote_record():
             with ZMQContextManager(zmq.DEALER,
                                    connect=self.database) as socket:
@@ -323,19 +404,19 @@ class Record():
                 })
                 ret = socket.recv_pyobj()
                 if isinstance(ret, BufferList):
-                    socket.send_pyobj({
-                        'method': 'bufferlist_slice',
-                        'record_id': self.id,
-                        'key': key,
-                        'slice': slice
-                    })
-                    lst = socket.recv_pyobj()
-                    ret._list = lst
-                    ret._slice = slice
                     if buffer_to_array:
+                        socket.send_pyobj({
+                            'method': 'bufferlist_slice',
+                            'record_id': self.id,
+                            'key': key,
+                            'slice': slice
+                        })
+                        lst = socket.recv_pyobj()
+                        ret._list = lst
+                        ret._slice = slice
                         return ret.array()
                     else:
-                        ret._database = self.database
+                        ret._data_id = self.database, self.id, key
                         return ret
                 else:
                     return ret

{qulab-2.0.9 → qulab-2.1.0}/qulab/scan/scan.py

@@ -8,6 +8,7 @@ import re
 import sys
 import uuid
 import warnings
+from concurrent.futures import ProcessPoolExecutor
 from graphlib import TopologicalSorter
 from pathlib import Path
 from types import MethodType
@@ -157,6 +158,11 @@ class Promise():
         return Promise(self.task, None, attr)
 
 
+def _run_function_in_process(buf):
+    func, args, kwds = dill.loads(buf)
+    return func(*args, **kwds)
+
+
 class Scan():
 
     def __new__(cls, *args, mixin=None, **kwds):
@@ -176,6 +182,8 @@ class Scan():
                  database: str | Path
                  | None = f'tcp://127.0.0.1:{default_record_port}',
                  dump_globals: bool = False,
+                 max_workers: int = 4,
+                 max_promise: int = 100,
                  mixin=None):
         self.id = task_uuid()
         self.record = None
@@ -207,23 +215,21 @@ class Scan():
         self._variables = {}
         self._main_task = None
         self._sock = None
-        self._sem = asyncio.Semaphore(100)
+        self._sem = asyncio.Semaphore(max_promise + 1)
         self._bar: dict[int, tqdm] = {}
         self._hide_pattern_re = re.compile('|'.join(self.description['hiden']))
-        self._task_queue = asyncio.Queue()
-        self._task_pool = []
+        self._msg_queue = asyncio.Queue()
+        self._prm_queue = asyncio.Queue()
         self._single_step = True
+        self._max_workers = max_workers
+        self._max_promise = max_promise
+        self._executors = ProcessPoolExecutor(max_workers=max_workers)
 
     def __del__(self):
         try:
             self._main_task.cancel()
         except:
             pass
-        for task in self._task_pool:
-            try:
-                task.cancel()
-            except:
-                pass
 
     def __getstate__(self) -> dict:
         state = self.__dict__.copy()
@@ -231,9 +237,10 @@ class Scan():
         del state['_sock']
         del state['_main_task']
         del state['_bar']
-        del state['_task_queue']
-        del state['_task_pool']
+        del state['_msg_queue']
+        del state['_prm_queue']
         del state['_sem']
+        del state['_executors']
         return state
 
     def __setstate__(self, state: dict) -> None:
@@ -242,12 +249,20 @@ class Scan():
         self._sock = None
         self._main_task = None
         self._bar = {}
-        self._task_queue = asyncio.Queue()
-        self._task_pool = []
-        self._sem = asyncio.Semaphore(100)
+        self._prm_queue = asyncio.Queue()
+        self._msg_queue = asyncio.Queue()
+        self._sem = asyncio.Semaphore(self._max_promise + 1)
+        self._executors = ProcessPoolExecutor(max_workers=self._max_workers)
         for opt in self.description['optimizers'].values():
             opt.scanner = self
 
+    def __del__(self):
+        try:
+            self._main_task.cancel()
+        except:
+            pass
+        self._executors.shutdown()
+
     @property
     def current_level(self):
         return self._current_level
@@ -256,8 +271,8 @@ class Scan():
     def variables(self) -> dict[str, Any]:
         return self._variables
 
-    async def emit(self, current_level, step, position, variables: dict[str,
-                                                                        Any]):
+    async def _emit(self, current_level, step, position, variables: dict[str,
+                                                                         Any]):
         for key, value in list(variables.items()):
             if inspect.isawaitable(value) and not self.hiden(key):
                 variables[key] = await value
@@ -280,6 +295,11 @@ class Scan():
                 for k, v in variables.items() if not self.hiden(k)
             })
 
+    def emit(self, current_level, step, position, variables: dict[str, Any]):
+        self._msg_queue.put_nowait(
+            asyncio.create_task(
+                self._emit(current_level, step, position, variables.copy())))
+
     def hide(self, name: str):
         self.description['hiden'].append(name)
         self._hide_pattern_re = re.compile('|'.join(self.description['hiden']))
@@ -289,11 +309,11 @@ class Scan():
 
     async def _filter(self, variables: dict[str, Any], level: int = 0):
         try:
-            return all([
-                await call_function(fun, variables) for fun in itertools.chain(
+            return all(await asyncio.gather(*[
+                call_function(fun, variables) for fun in itertools.chain(
                     self.description['filters'].get(level, []),
                     self.description['filters'].get(-1, []))
-            ])
+            ]))
         except:
             return True
 
@@ -432,11 +452,15 @@ class Scan():
 
     async def _update_progress(self):
         while True:
-            task = await self._task_queue.get()
-            if isinstance(task, asyncio.Event):
-                task.set()
-            elif inspect.isawaitable(task):
-                await task
+            task = await self._prm_queue.get()
+            await task
+            self._prm_queue.task_done()
+
+    async def _send_msg(self):
+        while True:
+            task = await self._msg_queue.get()
+            await task
+            self._msg_queue.task_done()
 
     async def run(self):
         assymbly(self.description)
@@ -452,44 +476,41 @@ class Scan():
             await self._run()
 
     async def _run(self):
-        task = asyncio.create_task(self._update_progress())
-        self._task_pool.append(task)
+        send_msg = asyncio.create_task(self._send_msg())
+        update_progress_task = asyncio.create_task(self._update_progress())
+
         self._variables = {'self': self}
-        self._variables.update(self.description['consts'].copy())
+
+        await update_variables(self._variables, self.description['consts'],
+                               self.description['setters'])
         for level, total in self.description['total'].items():
             if total == np.inf:
                 total = None
             self._bar[level] = tqdm(total=total)
-        for group in self.description['order'].get(-1, []):
-            for name in group:
-                if name in self.description['functions']:
-                    self.variables[name] = await call_function(
-                        self.description['functions'][name], self.variables)
-                if name in self.description['setters']:
-                    coro = self.description['setters'][name](
-                        self.variables[name])
-                    if inspect.isawaitable(coro):
-                        await coro
+
+        updates = await call_many_functions(
+            self.description['order'].get(-1, []),
+            self.description['functions'], self.variables)
+        await update_variables(self.variables, updates,
+                               self.description['setters'])
+
         self.record = await self.create_record()
         await self.work()
         for level, bar in self._bar.items():
             bar.close()
 
-        while not self._task_queue.empty():
-            evt = self._task_queue.get_nowait()
-            if isinstance(evt, asyncio.Event):
-                evt.set()
-            elif inspect.isawaitable(evt):
-                await evt
         if self._single_step:
-            for group in self.description['order'].get(-1, []):
-                for name in group:
-                    if name in self.description['getters']:
-                        self.variables[name] = await call_function(
-                            self.description['getters'][name], self.variables)
-            await self.emit(0, 0, 0, self.variables.copy())
-            await self.emit(-1, 0, 0, {})
-        task.cancel()
+            self.variables.update(await call_many_functions(
+                self.description['order'].get(-1, []),
+                self.description['getters'], self.variables))
+
+            self.emit(0, 0, 0, self.variables)
+            self.emit(-1, 0, 0, {})
+
+        await self._prm_queue.join()
+        update_progress_task.cancel()
+        await self._msg_queue.join()
+        send_msg.cancel()
         return self.variables
 
     async def done(self):
@@ -536,8 +557,8 @@ class Scan():
             return
         step = 0
         position = 0
-        self._task_queue.put_nowait(
-            self._reset_progress_bar(self.current_level))
+        self._prm_queue.put_nowait(self._reset_progress_bar(
+            self.current_level))
         async for variables in _iter_level(
                 self.variables,
                 self.description['loops'].get(self.current_level, []),
@@ -548,23 +569,18 @@ class Scan():
             if await self._filter(variables, self.current_level - 1):
                 yield variables
                 self._single_step = False
-                asyncio.create_task(
-                    self.emit(self.current_level - 1, step, position,
-                              variables.copy()))
+                self.emit(self.current_level - 1, step, position, variables)
                 step += 1
             position += 1
             self._current_level -= 1
-            self._task_queue.put_nowait(
+            self._prm_queue.put_nowait(
                 self._update_progress_bar(self.current_level, 1))
         if self.current_level == 0:
-            await self.emit(self.current_level - 1, 0, 0, {})
+            self.emit(self.current_level - 1, 0, 0, {})
             for name, value in self.variables.items():
                 if inspect.isawaitable(value):
                     self.variables[name] = await value
-            while not self._task_queue.empty():
-                task = self._task_queue.get_nowait()
-                if inspect.isawaitable(task):
-                    await task
+            await self._prm_queue.join()
 
     async def work(self, **kwds):
         if self.current_level in self.description['actions']:
@@ -589,7 +605,8 @@ class Scan():
         """
         self.description['actions'][level] = action
 
-    async def promise(self, awaitable: Awaitable) -> Promise:
+    async def promise(self, awaitable: Awaitable | Callable, *args,
+                      **kwds) -> Promise:
         """
         Promise to calculate asynchronous function and return the result in future.
 
@@ -602,8 +619,19 @@ class Scan():
         if inspect.isawaitable(awaitable):
             async with self._sem:
                 task = asyncio.create_task(self._await(awaitable))
-                self._task_queue.put_nowait(task)
+                self._prm_queue.put_nowait(task)
                 return Promise(task)
+        elif inspect.iscoroutinefunction(awaitable):
+            return await self.promise(awaitable(*args, **kwds))
+        elif callable(awaitable):
+            try:
+                buf = dill.dumps((awaitable, args, kwds))
+                task = asyncio.get_running_loop().run_in_executor(
+                    self._executors, _run_function_in_process, buf)
+                self._prm_queue.put_nowait(task)
+                return Promise(task)
+            except:
+                return awaitable(*args, **kwds)
         else:
             return awaitable
 
@@ -785,13 +813,17 @@ def assymbly(description):
     return description
 
 
-async def _update_variables(variables, updates, setters):
+async def update_variables(variables: dict[str, Any], updates: dict[str, Any],
+                           setters: dict[str, Callable]):
+    coros = []
     for name, value in updates.items():
         if name in setters:
             coro = setters[name](value)
             if inspect.isawaitable(coro):
-                await coro
+                coros.append(coro)
         variables[name] = value
+    if coros:
+        await asyncio.gather(*coros)
 
 
 async def _iter_level(variables,
@@ -824,31 +856,23 @@ async def _iter_level(variables,
         maxiter = min(maxiter, opt_cfg.maxiter)
 
     async for args in async_zip(*iters_d.values(), range(maxiter)):
-        await _update_variables(variables, dict(zip(iters_d.keys(),
-                                                    args[:-1])), setters)
+        await update_variables(variables, dict(zip(iters_d.keys(), args[:-1])),
+                               setters)
         for name, opt in opts.items():
             args = opt.ask()
             opt_cfg = optimizers[name]
-            await _update_variables(variables, {
+            await update_variables(variables, {
                 n: v
                 for n, v in zip(opt_cfg.dimensions.keys(), args)
             }, setters)
 
-        for group in order:
-            for name in group:
-                if name in functions:
-                    await _update_variables(variables, {
-                        name:
-                        await call_function(functions[name], variables)
-                    }, setters)
+        await update_variables(
+            variables, await call_many_functions(order, functions, variables),
+            setters)
 
         yield variables
 
-        for group in order:
-            for name in group:
-                if name in getters:
-                    variables[name] = await call_function(
-                        getters[name], variables)
+        variables.update(await call_many_functions(order, getters, variables))
 
         for name, opt in opts.items():
             opt_cfg = optimizers[name]
@@ -866,10 +890,28 @@ async def _iter_level(variables,
     for name, opt in opts.items():
         opt_cfg = optimizers[name]
         result = opt.get_result()
-        variables.update({
-            n: v
-            for n, v in zip(opt_cfg.dimensions.keys(), result.x)
-        })
+        await update_variables(
+            variables, {
+                name: value
+                for name, value in zip(opt_cfg.dimensions.keys(), result.x)
+            }, setters)
         variables[name] = result.fun
     if opts:
         yield variables
+
+
+async def call_many_functions(order: list[list[str]],
+                              functions: dict[str, Callable],
+                              variables: dict[str, Any]) -> dict[str, Any]:
+    ret = {}
+    for group in order:
+        waited = []
+        coros = []
+        for name in group:
+            if name in functions:
+                waited.append(name)
+                coros.append(call_function(functions[name], variables | ret))
+        if coros:
+            results = await asyncio.gather(*coros)
+            ret.update(dict(zip(waited, results)))
+    return ret

qulab-2.1.0/qulab/version.py

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

qulab-2.0.9/qulab/version.py

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