ti-runtime 0.2.1__tar.gz

ti_runtime-0.2.1/.gitignore

@@ -0,0 +1,75 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+*.egg
+dist/
+build/
+*.whl
+.eggs/
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.mypy_cache/
+.ruff_cache/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Taichi
+*.tcb
+ti_cache/
+
+# Artifacts
+*.artifact.json
+*.artifact.yaml
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Claude Code
+.claude/settings.local.json
+
+# Claude CoWork
+/MechDsl/
+*.npz
+# Golden-file regression snapshots must be tracked
+!packages/mechdsl-core/tests/golden/*.npz
+.gitnexus
+
+# Local working artifacts (orchestra runtime, plan-edit backups)
+.orchestra/
+**/.orchestra/
+*.original.md
+*.original.[0-9].md
+.claude/worktrees/
+.publisher-worktrees/
+.sesskey
+
+# Public-release staging tree (curated copy pushed to CEmM2/MechDSL)
+/dev/MechDSL/
+
+# MkDocs build output
+/site/
+
+# Generated wiki / code-intelligence index (logic-loom / akms tooling)
+.repo_wiki/
+
+# comment-sweep working tree
+/.comment-review/

ti_runtime-0.2.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Shmuel Osovski
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ti_runtime-0.2.1/PKG-INFO

@@ -0,0 +1,58 @@
+Metadata-Version: 2.4
+Name: ti-runtime
+Version: 0.2.1
+Summary: Neutral Taichi runtime: vector primitives, Tier-1 @ti.func helpers, and solver/operator injection seams for MechDSL-generated code
+Project-URL: Repository, https://github.com/CEmM2/MechDSL
+Author: Shmuel Osovski
+License-Expression: MIT
+License-File: LICENSE
+Keywords: fem,kernels,matrix-free,runtime,taichi
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: <3.14,>=3.11
+Requires-Dist: taichi>=1.7
+Description-Content-Type: text/markdown
+
+# ti-runtime
+
+Neutral Taichi runtime for **MechDSL** (PlanJune14, PJ-0).
+
+This package is the *seam + primitive* layer of the "Seams & Bodies" architecture:
+
+> MechDSL owns the **seams + primitives** (this package); algo2code generates the
+> **bodies** (solvers, preconditioners, constitutive updates, time integrators)
+> from LaTeX, injected into these seams.
+
+It contains only stable infrastructure — nothing algorithmic is hardcoded here:
+
+| Module | Contents |
+|---|---|
+| `vector_ops` | `@ti.kernel` vector primitives: `copy / axpy / xpay / scal / zero / dot / norm2` |
+| `tensor_ti`  | Tier-1 `@ti.func` helpers: `det3 / inv3 / F→C→E→J / Voigt / deviatoric / von_mises` |
+| `seams`      | injection plumbing: `Operator`, `PreconditionerBase`/`Identity`/`Diagonal`, `Solver`, `LinearSolveContext` |
+| `fields`     | `ti.init` + field-allocation boilerplate |
+| `hex8`       | Hex8 shape functions, natural-coord gradients, 2×2×2 Gauss quadrature (test harness / element operators) |
+
+## Provenance & invariants
+
+- **One-time harvest** (PlanJune14 D-B) — adapted from NumerixWeave
+  (`libs/tisolvers`, `libs/ticonstit`, `apps/tifem`); no ongoing sync.
+- **Portable output** (D-D) — generated artifacts depend on `ti_runtime`, never on
+  `mechdsl`, so they can later feed NumerixWeave / MOOSE / MFEM.
+- **Conventions** follow MechDSL `dev/design_docs/07-CONVENTIONS.md`: Voigt order
+  `[xx, yy, zz, xy, xz, yz]`, tensorial (unscaled shears), metric `diag(1,1,1,2,2,2)`.
+
+## Operator / solver contract
+
+A matrix-free operator is an in-place callable `apply(out, x) -> None` computing
+`out = A @ x` over Taichi fields. Inject it and run a (generated) solver against it:
+
+```python
+ctx = LinearSolveContext()
+ctx.set_operator(my_tangent_matvec)        # out = K(u) @ x, matrix-free @ti.kernel
+ctx.set_preconditioner(DiagonalPreconditioner(diag))
+# a generated PCG body calls ctx.apply_A / ctx.apply_preconditioner + vector_ops
+```

ti_runtime-0.2.1/README.md

@@ -0,0 +1,40 @@
+# ti-runtime
+
+Neutral Taichi runtime for **MechDSL** (PlanJune14, PJ-0).
+
+This package is the *seam + primitive* layer of the "Seams & Bodies" architecture:
+
+> MechDSL owns the **seams + primitives** (this package); algo2code generates the
+> **bodies** (solvers, preconditioners, constitutive updates, time integrators)
+> from LaTeX, injected into these seams.
+
+It contains only stable infrastructure — nothing algorithmic is hardcoded here:
+
+| Module | Contents |
+|---|---|
+| `vector_ops` | `@ti.kernel` vector primitives: `copy / axpy / xpay / scal / zero / dot / norm2` |
+| `tensor_ti`  | Tier-1 `@ti.func` helpers: `det3 / inv3 / F→C→E→J / Voigt / deviatoric / von_mises` |
+| `seams`      | injection plumbing: `Operator`, `PreconditionerBase`/`Identity`/`Diagonal`, `Solver`, `LinearSolveContext` |
+| `fields`     | `ti.init` + field-allocation boilerplate |
+| `hex8`       | Hex8 shape functions, natural-coord gradients, 2×2×2 Gauss quadrature (test harness / element operators) |
+
+## Provenance & invariants
+
+- **One-time harvest** (PlanJune14 D-B) — adapted from NumerixWeave
+  (`libs/tisolvers`, `libs/ticonstit`, `apps/tifem`); no ongoing sync.
+- **Portable output** (D-D) — generated artifacts depend on `ti_runtime`, never on
+  `mechdsl`, so they can later feed NumerixWeave / MOOSE / MFEM.
+- **Conventions** follow MechDSL `dev/design_docs/07-CONVENTIONS.md`: Voigt order
+  `[xx, yy, zz, xy, xz, yz]`, tensorial (unscaled shears), metric `diag(1,1,1,2,2,2)`.
+
+## Operator / solver contract
+
+A matrix-free operator is an in-place callable `apply(out, x) -> None` computing
+`out = A @ x` over Taichi fields. Inject it and run a (generated) solver against it:
+
+```python
+ctx = LinearSolveContext()
+ctx.set_operator(my_tangent_matvec)        # out = K(u) @ x, matrix-free @ti.kernel
+ctx.set_preconditioner(DiagonalPreconditioner(diag))
+# a generated PCG body calls ctx.apply_A / ctx.apply_preconditioner + vector_ops
+```

