conftest.py

"""Shared fixtures and helpers for regression tests."""

import os
import re
import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest


TESTS_DIR = Path(__file__).parent
EXPECTED_OUTPUT_DIR = TESTS_DIR / 'expected_output'
ACTUAL_OUTPUT_DIR = TESTS_DIR / 'actual_output'
CODE_DIR = TESTS_DIR.parent / 'code'

RTOL = 1e-14
ATOL = 1e-14

# LaTeX commands whose values depend on wall-clock timing (skip in comparison)
DONT_COMPARE = frozenset({
    'DMDTimeDMD',
    'DMDTimeSnapshot',
    'MedianTimeFOM',
    'MedianTimeROM',
    'MedianTimeSpeedup',
    'TimeReduction',
    'TimeSnapshots',
    'TimeTraining',
})


def test_env():
    """Environment variables for running scripts in CI mode."""
    env = os.environ.copy()
    env['CI'] = '1'
    env['OMP_NUM_THREADS'] = '1'
    env['OPENBLAS_NUM_THREADS'] = '1'
    env['MKL_NUM_THREADS'] = '1'
    env['MKL_CBWR'] = 'COMPATIBLE'
    env['MPLBACKEND'] = 'Agg'
    env['OUTPUT_DIR'] = ACTUAL_OUTPUT_DIR
    return env


def _count_header_rows(path: Path) -> int:
    """Return the number of header rows (metadata comments + header line)."""
    count = 0
    with open(path) as f:
        for line in f:
            line = line.strip()
            # Count comment lines (metadata)
            if not line or line.startswith('#'):
                count += 1
            # Count header line (non-numeric, non-comment)
            elif not _is_numeric_line(line):
                count += 1
            else:
                break
    return count


def _is_numeric_line(line: str) -> bool:
    """Check if a line contains numeric data (first token is a number)."""
    try:
        parts = line.split()
        if not parts:
            return False
        float(parts[0])
        return True
    except ValueError:
        return False


def _compute_min_tolerances(actual: np.ndarray, expected: np.ndarray) -> tuple[float, float]:
    """Compute minimum rtol/atol that would pass, given the other fixed at a baseline.

    Returns (min_rtol, min_atol) where:
    - min_rtol: minimum relative tolerance (assuming atol=1e-15)
    - min_atol: minimum absolute tolerance (assuming rtol=0)
    """
    abs_diff = np.abs(actual - expected)
    if abs_diff.size == 0:
        return 0.0, 0.0, ()

    worst_idx = np.unravel_index(np.argmax(abs_diff), abs_diff.shape)
    max_abs = abs_diff[worst_idx]
    exp_val = expected[worst_idx]

    # rtol needed if atol is tiny (1e-15 baseline)
    baseline_atol = 1e-15
    if np.abs(exp_val) > 1e-15:
        min_rtol = (max_abs - baseline_atol) / np.abs(exp_val)
    else:
        min_rtol = max_abs  # Fall back to absolute-like behavior

    # atol needed if rtol=0
    min_atol = max_abs

    return max(0, min_rtol), min_atol


def assert_dat_files_equal(actual_path: Path, expected_path: Path,
                           rtol: float = RTOL, atol: float = ATOL):
    """Assert that two .dat files (numpy savetxt format) are numerically equal.

    On failure, reports the minimum rtol/atol that would have passed this test.
    """
    actual = np.loadtxt(actual_path, skiprows=_count_header_rows(actual_path))
    expected = np.loadtxt(expected_path, skiprows=_count_header_rows(expected_path))

    try:
        np.testing.assert_allclose(
            actual, expected, rtol=rtol, atol=atol,
            err_msg=f'Mismatch in {actual_path.name} (comparing {actual_path} vs {expected_path})',
        )
    except AssertionError as err:
        # Compute what tolerances would have passed
        min_rtol, min_atol = _compute_min_tolerances(actual, expected)
        suggestion = (
            f'\n\n[TOLERANCE SUGGESTION] To pass this test, use at least one of:\n'
            f'  rtol={min_rtol:.6e} (with atol={atol})\n'
            f'  atol={min_atol:.6e} (with rtol={rtol})'
        )
        raise AssertionError(f'{err}\n{suggestion}') from None


