QuLab 2.10.10__cp313-cp313-win_amd64.whl

qulab/executor/template.py

@@ -0,0 +1,259 @@
+import ast
+import base64
+import hashlib
+import lzma
+import pickle
+import re
+import string
+import textwrap
+from typing import Any
+
+_notset = object()
+
+
+def VAR(name: str, /, *, default: Any = _notset) -> Any:
+    return name
+
+
+def encode_mapping(mapping):
+    mapping_bytes = lzma.compress(pickle.dumps(mapping))
+    hash_str = hashlib.md5(mapping_bytes).hexdigest()[:8]
+    mappping_code = '\n'.join(
+        textwrap.wrap(base64.b64encode(mapping_bytes).decode(),
+                      90,
+                      initial_indent='    ',
+                      subsequent_indent='    '))
+    return hash_str, mappping_code
+
+
+def decode_mapping(hash_str, mappping_code):
+    mapping_bytes = base64.b64decode(mappping_code)
+    if hash_str != hashlib.md5(mapping_bytes).hexdigest()[:8]:
+        raise ValueError("Hash does not match")
+    mapping = pickle.loads(lzma.decompress(mapping_bytes))
+    return mapping
+
+
+class TemplateTypeError(TypeError):
+    pass
+
+
+class TemplateKeyError(KeyError):
+    pass
+
+
+class TemplateVarExtractor(ast.NodeVisitor):
+
+    def __init__(self, fname, mapping):
+        self.var_func_def = (0, 0)
+        self.variables = set()
+        self.str_variables = set()
+        self.replacements = {}
+        self.fname = fname
+        self.mapping = mapping
+
+    def visit_Constant(self, node):
+        if isinstance(node.value, str):
+            self._process_string(node.value, node.lineno, node.col_offset,
+                                 node.end_lineno, node.end_col_offset)
+
+    def visit_JoinedStr(self, node):
+        for value in node.values:
+            if isinstance(value, ast.Constant) and isinstance(
+                    value.value, str):
+                self._process_string(value.value, value.lineno,
+                                     value.col_offset, value.end_lineno,
+                                     value.end_col_offset)
+        self.generic_visit(node)
+
+    def visit_FunctionDef(self, node):
+        if node.name == 'VAR':
+            self.var_func_def = (node.lineno, node.end_lineno)
+        self.generic_visit(node)
+
+    def visit_Call(self, node):
+        if isinstance(node.func, ast.Name) and node.func.id == 'VAR':
+            arg = node.args[0]
+            if not isinstance(arg, ast.Constant) or not isinstance(
+                    arg.value, str):
+                raise SyntaxError(
+                    f"Argument of VAR function must be a string."
+                    f" {self.fname}:{node.lineno}"
+                )
+            if len(node.args) != 1:
+                raise SyntaxError(
+                    f"VAR function only accept one argument."
+                    f" {self.fname}:{node.lineno}"
+                )
+            default = _notset
+            for k in node.keywords:
+                if k.arg == 'default':
+                    default = k.value
+                    # if isinstance(k.value, ast.Constant):
+                    #     # default = k.value.value
+                    #     default = k.value
+                    # else:
+                    #     default = k.value
+                    #     # raise SyntaxError(
+                    #     #     f"Argument of 'default' must be a constant. {self.fname}:{node.lineno}"
+                    #     # )
+                else:
+                    raise SyntaxError(
+                        f"VAR function only accept keyword argument 'default'."
+                        f" {self.fname}:{node.lineno}"
+                    )
+
+            if default is _notset:
+                # new_node = ast.Subscript(value=ast.Name(id="__VAR",
+                #                                         ctx=ast.Load()),
+                #                          slice=ast.Constant(value=arg.value),
+                #                          ctx=ast.Load())
+                if arg.value not in self.mapping:
+                    raise TemplateKeyError(
+                        f"The variable '{arg.value}' is not provided in mapping."
+                        f" {self.fname}:{node.lineno}"
+                    )
+                self.replacements[(node.lineno, node.end_lineno,
+                                   node.col_offset,
+                                   node.end_col_offset)] = ('VAR', arg.value,
+                                                            None, None)
+            else:
+                # new_node = ast.Call(
+                #     func=ast.Attribute(value=ast.Name(id='__VAR',
+                #                                       ctx=ast.Load()),
+                #                        attr='get',
+                #                        ctx=ast.Load()),
+                #     args=[ast.Constant(value=arg.value)],
+                #     keywords=[
+                #         ast.keyword(arg='default',
+                #                     value=value)
+                #     ],
+                #     ctx=ast.Load())
+                self.replacements[(node.lineno, node.end_lineno,
+                                   node.col_offset,
+                                   node.end_col_offset)] = ('VAR', arg.value,
+                                                            None, None)
+            # ast.fix_missing_locations(new_node)
+            # new_source = ast.unparse(new_node)
+            # print(new_source)
+
+            self.variables.add(arg.value)
+
+        self.generic_visit(node)
+
+    def _process_string(self, s: str, lineno: int, col_offset: int,
+                        end_lineno: int, end_col_offset: int):
+        """解析字符串内容，提取模板变量"""
+        lines = s.split('\n')
+        for offset, line in enumerate(lines):
+            current_lineno = lineno + offset
+            template = string.Template(line)
+            for var_name in template.get_identifiers():
+                if var_name not in self.mapping:
+                    raise TemplateKeyError(
+                        f"The variable '{var_name}' is not provided in mapping."
+                        f" {self.fname}:{current_lineno}")
+                if not isinstance(self.mapping[var_name], str):
+                    raise TemplateTypeError(
+                        f"Mapping value for '{var_name}' must be a string,"
+                        f" but got {type(self.mapping[var_name])}."
+                        f" {self.fname}:{current_lineno}")
+                self.str_variables.add(var_name)
+                start, stop = 0, len(line)
+                if current_lineno == lineno:
+                    start = col_offset
+                if current_lineno == end_lineno:
+                    stop = end_col_offset
+                self.replacements[(current_lineno, current_lineno, start,
+                                   stop)] = ('STR', var_name, None, None)
+
+
+def inject_mapping(source: str, mapping: dict[str, Any],
+                   fname: str) -> list[tuple[str, int]]:
+    hash_str, mapping_code = encode_mapping(mapping)
+
+    tree = ast.parse(source)
+
+    lines = source.splitlines()
+    lines_offset = [0 for _ in range(len(lines))]
+
+    extractor = TemplateVarExtractor(fname, mapping)
+    extractor.visit(tree)
+
+    # remove VAR function definition
+    if extractor.var_func_def != (0, 0):
+        for i in range(extractor.var_func_def[0] - 1,
+                       extractor.var_func_def[1]):
+            lines[i] = ''
+
+    for (lineno, end_lineno, col_offset,
+         end_col_offset), (kind, name, old_source,
+                           new_source) in extractor.replacements.items():
+        head = lines[lineno - 1][:col_offset + lines_offset[lineno - 1]]
+        tail = lines[end_lineno - 1][end_col_offset +
+                                     lines_offset[end_lineno - 1]:]
+        length_of_last_line = len(lines[end_lineno - 1])
+        content = lines[lineno - 1:end_lineno]
+        content[0] = content[0].removeprefix(head)
+        content[-1] = content[-1].removesuffix(tail)
+        content = '\n'.join(content)
+
+        if kind == 'STR':
+            template = string.Template(content)
+            formated_lines = template.substitute(mapping).splitlines()
+            formated_lines[0] = head + formated_lines[0]
+            formated_lines[-1] = formated_lines[-1] + tail
+            if len(formated_lines) == 1:
+                lines[lineno - 1] = formated_lines[0]
+                lines_offset[lineno -
+                             1] += len(lines[lineno - 1]) - length_of_last_line
+            else:
+                lines[lineno - 1:end_lineno] = formated_lines
+                lines_offset[end_lineno - 1] += len(
+                    lines[end_lineno - 1]) - length_of_last_line
+        else:
+            pattern = re.compile(
+                r'VAR\s*\(\s*'                # VAR(
+                r'(["\'])(\w+)\1'             # 捕获引号和变量名
+                r'(?:\s*,\s*default\s*=\s*'   # 匹配 default 参数
+                r'([^),]*)'                   # 捕获 default 值（排除逗号和括号）
+                r')?'                         # 可选参数组
+                r'\s*\)',                     # 闭合括号
+                re.DOTALL
+            ) # yapf: disable
+
+            def replacement(match):
+                quote = match.group(1)
+                var_name = match.group(2)
+                default_value = match.group(3)
+
+                base = f'__VAR_{hash_str}'
+                if default_value is not None:
+                    # 清除 default 值前后的空格
+                    clean_value = default_value.strip()
+                    return f'{base}.get({quote}{var_name}{quote}, {clean_value})'
+                else:
+                    return f'{base}[{quote}{var_name}{quote}]'
+
+            new_content = re.sub(pattern, replacement, content)
+
+            if lineno == end_lineno:
+                lines[lineno - 1] = head + new_content + tail
+                lines_offset[lineno -
+                             1] += len(lines[lineno - 1]) - length_of_last_line
+            else:
+                lines[lineno - 1] = head + new_content[:-1]
+                for i in range(lineno, end_lineno - 1):
+                    lines[i] = ''
+                lines[end_lineno - 1] = new_content[-1] + tail
+                lines_offset[end_lineno - 1] += len(
+                    lines[end_lineno - 1]) - length_of_last_line
+
+    injected_code = '\n'.join([
+        f"__QULAB_TEMPLATE__ = \"{fname}\"",
+        f"from qulab.executor.template import decode_mapping as __decode_{hash_str}",
+        f"__VAR_{hash_str} = __decode_{hash_str}(\"{hash_str}\", \"\"\"",
+        mapping_code, "    \"\"\")", *lines
+    ])
+
+    return injected_code, hash_str

