algo2code 0.2.1__tar.gz → 0.2.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {algo2code-0.2.1 → algo2code-0.2.2}/PKG-INFO +9 -2
- {algo2code-0.2.1 → algo2code-0.2.2}/README.md +8 -1
- {algo2code-0.2.1 → algo2code-0.2.2}/pyproject.toml +1 -1
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/__init__.py +11 -6
- algo2code-0.2.2/src/algo2code/backends/numpy_codegen.py +154 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/backends/taichi_codegen.py +13 -9
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/expr_parser.py +1 -2
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/__init__.py +11 -0
- algo2code-0.2.2/src/algo2code/library/algorithms/jacobi.tex +51 -0
- algo2code-0.2.2/src/algo2code/library/algorithms/newmark.tex +99 -0
- algo2code-0.2.2/src/algo2code/library/algorithms/pcg.tex +103 -0
- algo2code-0.2.2/src/algo2code/library/algorithms/radial_return_j2.tex +79 -0
- algo2code-0.2.2/src/algo2code/library/algorithms/radial_return_j2_kinematic.tex +82 -0
- algo2code-0.2.2/src/algo2code/library/algorithms/radial_return_j2_mixed.tex +114 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/pcg.py +3 -6
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/radial_return_j2.py +10 -28
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/radial_return_j2_kinematic.py +7 -25
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/radial_return_j2_mixed.py +7 -25
- algo2code-0.2.2/src/algo2code/library/sources.py +90 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_end_to_end.py +1 -1
- algo2code-0.2.2/tests/test_library_sources.py +46 -0
- algo2code-0.2.2/tests/test_numpy_codegen.py +147 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_smoke.py +3 -1
- algo2code-0.2.1/src/algo2code/backends/numpy_codegen.py +0 -1
- {algo2code-0.2.1 → algo2code-0.2.2}/.gitignore +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/LICENSE +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/__init__.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/algo_parser.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/ast_nodes.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/demo.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/expr_parser.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/main.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/pcg.tex +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/pcg_generated.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/taichi_codegen.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/test_algo2code.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/type_inference.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/algo_parser.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/ast_nodes.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/backends/__init__.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/backends/c_petsc_codegen.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/errors.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/.gitkeep +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/type_inference.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/conftest.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_algo_parser.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_deferrals.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_expr_parser.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_fail_loud.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_radial_return_codegen.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_radial_return_j2_kinematic_codegen.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_radial_return_j2_mixed_codegen.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_taichi_codegen.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_transpose_alias.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_type_inference.py +0 -0
- {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_vector_lowering.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: algo2code
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.2
|
|
4
4
|
Summary: Transpile LaTeX algorithm boxes (algpseudocode) to executable Taichi/NumPy/C code
|
|
5
5
|
Project-URL: Repository, https://github.com/CEmM2/MechDSL
|
|
6
6
|
Author: Shmuel Osovski
|
|
@@ -17,7 +17,14 @@ Description-Content-Type: text/markdown
|
|
|
17
17
|
|
|
18
18
|
# algo2code
|
|
19
19
|
|
|
20
|
-
Transpile LaTeX algorithm boxes (`algpseudocode`) to executable code targeting Taichi
|
|
20
|
+
Transpile LaTeX algorithm boxes (`algpseudocode`) to executable code targeting Taichi or NumPy (C/PETSc planned).
|
|
21
|
+
|
|
22
|
+
```python
|
|
23
|
+
from algo2code import transpile
|
|
24
|
+
|
|
25
|
+
transpile(source, backend="taichi") # Taichi kernels + Python driver
|
|
26
|
+
transpile(source, backend="numpy") # plain function that imports only NumPy
|
|
27
|
+
```
|
|
21
28
|
|
|
22
29
|
Parses `\begin{algorithmic}...\end{algorithmic}` environments with type-directed code generation.
|
|
23
30
|
Zero runtime dependencies — standard library only.
|
|
@@ -1,6 +1,13 @@
|
|
|
1
1
|
# algo2code
|
|
2
2
|
|
|
3
|
-
Transpile LaTeX algorithm boxes (`algpseudocode`) to executable code targeting Taichi
|
|
3
|
+
Transpile LaTeX algorithm boxes (`algpseudocode`) to executable code targeting Taichi or NumPy (C/PETSc planned).
|
|
4
|
+
|
|
5
|
+
```python
|
|
6
|
+
from algo2code import transpile
|
|
7
|
+
|
|
8
|
+
transpile(source, backend="taichi") # Taichi kernels + Python driver
|
|
9
|
+
transpile(source, backend="numpy") # plain function that imports only NumPy
|
|
10
|
+
```
|
|
4
11
|
|
|
5
12
|
Parses `\begin{algorithmic}...\end{algorithmic}` environments with type-directed code generation.
|
|
6
13
|
Zero runtime dependencies — standard library only.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
"""algo2code — Transpile LaTeX algorithm boxes (algpseudocode) to executable code.
|
|
2
2
|
|
|
3
|
-
Targets: Taichi
|
|
3
|
+
Targets: Taichi, NumPy (C/PETSc planned).
|
|
4
4
|
|
|
5
5
|
Usage:
|
|
6
6
|
from algo2code import transpile
|
|
@@ -8,10 +8,11 @@ Usage:
|
|
|
8
8
|
taichi_code = transpile(latex_source, backend='taichi')
|
|
9
9
|
"""
|
|
10
10
|
|
|
11
|
-
__version__ = "0.2.
|
|
11
|
+
__version__ = "0.2.2"
|
|
12
12
|
|
|
13
13
|
from .algo_parser import parse_algorithm
|
|
14
14
|
from .ast_nodes import Algorithm, VarType
|
|
15
|
+
from .backends.numpy_codegen import generate_numpy
|
|
15
16
|
from .backends.taichi_codegen import (
|
|
16
17
|
RUNTIME_INLINE,
|
|
17
18
|
RUNTIME_TI_RUNTIME,
|
|
@@ -33,9 +34,11 @@ def transpile(source: str, backend: str = "taichi", runtime: str = "inline") ->
|
|
|
33
34
|
LaTeX source containing \\begin{algorithmic} ... \\end{algorithmic}
|
|
34
35
|
with optional % directive comments.
|
|
35
36
|
backend : str
|
|
36
|
-
Target backend
|
|
37
|
+
Target backend: ``"taichi"`` (default) or ``"numpy"``. The NumPy
|
|
38
|
+
backend emits a plain function that imports only ``numpy``, so the
|
|
39
|
+
result runs without Taichi installed.
|
|
37
40
|
runtime : str
|
|
38
|
-
Import mode for the Taichi backend. ``"inline"`` (default) emits
|
|
41
|
+
Import mode for the Taichi backend (ignored by ``"numpy"``). ``"inline"`` (default) emits
|
|
39
42
|
private ``@ti.kernel`` definitions inline. ``"ti_runtime"`` emits
|
|
40
43
|
``from ti_runtime import vector_ops as _v`` and calls the shared
|
|
41
44
|
primitives for dot/norm/copy. algo2code itself never imports
|
|
@@ -51,8 +54,9 @@ def transpile(source: str, backend: str = "taichi", runtime: str = "inline") ->
|
|
|
51
54
|
|
|
52
55
|
if backend == "taichi":
|
|
53
56
|
return generate_taichi(algo, runtime=runtime)
|
|
54
|
-
|
|
55
|
-
|
|
57
|
+
if backend == "numpy":
|
|
58
|
+
return generate_numpy(algo)
|
|
59
|
+
raise ValueError(f"Unknown backend: {backend!r}. Supported: 'taichi', 'numpy'")
|
|
56
60
|
|
|
57
61
|
|
|
58
62
|
__all__ = [
|
|
@@ -63,6 +67,7 @@ __all__ = [
|
|
|
63
67
|
"Algorithm",
|
|
64
68
|
"UnsupportedConstructError",
|
|
65
69
|
"VarType",
|
|
70
|
+
"generate_numpy",
|
|
66
71
|
"generate_taichi",
|
|
67
72
|
"get_pcg_algorithm_latex",
|
|
68
73
|
"infer_types",
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
"""
|
|
2
|
+
NumPy code generator.
|
|
3
|
+
|
|
4
|
+
Takes the same typed Algorithm AST as the Taichi backend and emits a plain
|
|
5
|
+
Python function that depends on NumPy only. The generated module never
|
|
6
|
+
imports Taichi, so algorithms transpiled with this backend run on a lean
|
|
7
|
+
install (no ``[verify]`` extra).
|
|
8
|
+
|
|
9
|
+
The lowering mirrors :class:`~algo2code.backends.taichi_codegen.TaichiEmitter`
|
|
10
|
+
statement for statement, so the two backends agree on semantics:
|
|
11
|
+
|
|
12
|
+
- Scalar arithmetic and control flow are emitted exactly as in the Taichi
|
|
13
|
+
driver (plain Python on floats), so a scalar-only algorithm produces the
|
|
14
|
+
same function body under both backends and yields bit-identical results.
|
|
15
|
+
- Vector/matrix variables are ``numpy.ndarray`` buffers. Local working
|
|
16
|
+
vectors are allocated once with ``np.zeros(n)`` and every vector operation
|
|
17
|
+
writes *in place* (``out[...] = ...``), matching the Taichi field model:
|
|
18
|
+
a vector argument updated by the algorithm is updated for the caller too.
|
|
19
|
+
- Callables keep the in-place seam contract of the Taichi backend:
|
|
20
|
+
``M_inv(r, z)`` writes into ``z``, and a matrix-free operator declared
|
|
21
|
+
``% type A callable`` is applied as ``A(out, p)``.
|
|
22
|
+
- ``dot`` lowers to ``float(np.dot(a, b))``, ``norm`` to
|
|
23
|
+
``float(np.linalg.norm(a))``, a stored-matrix product to ``A @ x``.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
from __future__ import annotations
|
|
27
|
+
|
|
28
|
+
from ..ast_nodes import Algorithm, Assign, BinOp, Expr, FuncCall, UnaryOp, Var, VarType
|
|
29
|
+
from ..errors import UnsupportedConstructError
|
|
30
|
+
from .taichi_codegen import (
|
|
31
|
+
RUNTIME_INLINE,
|
|
32
|
+
TaichiEmitter,
|
|
33
|
+
_collect_arg_names,
|
|
34
|
+
_collect_vector_vars,
|
|
35
|
+
_var_name,
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class NumpyEmitter(TaichiEmitter):
|
|
40
|
+
"""Emit a NumPy-only Python function from a typed Algorithm AST."""
|
|
41
|
+
|
|
42
|
+
BACKEND_DOC = "Backend: NumPy (CPU)"
|
|
43
|
+
|
|
44
|
+
def __init__(self, algo: Algorithm):
|
|
45
|
+
super().__init__(algo, runtime=RUNTIME_INLINE)
|
|
46
|
+
|
|
47
|
+
def emit(self) -> str:
|
|
48
|
+
"""Generate the complete NumPy source."""
|
|
49
|
+
self._arg_names = _collect_arg_names(self.algo)
|
|
50
|
+
self._vector_vars = _collect_vector_vars(self.algo)
|
|
51
|
+
self._pre_scan_temps(self.algo.body)
|
|
52
|
+
|
|
53
|
+
parts = ["import numpy as np\n", "# ── Solver driver " + "─" * 57, self._emit_driver()]
|
|
54
|
+
return "\n\n".join(parts) + "\n"
|
|
55
|
+
|
|
56
|
+
# ── Allocation ───────────────────────────────────────────────────────
|
|
57
|
+
|
|
58
|
+
def _vector_field_alloc(self, vec_arg: str | None) -> str:
|
|
59
|
+
return "np.zeros(n)"
|
|
60
|
+
|
|
61
|
+
# ── Expressions ──────────────────────────────────────────────────────
|
|
62
|
+
|
|
63
|
+
def _emit_assign(self, stmt: Assign) -> None:
|
|
64
|
+
target = _var_name(stmt.target)
|
|
65
|
+
target_type = self.algo.type_annotations.get(stmt.target.display, stmt.target.inferred_type)
|
|
66
|
+
if (
|
|
67
|
+
target_type == VarType.VECTOR
|
|
68
|
+
and isinstance(stmt.value, Var)
|
|
69
|
+
and stmt.value.inferred_type == VarType.VECTOR
|
|
70
|
+
):
|
|
71
|
+
self._write(f"{target}[...] = {_var_name(stmt.value)}")
|
|
72
|
+
return
|
|
73
|
+
super()._emit_assign(stmt)
|
|
74
|
+
|
|
75
|
+
def _emit_unary(self, expr: UnaryOp) -> str:
|
|
76
|
+
if expr.op == "norm" and expr.operand.inferred_type != VarType.SCALAR:
|
|
77
|
+
return f"float(np.linalg.norm({self._emit_expr_str(expr.operand)}))"
|
|
78
|
+
return super()._emit_unary(expr)
|
|
79
|
+
|
|
80
|
+
def _emit_dot(self, expr: BinOp) -> str:
|
|
81
|
+
left = expr.left
|
|
82
|
+
if isinstance(left, UnaryOp) and left.op == "transpose":
|
|
83
|
+
left = left.operand
|
|
84
|
+
return f"float(np.dot({self._emit_expr_str(left)}, {self._emit_expr_str(expr.right)}))"
|
|
85
|
+
|
|
86
|
+
def _emit_func_call_name(self, expr: FuncCall) -> str:
|
|
87
|
+
if isinstance(expr.func, Var) and _var_name(expr.func) == "sqrt":
|
|
88
|
+
return "np.sqrt"
|
|
89
|
+
return super()._emit_func_call_name(expr)
|
|
90
|
+
|
|
91
|
+
# ── Vector lowering (in-place writes into ndarray buffers) ───────────
|
|
92
|
+
|
|
93
|
+
def _lower_vector_expr(self, expr: Expr, dest: str | None) -> str:
|
|
94
|
+
if isinstance(expr, Var):
|
|
95
|
+
name = _var_name(expr)
|
|
96
|
+
if dest is not None and dest != name:
|
|
97
|
+
self._write(f"{dest}[...] = {name}")
|
|
98
|
+
return dest
|
|
99
|
+
return name
|
|
100
|
+
|
|
101
|
+
if isinstance(expr, FuncCall):
|
|
102
|
+
out = dest if dest is not None else self._get_temp()
|
|
103
|
+
func = self._emit_func_call_name(expr)
|
|
104
|
+
args = ", ".join(self._emit_expr_str(a) for a in expr.args)
|
|
105
|
+
self._write(f"{func}({args}, {out})")
|
|
106
|
+
return out
|
|
107
|
+
|
|
108
|
+
if isinstance(expr, UnaryOp) and expr.op == "neg":
|
|
109
|
+
vec = self._lower_to_field(expr.operand)
|
|
110
|
+
out = dest if dest is not None else self._get_temp()
|
|
111
|
+
self._write(f"{out}[...] = -{vec}")
|
|
112
|
+
return out
|
|
113
|
+
|
|
114
|
+
if isinstance(expr, BinOp):
|
|
115
|
+
if expr.op == "matvec":
|
|
116
|
+
out = dest if dest is not None else self._get_temp()
|
|
117
|
+
vec = self._lower_to_field(expr.right)
|
|
118
|
+
if expr.left.inferred_type == VarType.CALLABLE:
|
|
119
|
+
self._write(f"{self._emit_expr_str(expr.left)}({out}, {vec})")
|
|
120
|
+
return out
|
|
121
|
+
self._write(f"{out}[...] = {self._emit_expr_str(expr.left)} @ {vec}")
|
|
122
|
+
return out
|
|
123
|
+
|
|
124
|
+
if expr.op in ("scale", "*"):
|
|
125
|
+
scalar, vec_expr = self._split_scale(expr)
|
|
126
|
+
if scalar is not None and vec_expr is not None:
|
|
127
|
+
vec = self._lower_to_field(vec_expr)
|
|
128
|
+
out = dest if dest is not None else self._get_temp()
|
|
129
|
+
self._write(f"{out}[...] = {scalar} * {vec}")
|
|
130
|
+
return out
|
|
131
|
+
|
|
132
|
+
if expr.op in ("+", "-"):
|
|
133
|
+
left = self._lower_to_field(expr.left)
|
|
134
|
+
coeff, right = self._lower_addend(expr.right)
|
|
135
|
+
beta = self._negate(coeff) if expr.op == "-" else coeff
|
|
136
|
+
out = dest if dest is not None else self._get_temp()
|
|
137
|
+
self._write(f"{out}[...] = {left} + {beta} * {right}")
|
|
138
|
+
return out
|
|
139
|
+
|
|
140
|
+
raise UnsupportedConstructError(
|
|
141
|
+
f"cannot lower vector expression to NumPy ops: {expr!r}. "
|
|
142
|
+
f"Supported vector forms are copy, callable application, matvec, "
|
|
143
|
+
f"scalar*vector, and vector +/- vector."
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def generate_numpy(algo: Algorithm) -> str:
|
|
148
|
+
"""Generate NumPy source code from a typed Algorithm AST.
|
|
149
|
+
|
|
150
|
+
The emitted module imports only ``numpy`` and defines one function named
|
|
151
|
+
after the algorithm, with the same signature and return values as the
|
|
152
|
+
Taichi backend's driver.
|
|
153
|
+
"""
|
|
154
|
+
return NumpyEmitter(algo).emit()
|
|
@@ -18,9 +18,9 @@ Import modes (``runtime`` parameter of :func:`generate_taichi`):
|
|
|
18
18
|
each generated file — backward-compatible, no runtime dependency.
|
|
19
19
|
- ``"ti_runtime"``: emit ``from ti_runtime import vector_ops as _v`` and call
|
|
20
20
|
the shared primitives for ``dot``/``norm``/``copy``/``vec_add``; only
|
|
21
|
-
``_matvec`` remains inlined (the matrix-free operator
|
|
22
|
-
|
|
23
|
-
the import line.
|
|
21
|
+
``_matvec`` remains inlined (the matrix-free operator is injected through
|
|
22
|
+
the seam by the consumer). algo2code itself never imports ``ti_runtime`` —
|
|
23
|
+
it only *emits* the import line.
|
|
24
24
|
|
|
25
25
|
Compatibility note (``"ti_runtime"`` mode):
|
|
26
26
|
``ti_runtime.vector_ops.dot`` uses ``x[I].dot(y[I])`` which requires
|
|
@@ -29,7 +29,7 @@ Compatibility note (``"ti_runtime"`` mode):
|
|
|
29
29
|
the consumer passes scalar fields. Callers that use vector fields (e.g.
|
|
30
30
|
mechdsl FEM drivers) are unaffected. This is a known field-model mismatch;
|
|
31
31
|
resolving it (either by changing the emitted field model to ``ti.Vector.field``
|
|
32
|
-
or by adding scalar variants to ti_runtime) is out of scope for
|
|
32
|
+
or by adding scalar variants to ti_runtime) is out of scope for this backend.
|
|
33
33
|
"""
|
|
34
34
|
|
|
35
35
|
from __future__ import annotations
|
|
@@ -177,7 +177,7 @@ class KernelCollector:
|
|
|
177
177
|
In ``"ti_runtime"`` mode the shared primitives (dot/norm/copy/vec_add)
|
|
178
178
|
are called via the ``_v`` alias imported from ``ti_runtime.vector_ops``.
|
|
179
179
|
Only ``_matvec`` has no ti_runtime equivalent and remains inlined
|
|
180
|
-
(the matrix-free operator seam
|
|
180
|
+
(the matrix-free operator is injected through the seam by the consumer).
|
|
181
181
|
"""
|
|
182
182
|
lines = []
|
|
183
183
|
if runtime == RUNTIME_TI_RUNTIME:
|
|
@@ -266,6 +266,9 @@ class TaichiEmitter:
|
|
|
266
266
|
docstring for the full description of each mode.
|
|
267
267
|
"""
|
|
268
268
|
|
|
269
|
+
#: Backend line written into the generated driver's docstring.
|
|
270
|
+
BACKEND_DOC = "Backend: Taichi (GPU)"
|
|
271
|
+
|
|
269
272
|
def __init__(self, algo: Algorithm, runtime: str = RUNTIME_INLINE):
|
|
270
273
|
self.algo = algo
|
|
271
274
|
self.runtime = runtime
|
|
@@ -341,7 +344,7 @@ class TaichiEmitter:
|
|
|
341
344
|
|
|
342
345
|
self._write('"""')
|
|
343
346
|
self._write("Auto-generated from LaTeX algorithmic environment.")
|
|
344
|
-
self._write(
|
|
347
|
+
self._write(self.BACKEND_DOC)
|
|
345
348
|
self._write('"""')
|
|
346
349
|
|
|
347
350
|
vec_arg = self._first_vector_arg()
|
|
@@ -394,8 +397,8 @@ class TaichiEmitter:
|
|
|
394
397
|
locals matching the *vector argument's* component layout (``<arg>.n``):
|
|
395
398
|
the ``ti_runtime`` ``dot``/``norm2`` primitives use ``x[I].dot(y[I])``,
|
|
396
399
|
which requires vector fields. This is what lets the matrix-free seam
|
|
397
|
-
|
|
398
|
-
mismatch
|
|
400
|
+
run end-to-end against ``ti_runtime`` (resolving the field-model
|
|
401
|
+
mismatch documented above). With no vector argument the component count
|
|
399
402
|
is unknown, so fall back to the scalar layout.
|
|
400
403
|
"""
|
|
401
404
|
if self.runtime == RUNTIME_TI_RUNTIME and vec_arg is not None:
|
|
@@ -847,7 +850,8 @@ def generate_taichi(algo: Algorithm, runtime: str = RUNTIME_INLINE) -> str:
|
|
|
847
850
|
``"ti_runtime"``
|
|
848
851
|
Emit ``from ti_runtime import vector_ops as _v`` and call the
|
|
849
852
|
shared primitives for dot/norm/copy/vec_add. Only ``_matvec``
|
|
850
|
-
remains inlined (the matrix-free operator
|
|
853
|
+
remains inlined (the matrix-free operator is injected through the
|
|
854
|
+
seam by the consumer).
|
|
851
855
|
algo2code itself never imports ti_runtime — it only emits the
|
|
852
856
|
import line.
|
|
853
857
|
|
|
@@ -225,8 +225,7 @@ class ExprParser:
|
|
|
225
225
|
|
|
226
226
|
Division at the term level shares precedence with multiplication
|
|
227
227
|
(left-associative) so ``a + b / c`` parses as ``a + (b / c)``
|
|
228
|
-
rather than dropping the divisor.
|
|
229
|
-
parser fix.
|
|
228
|
+
rather than dropping the divisor.
|
|
230
229
|
"""
|
|
231
230
|
left = self.parse_signed_factor()
|
|
232
231
|
while True:
|
|
@@ -6,14 +6,25 @@ small helper that returns it (and, where the parser supports it, the parsed
|
|
|
6
6
|
here instead of re-typing or re-inventing the algorithm text — this is the
|
|
7
7
|
canonical hand-off seam between ``algo2code`` and any consumer that wants a
|
|
8
8
|
reference implementation of a textbook iterative solver.
|
|
9
|
+
|
|
10
|
+
The verbatim ``.tex`` sources live in ``algo2code/library/algorithms/`` and
|
|
11
|
+
ship as package data; :func:`algorithm_source` reads any of them by name.
|
|
9
12
|
"""
|
|
10
13
|
|
|
11
14
|
from algo2code.library.pcg import (
|
|
12
15
|
PCG_ALGORITHM_LATEX,
|
|
13
16
|
get_pcg_algorithm_latex,
|
|
14
17
|
)
|
|
18
|
+
from algo2code.library.sources import (
|
|
19
|
+
ALGORITHM_SOURCES,
|
|
20
|
+
algorithm_source,
|
|
21
|
+
algorithm_source_path,
|
|
22
|
+
)
|
|
15
23
|
|
|
16
24
|
__all__ = [
|
|
25
|
+
"ALGORITHM_SOURCES",
|
|
17
26
|
"PCG_ALGORITHM_LATEX",
|
|
27
|
+
"algorithm_source",
|
|
28
|
+
"algorithm_source_path",
|
|
18
29
|
"get_pcg_algorithm_latex",
|
|
19
30
|
]
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
% Point-Jacobi preconditioner apply.
|
|
2
|
+
%
|
|
3
|
+
% Computes the Jacobi (diagonal) preconditioned vector:
|
|
4
|
+
%
|
|
5
|
+
% z_I = r_I / max(d_I, eps)
|
|
6
|
+
%
|
|
7
|
+
% where ``d`` is the diagonal of the system matrix (point-Jacobi: M = diag(d)),
|
|
8
|
+
% ``r`` is the residual vector, ``z`` is the output (M^{-1} r), and ``eps`` is a
|
|
9
|
+
% small positive guard against near-zero diagonal entries. All operands are
|
|
10
|
+
% degree-of-freedom vectors (one value per node or per component).
|
|
11
|
+
%
|
|
12
|
+
% Grammar note:
|
|
13
|
+
% -------------
|
|
14
|
+
% algo2code currently lowers vector/vector division (`z = r / d`) to the Python
|
|
15
|
+
% expression `z = (r / d)`, which is not runnable for `ti.Vector.field` operands
|
|
16
|
+
% (Taichi raises TypeError: unsupported operand type(s) for /: MatrixField /
|
|
17
|
+
% MatrixField). The grammar supports scalar operations and vector +/-/scale/
|
|
18
|
+
% callable-matvec, but has no `vec_ediv` (elementwise divide) primitive.
|
|
19
|
+
%
|
|
20
|
+
% The clean fix is a new `vec_ediv(z, r, d, eps)` primitive in ti_runtime
|
|
21
|
+
% vector_ops. The generated body below is therefore
|
|
22
|
+
% realised as a thin `@ti.data_oriented` adapter
|
|
23
|
+
% (`mechdsl.solver.jacobi_preconditioner.GeneratedJacobiPreconditioner`) whose
|
|
24
|
+
% `_apply` kernel calls `ti_runtime.vector_ops.ediv(z, r, d, eps)` — a
|
|
25
|
+
% generated body in the sense that the algorithm is expressed in LaTeX and
|
|
26
|
+
% realised via a structured adapter, not hand-coded heuristics.
|
|
27
|
+
%
|
|
28
|
+
% When algo2code gains a `vec_ediv` primitive (or scalar-field elementwise ops),
|
|
29
|
+
% this box can be re-transpiled to a fully auto-generated driver with no adapter
|
|
30
|
+
% shim.
|
|
31
|
+
%
|
|
32
|
+
% Inputs
|
|
33
|
+
% ------
|
|
34
|
+
% r : residual vector (numerator)
|
|
35
|
+
% d : diagonal vector (denominator)
|
|
36
|
+
% eps : small positive guard (scalar) -- assumed positive (SPD diagonal guard)
|
|
37
|
+
%
|
|
38
|
+
% Output
|
|
39
|
+
% ------
|
|
40
|
+
% z : preconditioned vector = M^{-1} r
|
|
41
|
+
|
|
42
|
+
% algorithm jacobi_apply
|
|
43
|
+
% backend taichi
|
|
44
|
+
% args r:vector, d:vector, eps:scalar
|
|
45
|
+
|
|
46
|
+
% type z vector
|
|
47
|
+
|
|
48
|
+
\begin{algorithmic}
|
|
49
|
+
\State $z_I = r_I \, / \, \max(d_I,\, \varepsilon)$ \Comment{elementwise; for all DOF indices $I$}
|
|
50
|
+
\Return $z$
|
|
51
|
+
\end{algorithmic}
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
% Newmark-beta implicit time integration — single step.
|
|
2
|
+
%
|
|
3
|
+
% The canonical Newmark-beta step authored for the **all-Taichi on-device
|
|
4
|
+
% time-integration seam path** (the temporal analogue of the matrix-free PCG in
|
|
5
|
+
% ``pcg.tex``): the per-step acceleration solve ``solve_a`` is a *callable*
|
|
6
|
+
% (matrix-free), so this box transpiles via
|
|
7
|
+
% ``algo2code.transpile(..., runtime='ti_runtime')`` into a body that advances a
|
|
8
|
+
% dynamic state ``(u, v, a)`` one step over ``ti.Vector.field`` DOF vectors
|
|
9
|
+
% **without NumPy in the hot path** — every update is a ``ti_runtime.vector_ops``
|
|
10
|
+
% AXPBY/copy call and the only non-primitive call is the injected ``solve_a``.
|
|
11
|
+
%
|
|
12
|
+
% Scheme (Hughes 2000, Box 9.1; Newmark 1959)
|
|
13
|
+
% -------------------------------------------
|
|
14
|
+
% For the semidiscrete equation of motion ``M a + C v + K u = F`` the Newmark
|
|
15
|
+
% family advances ``(u_n, v_n, a_n) -> (u_{n+1}, v_{n+1}, a_{n+1})`` via the
|
|
16
|
+
% predictor/corrector split:
|
|
17
|
+
%
|
|
18
|
+
% predictors (known from step n):
|
|
19
|
+
% u_pred = u_n + dt*v_n + dt^2*(1/2 - beta)*a_n
|
|
20
|
+
% v_pred = v_n + dt*(1 - gamma)*a_n
|
|
21
|
+
%
|
|
22
|
+
% acceleration solve (the only system solve — folded into the callable):
|
|
23
|
+
% a_{n+1} = solve_a(u_pred, v_pred)
|
|
24
|
+
% i.e. solve (M + gamma*dt*C + beta*dt^2*K) a_{n+1}
|
|
25
|
+
% = F_{n+1} - C*v_pred - K*u_pred
|
|
26
|
+
%
|
|
27
|
+
% correctors:
|
|
28
|
+
% u_{n+1} = u_pred + beta*dt^2 * a_{n+1}
|
|
29
|
+
% v_{n+1} = v_pred + gamma*dt * a_{n+1}
|
|
30
|
+
%
|
|
31
|
+
% gamma = 1/2 gives second-order accuracy; beta = 1/4 (with gamma = 1/2) is the
|
|
32
|
+
% average-acceleration / trapezoidal rule — unconditionally stable, no numerical
|
|
33
|
+
% damping. HHT-alpha is a one-parameter generalisation (alpha-weighted residual);
|
|
34
|
+
% it reuses the same predictor/corrector skeleton and is intentionally NOT
|
|
35
|
+
% authored here — Newmark-beta is the deliverable, HHT is optional/next.
|
|
36
|
+
%
|
|
37
|
+
% Seam contract (mirrors ti_runtime.seams.TimeIntegrationContext)
|
|
38
|
+
% ---------------------------------------------------------------
|
|
39
|
+
% solve_a(u_pred, v_pred, a_out) — acceleration solve: a_out = a_{n+1}
|
|
40
|
+
% (out LAST; the in-place callable convention,
|
|
41
|
+
% same shape as the PCG operator A(out, x)).
|
|
42
|
+
% The callable owns the mass / damping / stiffness / external-force data; the box
|
|
43
|
+
% is mass-/material-agnostic, exactly like ``pcg.tex`` is operator-agnostic.
|
|
44
|
+
%
|
|
45
|
+
% Grammar note
|
|
46
|
+
% ------------
|
|
47
|
+
% This box transpiles **directly** through algo2code in ``runtime='ti_runtime'``
|
|
48
|
+
% mode with no grammar gap: the predictor/corrector updates are all
|
|
49
|
+
% ``scalar*vector`` and ``vector +/- vector`` (fused ``_v.vec_add`` AXPBY calls),
|
|
50
|
+
% the scalar coefficients ``c0..c3`` are plain scalar arithmetic, and the
|
|
51
|
+
% acceleration solve is an in-place ``callable`` (``solve_a``) — the matrix-free
|
|
52
|
+
% seam form (11-ALGO2CODE §8.3). No NumPy, no dense ``_matvec``. (Verified: the
|
|
53
|
+
% emitted ``newmark_step`` body contains only ``_v.vec_add`` / ``_v.copy`` /
|
|
54
|
+
% ``solve_a(...)``.)
|
|
55
|
+
%
|
|
56
|
+
% The two predictor scalar coefficients are precomputed into named scalars
|
|
57
|
+
% (``c0``, ``c1``) and the corrector coefficients into (``c2``, ``c3``) so each
|
|
58
|
+
% vector update fuses to a single AXPBY; ``dt`` itself is the predictor velocity
|
|
59
|
+
% coefficient (``u + dt*v``).
|
|
60
|
+
%
|
|
61
|
+
% Inputs
|
|
62
|
+
% ------
|
|
63
|
+
% u : displacement vector at step n (carries the state; updated in place)
|
|
64
|
+
% v : velocity vector at step n (updated in place)
|
|
65
|
+
% a : acceleration vector at step n (updated in place)
|
|
66
|
+
% solve_a : acceleration solve (callable, in-place ``solve_a(u_pred, v_pred, a_out)``)
|
|
67
|
+
% dt : time step (scalar)
|
|
68
|
+
% beta : Newmark beta (scalar; 1/4 = average acceleration)
|
|
69
|
+
% gamma : Newmark gamma (scalar; 1/2 = second order, no algorithmic damping)
|
|
70
|
+
%
|
|
71
|
+
% Output
|
|
72
|
+
% ------
|
|
73
|
+
% u, v, a — the advanced state (u_{n+1}, v_{n+1}, a_{n+1})
|
|
74
|
+
|
|
75
|
+
% algorithm newmark_step
|
|
76
|
+
% backend taichi
|
|
77
|
+
% args u:vector, v:vector, a:vector, solve_a:callable, dt:scalar, beta:scalar, gamma:scalar
|
|
78
|
+
|
|
79
|
+
% type u_pred vector
|
|
80
|
+
% type v_pred vector
|
|
81
|
+
% type a_new vector
|
|
82
|
+
% type c_0 scalar
|
|
83
|
+
% type c_1 scalar
|
|
84
|
+
% type c_2 scalar
|
|
85
|
+
% type c_3 scalar
|
|
86
|
+
|
|
87
|
+
\begin{algorithmic}
|
|
88
|
+
\State $c_0 = dt \cdot dt \cdot (0.5 - \beta)$ % scalar
|
|
89
|
+
\State $c_1 = dt \cdot (1 - \gamma)$ % scalar
|
|
90
|
+
\State $c_2 = \beta \cdot dt \cdot dt$ % scalar
|
|
91
|
+
\State $c_3 = \gamma \cdot dt$ % scalar
|
|
92
|
+
\State $u_{\text{pred}} = u + dt \cdot v + c_0 \cdot a$ % vector
|
|
93
|
+
\State $v_{\text{pred}} = v + c_1 \cdot a$ % vector
|
|
94
|
+
\State $a_{\text{new}} = \text{solve\_a}(u_{\text{pred}}, v_{\text{pred}})$ % vector
|
|
95
|
+
\State $u = u_{\text{pred}} + c_2 \cdot a_{\text{new}}$ % vector
|
|
96
|
+
\State $v = v_{\text{pred}} + c_3 \cdot a_{\text{new}}$ % vector
|
|
97
|
+
\State $a = a_{\text{new}}$ % vector
|
|
98
|
+
\Return $u, v, a$
|
|
99
|
+
\end{algorithmic}
|
|
@@ -0,0 +1,103 @@
|
|
|
1
|
+
% Matrix-free Preconditioned Conjugate Gradient (PCG).
|
|
2
|
+
%
|
|
3
|
+
% The canonical PCG algorithm authored for the **all-Taichi on-device seam path**:
|
|
4
|
+
% the operator ``A`` and the preconditioner ``M_inv`` are both
|
|
5
|
+
% *callables* (matrix-free), so this box transpiles via
|
|
6
|
+
% ``algo2code.transpile(..., runtime='ti_runtime')`` into a body that solves
|
|
7
|
+
% ``A x = b`` over ``ti.Vector.field`` DOF vectors without ever forming or storing
|
|
8
|
+
% a matrix and without NumPy in the hot path.
|
|
9
|
+
%
|
|
10
|
+
% Provenance
|
|
11
|
+
% ----------
|
|
12
|
+
% Identical in algorithm to the canonical PCG LaTeX
|
|
13
|
+
% (``algo2code.library.pcg.PCG_ALGORITHM_LATEX`` / the hand twin
|
|
14
|
+
% ``mechdsl.solver.import_adapter.Algo2CodePCGSolver``) and to the matrix-free
|
|
15
|
+
% ``_CALLABLE_PCG_LATEX`` proven runnable in runtime mode. The *only* surface change vs
|
|
16
|
+
% the canonical box is the operator type:
|
|
17
|
+
%
|
|
18
|
+
% canonical: ``% args A:matrix, ... apply_M_inv:callable, ...``
|
|
19
|
+
% here: ``% args A:callable, ... M_inv:callable, ...``
|
|
20
|
+
%
|
|
21
|
+
% A matrix ``A`` lowers to a dense scalar-indexed ``_matvec`` (incompatible with
|
|
22
|
+
% the ``ti.Vector.field`` layout the seam uses); a callable ``A`` lowers to the
|
|
23
|
+
% in-place ``A(out, p)`` operator call — the ``ti_runtime`` ``apply_A(out, x)``
|
|
24
|
+
% seam contract. The runtime-mode guard requires exactly this: a callable
|
|
25
|
+
% operator + a vector argument.
|
|
26
|
+
%
|
|
27
|
+
% Seam contract (mirrors ti_runtime.seams.LinearSolveContext)
|
|
28
|
+
% -----------------------------------------------------------
|
|
29
|
+
% A(out, x) — operator: out = A @ x (out FIRST)
|
|
30
|
+
% M_inv(r, z) — preconditioner: z = M^{-1} r (out LAST)
|
|
31
|
+
%
|
|
32
|
+
% This box productionizes the spike's hand-written ``pcg`` body
|
|
33
|
+
% (``packages/mechdsl-core/tests/spike/svk_hex8_taichi.py``): the generated body
|
|
34
|
+
% maps line-for-line to that hand-written PCG, but is *derived from LaTeX* rather
|
|
35
|
+
% than hand-coded.
|
|
36
|
+
%
|
|
37
|
+
% Inputs
|
|
38
|
+
% ------
|
|
39
|
+
% A : matrix-free operator (callable, in-place ``A(out, x)``)
|
|
40
|
+
% b : right-hand side vector
|
|
41
|
+
% x : solution vector (carries the initial guess; solved in place)
|
|
42
|
+
% M_inv : preconditioner apply (callable, in-place ``M_inv(r, z)``)
|
|
43
|
+
% tol : relative convergence tolerance (scalar)
|
|
44
|
+
% maxiter : maximum iteration count (scalar)
|
|
45
|
+
%
|
|
46
|
+
% Output
|
|
47
|
+
% ------
|
|
48
|
+
% x, iterations, residual_norm, converged
|
|
49
|
+
%
|
|
50
|
+
% ``converged`` is a hard non-convergence flag: literal ``1`` on the two
|
|
51
|
+
% success returns (the ``r_0 = 0`` early-out and the in-loop tolerance hit),
|
|
52
|
+
% literal ``0`` on maxiter exhaustion and on the ``|pq| < 1e-300`` breakdown
|
|
53
|
+
% (which ``break``s out of the loop and falls through to the maxiter return).
|
|
54
|
+
% A caller that sees ``converged == 0`` must NOT consume ``x`` -- the inner
|
|
55
|
+
% solve failed and ``x`` is garbage (see seam_solve.make_seam_solver, which
|
|
56
|
+
% raises rather than letting a Newton step advance on it).
|
|
57
|
+
|
|
58
|
+
% algorithm pcg
|
|
59
|
+
% backend taichi
|
|
60
|
+
% args A:callable, b:vector, x:vector, M_inv:callable, tol:scalar, maxiter:scalar
|
|
61
|
+
|
|
62
|
+
% type r vector
|
|
63
|
+
% type z vector
|
|
64
|
+
% type p vector
|
|
65
|
+
% type q vector
|
|
66
|
+
% type rho scalar
|
|
67
|
+
% type rho_new scalar
|
|
68
|
+
% type alpha scalar
|
|
69
|
+
% type beta scalar
|
|
70
|
+
% type pq scalar
|
|
71
|
+
% type r0_norm scalar
|
|
72
|
+
% type r_norm scalar
|
|
73
|
+
|
|
74
|
+
\begin{algorithmic}
|
|
75
|
+
\State $r = b - A \cdot x$ % vector
|
|
76
|
+
\State $r_0 = \lVert r \rVert_2$ % scalar
|
|
77
|
+
\If{$r_0 = 0$}
|
|
78
|
+
\Return $x, 0, 0, 1$
|
|
79
|
+
\EndIf
|
|
80
|
+
\State $z = M^{-1}(r)$ % vector
|
|
81
|
+
\State $p = z$ % vector
|
|
82
|
+
\State $\rho = r^\top z$ % scalar
|
|
83
|
+
\For{$k = 1, 2, \ldots, \text{maxiter}$}
|
|
84
|
+
\State $q = A \cdot p$ % vector
|
|
85
|
+
\State $pq = p^\top q$ % scalar
|
|
86
|
+
\If{$|pq| < 10^{-300}$}
|
|
87
|
+
\State \textbf{break}
|
|
88
|
+
\EndIf
|
|
89
|
+
\State $\alpha = \frac{\rho}{pq}$ % scalar
|
|
90
|
+
\State $x = x + \alpha \, p$ % vector
|
|
91
|
+
\State $r = r - \alpha \, q$ % vector
|
|
92
|
+
\State $r_n = \lVert r \rVert_2$ % scalar
|
|
93
|
+
\If{$r_n < \text{tol} \cdot r_0$}
|
|
94
|
+
\Return $x, k, r_n, 1$
|
|
95
|
+
\EndIf
|
|
96
|
+
\State $z = M^{-1}(r)$ % vector
|
|
97
|
+
\State $\rho_{\text{new}} = r^\top z$ % scalar
|
|
98
|
+
\State $\beta = \frac{\rho_{\text{new}}}{\rho}$ % scalar
|
|
99
|
+
\State $p = z + \beta \, p$ % vector
|
|
100
|
+
\State $\rho = \rho_{\text{new}}$ % scalar
|
|
101
|
+
\EndFor
|
|
102
|
+
\Return $x, \text{maxiter}, \lVert r \rVert_2, 0$
|
|
103
|
+
\end{algorithmic}
|