COSY Infinity Backend Architecture

The COSY Backend allows sandalwood to harness the raw speed and high-order capabilities of the COSY Infinity differential algebra library (written in Fortran 77).

Unlike standard Python bindings (which often copy data excessively), this backend uses a custom Direct Memory Bridge designed for High-Performance Computing (HPC).

Architecture Overview

The backend consists of three layers:

  1. Fortran Core (`libcosy.so`): The compiled COSY Infinity library, patched with wrapper.f to expose a C-compatible ABI.

  2. C-Types Bridge (`cosy_backend.py`): A thin Python layer that marshals pointers and integers directly to the shared library.

  3. High-Level Wrapper (`CosyDA`): A Python class that manages the lifecycle of COSY variables using integer Indices, shielding the user from manual memory management. Ownership of an index can be safely transferred via CosyDA.transfer_ownership().

Memory Management Architecture

Warning

The COSY Fortran core is not thread-safe. COSY Infinity uses global STATIC memory (Fortran COMMON blocks) for all calculations. The underlying Fortran routines are strictly single-threaded.

Sandalwood’s Python layer has been hardened for concurrent access:

  • CosyIndexPool.acquire and release are protected by a threading.RLock, preventing duplicate-index allocation races when multiple threads hit the pool concurrently.

  • MultivariateTaylorFunction.initialize_mtf is protected by a module-level threading.RLock (_INIT_LOCK), ensuring class-state mutations are atomic.

However, concurrent calls into the Fortran library itself (e.g., two threads each computing da * da through COSY at the same time) will still cause memory corruption. Only one thread should invoke COSY arithmetic at any given moment. Use a process-level lock or multiprocessing worker isolation if you need parallelism over COSY computations.

Effective memory management is critical when bridging Python’s dynamic environment with COSY’s static Fortran roots.

CosyIndexPool: Object Recycling To prevent the overhead of frequent malloc/free cycles in the underlying stack, Sandalwood implements a Robust Memory Pooling (Object Pool Pattern). The CosyIndexPool maintains a list of available COSY variable indices. When a calculation needs a temporary variable, it “acquires” an index from the pool in O(1) time. Once the calculation is complete, the index is “released” back to the pool for future reuse. This prevents “stack thrashing” and significantly improves performance in iterative batch operations.

Since v0.1.3 (feat/backend-hardening), both acquire and release are protected by an internal threading.RLock so that the pool itself is safe to access from multiple Python threads simultaneously. The Fortran core remains single-threaded (see warning above).

CosyDA Ownership Transfer When a COSY operation produces a new CosyDA object, ownership of the underlying Fortran index must be transferred to the result wrapper exactly once — transferring twice would cause a double-free; failing to transfer would leak the index.

The transfer_ownership() method on CosyDA provides a single, explicit handoff:

res_da = some_operation()          # owned=True (will free on __del__)
idx = res_da.transfer_ownership() # res_da.owned now False — will NOT free
result = CosyMtfData(dimension, idx=idx, owned=True)  # new owner

This replaces the previous pattern of manually setting res_da.owned = False, which was easy to forget or apply twice.

Numerical Hygiene: DA_RESET Reusing memory indices requires strict hygiene. Before an index is acquired from the pool, it is subjected to a Hard Reset via the DA_RESET mechanism. This Fortran-level routine explicitly:

  • Sets the variable to a constant 0.0 using DACON.

  • Re-initializes the internal type metadata (NTYP) to the base Differential Algebra type (NDA).

  • Clears any stale high-order coefficients.

This ensures that every “new” variable is numerically clean, preventing intermediate results from leaking into subsequent calculations.

Configuration: Pool Sizing The memory pool size can be tuned via the environment variable:

export SANDALWOOD_COSY_POOL_SIZE=1024

The default value is 1024. For massive batch operations or high-order compositions involving millions of intermediate terms, increasing this value (e.g., to 4096) can reduce acquisition latency, at the cost of higher static memory usage.

