Performance Optimization in Sandalwood

One of the central challenges in Differential Algebra (DA) is the Combinatorial Explosion. As you increase the order of a Taylor series, the number of terms grows exponentially. Calculating (x + y)^10 involves manipulating hundreds of coefficients; calculating (x + y + z)^20 involves thousands.

Historically, this forced scientists to use rigid, compiled languages like Fortran (e.g., COSY Infinity) because Python’s overhead was too high. sandalwood bridges this gap by implementing advanced high-performance computing (HPC) techniques directly in the Python backend.

This guide details the optimizations that allow sandalwood to perform within a factor of 5-10x of optimized C++ code, while retaining Python’s flexibility.

The “Two-Language” Problem

In standard Python, every number is an “Object” (a box containing data + type info). Adding two numbers involves unboxing them, checking types, adding, and re-boxing the result.

  • The Bottleneck: In a Taylor series multiplication, we perform millions of tiny additions. Doing this loop in standard Python would be 1000x slower than C++.

  • The Solution: We must move the “Hot Loops” (the parts of code running millions of times) out of Python and into machine code, while keeping the high-level logic in Python.

1. “Dense Mode” Architecture

Standard symbolic libraries use a Sparse representation, storing terms as a list of (exponent, coefficient) pairs. To multiply two functions, they must generate every combination of exponents and then sort/merge them to combine like terms.

  • Why it’s slow: Sorting is an $O(N log N)$ operation. As the polynomial grows, the time spent sorting exponents dominates the calculation.

  • Our Optimization: When initialize_mtf is called, sandalwood pre-computes a Multiplication Table.

We map every possible exponent tuple (e.g., x^2 y^1) to a unique integer index. The multiplication logic becomes a simple “Lookup”:

# Naive Sparse Approach (Slow)
result = []
for term_a in A:
    for term_b in B:
         new_exp = term_a.exp + term_b.exp  # Tuple addition
         new_coeff = term_a.coeff * term_b.coeff
         result.append((new_exp, new_coeff))
result = sort_and_merge(result)  # Heavy sorting!

# Sandalwood Dense Approach (Fast)
# The 'Table' tells us exactly where the result goes. No sorting needed.
target_index = Table[index_a, index_b]
ResultArray[target_index] += CoeffA * CoeffB

This reduces the complexity from $O(N log N)$ to $O(1)$ per operation.

2. Numba JIT Compilation

Even with the Dense Table, iterating through arrays in Python is slower than C. To fix this, sandalwood uses Numba, a Just-In-Time (JIT) compiler.

  • What it does: At runtime, Numba translates our multiplication functions into Optimized Machine Code (LLVM IR), similar to what a C++ compiler produces.

  • Parallelization Strategy: We implementing a Map-Reduce pattern with thread-local buffers to solve the race condition problem:

    1. Allocation: Use get_num_threads() to allocate a (num_threads, output_size) buffer.

    2. Manual Chunking: We explicitly partition the outer loop index idx_a into chunks [start, end) for each thread using prange.

    3. Thread Isolation: Thread t writes only to buffer[t, :]. This eliminates the need for atomic locks.

    4. Reduction: A final np.sum(buffer, axis=0) merges the partial results.

Code Insight: Use of fastmath=True allows the compiler to re-associate floating point operations for SIMD vectorization.

3. Lazy Materialization

A major hidden cost in Python libraries is Object Creation. Creating a new Python object for every intermediate step in a calculation like x = (a + b) * (c + d) is expensive.

  • The Optimization: sandalwood uses Lazy Evaluation.

  • How it works: When you multiply two MTFs, we compute the result as a raw array of numbers (Dense format) but do not convert it back to the user-friendly format (Exponents/Coefficients) immediately.

  • Benefit: If you perform a chain of 100 operations, we skip the expensive conversion 99 times. We only “materialize” the full object when you actually ask to see the coefficients (e.g., printing the function).

4. Fast Evaluation with Power Caching

