QuLab 2.1.0__cp312-cp312-macosx_10_9_universal2.whl → 2.1.2__cp312-cp312-macosx_10_9_universal2.whl

qulab/scan/scan.py

@@ -1,35 +1,56 @@
 import asyncio
 import copy
-import datetime
 import inspect
 import itertools
 import os
 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
-from typing import Any, Awaitable, Callable, Iterable, Type
+from typing import Any, Awaitable, Callable, Iterable
 
 import dill
 import numpy as np
-import skopt
 import zmq
-from skopt.space import Categorical, Integer, Real
-from tqdm.notebook import tqdm
 
 from ..sys.rpc.zmq_socket import ZMQContextManager
 from .expression import Env, Expression, Symbol
 from .optimize import NgOptimizer
-from .recorder import Record, default_record_port
-from .utils import async_zip, call_function
+from .record import Record
+from .recorder import default_record_port
+from .space import Optimizer, OptimizeSpace, Space
+from .utils import async_zip, call_function, dump_globals
+
+try:
+    from tqdm.notebook import tqdm
+except:
+
+    class tqdm():
+
+        def update(self, n):
+            pass
+
+        def close(self):
+            pass
+
+        def reset(self):
+            pass
+
 
 __process_uuid = uuid.uuid1()
 __task_counter = itertools.count()
 
+if os.getenv('QULAB_SERVER'):
+    default_server = os.getenv('QULAB_SERVER')
+else:
+    default_server = f'tcp://127.0.0.1:{default_record_port}'
+if os.getenv('QULAB_EXECUTOR'):
+    default_executor = os.getenv('QULAB_EXECUTOR')
+else:
+    default_executor = default_server
+
 
 def task_uuid():
     return uuid.uuid3(__process_uuid, str(next(__task_counter)))
@@ -55,79 +76,6 @@ def _get_depends(func: Callable):
     return args
 
 
-class OptimizeSpace():
-
-    def __init__(self, optimizer: 'Optimizer', space):
-        self.optimizer = optimizer
-        self.space = space
-        self.name = None
-
-    def __len__(self):
-        return self.optimizer.maxiter
-
-
-class Optimizer():
-
-    def __init__(self,
-                 scanner: 'Scan',
-                 name: str,
-                 level: int,
-                 method: str | Type = skopt.Optimizer,
-                 maxiter: int = 1000,
-                 minimize: bool = True,
-                 **kwds):
-        self.scanner = scanner
-        self.method = method
-        self.maxiter = maxiter
-        self.dimensions = {}
-        self.name = name
-        self.level = level
-        self.kwds = kwds
-        self.minimize = minimize
-
-    def create(self):
-        return self.method(list(self.dimensions.values()), **self.kwds)
-
-    def Categorical(self,
-                    categories,
-                    prior=None,
-                    transform=None,
-                    name=None) -> OptimizeSpace:
-        return OptimizeSpace(self,
-                             Categorical(categories, prior, transform, name))
-
-    def Integer(self,
-                low,
-                high,
-                prior="uniform",
-                base=10,
-                transform=None,
-                name=None,
-                dtype=np.int64) -> OptimizeSpace:
-        return OptimizeSpace(
-            self, Integer(low, high, prior, base, transform, name, dtype))
-
-    def Real(self,
-             low,
-             high,
-             prior="uniform",
-             base=10,
-             transform=None,
-             name=None,
-             dtype=float) -> OptimizeSpace:
-        return OptimizeSpace(
-            self, Real(low, high, prior, base, transform, name, dtype))
-
-    def __getstate__(self) -> dict:
-        state = self.__dict__.copy()
-        del state['scanner']
-        return state
-
-    def __setstate__(self, state: dict) -> None:
-        self.__dict__.update(state)
-        self.scanner = None
-
-
 class Promise():
     __slots__ = ['task', 'key', 'attr']
 
@@ -180,7 +128,7 @@ class Scan():
                  app: str = 'task',
                  tags: tuple[str] = (),
                  database: str | Path
-                 | None = f'tcp://127.0.0.1:{default_record_port}',
+                 | None = default_server,
                  dump_globals: bool = False,
                  max_workers: int = 4,
                  max_promise: int = 100,
@@ -192,6 +140,7 @@ class Scan():
             'app': app,
             'tags': tags,
             'loops': {},
+            'intrinsic_loops': {},
             'consts': {},
             'functions': {},
             'getters': {},
@@ -201,6 +150,8 @@ class Scan():
             'actions': {},
             'dependents': {},
             'order': {},
+            'axis': {},
+            'independent_variables': set(),
             'filters': {},
             'total': {},
             'database': database,
@@ -341,10 +292,10 @@ class Scan():
         else:
             return Symbol(name)
 