ti_runtime-0.2.1/pyproject.toml

@@ -0,0 +1,40 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "ti-runtime"
+version = "0.2.1"
+description = "Neutral Taichi runtime: vector primitives, Tier-1 @ti.func helpers, and solver/operator injection seams for MechDSL-generated code"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.11,<3.14"
+authors = [
+    { name = "Shmuel Osovski" },
+]
+keywords = ["taichi", "fem", "matrix-free", "runtime", "kernels"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+]
+
+# Only depends on Taichi. MechDSL-generated artifacts import THIS package
+# (`from ti_runtime import ...`), never `mechdsl` — so generated output stays
+# portable (PlanJune14 D-D). Harvested one-time from NumerixWeave (D-B); MechDSL
+# then diverges freely.
+dependencies = [
+    "taichi>=1.7",
+]
+
+[tool.uv]
+# Supply-chain safety: ignore any package version uploaded after this date.
+exclude-newer = "2026-02-20T00:00:00Z"
+
+[project.urls]
+Repository = "https://github.com/CEmM2/MechDSL"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/ti_runtime"]

ti_runtime-0.2.1/src/ti_runtime/__init__.py

@@ -0,0 +1,49 @@
+"""ti-runtime — neutral Taichi runtime for MechDSL-generated code (PlanJune14 PJ-0).
+
+Seams + primitives only; algorithmic bodies are generated by algo2code and injected.
+See the module-level READMEs for the Seams & Bodies contract.
+"""
+
+from __future__ import annotations
+
+__version__ = "0.2.1"
+
+from . import fields, hex8, seams, tensor_ti, vector_ops
+from .seams import (
+    AccelSolve,
+    DiagonalPreconditioner,
+    IdentityPreconditioner,
+    Integrator,
+    LinearSolveContext,
+    Operator,
+    PreconditionerBase,
+    Solver,
+    TimeIntegrationContext,
+)
+from .vector_ops import axpy, copy, dot, ediv, norm2, scal, vec_add, xpay, zero
+
+__all__ = [
+    "AccelSolve",
+    "DiagonalPreconditioner",
+    "IdentityPreconditioner",
+    "Integrator",
+    "LinearSolveContext",
+    "Operator",
+    "PreconditionerBase",
+    "Solver",
+    "TimeIntegrationContext",
+    "axpy",
+    "copy",
+    "dot",
+    "ediv",
+    "fields",
+    "hex8",
+    "norm2",
+    "scal",
+    "seams",
+    "tensor_ti",
+    "vec_add",
+    "vector_ops",
+    "xpay",
+    "zero",
+]

ti_runtime-0.2.1/src/ti_runtime/fields.py

@@ -0,0 +1,44 @@
+"""Field allocation + ``ti.init`` boilerplate (PlanJune14 PJ-0).
+
+The minimal runtime preamble generated code (and the test harness) targets:
+choose a backend, allocate the standard field kinds. Kept deliberately thin —
+heavy field-registry / domain-manager machinery from NumerixWeave is *not*
+harvested (PlanJune14 "avoid" list).
+"""
+
+from __future__ import annotations
+
+import taichi as ti
+
+_ARCH = {
+    "cpu": ti.cpu,
+    "gpu": ti.gpu,
+    "cuda": ti.cuda,
+    "metal": ti.metal,
+    "vulkan": ti.vulkan,
+}
+
+
+def init(arch: str = "cpu", default_fp=ti.f64, **kwargs) -> None:
+    """Initialise Taichi for a named backend (``cpu``/``gpu``/``cuda``/``metal``)."""
+    if arch not in _ARCH:
+        raise ValueError(f"Unknown arch {arch!r}; expected one of {sorted(_ARCH)}.")
+    ti.init(arch=_ARCH[arch], default_fp=default_fp, **kwargs)
+
+
+def vector_field(dim: int, n: int, dtype=ti.f64):
+    """``n`` nodes × ``dim`` dof (e.g. displacement, force, residual)."""
+    return ti.Vector.field(dim, dtype=dtype, shape=n)
+
+
+def scalar_field(n: int, dtype=ti.f64):
+    return ti.field(dtype=dtype, shape=n)
+
+
+def matrix_field(rows: int, cols: int, shape, dtype=ti.f64):
+    return ti.Matrix.field(rows, cols, dtype=dtype, shape=shape)
+
+
+def index_field(shape, dtype=ti.i32):
+    """Integer field, e.g. element→node connectivity."""
+    return ti.field(dtype=dtype, shape=shape)

ti_runtime-0.2.1/src/ti_runtime/hex8.py

