algo2code 0.2.1__py3-none-any.whl

algo2code/__init__.py

@@ -0,0 +1,72 @@
+"""algo2code — Transpile LaTeX algorithm boxes (algpseudocode) to executable code.
+
+Targets: Taichi (MVP), NumPy, C/PETSc.
+
+Usage:
+    from algo2code import transpile
+
+    taichi_code = transpile(latex_source, backend='taichi')
+"""
+
+__version__ = "0.2.1"
+
+from .algo_parser import parse_algorithm
+from .ast_nodes import Algorithm, VarType
+from .backends.taichi_codegen import (
+    RUNTIME_INLINE,
+    RUNTIME_TI_RUNTIME,
+    generate_taichi,
+)
+from .errors import Algo2CodeError, UnsupportedConstructError
+from .expr_parser import parse_latex_expr
+from .library import PCG_ALGORITHM_LATEX, get_pcg_algorithm_latex
+from .type_inference import infer_types
+
+
+def transpile(source: str, backend: str = "taichi", runtime: str = "inline") -> str:
+    """
+    Full pipeline:  LaTeX source → parsed AST → type inference → code generation.
+
+    Parameters
+    ----------
+    source : str
+        LaTeX source containing \\begin{algorithmic} ... \\end{algorithmic}
+        with optional % directive comments.
+    backend : str
+        Target backend. Currently only 'taichi' is supported.
+    runtime : str
+        Import mode for the Taichi backend.  ``"inline"`` (default) emits
+        private ``@ti.kernel`` definitions inline.  ``"ti_runtime"`` emits
+        ``from ti_runtime import vector_ops as _v`` and calls the shared
+        primitives for dot/norm/copy.  algo2code itself never imports
+        ti_runtime regardless of this setting.
+
+    Returns
+    -------
+    str
+        Generated source code.
+    """
+    algo = parse_algorithm(source)
+    infer_types(algo)
+
+    if backend == "taichi":
+        return generate_taichi(algo, runtime=runtime)
+    else:
+        raise ValueError(f"Unknown backend: {backend!r}. Supported: 'taichi'")
+
+
+__all__ = [
+    "PCG_ALGORITHM_LATEX",
+    "RUNTIME_INLINE",
+    "RUNTIME_TI_RUNTIME",
+    "Algo2CodeError",
+    "Algorithm",
+    "UnsupportedConstructError",
+    "VarType",
+    "generate_taichi",
+    "get_pcg_algorithm_latex",
+    "infer_types",
+    "parse_algorithm",
+    "parse_latex_expr",
+    "transpile",
+]

algo2code/algo_parser.py

