pytransc.utils.forward_context

  1"""Forward pool context for log_likelihood function access.
  2
  3This module provides a mechanism for log_likelihood functions to access forward pools
  4that were passed to the samplers. It uses global variables that are automatically
  5process-local in multiprocessing environments.
  6
  7The typical usage pattern is:
  81. User passes forward_pool to a sampler function
  92. Sampler sets the pool before calling log_posterior
 103. Log_posterior function gets the pool via get_forward_pool()
 114. Sampler clears the pool after log_posterior call
 12
 13Examples
 14--------
 15In a log_posterior function:
 16
 17>>> from pytransc.utils.forward_context import get_forward_pool
 18>>> def log_posterior(params, state):
 19...     forward_pool = get_forward_pool()
 20...     if forward_pool is not None:
 21...         # Use pool for parallel forward solver calls
 22...         results = forward_pool.map(forward_solver, param_chunks)
 23...     else:
 24...         # Sequential execution
 25...         results = [forward_solver(p) for p in param_chunks]
 26"""
 27
 28# Global variable for forward pool access (process-local in multiprocessing)
 29_forward_pool = None
 30
 31
 32def set_forward_pool(pool):
 33    """Set forward pool in current process.
 34
 35    Parameters
 36    ----------
 37    pool : Any
 38        Pool object with map() method. Typically ProcessPoolExecutor,
 39        ThreadPoolExecutor, or schwimmbad pools.
 40
 41    Raises
 42    ------
 43    ValueError
 44        If pool does not have a map() method.
 45    RuntimeError
 46        If setting the pool fails for any reason.
 47
 48    Examples
 49    --------
 50    >>> from concurrent.futures import ProcessPoolExecutor
 51    >>> with ProcessPoolExecutor(max_workers=4) as pool:
 52    ...     set_forward_pool(pool)
 53    """
 54    global _forward_pool
 55
 56    # Validate pool has map() method
 57    if pool is not None and not hasattr(pool, 'map'):
 58        raise ValueError("Forward pool must have a 'map' method")
 59
 60    try:
 61        _forward_pool = pool
 62    except Exception as e:
 63        raise RuntimeError(f"Failed to set forward pool: {e}") from e
 64
 65
 66def get_forward_pool():
 67    """Get forward pool from current process.
 68
 69    Returns
 70    -------
 71    Pool or None
 72        The forward pool if available, None otherwise.
 73
 74    Examples
 75    --------
 76    >>> forward_pool = get_forward_pool()
 77    >>> if forward_pool is not None:
 78    ...     # Use pool for parallel execution
 79    ...     results = forward_pool.map(my_function, data_chunks)
 80    """
 81    return _forward_pool
 82
 83
 84def clear_forward_pool():
 85    """Clear forward pool from current process.
 86
 87    This should be called after log_posterior evaluation to clean up
 88    the pool reference and prevent memory leaks.
 89
 90    Raises
 91    ------
 92    RuntimeError
 93        If clearing the pool fails for any reason.
 94
 95    Examples
 96    --------
 97    >>> clear_forward_pool()  # Pool reference is cleared
 98    >>> assert get_forward_pool() is None
 99    """
100    global _forward_pool
101    try:
102        _forward_pool = None
103    except Exception as e:
104        raise RuntimeError(f"Failed to clear forward pool: {e}") from e