@@ -0,0 +1,53 @@
+"""Hex8 reference element: shape functions, gradients, quadrature (PlanJune14 PJ-0).
+
+Trilinear 8-node hexahedron with 2×2×2 Gauss quadrature — the minimal element
+the SVK spike (PJ-1) and the matrix-free tangent operator build on. Shape
+functions/gradients are ``@ti.func`` (call from a kernel); quadrature tables are
+Python constants iterated with ``ti.static``. Standard node ordering (corner
+signs below). Adapted from NumerixWeave ``apps/tifem`` ``Ref_elements/HEX8``.
+"""
+
+import taichi as ti
+
+N_NODES = 8
+DIM = 3
+N_QP = 8
+
+# Natural-coordinate corner signs (standard Hex8 node ordering).
+_CORNERS = (
+    (-1.0, -1.0, -1.0),
+    (1.0, -1.0, -1.0),
+    (1.0, 1.0, -1.0),
+    (-1.0, 1.0, -1.0),
+    (-1.0, -1.0, 1.0),
+    (1.0, -1.0, 1.0),
+    (1.0, 1.0, 1.0),
+    (-1.0, 1.0, 1.0),
+)
+
+_G = 1.0 / (3.0**0.5)
+# 2×2×2 Gauss points (corner sign pattern, scaled by 1/√3) and unit weights.
+QUAD_POINTS = tuple((sx * _G, sy * _G, sz * _G) for (sx, sy, sz) in _CORNERS)
+QUAD_WEIGHTS = (1.0,) * N_QP
+
+
+@ti.func
+def shape(xi, eta, zeta):
+    """Trilinear shape functions at natural coords → ``ti.Vector(8)``."""
+    N = ti.Vector.zero(ti.f64, 8)
+    for a in ti.static(range(8)):
+        sx, sy, sz = ti.static(_CORNERS[a])
+        N[a] = 0.125 * (1.0 + sx * xi) * (1.0 + sy * eta) * (1.0 + sz * zeta)
+    return N
+
+
+@ti.func
+def shape_grad_natural(xi, eta, zeta):
+    """``∂N_a/∂(ξ,η,ζ)`` → ``ti.Matrix(8, 3)``."""
+    dN = ti.Matrix.zero(ti.f64, 8, 3)
+    for a in ti.static(range(8)):
+        sx, sy, sz = ti.static(_CORNERS[a])
+        dN[a, 0] = 0.125 * sx * (1.0 + sy * eta) * (1.0 + sz * zeta)
+        dN[a, 1] = 0.125 * sy * (1.0 + sx * xi) * (1.0 + sz * zeta)
+        dN[a, 2] = 0.125 * sz * (1.0 + sx * xi) * (1.0 + sy * eta)
+    return dN

ti_runtime-0.2.1/src/ti_runtime/seams.py