@@ -0,0 +1,458 @@
+r"""
+Parser for LaTeX `algpseudocode` environments.
+
+Recognises:
+  \begin{algorithmic}  ...  \end{algorithmic}
+  \State $lhs = rhs$       % optional_comment
+  \For{$k = 0, 1, \ldots, N$}  ...  \EndFor
+  \While{$cond$}            ...  \EndWhile
+  \If{$cond$}               ...  \ElsIf{$cond$}  ...  \Else  ...  \EndIf
+  \Return $expr$
+  \State \textbf{break}
+
+Directive comments (lines starting with `%` OUTSIDE algorithmic):
+  % algorithm <name>
+  % backend <taichi|numpy|petsc>
+  % args <name>:<type>, ...
+  % type <varname> <scalar|vector|matrix|callable>
+
+The parser delegates all $...$ math fragments to expr_parser.
+"""
+
+from __future__ import annotations
+
+import re
+
+from .ast_nodes import (
+    Algorithm,
+    Assign,
+    Branch,
+    Break,
+    ForLoop,
+    Return,
+    Stmt,
+    Var,
+    VarType,
+    WhileLoop,
+)
+from .expr_parser import parse_assignment, parse_condition, parse_latex_expr
+
+# ── Directive parsing ────────────────────────────────────────────────────────
+
+_TYPE_MAP = {
+    "scalar": VarType.SCALAR,
+    "vector": VarType.VECTOR,
+    "matrix": VarType.MATRIX,
+    "callable": VarType.CALLABLE,
+    "matvec": VarType.MATRIX,  # alias
+}
+
+
+def _parse_directives(lines: list[str]) -> dict:
+    """Parse % directive comments before the algorithmic block."""
+    directives: dict = {
+        "name": "algorithm",
+        "backend": "taichi",
+        "args": [],
+        "types": {},
+    }
+
+    for line in lines:
+        line = line.strip()
+        if not line.startswith("%"):
+            continue
+        line = line[1:].strip()
+
+        if line.startswith("algorithm "):
+            directives["name"] = line.split(None, 1)[1].strip()
+
+        elif line.startswith("backend "):
+            directives["backend"] = line.split(None, 1)[1].strip()
+
+        elif line.startswith("args "):
+            arg_str = line.split(None, 1)[1]
+            for arg in arg_str.split(","):
+                arg = arg.strip()
+                if ":" in arg:
+                    name, typ = arg.split(":", 1)
+                    directives["args"].append(
+                        (name.strip(), _TYPE_MAP.get(typ.strip(), VarType.UNKNOWN))
+                    )
+                else:
+                    directives["args"].append((arg, VarType.UNKNOWN))
+
+        elif line.startswith("type "):
+            parts = line.split()
+            if len(parts) >= 3:
+                varname = parts[1]
+                vtype = _TYPE_MAP.get(parts[2], VarType.UNKNOWN)
+                directives["types"][varname] = vtype
+
+    return directives
+
+
+# ── Main parser ──────────────────────────────────────────────────────────────
+
+
+class AlgPseudocodeParser:
+    """
+    Parse a complete LaTeX source containing an algorithmic environment.
+
+    Usage:
+        algo = AlgPseudocodeParser(latex_string).parse()
+    """
+
+    def __init__(self, source: str):
+        self.source = source
+        self.lines: list[str] = []
+        self.pos = 0
+
+    def parse(self) -> Algorithm:
+        """Parse the full source and return an Algorithm AST."""
+        pre_lines, body_lines = self._split_sections()
+        directives = _parse_directives(pre_lines)
+
+        type_annotations = dict(directives["types"])
+        self._collect_inline_types(body_lines, type_annotations)
+
+        self.lines = body_lines
+        self.pos = 0
+        stmts = self._parse_block(terminators=[])
+
+        return Algorithm(
+            name=directives["name"],
+            backend=directives["backend"],
+            args=directives["args"],
+            body=stmts,
+            type_annotations=type_annotations,
+        )
+
+    def _split_sections(self) -> tuple[list[str], list[str]]:
+        """Split source into pre-algorithmic lines and body lines."""
+        all_lines = self.source.split("\n")
+        pre_lines = []
+        body_lines = []
+        in_body = False
+
+        for line in all_lines:
+            stripped = line.strip()
+
+            if re.match(r"\\begin\{algorithmic\}", stripped):
+                in_body = True
+                continue
+            if re.match(r"\\end\{algorithmic\}", stripped):
+                in_body = False
+                continue
+            if re.match(r"\\begin\{algorithm\}", stripped):
+                continue
+            if re.match(r"\\end\{algorithm\}", stripped):
+                continue
+            if re.match(r"\\caption\{", stripped):
+                continue
+
+            if in_body:
+                if stripped:  # skip blank lines
+                    body_lines.append(line)
+            else:
+                pre_lines.append(line)
+
+        return pre_lines, body_lines
+
+    def _collect_inline_types(self, lines: list[str], types: dict):
+        """Extract inline type annotations from % comments on \\State lines."""
+        for line in lines:
+            m = re.search(r"%\s*(\w+)\s*$", line)
+            if m:
+                hint = m.group(1).lower()
+                if hint in _TYPE_MAP:
+                    lhs_match = re.search(r"\\State\s+\$\s*([^=$]+?)\s*=", line)
+                    if lhs_match:
+                        var_name = self._extract_var_name(lhs_match.group(1))
+                        if var_name:
+                            types[var_name] = _TYPE_MAP[hint]
+
+    def _extract_var_name(self, lhs_latex: str) -> str | None:
+        """Extract a clean variable name from a LaTeX LHS fragment."""
+        lhs = re.sub(r"\\(?:mathbf|boldsymbol|bm)\{([^}]*)\}", r"\1", lhs_latex)
+        lhs = re.sub(r"_\{\\text\{([^}]*)\}\}", r"_\1", lhs)
+        lhs = re.sub(r"\\(\w+)", r"\1", lhs)
+        lhs = lhs.strip().replace(" ", "")
+        return lhs if lhs else None
+
+    # ── Statement parsing ────────────────────────────────────────────────
+
+    def _current_line(self) -> str | None:
+        if self.pos < len(self.lines):
+            return self.lines[self.pos].strip()
+        return None
+
+    def _advance(self):
+        self.pos += 1
+
+    def _parse_block(self, terminators: list[str]) -> list[Stmt]:
+        """Parse statements until a terminator command is found."""
+        stmts = []
+        while self.pos < len(self.lines):
+            line = self._current_line()
+            if line is None:
+                break
+
+            stripped = self._strip_comment(line)
+            if any(stripped.startswith(t) for t in terminators):
+                break
+
+            stmt = self._parse_statement()
+            if stmt is not None:
+                stmts.append(stmt)
+
+        return stmts
+
+    def _strip_comment(self, line: str) -> str:
+        """Remove trailing % comment but preserve % inside $...$."""
+        in_math = False
+        for i, ch in enumerate(line):
+            if ch == "$":
+                in_math = not in_math
+            elif ch == "%" and not in_math:
+                return line[:i].strip()
+        return line.strip()
+
+    def _extract_inline_comment(self, line: str) -> str:
+        """Extract the % comment portion."""
+        in_math = False
+        for i, ch in enumerate(line):
+            if ch == "$":
+                in_math = not in_math
+            elif ch == "%" and not in_math:
+                return line[i + 1 :].strip()
+        return ""
+
+    def _parse_statement(self) -> Stmt | None:
+        """Parse a single statement from the current line."""
+        line = self._current_line()
+        if line is None:
+            return None
+
+        stripped = self._strip_comment(line)
+        comment = self._extract_inline_comment(line)
+
+        if stripped.startswith("\\For"):
+            return self._parse_for(stripped)
+
+        if stripped.startswith("\\While"):
+            return self._parse_while(stripped)
+
+        if stripped.startswith("\\If"):
+            return self._parse_if(stripped)
+
+        if stripped.startswith("\\Return") or stripped.startswith("\\State \\Return"):
+            self._advance()
+            return self._parse_return(stripped)
+
+        if re.search(r"\\textbf\{break\}|\\Break|\\textbf\{Break\}", stripped):
+            self._advance()
+            return Break()
+
+        if stripped.startswith("\\State"):
+            self._advance()
+            return self._parse_state(stripped, comment)
+
+        # Skip unrecognized lines
+        self._advance()
+        return None
+
+    def _extract_math(self, text: str) -> str:
+        """Extract content between $ delimiters."""
+        m = re.search(r"\$(.+?)\$", text)
+        if m:
+            return m.group(1).strip()
+        return text.strip()
+
+    def _extract_brace_arg(self, text: str, command: str) -> str:
+        """Extract the {argument} after a \\Command."""
+        idx = text.find(command)
+        if idx < 0:
+            return ""
+        rest = text[idx + len(command) :]
+
+        brace_start = rest.find("{")
+        if brace_start < 0:
+            return ""
+
+        depth = 0
+        start = brace_start
+        for i in range(brace_start, len(rest)):
+            if rest[i] == "{":
+                depth += 1
+            elif rest[i] == "}":
+                depth -= 1
+                if depth == 0:
+                    return rest[start + 1 : i].strip()
+        return rest[start + 1 :].strip()
+
+    # ── For loop ─────────────────────────────────────────────────────────
+
+    def _parse_for(self, line: str) -> ForLoop:
+        r"""Parse \\For{$k = 0, 1, \\ldots, N$}  body  \\EndFor"""
+        arg = self._extract_brace_arg(line, "\\For")
+        arg = arg.strip("$ ")
+
+        var, start, end_expr = self._parse_for_range(arg)
+
+        self._advance()  # past the \For line
+        body = self._parse_block(terminators=["\\EndFor"])
+
+        cur_line = self._current_line()
+        if cur_line and self._strip_comment(cur_line).startswith("\\EndFor"):
+            self._advance()
+
+        return ForLoop(var=var, start=start, end_expr=end_expr, body=body)
+
+    def _parse_for_range(self, arg: str) -> tuple[str, int, str]:
+        r"""
+        Parse for-loop range specifications:
+          k = 0, 1, ..., N     → var='k', start=0, end='N'
+          k = 0, 1, 2, ...     → var='k', start=0, end=''
+          k = 1 to N           → var='k', start=1, end='N'
+        """
+        m = re.match(
+            r"([a-zA-Z]\w*)\s*=\s*(\d+)\s*,\s*\d+\s*,?\s*"
+            r"(?:\\ldots|\\dots|\\cdots|\.\.\.)\s*(?:,\s*)?"
+            r"(?:\\(?:text|mathrm)\{(\w+)\}|([a-zA-Z]\w*))?",
+            arg,
+        )
+        if m:
+            var = m.group(1)
+            start = int(m.group(2))
+            end_expr = m.group(3) or m.group(4) or ""
+            return var, start, end_expr
+
+        m = re.match(r"([a-zA-Z]\w*)\s*=\s*(\d+)\s+(?:to|\\to)\s+(\w+)", arg)
+        if m:
+            return m.group(1), int(m.group(2)), m.group(3)
+
+        m = re.match(r"([a-zA-Z]\w*)", arg)
+        var = m.group(1) if m else "k"
+        return var, 0, ""
+
+    # ── While loop ───────────────────────────────────────────────────────
+
+    def _parse_while(self, line: str) -> WhileLoop:
+        arg = self._extract_brace_arg(line, "\\While")
+        arg = arg.strip("$ ")
+        condition = parse_condition(arg)
+
+        self._advance()
+        body = self._parse_block(terminators=["\\EndWhile"])
+
+        cur_line = self._current_line()
+        if cur_line and self._strip_comment(cur_line).startswith("\\EndWhile"):
+            self._advance()
+
+        return WhileLoop(condition=condition, body=body)
+
+    # ── If / ElsIf / Else ────────────────────────────────────────────────
+
+    def _parse_if(self, line: str) -> Branch:
+        arg = self._extract_brace_arg(line, "\\If")
+        arg = arg.strip("$ ")
+        condition = parse_condition(arg)
+
+        self._advance()
+        if_body = self._parse_block(terminators=["\\EndIf", "\\ElsIf", "\\Else"])
+
+        elif_branches = []
+        else_body: list[Stmt] = []
+
+        while True:
+            cur_line = self._current_line()
+            if not cur_line:
+                break
+            cur = self._strip_comment(cur_line)
+            if cur.startswith("\\ElsIf"):
+                elif_arg = self._extract_brace_arg(cur, "\\ElsIf")
+                elif_arg = elif_arg.strip("$ ")
+                elif_cond = parse_condition(elif_arg)
+                self._advance()
+                elif_body = self._parse_block(terminators=["\\EndIf", "\\ElsIf", "\\Else"])
+                elif_branches.append((elif_cond, elif_body))
+            elif cur.startswith("\\Else"):
+                self._advance()
+                else_body = self._parse_block(terminators=["\\EndIf"])
+                break
+            else:
+                break
+
+        cur_line = self._current_line()
+        if cur_line and self._strip_comment(cur_line).startswith("\\EndIf"):
+            self._advance()
+
+        return Branch(
+            condition=condition,
+            if_body=if_body,
+            elif_branches=elif_branches,
+            else_body=else_body,
+        )
+
+    # ── Return ───────────────────────────────────────────────────────────
+
+    def _parse_return(self, line: str) -> Return:
+        m = re.search(r"\\Return\s*(.*)", self._strip_comment(line))
+        if not m:
+            return Return(values=[])
+
+        rest = m.group(1).strip().strip("$").strip()
+        if not rest:
+            return Return(values=[])
+
+        values = []
+        for part in self._split_top_level(rest, ","):
+            part = part.strip()
+            if part:
+                values.append(parse_latex_expr(part))
+
+        return Return(values=values)
+
+    def _split_top_level(self, text: str, sep: str) -> list[str]:
+        """Split text by separator, respecting brace depth."""
+        parts = []
+        depth = 0
+        current: list[str] = []
+        for ch in text:
+            if ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+            elif ch == sep and depth == 0:
+                parts.append("".join(current))
+                current = []
+                continue
+            current.append(ch)
+        parts.append("".join(current))
+        return parts
+
+    # ── State (assignment) ───────────────────────────────────────────────
+
+    def _parse_state(self, line: str, comment: str) -> Stmt | None:
+        r"""Parse \\State $lhs = rhs$"""
+        stripped = self._strip_comment(line)
+        stripped = re.sub(r"^\\State\s*", "", stripped).strip()
+        math = self._extract_math(stripped)
+        if not math:
+            return None
+
+        result = parse_assignment(math)
+        if result is None:
+            expr = parse_latex_expr(math)
+            return Assign(target=Var(name="_"), value=expr, comment=comment)
+
+        target, value = result
+        return Assign(target=target, value=value, comment=comment)
+
+
+# ── Public API ───────────────────────────────────────────────────────────────
+
+
+def parse_algorithm(source: str) -> Algorithm:
+    """Parse a LaTeX source containing an algorithmic environment."""
+    return AlgPseudocodeParser(source).parse()

