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:**

.. code-block:: python

   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:**

.. code-block:: python

   # 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 :meth:`~pyralysis.optimization.optimizer.Optimizer._finalize_iteration_state`
after each iteration. Custom optimizers **must** do the same for every carried
``Parameter``; see :doc:`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:**

.. code-block:: python

   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:

.. code-block:: python

   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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   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:**

.. code-block:: python

   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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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
:doc:`pipelines_distributed`.

**GPU Acceleration (Future):**

.. code-block:: python

   # 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:**

.. code-block:: python

   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:**

.. code-block:: python

   # 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**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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:**

.. code-block:: python

   # 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

----

:doc:`image_processing` | :doc:`testing`