@@ -0,0 +1,303 @@
+"""Injection plumbing — the *seams* generated bodies plug into (PlanJune14 PJ-0).
+
+MechDSL owns the stable wrappers; algo2code generates the bodies. The same
+``LinearSolveContext`` applies *whatever* operator / preconditioner / solver was
+injected, so the plumbing is algorithm-agnostic.
+
+A matrix-free operator is an in-place callable ``apply(out, x) -> None`` computing
+``out = A @ x`` over Taichi fields (typically a generated ``@ti.kernel``).
+Preconditioner protocol mirrors NumerixWeave ``tisolvers``: ``apply(z, r)`` sets
+``z = M^{-1} r``.
+"""
+
+import math
+from collections.abc import Callable
+
+import taichi as ti
+
+from . import vector_ops as vops
+
+# out = A @ x, in place.
+OperatorApply = Callable[..., None]
+
+
+@ti.data_oriented
+class PreconditionerBase:
+    """SPD preconditioner interface: ``apply(z, r)`` sets ``z = M^{-1} r``."""
+
+    def apply(self, z, r) -> None:
+        raise NotImplementedError
+
+    def assemble(self) -> None:
+        """Rebuild internal state from the current operator/tangent.
+
+        No-op for stateless preconditioners (Identity, point-Jacobi); stateful
+        ones (block-Jacobi) override.
+        """
+
+
+@ti.data_oriented
+class IdentityPreconditioner(PreconditionerBase):
+    """``M = I`` — unpreconditioned."""
+
+    def apply(self, z, r) -> None:
+        vops.copy(z, r)
+
+
+@ti.data_oriented
+class DiagonalPreconditioner(PreconditionerBase):
+    """Point-Jacobi: ``M = diag(d)``, ``M^{-1} = diag(1/d)``.
+
+    ``diag`` is a Taichi field of the same layout as the vectors; division is
+    guarded by ``max(d, eps)`` to avoid NaN on a zero/near-zero diagonal.
+    """
+
+    def __init__(self, diag: ti.template(), eps: float = 1e-12):
+        self.diag = diag
+        self.eps = eps
+
+    @ti.kernel
+    def _apply(self, z: ti.template(), r: ti.template(), eps: float):
+        for I in ti.grouped(r):
+            z[I] = r[I] / ti.max(self.diag[I], eps)
+
+    def apply(self, z, r) -> None:
+        self._apply(z, r, self.eps)
+
+
+class Operator:
+    """Holds the injected matrix-free operator ``apply(out, x): out = A @ x``."""
+
+    def __init__(self) -> None:
+        self._apply: OperatorApply | None = None
+
+    def set_apply(self, fn: OperatorApply) -> "Operator":
+        self._apply = fn
+        return self
+
+    def apply(self, out, x) -> None:
+        if self._apply is None:
+            raise RuntimeError("Operator has no body injected; call set_apply(...) first.")
+        self._apply(out, x)
+
+
+class Solver:
+    """Holds an injected (generated) solver body ``solve(ctx, b, x, ...)``."""
+
+    def __init__(self) -> None:
+        self._solve: Callable[..., object] | None = None
+
+    def set_solve(self, fn: Callable[..., object]) -> "Solver":
+        self._solve = fn
+        return self
+
+    def solve(self, *args, **kwargs):
+        if self._solve is None:
+            raise RuntimeError("Solver has no body injected; call set_solve(...) first.")
+        return self._solve(*args, **kwargs)
+
+
+class LinearSolveContext:
+    """The plumbing a generated linear solver targets.
+
+    Bundles the injected matrix-free operator and preconditioner; the generated
+    solver body calls :meth:`apply_A` / :meth:`apply_preconditioner` plus the
+    :mod:`ti_runtime.vector_ops` primitives. Defaults to the identity
+    preconditioner (unpreconditioned).
+    """
+
+    def __init__(self) -> None:
+        self.operator = Operator()
+        self.preconditioner: PreconditionerBase = IdentityPreconditioner()
+        self.solver = Solver()
+
+    def set_operator(self, fn: OperatorApply) -> "LinearSolveContext":
+        self.operator.set_apply(fn)
+        return self
+
+    def set_preconditioner(self, precond: PreconditionerBase) -> "LinearSolveContext":
+        self.preconditioner = precond
+        return self
+
+    def set_solver(self, fn: Callable[..., object]) -> "LinearSolveContext":
+        self.solver.set_solve(fn)
+        return self
+
+    def apply_A(self, out, x) -> None:
+        self.operator.apply(out, x)
+
+    def apply_preconditioner(self, z, r) -> None:
+        self.preconditioner.apply(z, r)
+
+
+# ── Time-integration seam ────────────────────────────────────────────────────
+#
+# The temporal analogue of LinearSolveContext: a generated *time integrator*
+# (e.g. a Newmark-beta step transpiled from its LaTeX algorithm spec) plugs
+# into a stable wrapper here, exactly as a generated linear solver plugs into
+# set_solver. The wrapper is integrator-agnostic — it applies *whatever* step
+# body was injected (Newmark-beta, central difference, HHT, ...).
+#
+# An integrator step is an in-place callable advancing the dynamic state one
+# step. The seam-injected acceleration solve mirrors the matrix-free operator:
+# accel_solve(u_pred, v_pred, a_out) sets a_out = a_{n+1} (out LAST), folding the
+# mass/damping/stiffness/external-force data — the box itself stays
+# mass-/material-agnostic, just as the PCG box is operator-agnostic.
+
+# (u, v, a, accel_solve, dt, beta, gamma) -> (u, v, a); advances state in place.
+IntegratorStep = Callable[..., object]
+
+# accel_solve(u_pred, v_pred, a_out): a_out = a_{n+1}, in place (out LAST).
+# Returns an optional status: ``None`` for a solve that cannot fail
+# (e.g. an elementwise SDOF/diagonal-mass solve), or a convergence flag for an
+# iterative solve (a falsy flag == not converged). TimeIntegrationContext.step
+# consumes it to roll back and fail loud rather than advancing on a bad solve.
+AccelSolveApply = Callable[..., object]
+
+
+class Integrator:
+    """Holds an injected (generated) time-integrator step body.
+
+    The body advances ``(u, v, a)`` one step in place, calling the injected
+    acceleration solve plus the :mod:`ti_runtime.vector_ops` primitives — the
+    shape ``dev/algorithms/newmark.tex`` transpiles to.
+    """
+
+    def __init__(self) -> None:
+        self._step: IntegratorStep | None = None
+
+    def set_step(self, fn: IntegratorStep) -> "Integrator":
+        self._step = fn
+        return self
+
+    def step(self, *args, **kwargs):
+        if self._step is None:
+            raise RuntimeError("Integrator has no body injected; call set_integrator(...) first.")
+        return self._step(*args, **kwargs)
+
+
+class AccelSolve:
+    """Holds the injected acceleration solve ``apply(u_pred, v_pred, a_out)``.
+
+    Sets ``a_out = a_{n+1}`` — the solution of
+    ``(M + gamma*dt*C + beta*dt^2*K) a_{n+1} = F_{n+1} - C*v_pred - K*u_pred``.
+    For an SDOF / diagonal-mass problem this is an elementwise solve; for a full
+    FEM system it is itself a (linear) solve. Matrix-free: the wrapper applies
+    whatever callable was injected.
+    """
+
+    def __init__(self) -> None:
+        self._apply: AccelSolveApply | None = None
+
+    def set_apply(self, fn: AccelSolveApply) -> "AccelSolve":
+        self._apply = fn
+        return self
+
+    def apply(self, u_pred, v_pred, a_out) -> object:
+        if self._apply is None:
+            raise RuntimeError("AccelSolve has no body injected; call set_accel_solve(...) first.")
+        # Forward the injected solve's return value: a convergence status
+        # when the solve is iterative, else ``None``. The caller (step) decides.
+        return self._apply(u_pred, v_pred, a_out)
+
+
+class TimeIntegrationContext:
+    """The plumbing a generated time integrator targets (PlanJune14 P6-1).
+
+    The temporal twin of :class:`LinearSolveContext`. Bundles the injected
+    acceleration solve and the injected integrator step; the generated step body
+    calls :meth:`apply_accel_solve` (the matrix-free seam) plus the
+    :mod:`ti_runtime.vector_ops` primitives. :meth:`step` advances the dynamic
+    state ``(u, v, a)`` one step in place by applying the injected integrator,
+    with no NumPy in the hot path.
+
+    Parameters ``dt`` / ``beta`` / ``gamma`` are the Newmark step controls; they
+    default to the average-acceleration scheme (``beta = 1/4``, ``gamma = 1/2``)
+    — unconditionally stable, second-order, no algorithmic damping.
+    """
+
+    def __init__(self, dt: float = 1.0, beta: float = 0.25, gamma: float = 0.5) -> None:
+        self.integrator = Integrator()
+        self.accel = AccelSolve()
+        self.dt = dt
+        self.beta = beta
+        self.gamma = gamma
+
+    def set_integrator(self, fn: IntegratorStep) -> "TimeIntegrationContext":
+        self.integrator.set_step(fn)
+        return self
+
+    def set_accel_solve(self, fn: AccelSolveApply) -> "TimeIntegrationContext":
+        self.accel.set_apply(fn)
+        return self
+
+    def apply_accel_solve(self, u_pred, v_pred, a_out) -> object:
+        return self.accel.apply(u_pred, v_pred, a_out)
+
+    def step(self, u, v, a):
+        """Advance the state ``(u, v, a)`` one step via the injected integrator.
+
+        Wires the seam conventions to the generated body's callable arguments:
+        the integrator's ``solve_a(u_pred, v_pred, a_out)`` argument is bound to
+        :meth:`apply_accel_solve` (the injected acceleration solve), and the
+        ``(dt, beta, gamma)`` controls come from this context. Returns whatever
+        the generated step returns (typically ``(u, v, a)``).
+
+        Fail-loud contract (WI-3). The step is committed only if the injected
+        acceleration solve succeeded and the advanced state is finite. Before
+        stepping we snapshot ``(u, v, a)``; if the injected solve **raises**
+        (e.g. the WI-2 seam PCG on non-convergence), **reports** a falsy
+        convergence status, or the advanced acceleration is **non-finite**, we
+        restore the snapshot and raise ``RuntimeError`` rather than leaving the
+        dynamic state half-advanced on a bad solve. This mirrors the host
+        Newton driver's snapshot/rollback + isfinite guard (newton.py). P6-1 is
+        not yet production-wired, so this establishes the contract without
+        redesigning the matrix-free body; the per-step snapshot is a
+        step-boundary archival copy (the same ``.to_numpy()`` round-trip the
+        P6-1 tests use at step boundaries), not a hot-path operation.
+        """
+        # Step-boundary snapshot for rollback (not the generated hot path).
+        u_snap = u.to_numpy()
+        v_snap = v.to_numpy()
+        a_snap = a.to_numpy()
+
+        def _restore() -> None:
+            u.from_numpy(u_snap)
+            v.from_numpy(v_snap)
+            a.from_numpy(a_snap)
+
+        status: dict[str, object] = {}
+
+        def accel_solve(u_pred, v_pred, a_out):
+            status["accel"] = self.apply_accel_solve(u_pred, v_pred, a_out)
+
+        try:
+            result = self.integrator.step(u, v, a, accel_solve, self.dt, self.beta, self.gamma)
+        except Exception:
+            # A raising solve (e.g. the seam-injected PCG) may have left state
+            # half-advanced -- restore before propagating.
+            _restore()
+            raise
+
+        # A reported non-convergence status (a falsy flag that is not ``None``;
+        # ``None`` means "no status / cannot fail", e.g. an elementwise solve).
+        accel_status = status.get("accel")
+        if accel_status is not None and not accel_status:
+            _restore()
+            raise RuntimeError(
+                "Time-integration acceleration solve did not converge "
+                f"(status={accel_status!r}); restored the pre-step state rather "
+                "than advancing (u, v, a) on a non-converged acceleration."
+            )
+
+        # Non-finite advanced acceleration poisons all subsequent steps; catch
+        # it here (mirrors the Newton driver's isfinite residual guard).
+        if not math.isfinite(vops.norm2(a)):
+            _restore()
+            raise RuntimeError(
+                "Time-integration step produced a non-finite acceleration "
+                "(||a|| is NaN/Inf); restored the pre-step state rather than "
+                "advancing on a corrupt acceleration."
+            )
+
+        return result