Specialized Physics Kernels

The COSY backend includes specialized kernels for high-performance physics simulations, accessible via Hybrid Dispatch.

Biot-Savart Solver The biot_savart_batch function automatically dispatches to the most efficient kernel based on input types:

  • Fast Path (Discrete Mode): If inputs are float64 arrays, the backend dispatches to a raw Fortran kernel (COMPUTE_BIOT_SAVART_BATCH_FAST). This bypasses the DA layer entirely, offering speeds comparable to compiled C/C++ code (~100x faster than Python loops).

  • Parametric Path (MTF Mode): If inputs are DA objects, the backend uses the high-order DA kernel. This allows for computing derivatives and Taylor parametric equations and high-order derivatives.

[ Python User ]  <--  sandalwood.mtf
       |
[ Python Backend ]  (cosy_backend.py)
       |  (ctypes / pointers)
       v
[ Fortran Wrapper ] (wrapper.f / ISO_C_BINDING)
       |  (Common Blocks / Arrays)
       v
[ COSY Core ]       (Static Memory Pool)

Memory Management: The “CosyScope”

COSY Infinity uses a Global Static Memory Pool (Common Blocks /DACOM/). It allocates variables sequentially on a stack. This poses a challenge for Python, which uses a dynamic Garbage Collector.

If Python objects are deleted out of order, they would leave “holes” in the COSY stack, leading to memory fragmentation or corruption.

The Solution: CosyScope

We implemented a custom context manager, CosyScope, which treats memory allocation like a stack frame in a compiler.

from sandalwood.backends.cosy.cosy_backend import CosyScope

# State: IVAR=10, IMEM=500
with CosyScope():
    # Allocation happens here
    x = CosyBackend.var(0)
    y = x * x + 2
    # State: IVAR=12, IMEM=550

# Exit Scope:
# The backend strictly "rewinds" the internal pointers.
# State: IVAR=10, IMEM=500
# All temporary variables (x, y) are effectively freed instantly.

This approach is O(1) (constant time) for freeing memory, regardless of how many millions of temporary objects were created inside the loop.

Performance Optimizations

1. Batch Evaluation with OpenMP

Evaluating a high-order polynomial at millions of points is the most expensive operation in tracking simulations. We offload this entirely to Fortran.

  • Function: eval_da_batch (in wrapper.f)

  • Technique: It flattens the pointer structure and uses OpenMP directives (!$OMP PARALLEL DO) to distribute the workload across all CPU cores.

  • Speedup: ~50-100x compared to NumPy broadcasting for high orders.

2. Direct “Flat” Data Transfer

Standard DA libraries often serialize coefficients into complex objects. Sandalwood uses specialized subroutines (get_all_coeffs_flat and cosy_set_coeffs) that copy internal COSY memory blocks directly to/from pre-allocated NumPy array buffers.

  • Benefit: Zero-copy overhead for large coefficient transfers (bidirectional).

  • Format: The wrapper flattens the multi-dimensional exponent array into a 1D C-integer array, minimizing marshalling and transition costs.

Memory Configuration

For advanced users, the COSY backend’s internal memory limits (e.g., the size of the storage stack or the maximum number of variables) can be customized without modifying the core Fortran source files.

A configuration file is provided at src/sandalwood/backends/cosy/cosy_config.env. You can modify the following parameters:

  • COSY_LMEM: Length of the main storage stack (Default: 140,000,000).

  • COSY_LVAR: Maximum number of variables (Default: 10,000,000).

  • COSY_LEA: Maximum number of monomials. Increase this for extremely high-order calculations.

  • COSY_LNO: Maximum supported Taylor order (Default: 99).

  • COSY_LNV: Maximum number of variables (Default: 40).

### Compiler Support (Linux)

