Source code for llg3d.post.extract

"""
Extract scalar values from .npz result files.

Use the ``llg3d.extract`` command line tool to extract scalar values from .npz result
files:

.. command-output:: llg3d.extract --help


Extract the total execution time from a ``run.npz`` file:

.. command-output:: llg3d.extract run.npz results/metrics/total_time
    :cwd: ../execute


Extract both the total execution time and the time per iteration:

.. command-output:: llg3d.extract run.npz results/metrics/total_time \
    results/metrics/time_per_ite
    :cwd: ../execute
"""

import argparse
from pathlib import Path

import numpy as np

from llg3d.io import RunResults, load_results




def _navigate(value: object, levels: list[str]):
    """
    Recursively navigate through nested structures using keys.

    Tries subscript access (dicts/arrays) first, then attribute access (objects).

    Args:
        value: The current value to navigate.
        levels: List of keys/attributes to navigate through.

    Returns:
        The final value after navigating through all levels.
    """
    if not levels:
        return value

    k = levels[0]
    try:
        # Try dict/array subscript access
        next_value = value[k]  # type: ignore
    except (TypeError, KeyError, IndexError):
        # Fallback to attribute access
        next_value = getattr(value, k)

    return _navigate(next_value, levels[1:])



ExtractedValue = float | int | str | bool | np.integer | np.floating




def extract_values(input: Path | RunResults, *keys: str) -> list[ExtractedValue]:
    """
    Extract scalar values from a .npz file or a RunResults object.

    Args:
        input: Path to the .npz result file or a RunResults object
        *keys: tuple of keys of the scalar values to extract
            (slash-separated for nested keys)

    Returns:
        The list of extracted scalar values

    Raises:
        ValueError: If one of the extracted values is not a scalar
    """
    if isinstance(input, Path):
        results = load_results(input)
    else:
        results = input
    values: list[ExtractedValue] = []
    for key in keys:
        # Split path into levels
        levels: list[str] = key.split("/")
        raw_value = _navigate(results, levels)

        # Handle numpy scalars and arrays with shape ()
        if isinstance(raw_value, np.ndarray):
            if raw_value.shape != ():
                raise ValueError(f"Extracted value for key '{key}' is not a scalar.")
            value = raw_value.item()  # Convert numpy scalar to Python scalar
        else:
            value = raw_value

        if not isinstance(value, (float, int, str, bool, np.integer, np.floating)):
            raise ValueError(f"Extracted value for key '{key}' is not a scalar.")
        values.append(value)
    return values





def main():  # pragma: no cover
    """Parse command line arguments and print simulation info."""
    parser = argparse.ArgumentParser(
        description="Extract scalar values from .npz result files.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("filename", help="Path to the .npz result file", type=Path)
    parser.add_argument(
        "keys",
        nargs="+",
        help=(
            "Key(s) of the scalar value(s) to extract (slash-separated for nested keys)"
        ),
    )
    args = parser.parse_args()

    values = extract_values(args.filename, *args.keys)
    print(" ".join(map(str, values)))