Source code for ptyrad.runtime.seed

"""
Random seed resolution and initialization.

This module provides utilities to ensure reproducibility across PtyRAD 
reconstructions by managing random seeds for Python, NumPy, and PyTorch. 
It includes logic to prioritize user-provided seeds and automatically 
enforce seeding during distributed (multi-GPU) runs to prevent divergent 
model initializations.
"""

import logging
from typing import Optional

logger = logging.getLogger(__name__)



def resolve_seed_priority(args_seed, params_seed, acc):
    """Determines the final random seed to use based on user inputs and runtime context.

    This function resolves the seed using the following priority hierarchy:
    
    1. Command Line Argument (`args_seed`)
    2. Configuration File (`params_seed`)
    3. Automatic Fallback (forces seed to 42 if running in multi-GPU mode to 
       ensure consistent probe/object initializations across processes).
    4. None (no fixed seed).

    Args:
        args_seed (int or None): The seed provided via command-line arguments.
        params_seed (int or None): The seed provided in the parameters YAML file.
        acc (accelerate.Accelerator or None): The active HuggingFace Accelerator 
            instance, used to check the number of active processes.

    Returns:
        int or None: The resolved integer seed, or None if no seeding is required.
    """
    
    if  args_seed is not None:
        seed = args_seed
        logger.info(f"Random seed: {seed} provided by CLI argument")
    elif  params_seed is not None:
        seed = params_seed
        logger.info(f"Random seed: {seed} provided by params file")
    elif acc is not None and acc.num_processes > 1:
        seed = 42 # seed is required otherwise the probe position with random displacement could cause objects with different shapes
        logger.info(f"Random seed: {seed} is set automatically because multi GPU is detected but no seed is provided")
    else:
        seed = None
    return seed




def set_random_seed(seed: Optional[int], deterministic: bool = False):
    """Sets the random seeds for Python, NumPy, and PyTorch operations.

    Ensures that pseudo-random number generators (PRNGs) across all relevant 
    libraries and devices (CPU and all available CUDA GPUs) are synchronized 
    for reproducible reconstructions.

    Args:
        seed (int, optional): The integer seed to apply. If None, this function 
            does nothing.
        deterministic (bool, optional): If True, forces PyTorch to use deterministic 
            algorithms. Note that this can significantly impact performance and 
            may cause crashes if certain operations do not have a deterministic 
            implementation. Defaults to False.
            
    """
    
    import random

    import numpy as np
    import torch
    
    if seed is not None:
        random.seed(seed)                  # Python's RNG
        np.random.seed(seed)               # NumPy RNG
        torch.manual_seed(seed)            # PyTorch CPU RNG
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed) # All GPUs
        
        if deterministic:
            # This would slow down the operation a bit: https://docs.pytorch.org/docs/stable/notes/randomness.html
            torch.use_deterministic_algorithms(True)