ti_runtime-0.2.1/src/ti_runtime/tensor_ti.py

@@ -0,0 +1,93 @@
+"""Tier-1 tensor / kinematics ``@ti.func`` helpers (PlanJune14 PJ-0).
+
+The building blocks generated constitutive / element code calls — kept as
+``@ti.func`` so they inline cheaply and don't burn JIT budget. Conventions follow
+MechDSL ``dev/design_docs/07-CONVENTIONS.md``: spatial/material 3×3 tensors, Voigt
+order ``[xx, yy, zz, xy, xz, yz]``, **tensorial** (unscaled shears), metric
+``G = diag(1,1,1,2,2,2)``. Harvested/adapted from NumerixWeave ``ticonstit``.
+
+These are ``@ti.func`` — call them from inside a ``@ti.kernel``.
+"""
+
+import taichi as ti
+
+mat3 = ti.types.matrix(3, 3, ti.f64)
+vec6 = ti.types.vector(6, ti.f64)
+
+
+@ti.func
+def identity3() -> mat3:
+    return ti.Matrix.identity(ti.f64, 3)
+
+
+@ti.func
+def det3(A) -> ti.f64:
+    """Determinant of a 3×3 matrix."""
+    return A.determinant()
+
+
+@ti.func
+def inv3(A) -> mat3:
+    """Inverse of a 3×3 matrix."""
+    return A.inverse()
+
+
+@ti.func
+def deformation_gradient(grad_u) -> mat3:
+    """``F = I + ∂u_i/∂X_J`` from the material displacement gradient (3×3)."""
+    return ti.Matrix.identity(ti.f64, 3) + grad_u
+
+
+@ti.func
+def jacobian(F) -> ti.f64:
+    """``J = det(F)``."""
+    return F.determinant()
+
+
+@ti.func
+def right_cauchy_green(F) -> mat3:
+    """``C = F^T F``."""
+    return F.transpose() @ F
+
+
+@ti.func
+def green_lagrange(F) -> mat3:
+    """``E = ½(F^T F − I)``."""
+    return 0.5 * (F.transpose() @ F - ti.Matrix.identity(ti.f64, 3))
+
+
+@ti.func
+def trace3(A) -> ti.f64:
+    return A[0, 0] + A[1, 1] + A[2, 2]
+
+
+@ti.func
+def deviatoric(A) -> mat3:
+    """Deviatoric part ``A − (tr A / 3) I``."""
+    p = (A[0, 0] + A[1, 1] + A[2, 2]) / 3.0
+    return A - p * ti.Matrix.identity(ti.f64, 3)
+
+
+@ti.func
+def von_mises(S) -> ti.f64:
+    """von Mises equivalent stress ``sqrt(3/2 · s:s)`` of a symmetric 3×3 ``S``."""
+    s = deviatoric(S)
+    return ti.sqrt(1.5 * (s * s).sum())
+
+
+@ti.func
+def to_voigt(S) -> vec6:
+    """Symmetric 3×3 → Voigt-6 ``[xx, yy, zz, xy, xz, yz]`` (unscaled shears)."""
+    return ti.Vector([S[0, 0], S[1, 1], S[2, 2], S[0, 1], S[0, 2], S[1, 2]], dt=ti.f64)
+
+
+@ti.func
+def from_voigt(v) -> mat3:
+    """Voigt-6 → symmetric 3×3 (off-diagonals not doubled — tensorial)."""
+    return ti.Matrix([[v[0], v[3], v[4]], [v[3], v[1], v[5]], [v[4], v[5], v[2]]], dt=ti.f64)
+
+
+@ti.func
+def double_contract(A, B) -> ti.f64:
+    """Full tensor double contraction ``A:B`` of two 3×3 tensors."""
+    return (A * B).sum()