def parse_tex_commands(path: Path) -> dict[str, str]:
    """Parse \\newcommand{\\Name}{value} entries from a results.tex file."""
    pattern = re.compile(r'\\newcommand\{\\(\w+)\}\{(.+?)\}')
    commands = {}
    for line in path.read_text().splitlines():
        m = pattern.match(line)
        if m:
            commands[m.group(1)] = m.group(2)
    return commands


def assert_tex_files_equal(actual_path: Path, expected_path: Path,
                           rtol: float = RTOL, atol: float = ATOL):
    """Assert that non-timing values in results.tex files match.

    On failure, reports the minimum rtol/atol that would have passed.
    """
    actual_cmds = parse_tex_commands(actual_path)
    expected_cmds = parse_tex_commands(expected_path)

    for name, expected_val in expected_cmds.items():
        if name in DONT_COMPARE:
            continue
        assert name in actual_cmds, f'Missing command \\{name} in {actual_path}'
        actual_val = actual_cmds[name]
        try:
            # Try numeric conversion ...
            actual_val = float(actual_val)
            expected_val = float(expected_val)
        except ValueError:
            # ... and fall back to exact string match for integers/strings
            assert actual_val == expected_val, (
                f'\\{name}: {actual_val!r} != {expected_val!r}'
            )
            return
        # Now that we're here we know we're dealing with floats.
        try:
            np.testing.assert_allclose(actual_val, expected_val, rtol=rtol, atol=atol,
                                        err_msg=f'\\{name}: {actual_val} != {expected_val}')
        except AssertionError as err:
            min_rtol, min_atol = _compute_min_tolerances(
                np.array(a), np.array(e)
            )
            suggestion = (
                f'\n\n[TOLERANCE SUGGESTION] To pass this test, use at least one of:\n'
                f'  rtol={min_rtol:.6e} (with atol={atol})\n'
                f'  atol={min_atol:.6e} (with rtol={rtol})'
            )
            raise AssertionError(f'{err}\n{suggestion}') from None


def assert_output_dir_matches(actual_dir: Path, expected_dir: Path,
                              rtol: float = RTOL, atol: float = ATOL):
    """Recursively compare all .dat and .tex files between two directory trees."""
    expected_files = sorted(expected_dir.rglob('*'))
    compared = 0

    for expected_file in expected_files:
        if not expected_file.is_file():
            continue
        rel = expected_file.relative_to(expected_dir)
        actual_file = actual_dir / rel

        assert actual_file.exists(), f'Missing output file: {actual_file}'

        if expected_file.suffix in '.dat' or expected_file.name.endswith('_limits.txt'):
            assert_dat_files_equal(actual_file, expected_file, rtol=rtol, atol=atol)
            compared += 1
        elif expected_file.suffix == '.tex':
            assert_tex_files_equal(actual_file, expected_file, rtol=rtol, atol=atol)
            compared += 1
        else:
            raise RuntimeError(f'I do not know how to compare {expected_file} with {actual_file}!')

    assert compared > 0, f'No comparable files found in {expected_dir}'


def run_script_and_compare(script: Path, output_subdirs: Path | list[Path],
                           rtol: float = RTOL, atol: float = ATOL):
    """Run a script in CI mode and compare its output directory to expected results."""
    if isinstance(output_subdirs, Path):
        output_subdirs = [output_subdirs]
    result = subprocess.run(
        [sys.executable, str(script)],
        cwd=str(script.parent),
        env=test_env(),
        capture_output=True,
        text=True,
    )
    assert result.returncode == 0, (
        f'{script.name} failed with exit code {result.returncode}\n'
        f'--- stdout ---\n{result.stdout[-2000:]}\n'
        f'--- stderr ---\n{result.stderr[-2000:]}'
    )
    for subdir in output_subdirs:
        assert_output_dir_matches(
            ACTUAL_OUTPUT_DIR / subdir,
            EXPECTED_OUTPUT_DIR / subdir,
            rtol=rtol, atol=atol,
        )