-    def _add_loop_var(self, name: str, level: int, range):
+    def _add_search_space(self, name: str, level: int, space):
         if level not in self.description['loops']:
             self.description['loops'][level] = []
-        self.description['loops'][level].append((name, range))
+        self.description['loops'][level].append((name, space))
 
     def add_depends(self, name: str, depends: list[str]):
         if isinstance(depends, str):
@@ -353,7 +304,7 @@ class Scan():
             self.description['dependents'][name] = set()
         self.description['dependents'][name].update(depends)
 
-    def add_filter(self, func: Callable, level: int):
+    def add_filter(self, func: Callable, level: int = -1):
         """
         Add a filter function to the scan.
 
@@ -365,7 +316,11 @@ class Scan():
             self.description['filters'][level] = []
         self.description['filters'][level].append(func)
 
-    def set(self, name: str, value, setter: Callable | None = None):
+    def set(self,
+            name: str,
+            value,
+            depends: Iterable[str] | None = None,
+            setter: Callable | None = None):
         try:
             dill.dumps(value)
         except:
@@ -374,31 +329,54 @@ class Scan():
             self.add_depends(name, value.symbols())
             self.description['functions'][name] = value
         elif callable(value):
-            self.add_depends(name, _get_depends(value))
-            self.description['functions'][name] = value
+            if depends:
+                self.add_depends(name, depends)
+                s = ','.join(depends)
+                self.description['functions'][f'_tmp_{name}'] = value
+                self.description['functions'][name] = eval(
+                    f"lambda self, {s}: self.description['functions']['_tmp_{name}']({s})"
+                )
+            else:
+                self.add_depends(name, _get_depends(value))
+                self.description['functions'][name] = value
         else:
+            try:
+                value = Space.fromarray(value)
+            except:
+                pass
             self.description['consts'][name] = value
         if setter:
             self.description['setters'][name] = setter
 
     def search(self,
                name: str,
-               range: Iterable | Expression | Callable | OptimizeSpace,
+               space: Iterable | Expression | Callable | OptimizeSpace,
                level: int | None = None,
-               setter: Callable | None = None):
+               setter: Callable | None = None,
+               intrinsic: bool = False):
         if level is not None:
-            assert level >= 0, 'level must be greater than or equal to 0.'
-        if isinstance(range, OptimizeSpace):
-            range.name = name
-            range.optimizer.dimensions[name] = range.space
-            self._add_loop_var(name, range.optimizer.level, range)
-            self.add_depends(range.optimizer.name, [name])
+            if not intrinsic:
+                assert level >= 0, 'level must be greater than or equal to 0.'
+        if intrinsic:
+            assert isinstance(space, (np.ndarray, list, tuple, range, Space)), \
+                'space must be an instance of np.ndarray, list, tuple, range or Space.'
+            self.description['intrinsic_loops'][name] = level
+            self.set(name, space)
+        elif isinstance(space, OptimizeSpace):
+            space.name = name
+            space.optimizer.dimensions[name] = space.space
+            self._add_search_space(name, space.optimizer.level, space)
+            self.add_depends(space.optimizer.name, [name])
         else:
             if level is None:
                 raise ValueError('level must be provided.')
-            self._add_loop_var(name, level, range)
-            if isinstance(range, Expression) or callable(range):
-                self.add_depends(name, range.symbols())
+            try:
+                space = Space.fromarray(space)
+            except:
+                pass
+            self._add_search_space(name, level, space)
+            if isinstance(space, Expression) or callable(space):
+                self.add_depends(name, space.symbols())
         if setter:
             self.description['setters'][name] = setter
 
@@ -481,7 +459,14 @@ class Scan():
 
         self._variables = {'self': self}
 
-        await update_variables(self._variables, self.description['consts'],
+        consts = {}
+        for k, v in self.description['consts'].items():
+            if isinstance(v, Space):
+                consts[k] = v.toarray()
+            else:
+                consts[k] = v
+
+        await update_variables(self._variables, consts,
                                self.description['setters'])
         for level, total in self.description['total'].items():
             if total == np.inf:
@@ -527,7 +512,7 @@ class Scan():
         import asyncio
         self._main_task = asyncio.create_task(self.run())
 
-    async def submit(self, server='tcp://127.0.0.1:6788'):
+    async def submit(self, server=default_executor):
         assymbly(self.description)
         async with ZMQContextManager(zmq.DEALER, connect=server) as socket:
             await socket.send_pyobj({
@@ -640,51 +625,6 @@ class Scan():
             return await awaitable
 
 
-class Unpicklable:
-
-    def __init__(self, obj):
-        self.type = str(type(obj))
-        self.id = id(obj)
-
-    def __repr__(self):
-        return f'<Unpicklable: {self.type} at 0x{id(self):x}>'
-
-
-class TooLarge:
-
-    def __init__(self, obj):
-        self.type = str(type(obj))
-        self.id = id(obj)
-
-    def __repr__(self):
-        return f'<TooLarge: {self.type} at 0x{id(self):x}>'
-
-
-def dump_globals(ns=None, *, size_limit=10 * 1024 * 1024, warn=False):
-    import __main__
-
-    if ns is None:
-        ns = __main__.__dict__
-
-    namespace = {}
-
-    for name, value in ns.items():
-        try:
-            buf = dill.dumps(value)
-        except:
-            namespace[name] = Unpicklable(value)
-            if warn:
-                warnings.warn(f'Unpicklable: {name} {type(value)}')
-        if len(buf) > size_limit:
-            namespace[name] = TooLarge(value)
-            if warn:
-                warnings.warn(f'TooLarge: {name} {type(value)}')
-        else:
-            namespace[name] = buf
-
-    return namespace
-
-
 def assymbly(description):
     import __main__
     from IPython import get_ipython
@@ -810,6 +750,38 @@ def assymbly(description):
             if ready:
                 description['order'][level].append(ready)
                 keys -= set(ready)
+
+    axis = {}
+    independent_variables = set(description['intrinsic_loops'].keys())
+
+    for name in description['consts']:
+        axis[name] = ()
+    for level, range_list in description['loops'].items():
+        for name, iterable in range_list:
+            if isinstance(iterable, OptimizeSpace):
+                axis[name] = tuple(range(level + 1))
+                continue
+            elif isinstance(iterable, (np.ndarray, list, tuple, range, Space)):
+                independent_variables.add(name)
+            axis[name] = (level, )
+
+    for level, group in description['order'].items():
+        for names in group:
+            for name in names:
+                if name not in description['dependents']:
+                    if name not in axis:
+                        axis[name] = (level, )
+                else:
+                    d = set()
+                    for n in description['dependents'][name]:
+                        d.update(axis[n])
+                    if name not in axis:
+                        axis[name] = tuple(sorted(d))
+                    else:
+                        axis[name] = tuple(sorted(set(axis[name]) | d))
+    description['axis'] = axis
+    description['independent_variables'] = independent_variables
+
     return description
 
 
@@ -845,6 +817,8 @@ async def _iter_level(variables,
                 opts[iter.optimizer.name] = iter.optimizer.create()
         elif isinstance(iter, Expression):
             iters_d[name] = iter.eval(env)
+        elif isinstance(iter, Space):
+            iters_d[name] = iter.toarray()
         elif callable(iter):
             iters_d[name] = await call_function(iter, variables)
         else:

qulab/scan/server.py

@@ -1,20 +1,18 @@
 import asyncio
 import pickle
-import sys
-import time
-import uuid
-from pathlib import Path
-from .scan import Scan
+
 import click
 import dill
-import numpy as np
 import zmq
 from loguru import logger
 
 from qulab.sys.rpc.zmq_socket import ZMQContextManager
 
+from .scan import Scan
+
 pool = {}
 
+
 class Request():
     __slots__ = ['sock', 'identity', 'msg', 'method']
 

qulab/scan/space.py

@@ -0,0 +1,172 @@
+from typing import Type
+
+import numpy as np
+import skopt
+from skopt.space import Categorical, Integer, Real
+
+
+class Space():
+
+    def __init__(self, function, *args, **kwds):
+        self.function = function
+        self.args = args
+        self.kwds = kwds
+
+    def __repr__(self):
+        if self.function == 'asarray':
+            return repr(self.args[0])
+        args = ', '.join(map(repr, self.args))
+        kwds = ', '.join(f'{k}={v!r}' for k, v in self.kwds.items())
+        return f"{self.function}({args}, {kwds})"
+
+    def __len__(self):
+        return len(self.toarray())
+
+    @classmethod
+    def fromarray(cls, array):
+        if isinstance(array, Space):
+            return array
+        if isinstance(array, (list, tuple)):
+            array = np.array(array)
+        try:
+            a = np.linspace(array[0], array[-1], len(array), dtype=array.dtype)
+            if np.allclose(a, array):
+                return cls('linspace',
+                           array[0],
+                           array[-1],
+                           len(array),
+                           dtype=array.dtype)
+        except:
+            pass
+        try:
+            a = np.logspace(np.log10(array[0]),
+                            np.log10(array[-1]),
+                            len(array),
+                            base=10,
+                            dtype=array.dtype)
+            if np.allclose(a, array):
+                return cls('logspace',
+                           np.log10(array[0]),
+                           np.log10(array[-1]),
+                           len(array),
+                           base=10,
+                           dtype=array.dtype)
+        except:
+            pass
+        try:
+            a = np.logspace(np.log2(array[0]),
+                            np.log2(array[-1]),
+                            len(array),
+                            base=2,
+                            dtype=array.dtype)
+            if np.allclose(a, array):
+                return cls('logspace',
+                           np.log2(array[0]),
+                           np.log2(array[-1]),
+                           len(array),
+                           base=2,
+                           dtype=array.dtype)
+        except:
+            pass
+        try:
+            a = np.geomspace(array[0],
+                             array[-1],
+                             len(array),
+                             dtype=array.dtype)
+            if np.allclose(a, array):
+                return cls('geomspace',
+                           array[0],
+                           array[-1],
+                           len(array),
+                           dtype=array.dtype)
+        except:
+            pass
+        return array
+
+    def toarray(self):
+        return getattr(np, self.function)(*self.args, **self.kwds)
+
+
+def logspace(start, stop, num=50, endpoint=True, base=10):
+    return Space('logspace', start, stop, num, endpoint=endpoint, base=base)
+
+
+def linspace(start, stop, num=50, endpoint=True):
+    return Space('linspace', start, stop, num, endpoint=endpoint)
+
+
+def geomspace(start, stop, num=50, endpoint=True):
+    return Space('geomspace', start, stop, num, endpoint=endpoint)
+
+
+class OptimizeSpace():
+
+    def __init__(self, optimizer: 'Optimizer', space):
+        self.optimizer = optimizer
+        self.space = space
+        self.name = None
+
+    def __len__(self):
+        return self.optimizer.maxiter
+
+
+class Optimizer():
+
+    def __init__(self,
+                 scanner,
+                 name: str,
+                 level: int,
+                 method: str | Type = skopt.Optimizer,
+                 maxiter: int = 1000,
+                 minimize: bool = True,
+                 **kwds):
+        self.scanner = scanner
+        self.method = method
+        self.maxiter = maxiter
+        self.dimensions = {}
+        self.name = name
+        self.level = level
+        self.kwds = kwds
+        self.minimize = minimize
+
+    def create(self):
+        return self.method(list(self.dimensions.values()), **self.kwds)
+
+    def Categorical(self,
+                    categories,
+                    prior=None,
+                    transform=None,
+                    name=None) -> OptimizeSpace:
+        return OptimizeSpace(self,
+                             Categorical(categories, prior, transform, name))
+
+    def Integer(self,
+                low,
+                high,
+                prior="uniform",
+                base=10,
+                transform=None,
+                name=None,
+                dtype=np.int64) -> OptimizeSpace:
+        return OptimizeSpace(
+            self, Integer(low, high, prior, base, transform, name, dtype))
+
+    def Real(self,
+             low,
+             high,
+             prior="uniform",
+             base=10,
+             transform=None,
+             name=None,
+             dtype=float) -> OptimizeSpace:
+        return OptimizeSpace(
+            self, Real(low, high, prior, base, transform, name, dtype))
+
+    def __getstate__(self) -> dict:
+        state = self.__dict__.copy()
+        del state['scanner']
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        self.scanner = None

qulab/scan/utils.py

@@ -1,11 +1,59 @@
 import ast
 import asyncio
 import inspect
+import warnings
 from typing import Any, Callable
 
+import dill
+
 from .expression import Env, Expression
 
 
+class Unpicklable:
+
+    def __init__(self, obj):
+        self.type = str(type(obj))
+        self.id = id(obj)
+
+    def __repr__(self):
+        return f'<Unpicklable: {self.type} at 0x{id(self):x}>'
+
+
+class TooLarge:
+
+    def __init__(self, obj):
+        self.type = str(type(obj))
+        self.id = id(obj)
+
+    def __repr__(self):
+        return f'<TooLarge: {self.type} at 0x{id(self):x}>'
+
+
+def dump_globals(ns=None, *, size_limit=10 * 1024 * 1024, warn=False):
+    import __main__
+
+    if ns is None:
+        ns = __main__.__dict__
+
+    namespace = {}
+
+    for name, value in ns.items():
+        try:
+            buf = dill.dumps(value)
+        except:
+            namespace[name] = Unpicklable(value)
+            if warn:
+                warnings.warn(f'Unpicklable: {name} {type(value)}')
+        if len(buf) > size_limit:
+            namespace[name] = TooLarge(value)
+            if warn:
+                warnings.warn(f'TooLarge: {name} {type(value)}')
+        else:
+            namespace[name] = buf
+
+    return namespace
+
+
 def is_valid_identifier(s: str) -> bool:
     """
     Check if a string is a valid identifier.

qulab/version.py

@@ -1 +1 @@
-__version__ = "2.1.0"
+__version__ = "2.1.2"