Evaluating a high-order polynomial at millions of points (e.g., particle tracking) implies computing x^p millions of times. pow(x, p) is slow.

  • Algorithm: We use a 2D Thread-Local Power Cache.

    1. For each point (in parallel), we pre-compute powers [x^0, x^1, … x^N] for each variable dimension.

    2. These are stored in a simple array cached_powers[dim][order].

    3. The evaluation loop simply looks up cached_powers[d][idx] instead of calling pow.

  • Memory Safety: The neval method automatically detects large datasets and processes them in Chunks (default: 10,000 points) to keep the working set within CPU L1/L2 cache.

5. Vectorized Object Instantiation

When working with large arrays of COSY results, standard Python loops for wrapping C pointers into Python objects can be a significant bottleneck.

Anti-Pattern (Slow):

# Slow loop creating objects one by one
objects = []
for idx in cosy_indices:
    objects.append(MultivariateTaylorFunction(idx))

Best Practice (Fast):

Use from_cosy_indices to perform bulk instantiation. This moves the loop to optimized C/Fortran or vectorized Python logic, reducing overhead.

# Fast bulk creation
objects = MultivariateTaylorFunction.from_cosy_indices(cosy_indices)

Summary of Speedups

Internal benchmarks show the impact of these optimizations:

Operation

Naive Python

Optimized Sandalwood

Multiplication (Ord 12)

~500 ms

~45 ms

Evaluation (1M pts)

Crash (Out of Memory)

Stable & Fast

6. Hybrid Dispatch & Fortran Fast Path

While COSY Infinity provides exact high-order derivatives, the overhead of Differential Algebra (DA) tracking can be excessive for simple scalar evaluations (e.g., calculating the magnetic field at millions of points for visualization).

  • The Problem: Using generic DA arithmetic for floating-point numbers incurs overhead from object wrapping, memory allocation, and sparse/dense array management.

  • The Solution: sandalwood implements a Hybrid Dispatch mechanism in CosyBackend.

    1. Inspection: The backend checks if the input coordinates are pure Python float/int or numpy.ndarray (scalars) versus MultivariateTaylorFunction objects.

    2. Fast Path (Scalars): If inputs are scalars, execution is routed to the new compute_biot_savart_batch_fast Fortran routine. * Implementation: A pure DOUBLE PRECISION implementation of the Biot-Savart law. * Parallelism: Uses OpenMP to parallelize the loop over target points. * Performance: Avoids all DA overhead, achieving near-native Fortran speed (orders of magnitude faster than DA evaluation for scalars).

    3. General Path (DA): If inputs are MTFs, execution routes to the standard DA-enabled Fortran routines (compute_biot_savart_batch) to preserve derivative information.

This ensures the best of both worlds: high-performance visualization (using Fast Path) and high-precision optimization (using DA Path).

7. Phase 1 Optimizations (Vectorization & Batching)

Recent updates (Phase 1) focused on core arithmetic and composition bottlenecks.

Vectorized Horner’s Method

The neval kernel (sandalwood.numba_kernels.evaluate_dense_kernel) was refactored to use a Block-Based Reduction strategy. * Block Processing: Points are processed in blocks of 256 to maximize cache locality. * Transposed Layout: Powers are pre-computed in a (dim, max_order, block_size) layout, ensuring that the inner loop over block_size accesses contiguous memory, enabling SIMD vectorization. * Result: Significant speedup in large-scale evaluations.

Batch Composition

TaylorMap.compose previously performed term-by-term addition, incurring high overhead from repeated object creation and merging. * Optimization: Terms are now accumulated in a list terms_to_sum. * Batch Add: A new _batch_add method (in taylor_function.py) efficiently sums these terms in one pass (or using optimized reductions), minimizing intermediate object churn.

Fast-Path Dispatch

Arithmetic operator dispatch (__add__, __mul__, etc.) uses a simple if self._IMPLEMENTATION == "cosy" branch inside each dunder method. This avoids any global class mutation at initialization time and is safe from concurrent threads. The branch is predictable (same condition throughout the lifetime of a session), so modern CPUs can speculate it with near-zero overhead.

COSY Scalar Operators

CosyMtfData (the data holder for COSY backend) now supports direct operator overloading (e.g., da + 1.0), removing the need for explicit constant wrapping in standard Python arithmetic.