qulab/executor/utils.py

@@ -0,0 +1,194 @@
+import inspect
+from pathlib import Path
+
+from ..cli.config import get_config_value
+from .load import load_workflow
+
+
+class Node:
+
+    def __init__(self, name: str):
+        self.name = name
+        self.dependents = []
+
+
+class Tree:
+
+    def __init__(self):
+        self.nodes = {}
+        self.heads = []
+
+    def add_node(self, node: str):
+        self.nodes[node] = Node(node)
+
+
+def dependent_tree(node: str, code_path: str | Path) -> dict[str, list[str]]:
+    '''
+    Returns a dict of nodes and their dependents.
+    '''
+    tree = {}
+    for n in load_workflow(node, code_path).depends():
+        tree[n] = dependent_tree(n, code_path)
+    return tree
+
+
+def workflow_template(workflow: str, deps: list[str]) -> str:
+    return f"""
+import numpy as np
+from loguru import logger
+
+from qulab import VAR
+from qulab.typing import Report
+
+
+# 多长时间应该检查一次校准实验，单位是秒。
+__timeout__ = 7*24*3600
+
+def depends():
+    return {deps!r}
+
+
+async def calibrate():
+    logger.info(f"running {workflow} ...")
+
+    # calibrate 是一个完整的校准实验，如power Rabi，Ramsey等。
+    # 你需要足够的扫描点，以使得后续的 analyze 可以拟合出合适的参数。
+
+    # 这里只是一个示例，实际上你需要在这里写上你的校准代码。
+    x = np.linspace(0, 2*np.pi, 101)
+    y = []
+    for i in x:
+        y.append(np.sin(i))
+
+    logger.info(f"running {workflow} ... finished!")
+    return x, y
+
+
+async def analyze(report: Report, history: Report | None = None) -> Report:
+    \"\"\"
+    分析校准结果。
+
+    report: Report
+        本次校准实验的数据。
+    history: Report | None
+        上次校准实验数据和分析结果，如果有的话。
+    \"\"\"
+    import random
+
+    # 这里添加你的分析过程，运行 calibrate 得到的数据，在 report.data 里
+    # 你可以得到校准的结果，然后根据这个结果进行分析。
+    x, y = report.data
+
+    # 完整校准后的状态有两种：OK 和 Bad，分别对应校准成功和校准失败。
+    # 校准失败是指出现坏数据，无法简单通过重新运行本次校准解决，需要
+    # 检查前置步骤。
+    report.state = random.choice(['OK', 'Bad'])
+
+    # 参数是一个字典，包含了本次校准得到的参数，后续会更新到config表中。
+    report.parameters = {{'gate.R.Q1.params.amp':1}}
+
+    # 其他信息可以是任何可序列化的内容，你可以将你想要记录的信息放在这里。
+    # 下次校准分析时，这些信息也会在 history 参数中一起传入，帮助你在下
+    # 次分析时对比参考。
+    report.other_infomation = {{}}
+
+    return report
+
+
+async def check():
+    logger.info(f"checking {workflow} ...")
+
+    # check 是一个快速检查实验，用于检查校准是否过时。
+    # 你只需要少数扫描点，让后续的 check_analyze 知道参数是否漂移，数据
+    # 坏没坏就够了，不要求拟合。
+
+    # 这里只是一个示例，实际上你需要在这里写上你的检查代码。
+    x = np.linspace(0, 2*np.pi, 5)
+    y = []
+    for i in x:
+        y.append(np.sin(i))
+
+    logger.info(f"checking {workflow} ... finished!")
+    return x, y
+
+
+async def check_analyze(report: Report, history: Report | None = None) -> Report:
+    \"\"\"
+    分析检查结果。
+
+    report: Report
+        本次检查实验的数据。
+    history: Report | None
+        上次检查实验数据和分析结果，如果有的话。
+    \"\"\"
+    import random
+
+    # 这里添加你的分析过程，运行 check 得到的数据，在 report.data 里
+    # 你可以得到校准的结果，然后根据这个结果进行分析。
+    x, y = report.data
+
+    # 状态有三种：Outdated, OK, Bad，分别对应过时、正常、坏数据。
+    # Outdated 是指数据过时，即参数漂了，需要重新校准。
+    # OK 是指数据正常，参数也没漂，不用重新校准。
+    # Bad 是指数据坏了，无法校准，需要检查前置步骤。
+    report.state = random.choice(['Outdated', 'OK', 'Bad'])
+
+    return report
+
+
+async def oracle(report: Report,
+           history: Report | None = None,
+           system_state: dict[str:str] | None = None) -> Report:
+    \"\"\"
+    谕示：指凭直觉或经验判断，改动某些配置，以期望下次校准成功。
+    
+    当校准失败时，根据 analyze 的结果，尝试改变某些配置再重新校准整个系统。
+    比如通常我们在死活测不到 rabi 或能谱时，会换一个 idle bias 再试试。这
+    里我们凭直觉设的那个 bias 值，就是一个谕示，可以通过 oracle 来设定。
+
+    该函数代入的参数 report 是 analyze 函数的返回值。
+    \"\"\"
+
+    # report.oracle['Q0.bias'] = 0.1
+    # report.oracle['Q1.bias'] = -0.03
+
+    return report
+"""
+
+
+async def debug_analyze(
+        report_index: int,
+        code_path: str | Path = get_config_value('code', Path),
+        data_path: str | Path = get_config_value('data', Path),
+) -> None:
+    from .storage import get_report_by_index
+
+    report = get_report_by_index(report_index, data_path)
+    if report is None:
+        raise ValueError(f'Invalid report index: {report_index}')
+    workflow = report.workflow
+    wf = load_workflow(workflow, code_path)
+    if wf is None:
+        raise ValueError(f'Invalid workflow: {workflow}')
+    if hasattr(wf, '__QULAB_TEMPLATE__'):
+        template_mtime = (Path(code_path) /
+                          wf.__QULAB_TEMPLATE__).stat().st_mtime
+        if template_mtime > wf.__mtime__:
+            for k in dir(wf):
+                if k.startswith('__VAR_') and len(k) == len('__VAR_17fb4dde'):
+                    var_dict = getattr(wf, k)
+                    break
+            else:
+                var_dict = {}
+            wf = load_workflow((wf.__QULAB_TEMPLATE__, workflow, var_dict),
+                               code_path)
+
+    report = wf.analyze(report, report.previous)
+    if inspect.isawaitable(report):
+        report = await report
+    if hasattr(wf, 'plot'):
+        if inspect.iscoroutinefunction(wf.plot):
+            await wf.plot(report)
+        else:
+            wf.plot(report)
+    return report