ti_runtime-0.2.1/src/ti_runtime/vector_ops.py

@@ -0,0 +1,88 @@
+"""Vector-primitive ``@ti.kernel`` operations (PlanJune14 PJ-0).
+
+Layout-agnostic primitives over Taichi *vector* fields (``ti.Vector.field(d)``),
+the shape FEM degrees-of-freedom use. Generated solvers call these instead of
+re-emitting their own (PlanJune14 PJ-2).
+
+The scalar coefficient ``a`` is a **runtime ``float`` argument**, not
+``ti.template()`` — annotating a Python scalar as ``ti.template()`` bakes the
+value into the compiled kernel, triggering a fresh JIT per distinct value
+(observed as 100k+ kernel variants in tight Krylov loops). Harvested from
+NumerixWeave ``tisolvers``.
+"""
+
+import taichi as ti
+
+
+@ti.kernel
+def copy(dst: ti.template(), src: ti.template()):
+    """``dst <- src`` (elementwise)."""
+    for I in ti.grouped(src):
+        dst[I] = src[I]
+
+
+@ti.kernel
+def axpy(y: ti.template(), a: float, x: ti.template()):
+    """``y <- y + a*x``."""
+    for I in ti.grouped(y):
+        y[I] += a * x[I]
+
+
+@ti.kernel
+def xpay(x: ti.template(), a: float, y: ti.template()):
+    """``x <- a*x + y`` (the CG search-direction update)."""
+    for I in ti.grouped(x):
+        x[I] = a * x[I] + y[I]
+
+
+@ti.kernel
+def scal(x: ti.template(), a: float):
+    """``x <- a*x``."""
+    for I in ti.grouped(x):
+        x[I] = a * x[I]
+
+
+@ti.kernel
+def _dot(x: ti.template(), y: ti.template()) -> ti.f64:
+    s = 0.0
+    for I in ti.grouped(x):
+        s += x[I].dot(y[I])
+    return s
+
+
+def dot(x, y) -> float:
+    """Euclidean inner product ``x . y`` over a vector field (returns a float)."""
+    return float(_dot(x, y))
+
+
+def norm2(x) -> float:
+    """Euclidean 2-norm ``||x||_2``."""
+    return float(_dot(x, x)) ** 0.5
+
+
+@ti.kernel
+def vec_add(out: ti.template(), a: float, x: ti.template(), b: float, y: ti.template()):
+    """``out[I] = a*x[I] + b*y[I]`` (AXPBY, elementwise over any field shape)."""
+    for I in ti.grouped(out):
+        out[I] = a * x[I] + b * y[I]
+
+
+def zero(x) -> None:
+    """Zero a field in place."""
+    x.fill(0.0)
+
+
+@ti.kernel
+def ediv(z: ti.template(), r: ti.template(), d: ti.template(), eps: float):
+    """Elementwise guarded divide: ``z[I] = r[I] / max(d[I], eps)``.
+
+    Point-Jacobi preconditioner kernel (PlanJune14 P4-1): applies M^{-1} r
+    where M = diag(d).  The ``eps`` guard prevents division by zero on a
+    near-zero diagonal; the caller is responsible for choosing a physically
+    appropriate value (default: 1e-12).
+
+    Layout-agnostic: works for any ``ti.Vector.field`` or scalar ``ti.field``
+    where element-wise max and division are defined.
+    """
+    for I in ti.grouped(r):
+        z[I] = r[I] / ti.max(d[I], eps)

ti_runtime-0.2.1/tests/conftest.py

@@ -0,0 +1,13 @@
+"""Shared fixtures for ti-runtime tests."""
+
+from __future__ import annotations
+
+import pytest
+import taichi as ti
+
+
+@pytest.fixture(autouse=True)
+def _ti_cpu():
+    """Re-init Taichi on CPU before each test so freshly-allocated fields are valid."""
+    ti.init(arch=ti.cpu, default_fp=ti.f64)
+    yield

ti_runtime-0.2.1/tests/test_hex8.py

@@ -0,0 +1,46 @@
+"""PJ-0 — Hex8 shape functions, gradients, quadrature."""
+
+import numpy as np
+import pytest
+import taichi as ti
+
+from ti_runtime import hex8
+
+pytestmark = pytest.mark.slow  # executes Taichi kernels (JIT)
+
+
+def test_quadrature_tables():
+    assert hex8.N_QP == 8
+    assert len(hex8.QUAD_POINTS) == 8
+    assert hex8.QUAD_WEIGHTS == (1.0,) * 8
+    # 2x2x2 Gauss weights sum to the reference cube volume (2^3 = 8).
+    assert sum(hex8.QUAD_WEIGHTS) == pytest.approx(8.0)
+
+
+def test_partition_of_unity_and_gradient_sum():
+    N = ti.Vector.field(8, ti.f64, shape=())
+    dN = ti.Matrix.field(8, 3, ti.f64, shape=())
+
+    @ti.kernel
+    def run(xi: ti.f64, eta: ti.f64, zeta: ti.f64):
+        N[None] = hex8.shape(xi, eta, zeta)
+        dN[None] = hex8.shape_grad_natural(xi, eta, zeta)
+
+    run(0.3, -0.6, 0.2)
+    # Shape functions partition unity; natural gradients sum to zero per direction.
+    assert N[None].to_numpy().sum() == pytest.approx(1.0, rel=1e-12)
+    np.testing.assert_allclose(dN[None].to_numpy().sum(axis=0), 0.0, atol=1e-12)
+
+
+def test_shape_is_nodal_delta_at_corners():
+    N = ti.Vector.field(8, ti.f64, shape=())
+
+    @ti.kernel
+    def run(xi: ti.f64, eta: ti.f64, zeta: ti.f64):
+        N[None] = hex8.shape(xi, eta, zeta)
+
+    for a, (sx, sy, sz) in enumerate(hex8._CORNERS):
+        run(float(sx), float(sy), float(sz))
+        expected = np.zeros(8)
+        expected[a] = 1.0
+        np.testing.assert_allclose(N[None].to_numpy(), expected, atol=1e-12)

