utils

Utility functions for the different SRC algorithms.

Functions

default_rng

default_rng(seed: int | None) -> np.random.Generator

Return a seeded NumPy Generator.

Always uses NumPy so that the same seed produces identical draws regardless of the device, and avoids CuPy Generator API differences (e.g. missing .normal()).

Source code in src/src_method/utils/_backend.py

def default_rng(
    seed: int | None,
) -> np.random.Generator:
    """Return a seeded NumPy ``Generator``.

    Always uses NumPy so that the same seed produces identical draws
    regardless of the device, and avoids CuPy ``Generator`` API
    differences (e.g. missing ``.normal()``).
    """
    return np.random.default_rng(seed)

get_xp

get_xp(device: str) -> ModuleType

Return the array module for the requested device.

PARAMETER	DESCRIPTION
device	"cpu" for numpy or "gpu" for cupy. TYPE: str

RETURNS	DESCRIPTION
ModuleType	The numpy or cupy module.

RAISES	DESCRIPTION
ValueError	If device is not recognised.
ImportError	If device="gpu" but cupy is not installed.

Source code in src/src_method/utils/_backend.py

def get_xp(device: str) -> ModuleType:
    """Return the array module for the requested device.

    Args:
        device: ``"cpu"`` for numpy or ``"gpu"`` for cupy.

    Returns:
        The numpy or cupy module.

    Raises:
        ValueError: If ``device`` is not recognised.
        ImportError: If ``device="gpu"`` but cupy is not installed.
    """
    if device == "cpu":
        return np
    if device == "gpu":
        import cupy  # noqa: PLC0415  (lazy: optional dependency)

        return cupy
    msg = f"Unknown device {device!r}; expected 'cpu' or 'gpu'."
    raise ValueError(msg)

setup_logging

setup_logging(default_level: str = 'INFO') -> None

Configure logging for the entire application.

PARAMETER	DESCRIPTION
default_level	Default logging level. TYPE: str DEFAULT: 'INFO'

The logging level can be overridden by setting the LOG_LEVEL_SRC environment. For example:

LOG_LEVEL_SRC=DEBUG python -m <your_module>
pytest -o log_cli_level=DEBUG tests/

Source code in src/src_method/utils/logging_config.py

def setup_logging(default_level: str = "INFO") -> None:
    """Configure logging for the entire application.

    Args:
        default_level: Default logging level.

    The logging level can be overridden by setting the `LOG_LEVEL_SRC`
    environment. For example:

    ```bash
    LOG_LEVEL_SRC=DEBUG python -m <your_module>
    pytest -o log_cli_level=DEBUG tests/
    ```
    """
    # Configure Python's standard logging
    level_name = os.environ.get("LOG_LEVEL_SRC", default_level)
    level = logging.getLevelName(level_name.upper())
    logging.basicConfig(
        level=level,
        format="%(message)s",
        stream=sys.stdout,
    )

    # Remove colors for slurm logs but not for pytest
    use_colors = "SLURM_JOB_ID" not in os.environ

    # Configure structlog to use the standard logger
    structlog.configure(
        processors=[
            structlog.contextvars.merge_contextvars,
            structlog.processors.add_log_level,
            structlog.processors.StackInfoRenderer(),
            structlog.dev.set_exc_info,
            structlog.processors.TimeStamper(fmt="%Y-%m-%d %H:%M:%S", utc=False),
            structlog.dev.ConsoleRenderer(colors=use_colors),
        ],
        logger_factory=structlog.stdlib.LoggerFactory(),
        wrapper_class=structlog.stdlib.BoundLogger,
        cache_logger_on_first_use=True,
    )

to_numpy

to_numpy(arr: NDArray) -> np.ndarray

Bring an array onto the host as a numpy array (no-op for numpy).

Source code in src/src_method/utils/_backend.py

def to_numpy(arr: NDArray) -> np.ndarray:
    """Bring an array onto the host as a numpy array (no-op for numpy)."""
    if isinstance(arr, np.ndarray):
        return arr
    # cupy.ndarray exposes .get(); fall back to np.asarray for other dispatchers.
    get = getattr(arr, "get", None)
    return get() if callable(get) else np.asarray(arr)

truncated_qr

truncated_qr(
    matrix: NDArray, cutoff: float, xp: ModuleType = np
) -> NDArray

QR with SVD-based rank truncation, returning only the isometry.

Decomposes matrix = Q @ R, then truncates via SVD on R, discarding singular values below cutoff * sigma_max.

When cutoff <= 0 this falls back to a plain QR (no truncation).

PARAMETER	DESCRIPTION
matrix	Input matrix of shape (m, n). TYPE: NDArray
cutoff	Relative singular-value threshold. Singular values satisfying s < cutoff * s_max are discarded. TYPE: float
xp	Array module (numpy or cupy); defaults to numpy. TYPE: ModuleType DEFAULT: numpy

RETURNS	DESCRIPTION
NDArray	The truncated isometry Q of shape (m, rank).

Source code in src/src_method/utils/linalg.py

def truncated_qr(matrix: NDArray, cutoff: float, xp: ModuleType = np) -> NDArray:
    """QR with SVD-based rank truncation, returning only the isometry.

    Decomposes ``matrix = Q @ R``, then truncates via SVD on R,
    discarding singular values below ``cutoff * sigma_max``.

    When ``cutoff <= 0`` this falls back to a plain QR (no truncation).

    Args:
        matrix: Input matrix of shape (m, n).
        cutoff: Relative singular-value threshold.  Singular values
            satisfying ``s < cutoff * s_max`` are discarded.
        xp: Array module (``numpy`` or ``cupy``); defaults to numpy.

    Returns:
        The truncated isometry Q of shape ``(m, rank)``.
    """
    m, n = matrix.shape
    transpose = m < n

    if cutoff <= 0:
        return xp.linalg.qr(matrix)[0]

    Q, R = xp.linalg.qr(matrix.T if transpose else matrix)
    R_np = R.get() if hasattr(R, "get") else np.asarray(R)
    U, S, _ = np.linalg.svd(R_np.T if transpose else R_np, full_matrices=False)
    rank = max(1, int((cutoff * S[0] <= S).sum()))

    if transpose:
        Q_trunc = U[:, :rank]
    else:
        U_trunc = xp.asarray(U[:, :rank]) if xp is not np else U[:, :rank]
        Q_trunc = Q @ U_trunc

    if xp is not np:
        return xp.asarray(Q_trunc)
    return Q_trunc