8. Phase 2 Optimizations (Solver & Batching)

The second phase targeted the integration layer between the high-level solvers and the backend kernels.

Structure of Arrays (SoA)

In em-simulation-platform, calculations were refactored to strictly enforce Structure of Arrays (SoA) memory layout. Instead of processing lists of FieldVector objects (Array of Structures), the solvers now manipulate contiguous arrays for Bx, By, and Bz components independently.

  • Benefit: Improved cache locality and SIMD vectorization potential during Biot-Savart integration.

  • Result: The calculate_b_field function now returns a SoA-optimized VectorField object directly.

Batch COSY Operations

To reduce the overhead of switching between Python and C for every single arithmetic operation, we implemented Batch Dispatch in CosyBackend.

  • `linear_combination`: Computes $sum c_i cdot x_i$ entirely within the C kernel. This is used for TaylorMap.compose and aggregating contributions from multiple source segments.

  • `batch_arithmetic`: Performs element-wise operations (Add, Sub, Mul, Div) on lists of DA objects in a single call.

  • Optimization Strategy: Due to stability issues with the native da_lin_comb COSY function, we implemented a robust Iterative Fallback at the C-interface level. This loops over inputs using direct low-level C functions (compute_da_add, compute_da_mul_const) without returning control to Python, preserving performance while ensuring correctness.

9. Final Benchmark Results

We benchmarked the RingCoil B-field calculation (Biot-Savart Law) to measure the cumulative impact of these optimizations. The comparison is between the optimized Python backend and the new SoA-optimized COSY backend.

Test Case: 1,000 source segments, Order 1 calculation.

Field Points

Python (s)

COSY (s)

Speedup

100

0.0157

0.0112

1.41x

10,000

1.1031

0.0074

148.68x

100,000

11.1953

0.0243

461.23x

Conclusion: The combination of SoA layout, Fast-Path dispatch, and Batch processing has yielded a ~460x speedup for large-scale simulations.

10. Map Composition Optimization (Phase 4)

Map composition (substituting one map into another: $F(G(x))$) is computationally expensive because it requires expanding high-order polynomials raised to high powers.

  • The Bottleneck: In sparse mode, computing $(a + b + dots)^n$ generates an explosion of intermediate terms before simplification.

  • The Solution: We implemented a Hybrid Strategy that combines high-level topological logic with Just-In-Time (JIT) compiled kernels using Numba.

Algorithm Design:

  1. Pre-computation: Before the main loop, we compute all necessary powers of the inner map components $(G_1(x)^p, G_2(x)^p, dots)$ and store them in a contiguous 3D dense array (dim, max_order+1, n_terms).

  2. Parallel Accumulation: We iterate over the terms of the outer map $F$ in parallel. * For each term $c cdot x_1^{p_1} x_2^{p_2} dots$, we fetch the pre-computed dense vectors. * We multiply them using the global Multiplication Table kernel (dense_mul). * The result is accumulated into Thread-Local Buffers to avoid race conditions.

  3. Result: The final dense array is converted back to a sparse MTF only once at the end.

Performance Impact: Benchmarks show a 3x - 5x speedup compared to the pure Python implementation for typical 2D and 3D maps at orders 4-8. This optimization is automatically engaged when the multiplication table is available (dense mode).

11. Phase 5 Optimizations (Python Structural)

Recent structural changes to the TaylorMap class significantly accelerated execution purely within Python:

  • List vs NumPy Object Arrays: TaylorMap.components previously used numpy.ndarray to hold objects. NumPy incurs severe overhead when doing operations on object arrays. These were replaced with standard Python lists and comprehensions, speeding up vector arithmetic by nearly 3x.

  • Linear Math in Inversion: TaylorMap.invert previously relied on generic map composition (which computes all cross terms) to invert the linear part of a map. This was optimized to use a specialized linear matrix-vector multiplication routine (MultivariateTaylorFunction._batch_add), avoiding polynomial composition overhead entirely.

  • Restored COSY Fast-Path: Re-enabled the COSY native backend to process TaylorMap.compose entirely within Fortran, avoiding the slow extraction and sparse evaluation cycle when components are tracked in COSY.