ti_runtime-0.2.1/tests/test_seams.py

@@ -0,0 +1,99 @@
+"""PJ-0 — injection seams + a mini-CG composition (PJ-1 preview).
+
+The composition test is the important one: it injects a matrix-free operator and
+runs a hand-written CG (the shape algo2code will *generate*) using only the seam
+(`apply_A`) and `vector_ops` primitives — proving the seams + primitives compose
+into a working solver with no NumPy in the operator/solve hot path.
+"""
+
+import numpy as np
+import pytest
+import taichi as ti
+
+from ti_runtime import vector_ops as v
+from ti_runtime.seams import (
+    DiagonalPreconditioner,
+    IdentityPreconditioner,
+    LinearSolveContext,
+    Operator,
+)
+
+pytestmark = pytest.mark.slow  # executes Taichi kernels (JIT)
+
+
+def _vfield(vals: np.ndarray):
+    vals = np.ascontiguousarray(vals, dtype=np.float64)
+    f = ti.Vector.field(vals.shape[1], ti.f64, shape=vals.shape[0])
+    f.from_numpy(vals)
+    return f
+
+
+def test_operator_requires_body():
+    op = Operator()
+    out = ti.Vector.field(3, ti.f64, shape=2)
+    x = ti.Vector.field(3, ti.f64, shape=2)
+    with pytest.raises(RuntimeError, match="no body injected"):
+        op.apply(out, x)
+
+
+def test_identity_preconditioner_copies():
+    r = _vfield(np.arange(6.0).reshape(2, 3))
+    z = ti.Vector.field(3, ti.f64, shape=2)
+    IdentityPreconditioner().apply(z, r)
+    np.testing.assert_allclose(z.to_numpy(), r.to_numpy())
+
+
+def test_diagonal_preconditioner_inverts():
+    rv = np.array([[2.0, 4.0, 6.0]])
+    dv = np.array([[2.0, 4.0, 8.0]])
+    r, d = _vfield(rv), _vfield(dv)
+    z = ti.Vector.field(3, ti.f64, shape=1)
+    DiagonalPreconditioner(d).apply(z, r)
+    np.testing.assert_allclose(z.to_numpy(), rv / dv, rtol=1e-12)
+
+
+# Block-diagonal SPD operator: out[i] = M @ x[i] with a fixed SPD 3x3 M.
+_M = np.array([[4.0, 1.0, 0.0], [1.0, 3.0, 1.0], [0.0, 1.0, 2.0]])
+
+
+@ti.kernel
+def _apply_M(out: ti.template(), x: ti.template()):
+    M = ti.Matrix([[4.0, 1.0, 0.0], [1.0, 3.0, 1.0], [0.0, 1.0, 2.0]], dt=ti.f64)
+    for i in out:
+        out[i] = M @ x[i]
+
+
+def _cg(ctx: LinearSolveContext, b, x, tol=1e-12, maxiter=100):
+    """Hand-written CG (the body algo2code generates) over the seam + primitives."""
+    n = b.shape[0]
+    r = ti.Vector.field(3, ti.f64, shape=n)
+    p = ti.Vector.field(3, ti.f64, shape=n)
+    ap = ti.Vector.field(3, ti.f64, shape=n)
+    v.copy(r, b)  # x0 = 0  ->  r = b - A x0 = b
+    v.copy(p, r)
+    rz = v.dot(r, r)
+    for _ in range(maxiter):
+        ctx.apply_A(ap, p)
+        alpha = rz / v.dot(p, ap)
+        v.axpy(x, alpha, p)
+        v.axpy(r, -alpha, ap)
+        rz_new = v.dot(r, r)
+        if rz_new**0.5 < tol:
+            break
+        v.xpay(p, rz_new / rz, r)  # p = beta*p + r
+        rz = rz_new
+    return x
+
+
+def test_cg_composition_solves_spd_system():
+    rng = np.random.default_rng(7)
+    n = 5
+    bv = rng.standard_normal((n, 3))
+    b = _vfield(bv)
+    x = ti.Vector.field(3, ti.f64, shape=n)  # zero initial guess
+
+    ctx = LinearSolveContext().set_operator(_apply_M)
+    _cg(ctx, b, x)
+
+    expected = np.linalg.solve(_M, bv.T).T  # per-node M^{-1} b
+    np.testing.assert_allclose(x.to_numpy(), expected, atol=1e-9)

ti_runtime-0.2.1/tests/test_tensor_ti.py

