Source code for llg3d.post.info

"""
Dump simulation results from a .npz file.

To browse the content of the ``run.npz`` file, use the ``llg3d.info`` command:

.. command-output:: llg3d.info run.npz
    :cwd: ../execute/

The numpy arrays cans be previewed with more detailed information using the
``--verbose`` option:

.. command-output:: llg3d.info run.npz --verbose
    :cwd: ../execute/

"""

import argparse
import textwrap
from pathlib import Path

import numpy as np

from ..io import Metrics, Records, RunResults, format_profiling_table, load_results

INDENT = 4 * " "




def get_array_memory(arr: np.ndarray) -> str:
    """Return a human-readable string of the array memory size."""
    nbytes: float | int = arr.nbytes
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if nbytes < 1024.0:
            return f"{nbytes:.2f} {unit}"
        nbytes /= 1024.0
    return f"{nbytes:.2f} PB"





def summarize_array(arr: np.ndarray) -> str:
    """Return a summary of the array."""
    s = f"""shape: {arr.shape}
dtype: {arr.dtype}
min: {arr.min():.6e}, max: {arr.max():.6e}
mean: {arr.mean():.6e}, std: {arr.std():.6e}
memory: {get_array_memory(arr)}
"""
    s += np.array2string(
        arr,
        max_line_width=120,
        threshold=50,
        edgeitems=2,
        formatter={"float": lambda x: f"{x:.3e}"},
    )
    return textwrap.indent(s, INDENT)





def dict_to_str(d: dict) -> str:
    """
    Convert a dictionary to a formatted string.

    Right-aligns the values for better readability.
    Skips empty dicts and empty lists.

    Args:
        d: The dictionary to convert.

    Returns:
        A formatted string representation of the dictionary.
    """
    # Filter out empty dicts and empty lists
    filtered = {
        k: v for k, v in d.items() if not (isinstance(v, (dict, list)) and not v)
    }

    if not filtered:
        return ""

    lines = []
    max_key_length = max(len(str(key)) for key in filtered.keys())
    for key, value in filtered.items():
        lines.append(f"{key}: {' ' * (max_key_length - len(str(key)))}{value}")
    return textwrap.indent("\n".join(lines), INDENT)





def _format_array(arr: np.ndarray, verbose: bool) -> str:
    """Format an array depending on verbosity."""
    if verbose:
        return summarize_array(arr)
    else:
        return textwrap.indent(f"shape: {arr.shape}, dtype: {arr.dtype}", INDENT)





def _format_records(
    records: Records | dict, verbose: bool, indent: int = 0
) -> list[str]:
    """Recursively format records section (nested dicts/arrays)."""
    lines: list[str] = []
    prefix = " " * indent
    for key, value in records.items():
        if isinstance(value, dict):
            # Nested dict: recurse
            lines.append(f"{prefix}{key}")
            lines.extend(_format_records(value, verbose, indent + 4))
        elif isinstance(value, np.ndarray):
            # Numpy array: format accordingly
            lines.append(f"{prefix}{key}")
            lines.append(
                textwrap.indent(_format_array(value, verbose), prefix + " " * 4)
            )
        else:
            # Other types: just print key and value
            lines.append(f"{prefix}{key}: {value}")
    return lines





def _format_metrics(metrics: Metrics) -> list[str]:
    """
    Format metrics data, handling profiling_stats specially.

    Args:
        metrics: Dictionary containing metrics data

    Returns:
        List of formatted lines to append
    """
    lines = []

    # Extract metrics without profiling_stats
    metrics_noprof = {k: v for k, v in metrics.items() if k != "profiling_stats"}
    formatted = dict_to_str(metrics_noprof)
    lines.append(formatted)

    # Format profiling_stats as a table if present
    if "profiling_stats" in metrics and metrics["profiling_stats"]:
        lines.append("profiling_stats")
        total_time = metrics.get("total_time")
        lines.append(
            textwrap.indent(
                format_profiling_table(metrics["profiling_stats"], total_time), INDENT
            )
        )

    return lines





def get_info(filename: Path, verbose: bool = False) -> str:
    """
    Returns the simulation parameters and results from a .npz file.

    Args:
        filename: Path to the .npz result file
        verbose: If True, includes detailed array information

    Returns:
        A formatted string with simulation information
    """
    run: RunResults = load_results(filename)

    lines: list[str] = []

    # Params
    lines.append("params")
    lines.append(dict_to_str(run.params.as_dict()))

    # Metrics
    lines.append("results/metrics")
    lines.extend(_format_metrics(run.results["metrics"]))

    # Observables
    if "observables" in run.results:
        observables = run.results["observables"]
        lines.append("results/observables")
        lines.append(dict_to_str(dict(observables)))

    # Records
    if "records" in run.results:
        records = run.results["records"]
        lines.append("results/records")
        lines.extend(_format_records(records, verbose))

    return "\n".join(lines)





def main():  # pragma: no cover
    """Parse command line arguments and print simulation info."""
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("filename", help="Path to the .npz result file", type=Path)
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    args = parser.parse_args()
    print(get_info(args.filename, args.verbose))