Performance & Scalability

Performance & Scalability#

Overview#

Pyralysis is designed for scalability and performance optimization, leveraging Dask for parallel processing, efficient memory management, and smart optimization strategies. This guide covers best practices for achieving optimal performance with both small and large datasets.

Note

Performance optimization in Pyralysis involves balancing memory usage, computation speed, and scalability across different hardware configurations.

Dask Integration#

Pyralysis leverages Dask for scalable parallel processing:

Key Benefits: - Out-of-core Processing: Handle datasets larger than available memory - Parallel Computation: Automatic parallelization across CPU cores - Lazy Evaluation: Defer computation until results are needed - Distributed Computing: Scale across multiple machines

Basic Dask Usage:

import dask.array as da
from pyralysis.reconstruction import Parameter

# Create large dask array
large_data = da.random.random((4096, 4096), chunks=(256, 256))

# Create parameter with dask backend
param = Parameter(
    data=large_data,
    cellsize=0.1*u.arcsec,
    chunks=(256, 256)
)

Chunking Strategies#

Optimal chunking is crucial for performance:

Chunk Size Guidelines:

  • Small Images (< 1024x1024): Use chunks similar to image size

  • Medium Images (1024x1024 - 4096x4096): Use 256x256 or 512x512 chunks

  • Large Images (> 4096x4096): Use 256x256 chunks for memory efficiency

Chunking Examples:

# Manual chunking
param = Parameter(
    data=data,
    chunks=(256, 256)  # Good for most use cases
)

# Automatic chunking
param = Parameter(
    data=data,
    chunks="auto"  # Let Dask choose optimal chunks
)

# Custom chunking for specific dimensions
param = Parameter(
    data=data,
    chunks={"dim_0": 256, "dim_1": 256, "dim_2": 50}
)

Gridding performance (DirtyMapper / weighting)#

Visibility gridding uses da.delayed tasks per row chunk of flattened u_pix / weights / data, then a pairwise tree combine of partial UV grids.

Visibility chunks (parallelism and memory)

  • Prefer fewer, larger row-aligned chunks when the padded grid is large (e.g. chunks={"row": 50000, "chan": 1, "corr": 4} in MS read).

  • Each chunk task allocates a full (v_pad, u_pad, n_corrs) accumulator; peak RAM scales with concurrent_tasks × 2 × v_pad × u_pad × n_corrs × dtype size.

  • Keep DirtyMapper.chunks for the output UV cube (FFT/normalization) separate from visibility chunking; "auto" or ~256 px spatial blocks are typical.

CPU vs GPU

  • CPU: Numba grid_block_unified_numba (serial per chunk).

  • GPU: Numba CUDA one thread per visibility with cuda.atomic.add; collection backend must match MS (array.backend=cupy). Tunable via PYRALYSIS_GRIDDING_CUDA_THREADS (default 256). Compare backends with tests/unit/utils/gridding/test_kernels_backend.py and examples/scripts/gridding_profile_chunk.py (HD163296 chunk sweeps).

Regression tests

  • pytest tests/unit/utils/gridding/test_reduction.py -m "not gpu"

  • pytest tests/unit/utils/gridding/test_performance.py -m slow (timing smoke)

Memory Management#

Efficient memory usage strategies:

Iterate state (required for Dask optimizers):

Each optimization iteration assigns new lazy arrays to the primary Parameter (and, for FISTA, the accelerated iterate). Without collapsing that lineage, Dask graph size grows O(iterations) and host or GPU memory can OOM even when individual forward models fit in RAM.

Built-in optimizers call _finalize_iteration_state() after each iteration. Custom optimizers must do the same for every carried Parameter; see Optimizers (Dask iterate state).

This is separate from the optional caches below: persist_gradient collapses the aggregate gradient, not the iterate; LBFGS persists correction pairs and search directions because they are read multiple times within one iteration. SDMM always finalizes auxiliary z/u arrays each outer iteration via _finalize_dask_arrays.

Gradient Persistence:

from pyralysis.optimization import ObjectiveFunction

# Enable gradient caching for repeated evaluations
obj_func = ObjectiveFunction(
    terms=terms,
    parameter=image,
    persist_gradient=True  # Cache gradients in memory
)

# Gradients are computed once and reused
for iteration in range(max_iterations):
    obj_func.calculate_gradient(iteration=iteration)

Proximal Operator Persistence:

For proximal optimization methods (FISTA, SDMM), you can enable persistence of proximal operator results to maintain consistent iteration times:

from pyralysis.optimization.terms.regularizers import L1Norm, Tikhonov

# Enable proximal persistence for specific regularizers
l1_term = L1Norm(penalization_factor=0.01, persist_proximal=True)
tikh_term = Tikhonov(penalization_factor=0.001, persist_proximal=True)

# Proximal results are persisted in memory after first computation
# This maintains consistent iteration times in FISTA/SDMM

Memory Optimization Techniques:

# Persist frequently accessed data
param.persist()

# Compute when needed
if param.is_dask():
    numpy_data = param.compute()  # Convert to numpy

# Monitor memory usage
print(f"Memory usage: {param.memory_usage()}")

# Use in-place operations
param *= 2.0  # In-place multiplication

Performance Optimization#

Strategies for optimal performance:

1. Efficient Term Processing:

# Process terms efficiently
obj_func = ObjectiveFunction(terms=terms, parameter=image)

# Enable smart term processing
obj_func.persist_gradient = True

# Use appropriate chunking
obj_func.chunks = (256, 256)

2. Multi-channel Optimization:

# Handle multi-channel data efficiently
if image.ndim > 2:
    # Process specific channels
    single_channel = image[:, :, 0]

    # Use parameter_index in regularizers
    l1_term = L1Norm(parameter_index=0, penalization_factor=0.01)

3. Masking for Performance:

from pyralysis.reconstruction.mask import Mask

# Apply operations only to relevant regions
mask = Mask(data=region_of_interest)

# Use mask in regularization
l1_term.function(mask=mask)  # Only processes masked pixels

Optimizer-Specific Performance Features#

L-BFGS Memory Management#

The optimized L-BFGS implementation includes several advanced performance features:

Efficient History Management:

from pyralysis.optimization.optimizer import LBFGS

# Standard memory-optimized configuration
optimizer = LBFGS(
    parameter=image,
    objective_function=obj_func,
    max_corrections=10,          # Max pairs (RAM + disk) for L-BFGS
    linesearch=ls
)

# For very large parameters with memory constraints
optimizer = LBFGS(
    parameter=large_image,
    objective_function=obj_func,
    max_corrections=10,
    memory_threshold_gb=50.0,   # Use sequential processing more aggressively
    use_disk_storage=True,      # Allow zarr spill when in_memory_limit is exceeded
    in_memory_limit=5,          # Keep only 5 pairs in RAM
    linesearch=ls
)

Key Memory Optimizations:

  • Lazy allocation: Correction pairs live in collections.deque structures; nothing is allocated until the first accepted pair. The total stored is capped by max_corrections (oldest dropped when full).

  • Hybrid Processing: Automatically switches between vectorized Dask operations (for speed) and sequential processing (for memory efficiency) based on estimated memory usage.

  • Memory Threshold: Configurable threshold (memory_threshold_gb, default: 100.0 GB) determines when to use sequential processing to avoid memory overflow.

  • Chunking Normalization: All correction pairs are automatically rechunked to match parameter chunking, reducing fragmentation and improving Dask performance.

  • Computation Caching: Gamma scaling factor and rho array conversions are cached to avoid redundant computations.

  • Optional disk storage: With use_disk_storage=True and an explicit in_memory_limit, oldest in-memory pairs can be written to zarr; total pairs still respect max_corrections.

Large-Scale Computation:

# Optimized for large images with complex gradients
large_image = Image(data=complex_data, cellsize=0.01*un.arcsec)
optimizer = LBFGS(
    parameter=large_image,
    memory_threshold_gb=50.0,   # Lower threshold for large arrays
    use_disk_storage=True,      # Allow spill to zarr
    in_memory_limit=5,          # Bound RAM; total pairs still <= max_corrections
    ...
)

# Uses lazy allocation and sequential processing when needed

Conjugate Gradient Robustness#

Enhanced CG variants with improved complexity handling:

Complex Parameter Support:

# All CG variants support complex numbers automatically
from pyralysis.optimization.optimizer import HagerZhang

optimizer = HagerZhang(
    parameter=complex_polarization_image,
    objective_function=obj_func,
    linesearch=ls
)

# Uses optimized scalar computation for beta parameters
# Improved numerical stability for complex gradients

Scalability Features#

Pyralysis scales from single machines to distributed clusters:

Single Machine Optimization:

# Configure for single machine
import dask.config

# Set number of threads
dask.config.set(scheduler='threads')

# Configure chunking
param.chunks = (512, 512)  # Larger chunks for single machine

Distributed Computing:

# Configure for distributed cluster
from dask.distributed import Client

# Connect to cluster
client = Client("scheduler-address:8786")

# Use distributed scheduler
import dask.config
dask.config.set(scheduler='distributed')

For simulation and imaging pipelines, you can drive the same idea through context configuration and the built-in setup/teardown steps instead of manual Client wiring in user code; see Pipelines and distributed Dask.

GPU Acceleration (Future):

# GPU support is in development
# Future: Automatic GPU acceleration for compatible operations

# Check GPU availability
try:
    import cupy as cp
    print("GPU acceleration available")
except ImportError:
    print("GPU acceleration not available")

Benchmarking Performance#

Monitor and optimize performance:

Performance Metrics:

import time

# Time function evaluation
start_time = time.time()
function_value = obj_func.calculate_function()
eval_time = time.time() - start_time
print(f"Function evaluation time: {eval_time:.2f}s")

# Time gradient computation
start_time = time.time()
obj_func.calculate_gradient(iteration=1)
grad_time = time.time() - start_time
print(f"Gradient computation time: {grad_time:.2f}s")

Memory Profiling:

# Monitor memory usage
print(f"Parameter memory: {param.memory_usage()}")
print(f"Image memory: {image.memory_usage()}")

# Check chunk information
if param.current_chunks:
    print(f"Current chunks: {param.current_chunks}")

Best Practices#

1. Chunking Strategy:

  • Start with automatic chunking (chunks=”auto”)

  • Adjust based on memory constraints and performance

  • Use smaller chunks for memory-constrained systems

  • Use larger chunks for single-machine setups

2. Memory Management:

  • Enable gradient persistence for repeated evaluations

  • Use appropriate chunking to balance memory and performance

  • Monitor memory usage during optimization

  • Use in-place operations when possible

3. Performance Monitoring:

  • Profile function and gradient evaluation times

  • Monitor memory usage patterns

  • Use appropriate chunking for your data size

  • Enable persistence for frequently accessed data

4. Scalability Planning:

  • Start with single-machine optimization

  • Scale to distributed computing for large datasets

  • Use appropriate chunking for your cluster size

  • Monitor cluster performance and resource usage

Example: Performance-Optimized Workflow

# Create parameter with optimal chunking
param = Parameter(
    data=large_data,
    cellsize=0.1*u.arcsec,
    chunks=(256, 256)  # Optimal for most systems
)

# Create objective function with performance features
obj_func = ObjectiveFunction(
    terms=terms,
    parameter=param,
    persist_gradient=True,  # Cache gradients
    chunks="auto"           # Automatic chunking
)

# Monitor performance
start_time = time.time()

# Optimization loop
for iteration in range(max_iterations):
    function_value = obj_func.calculate_function()
    obj_func.calculate_gradient(iteration=iteration)

    # Monitor progress
    if iteration % 10 == 0:
        elapsed = time.time() - start_time
        print(f"Iteration {iteration}: {elapsed:.2f}s elapsed")

total_time = time.time() - start_time
print(f"Total optimization time: {total_time:.2f}s")

Troubleshooting Performance Issues#

Common performance problems and solutions:

1. Memory Issues:

# Problem: Out of memory
# Solution: Reduce chunk size
param.chunks = (128, 128)  # Smaller chunks

# Problem: Excessive memory usage
# Solution: Enable persistence selectively
obj_func.persist_gradient = False  # Disable if memory is limited

# For proximal methods, disable proximal persistence if needed
l1_term.persist_proximal = False  # Disable proximal persistence

2. Slow Computation:

# Problem: Small chunks causing overhead
# Solution: Increase chunk size
param.chunks = (512, 512)  # Larger chunks

# Problem: No parallelization
# Solution: Check Dask configuration
import dask.config
print(f"Number of threads: {dask.config.get('num_workers')}")

3. Inefficient Operations:

# Problem: Repeated gradient computation
# Solution: Enable gradient persistence
obj_func.persist_gradient = True

# Problem: Large memory footprint
# Solution: Use appropriate chunking
param.chunks = (256, 256)  # Balance memory and performance

Image Processing | Testing & Quality Assurance