Architecture & Design

A single unified solver runs on CPU or GPU by injecting a MathBackend. Where an abstraction earns its keep it lives in a co-located base.py (HeatSolver, MathBackend, LaserPath); single-implementation interfaces are deliberately avoided.

Directory Structure

spectral_galerkin_heat/
├── pyproject.toml                  # Package definition and dependencies
├── README.md
├── src/
│   └── spectral_galerkin_heat/
│       ├── __init__.py             # Package version + public API
│       ├── runner.py               # StandaloneHeatRunner: simulation loop orchestrator
│       ├── core/
│       │   ├── parameters.py       # Data classes: SimulationContext, NumParams, MaterialParams, …
│       │   ├── laser.py            # LaserState dataclass + LaserPath ABC
│       │   └── vector.py           # Vec3: immutable (x, y, z) value type
│       ├── backends/
│       │   ├── base.py             # MathBackend: bundles array module (xp) + kernels
│       │   ├── numpy_backend.py    # NumpyBackend (CPU)
│       │   ├── cupy_backend.py     # CupyBackend (GPU)
│       │   ├── _registry.py        # backend registry: register_backend + get_backend
│       │   └── __init__.py         # get_backend(name); registers built-in backends
│       ├── solvers/
│       │   ├── base.py             # HeatSolver ABC
│       │   ├── spectral.py         # SpectralSolver (unified CPU/GPU, backend-injected)
│       │   ├── spectral_cpu_linear.py
│       │   └── __init__.py         # build_solver(context): solver selector
│       ├── physics/
│       │   ├── spectral_state.py        # Backend-agnostic state (SpectralGrid, FineMeshState, …)
│       │   ├── spectral_ops.py          # Backend-agnostic free functions
│       │   ├── spectral_cpu_kernels.py  # Low-level CPU spectral kernels (NumPy / Numba)
│       │   ├── spectral_gpu_kernels.py  # Low-level GPU spectral kernels (CuPy / CUDA)
│       │   └── spectral_helpers.py      # DCT reconstruction utilities
│       └── io_utils/
│           ├── spectral_fs_io.py   # LocalFSIOManager  (HDF5 + XDMF)
│           ├── xdmf_io.py          # XdmfBuilder + structured/unstructured readers
|           ...
├── simulations/
│   ├── main.py                     # CLI entry point: parses YAML, builds solver, runs simulation
│   └── examples/                   # Example YAML configs (increasing complexity) + orchestrator.py
│       └── paths/                  # G-code laser paths for the examples
└── docs/
    └── ARCHITECTURE.md

Component Descriptions

  • CLI (simulations/main.py): Parses a YAML config, builds the solver via build_solver(context), and delegates execution to StandaloneHeatRunner.

  • Runner (runner.py): StandaloneHeatRunner — the simulation loop orchestrator. Owns the while t < t_end loop; receives a HeatSolver and a LocalFSIOManager by injection (the I/O manager defaults to a fresh LocalFSIOManager), delegating all physics to the former and all I/O to the latter.

  • Solver selector (solvers/__init__.py): build_solver(context) — a single dispatch on backend that returns the configured HeatSolver. It replaces the former SimulationFactory hierarchy, which only ever differed in this one choice.

  • Solver base (solvers/base.py): HeatSolver ABC — declares initialize(), step(), and finalize(). The unified SpectralSolver runs on CPU or GPU purely through its injected MathBackend; SpectralSolverCPULinear is a separate linear variant.

  • Backends (backends/): MathBackend bundles an array module (xp, NumPy or CuPy) with its matching physics-kernel module. A name→factory registry (_registry.py) resolves backends: get_backend(name) returns NumpyBackend or CupyBackend, so the same solver code targets either device. Extra backends register themselves with the @register_backend("name") decorator, so adding one never requires editing get_backend.

  • I/O (io_utils/spectral_fs_io.py): LocalFSIOManager — the concrete filesystem I/O handler (HDF5 + XDMF). It owns the full I/O contract (initialize, process_step, process_end, finalize, save_step, …); save_step dispatches per output type to small _save_* helpers.

  • Laser (core/laser.py): LaserState dataclass and LaserPath ABC. Used by the solver (to query laser position/power) and io_utils/gcode_path.py.

  • Physics (physics/): Backend-agnostic state and free functions (spectral_state, spectral_ops) plus the pure numerical kernels (spectral_cpu_kernels Numba / spectral_gpu_kernels CuPy).