Sandalwood supports two main Fortran compilers on Linux:

  • Intel Fortran (ifx): The recommended default. Recommended for maximum performance on Intel hardware. It provides superior auto-vectorization and highly optimized OpenMP performance.

  • GNU Fortran (gfortran): The standard open-source fallback. Robust and widely available.

### Intel oneAPI Integration

If using ifx, you must ensure the Intel environment variables are loaded in your current shell before building. This is typically done by sourcing the setvars.sh script (Linux) or using the “Intel oneAPI Command Prompt” (Windows).

# Linux example (run in shell or add to .bashrc)
source /opt/intel/oneapi/setvars.sh

### Building the Backend

There are three ways to build the COSY backend components:

  1. Standard Python Install When you run pip install -e ., the setup.py script automatically detects your compiler, patches the memory limits, and builds the shared library.

  2. Standalone Shared Library Build Use this for development or after changing memory parameters:

    bash src/sandalwood/backends/cosy/compile_cosy.sh

  3. Raw COSY Binary Build (“cosy-raw”) For benchmarking against raw COSY Infinity or running legacy .FOX scripts, you can compile a standalone COSY executable:

    bash scripts/benchmarks/benchmark.sh help

    The benchmark script also respects the COSY_COMPILER setting from the config file.

The build script automatically creates a temporary build directory, patches the Fortran sources with your new limits, and links the updated libcosy.so. This ensures the original COSY source files in the repository remain un-modified.

3. Static “Scratchpad” Allocation

To avoid the overhead of finding free memory slots for every intermediate calculation (e.g., temp = a * b), the wrapper.f module pre-allocates a persistent Scratchpad (Common Block /DASCRATCH/) of 20 variables.

Intermediate operation wrappers (like compute_da_add_const) reuse these slots, eliminating the need to modify the global memory stack pointer (IVAR) for temporary arithmetic results. This reduces “stack churn” and improves cache locality.

4. Taylor Map Composition (POLVAL)

Composing two high-order maps is mathematically equivalent to substituting one polynomial into another. Sandalwood uses a specialized wrapper for the COSY POLVAL routine.

  • Technique: Instead of manual substitution in Python loops, we marshal the entire map to Fortran and use COSY’s highly optimized, recursive substitution algorithm.

  • Benefit: This is the most efficient way to perform symplectic tracking or map concatenation.

5. Vectorized Batch Arithmetic

For operations on large arrays of DA objects (e.g., adding two vectors of 10,000 Taylor series), we provide “Batch” variants.

  • Function: compute_da_add_batch, compute_da_mul_batch, etc.

  • Technique: These routines perform the loop over the array entirely within Fortran, avoiding 10,000 costly transitions between the Python interpreter and the shared library.

6. Novel: Stable Complex Arithmetic

COSY’s internal complex division and inversion can occasionally encounter stability issues with extremely small coefficients. Sandalwood includes Patched Implementation (SANDALWOOD_CDMUI and SANDALWOOD_CDDCD) that use a more robust normalization strategy.

  • Improvement: Ensures high-precision results even in numerically sensitive regions of the complex plane.

Since v0.1.3, the Python-level CosyMtfData.inverse() and divide() guards have been tightened from exact == 0 comparisons to tolerant abs(c0) < 1e-14 checks, preventing false positives from floating-point rounding near zero.

Extending the Backend

To add new COSY functions:

  1. Modify `wrapper.f`: Add a subroutine using ISO_C_BINDING.

    SUBROUTINE MY_FUNC(IDX_IN, IDX_OUT) BIND(C, NAME='my_func')
    USE ISO_C_BINDING
    INTEGER(C_INT) IDX_IN, IDX_OUT
    ! ... Call COSY internal ...
END SUBROUTINE

  2. Update `cosy_backend.py`: Bind the function signature.

    bind_cosy_func("my_func", [POINTER(c_int), POINTER(c_int)])

  3. Recompile: Run compile_cosy.sh.