Architecture and Codebase Structure

This page maps source ownership, dependencies, and solver execution in Tarang.jl.

Use the website's dev version for the structure on main; stable describes the latest tagged release. The published manual is built from docs/src/, with navigation and deployment configured in docs/make.jl.

Package layout

src/
├── Tarang.jl                 root module; declarative bootstrap only
├── dependencies.jl           package imports
├── load_order.jl             ordered implementation manifests (see below)
├── public_api.jl             checked supported-API registry (@public_api)
├── runtime_init.jl           MPI, FFTW, logging, and extension startup
├── api/
│   ├── public/               supported root exports, one file per capability
│   ├── namespaces.jl         Tarang.Fields / .Problems / .Solvers / ... facades
│   └── *.jl                  the facade bodies
├── core/
│   ├── architectures.jl, module_contracts.jl   CPU/GPU contract, ownership rules
│   ├── basis/                basis contracts, wavenumbers, product matrices
│   ├── boundary_conditions/  BC construction and types
│   ├── cartesian_operators/  Cartesian differential operator core, dispatch, eval
│   ├── distributor/          MPI layouts, decomposition convention, transposes
│   ├── field/                ScalarField/VectorField/TensorField
│   │   ├── field_data/       storage, copies, scales (dealiasing) — per-field data
│   │   └── field_layout/     :g/:c layout transitions and field arithmetic
│   ├── forcing/              stochastic + deterministic forcing (types, generation, application)
│   ├── les_models.jl         array-level Smagorinsky and AMD closures; shared CPU/GPU kernels
│   ├── nonlinear/            nonlinear products, 3/2 padding, dealiasing
│   ├── operators/            symbolic operator tree
│   │   ├── derivatives/      Fourier / polynomial derivatives, matrix apply
│   │   ├── matrices/         operator → sparse matrix builders
│   │   ├── operations/       integrate, interpolate, lift/convert
│   │   └── tensor/           grad/div/curl/Laplacian/fractional Laplacian
│   ├── problems/             equation parsing and EquationIR
│   │   └── problem_matrices/ EquationIR → global mass/linear blocks
│   ├── solvers/              solver construction, ExecutionPlan, lazy (compiled) RHS, stepping loop
│   ├── subsystems/           per-Fourier-mode subproblems, mode batching (+ KA kernels)
│   ├── timesteppers/         RK, multistep, diagonal-IMEX, ETD schemes and path selection
│   ├── transforms/           serial transforms, layout rules, GPU dispatch hooks
│   └── transpose/            TransposableField: MPI pencil transposes (pack/unpack, async)
├── tools/                    matrix solvers, NetCDF I/O, checkpoints, config, logging
│   └── temporal_filters/     temporal, wave/mean, and GQL filters
└── extras/
    └── flow_tools/           CFL, spectra, QG/streamfunction diagnostics, quick domains, plotting

ext/
├── TarangCUDAExt.jl          extension module and ordered CUDA includes
└── cuda/                    device architecture, cuFFT/DCT-I transforms, Chebyshev
                             derivatives, batched matsolvers, NCCL transposes, memory

test/
├── file_lists.jl             shared CPU, optional, GPU, MPI, and distributed GPU registry
├── runtests.jl               package test runner
├── run_gpu_ci.jl             CUDA test driver
└── run_mpi_ci.jl             launches each registered MPI test in its own process world

docs/
├── make.jl                  Documenter build, navigation, and deployment
└── src/                     published manual, tutorials, API pages, and assets

src/load_order.jl is the whole include order, as twelve manifests. Each one is a flat list of includes owning one slice; add implementation files to the owning manifest and never as a one-off include in src/Tarang.jl:

Entry points such as src/core/field.jl, src/core/operators/operators.jl, src/core/transforms.jl, and src/core/stochastic_forcing.jl aggregate implementation files from their respective directories. Forcing is split under src/core/forcing/.

Dependency direction

contracts and utilities
        ↓
fields, bases, distributors
        ↓
operators and problems
        ↓
compiled artifacts and subproblems
        ↓
solvers and timesteppers
        ↓
models, output, and extras

CUDA extension ──implements──► backend hooks declared by core
public API     ──exposes───► selected bindings from all layers

Core must load without CUDA. CUDA-specific module bindings for CUSOLVER and CUSPARSE are methods supplied by TarangCUDAExt; core owns only the solver contracts and backend-neutral orchestration.

Public API boundary

Supported root exports are declared with @public_api under src/api/public/. The macro both exports each name and registers it in the checked manifest returned by:

Tarang.public_api_names()
Tarang.is_public_api(:InitialValueSolver)  # true

Implementation files still contain compatibility exports from releases before the boundary existed. Treat those as legacy, not as permission to grow the root API. New supported names belong in one public capability file and, when appropriate, in a facade such as Tarang.Fields or Tarang.Solvers.

Problem compilation lifecycle

Problem construction has three distinct kinds of state:

EquationIR temporarily implements AbstractDict{String,Any} so downstream code using keys such as "M" continues to work. Internal code should prefer the named fields. Likewise, matrix and subproblem entries are mirrored into problem.parameters for compatibility, but runtime code reads problem.compiled as the canonical owner.

reset_compiled_problem! clears matrices, subproblems, and its RuntimeCacheContext before rebuilding. Per-problem caches therefore cannot leak through user parameters or be reused by an unrelated solver run.

Solver build and step path

For an InitialValueProblem, trace these files:

1. core/solvers/solver_types.jl resets compiled state, parses equations, assembles global compatibility matrices, builds subproblems, and compiles the RHS plan.
2. core/problems/problem_matrices/ converts each EquationIR into sparse mass and linear blocks.
3. core/subsystems/ groups Fourier modes, builds small coupled systems, applies valid-mode filtering, and owns per-mode runtime buffers.
4. core/solvers/lazy_rhs.jl translates explicit expressions into a type-specialized evaluation tree.
5. core/solvers/solver_execution_plan.jl records, once, the facts every later decision reads: architecture (:cpu/:gpu), distribution, spectral structure, and whether global matrices and subproblems were assembled.
6. core/solvers/solver_stepping.jl refreshes dynamic boundary conditions and calls the timestepper dispatcher (core/timesteppers/dispatch.jl), which checks stochastic-forcing and single-GPU implicit-operator compatibility.
7. core/timesteppers/step_selection.jl chooses the runtime path; the per-scheme step_*! functions then run one of the paths below.

The resulting flow is:

equation strings
    ↓ parse
EquationIR
    ↓ compile
CompiledProblem {global matrices, subproblems, caches}
    ↓ construct
InitialValueSolver {RHS policy, lazy plan, timestep state}
    ↓ step!
refresh BCs → evaluate RHS → per-mode solve → update fields

Timestepper runtime paths

Runtime selection uses the execution plan and timestepper capabilities. Unsupported configurations raise an error.

The user-facing consequences (which scheme runs where, and what refuses) are tabulated in Time Steppers.

core/timesteppers/types.jl declares const RK443_IMEX = RK443. Both names therefore use the same dispatch, workspace allocation, and backend capability traits. Changes to RK443 belong in the shared implementation. core/timesteppers/state.jl owns reusable stage fields and timestepper buffers; per-mode solve buffers and factorizations belong to subsystem caches.

Output and checkpoint ownership

Persistence is layered under src/tools/ so the field and solver core does not depend on file formats:

tools/load_output.jl loads scheduled output before the evaluator; tools/load_runtime.jl loads reconstruction and slab I/O before field, stochastic, and solver persistence. NetCDF writes stage device arrays to host memory explicitly. Restart validates metadata across all processor files before changing destination fields. See I/O for restart compatibility, including the Julia-version restriction for stochastic state.

LES model ownership

src/core/load_models.jl loads src/core/les_models.jl, which owns the SmagorinskyModel and AMDModel types, gradient validation, eddy viscosity and diffusivity calculations, SGS stress, and diagnostics. Both backends broadcast the same scalar kernels over architecture-specific arrays. There is no separate CUDA AMD implementation under ext/cuda/; the extension supplies device arrays and transfers through the architecture interface.

Callers supply grid-space gradients and apply the resulting stress to the RHS. See LES Models for AMD gradient normalization and the full variable-viscosity stress, and Testing for CPU/JLArrays/CUDA coverage.

RHS execution policy

rhs_fallback=:auto resolves per solver:

Use rhs_fallback=:strict to require compilation on serial CPU too. Use :interpreted only for a verified CPU or supported MPI compatibility case. GPU state rejects :interpreted explicitly, and distributed all-Fourier interpreted execution is rejected unconditionally because it is not correct.

This rule is broader than matrix-solver selection: a GPU field cannot select a CPU-only coupled solver, and :gpu never silently degrades to a CPU solver. NetCDF output is an explicit host I/O boundary, not a computational fallback.

GPU ownership

The core/extension split is:

Supported single-GPU IVPs are 2D/3D pure Fourier and mixed Fourier–Chebyshev layouts. Their transforms, RHS evaluation, and coupled subproblem solves remain device-resident after warm-up. Unsupported layouts raise an error.

MPI data movement

Per-mode linear solves are rank-local. Communication surrounds them:

• pure Fourier problems communicate inside distributed FFTs;
• mixed Fourier–Chebyshev problems additionally transpose between the FFT pencil and solve layout once per stage or step;
• diagnostics use collective reductions;
• output may gather or write rank-local files according to its handler.

Collectives must remain outside the per-subproblem loop and every rank must issue them in the same order.

Which axes are decomposed

One function answers this for the whole codebase:

decomposed_axes(dist, ndim)   # global axis indices that are split, ascending
mesh_axis_for(dist, ndim, axis)   # which mesh dimension splits `axis`, or nothing

The two conventions it encodes differ: with PencilArrays the last length(mesh) axes are decomposed, and with TransposableField (GPU+MPI) the first ones are. Both live in src/core/distributor/distributor_core.jl and nowhere else.

Use these helpers at call sites; test_decomposition_convention.jl detects independent copies of the axis arithmetic. Pass the field's dimensionality as ndim, which can differ from dist.dim.

Extension checklist

When adding a feature:

1. Put implementation in the owning core/tool/extension directory.
2. Keep dependency direction downward; do not make core depend on an API facade or on CUDA.
3. Store compiled or temporary state in CompiledProblem, RuntimeCacheContext, or a typed subsystem cache, not in user parameters.
4. Add a lazy-RHS translation or make unsupported execution fail explicitly.
5. Declare supported user-facing names with @public_api and update the relevant facade.
6. Register tests in test/file_lists.jl when adding a test file, and git add it: the inventory test fails on a registered file that is not tracked.
7. Update this page only when ownership or the runtime path changes.

See also

• Solvers
• Time Steppers
• GPU Computing
• Tau Method
• Testing