algo2code/ast_nodes.py

@@ -0,0 +1,171 @@
+"""
+AST node definitions for algorithmic pseudocode → code transpiler.
+
+Two-level AST:
+  1. Control-flow nodes (ForLoop, WhileLoop, Branch, Assign, Return, Break)
+  2. Expression nodes (BinOp, UnaryOp, FuncCall, Var, Scalar, Norm, Transpose)
+
+The expression nodes are *typed* (scalar / vector / matrix / callable)
+to drive correct code generation for Taichi kernels vs. Python-scope code.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum, auto
+
+# ── Type system ──────────────────────────────────────────────────────────────
+
+
+class VarType(Enum):
+    SCALAR = auto()
+    VECTOR = auto()
+    MATRIX = auto()
+    CALLABLE = auto()  # e.g. preconditioner M^{-1}
+    UNKNOWN = auto()
+
+
+# ── Expression nodes ─────────────────────────────────────────────────────────
+
+
+@dataclass
+class Expr:
+    """Base class for all expression nodes."""
+
+    inferred_type: VarType = field(default=VarType.UNKNOWN, repr=False)
+
+
+@dataclass
+class Var(Expr):
+    """A named variable:  x, r, A, \\alpha, \\rho_{\\text{new}}, etc."""
+
+    name: str = ""
+    subscript: str | None = None  # e.g. "new" from \rho_{\text{new}}
+
+    def __post_init__(self):
+        super().__init__()
+
+    @property
+    def display(self) -> str:
+        if self.subscript:
+            return f"{self.name}_{self.subscript}"
+        return self.name
+
+
+@dataclass
+class Number(Expr):
+    """A literal number."""
+
+    value: float = 0.0
+
+    def __post_init__(self):
+        super().__init__(inferred_type=VarType.SCALAR)
+
+
+@dataclass
+class BinOp(Expr):
+    """Binary operation:  a + b,  A*p,  r^T z,  a/b,  etc."""
+
+    op: str = ""  # '+', '-', '*', '/', 'dot', 'matvec', 'matmul'
+    left: Expr = field(default_factory=Expr)
+    right: Expr = field(default_factory=Expr)
+
+    def __post_init__(self):
+        super().__init__()
+
+
+@dataclass
+class UnaryOp(Expr):
+    """Unary operation:  -x, ‖r‖, etc."""
+
+    op: str = ""  # 'neg', 'norm', 'transpose'
+    operand: Expr = field(default_factory=Expr)
+
+    def __post_init__(self):
+        super().__init__()
+
+
+@dataclass
+class FuncCall(Expr):
+    """Function application:  M^{-1}(r), f(x,y), etc."""
+
+    func: Expr = field(default_factory=Expr)
+    args: list[Expr] = field(default_factory=list)
+
+    def __post_init__(self):
+        super().__init__()
+
+
+# ── Control-flow nodes ───────────────────────────────────────────────────────
+
+
+@dataclass
+class Stmt:
+    """Base class for all statement nodes."""
+
+    pass
+
+
+@dataclass
+class Assign(Stmt):
+    """Assignment:  x = expr"""
+
+    target: Var = field(default_factory=Var)
+    value: Expr = field(default_factory=Expr)
+    comment: str = ""  # inline annotation: "% vector", "% matvec"
+
+
+@dataclass
+class ForLoop(Stmt):
+    """For loop:  \\For{k = 0, 1, ..., maxiter}  body  \\EndFor"""
+
+    var: str = ""
+    start: int = 0
+    end_expr: str = ""  # e.g. "maxiter", or "" for indefinite
+    body: list[Stmt] = field(default_factory=list)
+
+
+@dataclass
+class WhileLoop(Stmt):
+    """While loop:  \\While{cond}  body  \\EndWhile"""
+
+    condition: Expr = field(default_factory=Expr)
+    body: list[Stmt] = field(default_factory=list)
+
+
+@dataclass
+class Branch(Stmt):
+    """If / ElsIf / Else:  \\If{cond}  body  \\EndIf"""
+
+    condition: Expr = field(default_factory=Expr)
+    if_body: list[Stmt] = field(default_factory=list)
+    elif_branches: list[tuple[Expr, list[Stmt]]] = field(default_factory=list)
+    else_body: list[Stmt] = field(default_factory=list)
+
+
+@dataclass
+class Return(Stmt):
+    """Return statement:  \\Return expr"""
+
+    values: list[Expr] = field(default_factory=list)
+
+
+@dataclass
+class Break(Stmt):
+    """Break statement."""
+
+    pass
+
+
+# ── Top-level algorithm ──────────────────────────────────────────────────────
+
+
+@dataclass
+class Algorithm:
+    """A complete parsed algorithm."""
+
+    name: str = ""
+    backend: str = "taichi"
+    args: list[tuple[str, VarType]] = field(default_factory=list)
+    body: list[Stmt] = field(default_factory=list)
+    type_annotations: dict[str, VarType] = field(default_factory=dict)  # var_name -> type

algo2code/backends/__init__.py

@@ -0,0 +1 @@
+"""Code generation backends."""

algo2code/backends/c_petsc_codegen.py

@@ -0,0 +1 @@
+"""C/PETSc backend (planned)."""