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.
Files changed (56) hide show
  1. {algo2code-0.2.1 → algo2code-0.2.2}/PKG-INFO +9 -2
  2. {algo2code-0.2.1 → algo2code-0.2.2}/README.md +8 -1
  3. {algo2code-0.2.1 → algo2code-0.2.2}/pyproject.toml +1 -1
  4. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/__init__.py +11 -6
  5. algo2code-0.2.2/src/algo2code/backends/numpy_codegen.py +154 -0
  6. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/backends/taichi_codegen.py +13 -9
  7. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/expr_parser.py +1 -2
  8. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/__init__.py +11 -0
  9. algo2code-0.2.2/src/algo2code/library/algorithms/jacobi.tex +51 -0
  10. algo2code-0.2.2/src/algo2code/library/algorithms/newmark.tex +99 -0
  11. algo2code-0.2.2/src/algo2code/library/algorithms/pcg.tex +103 -0
  12. algo2code-0.2.2/src/algo2code/library/algorithms/radial_return_j2.tex +79 -0
  13. algo2code-0.2.2/src/algo2code/library/algorithms/radial_return_j2_kinematic.tex +82 -0
  14. algo2code-0.2.2/src/algo2code/library/algorithms/radial_return_j2_mixed.tex +114 -0
  15. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/pcg.py +3 -6
  16. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/radial_return_j2.py +10 -28
  17. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/radial_return_j2_kinematic.py +7 -25
  18. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/radial_return_j2_mixed.py +7 -25
  19. algo2code-0.2.2/src/algo2code/library/sources.py +90 -0
  20. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_end_to_end.py +1 -1
  21. algo2code-0.2.2/tests/test_library_sources.py +46 -0
  22. algo2code-0.2.2/tests/test_numpy_codegen.py +147 -0
  23. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_smoke.py +3 -1
  24. algo2code-0.2.1/src/algo2code/backends/numpy_codegen.py +0 -1
  25. {algo2code-0.2.1 → algo2code-0.2.2}/.gitignore +0 -0
  26. {algo2code-0.2.1 → algo2code-0.2.2}/LICENSE +0 -0
  27. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/__init__.py +0 -0
  28. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/algo_parser.py +0 -0
  29. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/ast_nodes.py +0 -0
  30. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/demo.py +0 -0
  31. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/expr_parser.py +0 -0
  32. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/main.py +0 -0
  33. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/pcg.tex +0 -0
  34. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/pcg_generated.py +0 -0
  35. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/taichi_codegen.py +0 -0
  36. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/test_algo2code.py +0 -0
  37. {algo2code-0.2.1 → algo2code-0.2.2}/prototypes/type_inference.py +0 -0
  38. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/algo_parser.py +0 -0
  39. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/ast_nodes.py +0 -0
  40. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/backends/__init__.py +0 -0
  41. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/backends/c_petsc_codegen.py +0 -0
  42. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/errors.py +0 -0
  43. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/library/.gitkeep +0 -0
  44. {algo2code-0.2.1 → algo2code-0.2.2}/src/algo2code/type_inference.py +0 -0
  45. {algo2code-0.2.1 → algo2code-0.2.2}/tests/conftest.py +0 -0
  46. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_algo_parser.py +0 -0
  47. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_deferrals.py +0 -0
  48. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_expr_parser.py +0 -0
  49. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_fail_loud.py +0 -0
  50. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_radial_return_codegen.py +0 -0
  51. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_radial_return_j2_kinematic_codegen.py +0 -0
  52. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_radial_return_j2_mixed_codegen.py +0 -0
  53. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_taichi_codegen.py +0 -0
  54. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_transpose_alias.py +0 -0
  55. {algo2code-0.2.1 → algo2code-0.2.2}/tests/test_type_inference.py +0 -0
  56. {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.1
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, NumPy, or C/PETSc.
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, NumPy, or C/PETSc.
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.
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "algo2code"
7
- version = "0.2.1"
7
+ version = "0.2.2"
8
8
  description = "Transpile LaTeX algorithm boxes (algpseudocode) to executable Taichi/NumPy/C code"
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -1,6 +1,6 @@
1
1
  """algo2code — Transpile LaTeX algorithm boxes (algpseudocode) to executable code.
2
2
 
3
- Targets: Taichi (MVP), NumPy, C/PETSc.
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.1"
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. Currently only 'taichi' is supported.
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
- else:
55
- raise ValueError(f"Unknown backend: {backend!r}. Supported: 'taichi'")
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 seam is P2-2's
22
- concern). algo2code itself never imports ``ti_runtime`` — it only *emits*
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 P2-1.
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 is P2-2's concern).
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("Backend: Taichi (GPU)")
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
- (P2-2) run end-to-end against ``ti_runtime`` (resolving the field-model
398
- mismatch P2-1 documented). With no vector argument the component count
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 seam is P2-2's concern).
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. post_recovery_plan Phase 5
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}