@@ -0,0 +1,76 @@
+"""PJ-0 — Tier-1 tensor/kinematics @ti.func helpers."""
+
+import numpy as np
+import pytest
+import taichi as ti
+
+from ti_runtime import tensor_ti as T
+
+pytestmark = pytest.mark.slow  # executes Taichi kernels (JIT)
+
+
+def _mat_in(vals: np.ndarray):
+    f = ti.Matrix.field(3, 3, ti.f64, shape=())
+    f[None] = ti.Matrix(np.ascontiguousarray(vals, dtype=np.float64).tolist())
+    return f
+
+
+def test_det_and_kinematics():
+    rng = np.random.default_rng(0)
+    Fv = np.eye(3) + 0.1 * rng.standard_normal((3, 3))
+    Fin = _mat_in(Fv)
+    out = ti.field(ti.f64, shape=4)
+    Cout = ti.Matrix.field(3, 3, ti.f64, shape=())
+    Eout = ti.Matrix.field(3, 3, ti.f64, shape=())
+
+    @ti.kernel
+    def run():
+        F = Fin[None]
+        out[0] = T.det3(F)
+        out[1] = T.jacobian(F)
+        Cout[None] = T.right_cauchy_green(F)
+        Eout[None] = T.green_lagrange(F)
+
+    run()
+    assert out[0] == pytest.approx(np.linalg.det(Fv), rel=1e-12)
+    assert out[1] == pytest.approx(np.linalg.det(Fv), rel=1e-12)
+    np.testing.assert_allclose(Cout[None].to_numpy(), Fv.T @ Fv, rtol=1e-12)
+    np.testing.assert_allclose(
+        Eout[None].to_numpy(), 0.5 * (Fv.T @ Fv - np.eye(3)), rtol=1e-12, atol=1e-14
+    )
+
+
+def test_deformation_gradient_from_grad_u():
+    gu = np.array([[0.1, 0.0, 0.0], [0.0, -0.05, 0.0], [0.0, 0.0, 0.02]])
+    gin = _mat_in(gu)
+    Fout = ti.Matrix.field(3, 3, ti.f64, shape=())
+
+    @ti.kernel
+    def run():
+        Fout[None] = T.deformation_gradient(gin[None])
+
+    run()
+    np.testing.assert_allclose(Fout[None].to_numpy(), np.eye(3) + gu, rtol=1e-12)
+
+
+def test_voigt_roundtrip_and_von_mises():
+    s = np.array([[10.0, 4.0, 2.0], [4.0, -5.0, 1.0], [2.0, 1.0, 3.0]])  # symmetric
+    sin = _mat_in(s)
+    voigt = ti.Vector.field(6, ti.f64, shape=())
+    back = ti.Matrix.field(3, 3, ti.f64, shape=())
+    vm = ti.field(ti.f64, shape=())
+
+    @ti.kernel
+    def run():
+        v6 = T.to_voigt(sin[None])
+        voigt[None] = v6
+        back[None] = T.from_voigt(v6)
+        vm[None] = T.von_mises(sin[None])
+
+    run()
+    np.testing.assert_allclose(voigt[None].to_numpy(), [10.0, -5.0, 3.0, 4.0, 2.0, 1.0], rtol=1e-12)
+    np.testing.assert_allclose(back[None].to_numpy(), s, rtol=1e-12)
+    # reference von Mises: sqrt(3/2 dev:dev)
+    dev = s - np.trace(s) / 3.0 * np.eye(3)
+    ref = np.sqrt(1.5 * np.sum(dev * dev))
+    assert vm[None] == pytest.approx(ref, rel=1e-12)

ti_runtime-0.2.1/tests/test_vector_ops.py

@@ -0,0 +1,73 @@
+"""PJ-0 — vector-primitive kernels."""
+
+import numpy as np
+import pytest
+import taichi as ti
+
+from ti_runtime import vector_ops as v
+
+pytestmark = pytest.mark.slow  # executes Taichi kernels (JIT)
+
+
+def _vfield(vals: np.ndarray):
+    vals = np.ascontiguousarray(vals, dtype=np.float64)
+    f = ti.Vector.field(vals.shape[1], ti.f64, shape=vals.shape[0])
+    f.from_numpy(vals)
+    return f
+
+
+def test_copy():
+    rng = np.random.default_rng(0)
+    xv = rng.standard_normal((5, 3))
+    x = _vfield(xv)
+    y = ti.Vector.field(3, ti.f64, shape=5)
+    v.copy(y, x)
+    np.testing.assert_allclose(y.to_numpy(), xv, rtol=1e-12)
+
+
+def test_axpy():
+    rng = np.random.default_rng(1)
+    xv, yv = rng.standard_normal((4, 3)), rng.standard_normal((4, 3))
+    x, y = _vfield(xv), _vfield(yv)
+    v.axpy(y, 2.5, x)
+    np.testing.assert_allclose(y.to_numpy(), yv + 2.5 * xv, rtol=1e-12)
+
+
+def test_xpay():
+    rng = np.random.default_rng(2)
+    xv, yv = rng.standard_normal((4, 3)), rng.standard_normal((4, 3))
+    x, y = _vfield(xv), _vfield(yv)
+    v.xpay(x, -0.5, y)
+    np.testing.assert_allclose(x.to_numpy(), -0.5 * xv + yv, rtol=1e-12)
+
+
+def test_scal():
+    rng = np.random.default_rng(3)
+    xv = rng.standard_normal((6, 3))
+    x = _vfield(xv)
+    v.scal(x, 3.0)
+    np.testing.assert_allclose(x.to_numpy(), 3.0 * xv, rtol=1e-12)
+
+
+def test_dot_and_norm():
+    rng = np.random.default_rng(4)
+    xv, yv = rng.standard_normal((7, 3)), rng.standard_normal((7, 3))
+    x, y = _vfield(xv), _vfield(yv)
+    assert v.dot(x, y) == pytest.approx(float((xv * yv).sum()), rel=1e-12)
+    assert v.norm2(x) == pytest.approx(float(np.linalg.norm(xv)), rel=1e-12)
+
+
+def test_vec_add():
+    rng = np.random.default_rng(5)
+    xv, yv = rng.standard_normal((5, 3)), rng.standard_normal((5, 3))
+    x, y = _vfield(xv), _vfield(yv)
+    out = ti.Vector.field(3, ti.f64, shape=5)
+    a, b = 2.0, -0.5
+    v.vec_add(out, a, x, b, y)
+    np.testing.assert_allclose(out.to_numpy(), a * xv + b * yv, rtol=1e-12)
+
+
+def test_zero():
+    x = _vfield(np.ones((3, 3)))
+    v.zero(x)
+    np.testing.assert_allclose(x.to_numpy(), 0.0)