Adding a backend (PyTorch, JAX, …)

The solver doesn’t need a specific array library to run on. It only uses backend.xp (the array module) and backend.kernels (the math functions), and almost all the physics code already works with any NumPy-like xp. So a new backend is mostly plumbing — supplying four things:

1. An array module xp. It needs to behave like NumPy: zeros, arange, exp, sum, in-place writes (multiply(..., out=), a[0,0,0] = v), .astype, etc. NumPy and CuPy work out of the box. PyTorch and JAX could be used by writing wrappers that rename a few things (and JAX arrays can’t be modified in place, so those writes need adapting).

2. A kernels module. Most functions are reused as-is from spectral_ops and spectral_state — don’t rewrite them. You only need to provide the ones that truly depend on the library:

  • the cosine transform (DCT_II / IDCT_II);

  • the sub-pixel shift (_ndshift);

  • a couple of simple per-cell formulas (the evaporation and latent-heat source terms, and the Gaussian flux core.laser.super_gaussian_flux).

    Then subclass SpectralSolverState, bind your own xp, and give those primitives to BackendHooks. Copy spectral_cpu_kernels.py as a template.

3. A MathBackend and registration. Subclass MathBackend with your name, xp, and kernels, and add @register_backend("mybackend").

4. A config name. Hook the simulation.backend string up to your backend in solvers/__init__.py (build_solver).

Keep the kernels working for both float32 and float64, and check your backend against the CPU one with test_cpu_gpu_equivalence_e2e in tests/integration/test_pipeline.py.

Architecture Diagram

        flowchart TD
    %% --- Styling ---
    classDef client   fill:#ff7675,stroke:#d63031,stroke-width:2px,color:#000;
    classDef abstract fill:#81ecec,stroke:#00cec9,stroke-width:2px,stroke-dasharray:5 5,color:#000;
    classDef concrete fill:#55efc4,stroke:#00b894,stroke-width:2px,color:#000;
    classDef product  fill:#fdcb6e,stroke:#e17055,stroke-width:2px,color:#000;
    classDef core     fill:#a29bfe,stroke:#6c5ce7,stroke-width:2px,color:#000;

    %% --- Nodes ---
    CLI["simulations/main.py"]:::client
    Runner["runner.py\nStandaloneHeatRunner"]:::client
    Context["core/parameters.py\nSimulationContext"]:::core
    Laser["core/laser.py\nLaserState · LaserPath"]:::core

    Build["solvers/__init__.py\nbuild_solver()"]:::client
    ISolv["solvers/base.py\n≪abstract≫ HeatSolver"]:::abstract
    IBack["backends/base.py\n≪abstract≫ MathBackend"]:::abstract

    Solver["solvers/spectral.py\nSpectralSolver"]:::product
    CPUBack["backends/numpy_backend.py\nNumpyBackend"]:::concrete
    GPUBack["backends/cupy_backend.py\nCupyBackend"]:::concrete
    FSIO["io_utils/spectral_fs_io.py\nLocalFSIOManager"]:::product

    %% --- Relationships ---
    CLI -->|"creates context"| Context
    CLI -->|"build_solver(context)"| Build
    CLI -->|"instantiates with solver + io"| Runner
    Build -->|"returns"| Solver
    Runner -->|"calls"| ISolv
    Runner -->|"calls"| FSIO

    Solver -.->|"implements"| ISolv
    Solver -->|"injected with"| IBack
    IBack  -.->|"implemented by"| CPUBack
    IBack  -.->|"implemented by"| GPUBack

    Solver -->|"queries"| Laser