add mtt show

docs/src/dev-docs/cli/index.rst

@@ -2,14 +2,15 @@ CLI API
 =======
 
 This is the API for the command line interface ``cli`` functions for the ``train``,
-the ``eval`` and the ``export`` functions of ``metatrain``.
+the ``eval``, the ``export`` and the ``show`` functions of ``metatrain``.
 
 .. toctree::
    :maxdepth: 1
 
    train
    eval
    export
+   show
 
 We provide a custom formatter class for the formatting the help message of the
 ``argparse`` package.

docs/src/dev-docs/cli/show.rst

@@ -0,0 +1,7 @@
+Show
+####
+
+.. automodule:: metatrain.cli.show
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/src/getting-started/checkpoints.rst

@@ -37,6 +37,24 @@ of the current directory. If no checkpoint is found, the training will start
 from scratch.
 
 
+Inspecting a model
+------------------
+
+The ``show`` command prints a summary of a saved model without need to write any Python
+code. This shows quick info about targets and atomic types the model supports, the
+attached metadata and, for checkpoints, the architecture and training state.
+
+.. code-block:: bash
+
+    mtt show model.ckpt
+    mtt show model.pt
+
+As for the other sub-commands, a URL can be provided instead of a local file path to
+show a remote model.
+
+This is useful, for example, to find the target names of a pretrained model before
+fine-tuning it.
+
 Exporting models
 ----------------
 

src/metatrain/__main__.py

@@ -14,6 +14,7 @@
     _prepare_export_model_args,
     export_model,
 )
+from .cli.show import _add_show_model_parser, show_model
 from .cli.train import (
     _add_train_model_parser,
     _prepare_train_model_args,
@@ -59,6 +60,7 @@ def main():
     subparser = ap.add_subparsers(help="sub-command help")
     _add_eval_model_parser(subparser)
     _add_export_model_parser(subparser)
+    _add_show_model_parser(subparser)
     _add_train_model_parser(subparser)
 
     args = ap.parse_args()
@@ -102,6 +104,8 @@ def main():
             elif callable == "export_model":
                 _prepare_export_model_args(args)
                 export_model(**args.__dict__)
+            elif callable == "show_model":
+                show_model(**args.__dict__)
             elif callable == "train_model":
                 _prepare_train_model_args(args)
                 train_model(**args.__dict__)

src/metatrain/cli/show.py

@@ -0,0 +1,253 @@
+import argparse
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+from ase.data import chemical_symbols
+from metatomic.torch import (
+    AtomisticModel,
+    ModelMetadata,
+    ModelOutput,
+    load_atomistic_model,
+)
+
+from ..utils.data import TargetInfo
+from ..utils.io import is_exported_file, model_from_checkpoint, resolve_model_path
+from .formatter import CustomHelpFormatter
+
+
+def _add_show_model_parser(subparser: argparse._SubParsersAction) -> None:
+    """Add `show_model` parameters to an argparse (sub)-parser.
+
+    :param subparser: The argparse (sub)-parser to add the parameters to.
+    """
+
+    if show_model.__doc__ is not None:
+        description = show_model.__doc__.split(r":param")[0]
+    else:
+        description = None
+
+    parser = subparser.add_parser(
+        "show",
+        description=description,
+        formatter_class=CustomHelpFormatter,
+    )
+    parser.set_defaults(callable="show_model")
+
+    parser.add_argument(
+        "path",
+        type=str,
+        help="Model to show. Can be either a checkpoint (.ckpt) or an exported model "
+        "(.pt) as a local file or a URL.",
+    )
+    parser.add_argument(
+        "-e",
+        "--extensions-dir",
+        type=str,
+        required=False,
+        dest="extensions_directory",
+        default=None,
+        help=(
+            "path to a directory containing extensions required by an exported model"
+        ),
+    )
+
+
+def show_model(
+    path: Union[Path, str],
+    extensions_directory: Optional[Union[Path, str]] = None,
+) -> None:
+    """Show the contents of a saved model.
+
+    This prints a summary of a checkpoint (``.ckpt``) or exported model (``.pt``),
+    including architecture, targets that the model can predict, atomic types the model
+    supports and the attached metadata.
+
+    :param path: local or remote path to the model file (either a ``.ckpt`` checkpoint
+        or an exported ``.pt`` model)
+    :param extensions_directory: path to a directory containing all extensions required
+        by an exported model
+    """
+    if Path(path).suffix in [".yaml", ".yml"]:
+        raise ValueError(
+            f"path '{path}' seems to be a YAML option file and not a model"
+        )
+
+    local_path = resolve_model_path(path)
+
+    if is_exported_file(local_path):
+        model = load_atomistic_model(
+            local_path, extensions_directory=extensions_directory
+        )
+        lines = _describe_exported_model(model)
+    else:
+        if extensions_directory is not None:
+            logging.warning(
+                "the `--extensions-dir` option is only used for exported models and "
+                "will be ignored for checkpoints"
+            )
+        checkpoint = torch.load(local_path, weights_only=False, map_location="cpu")
+        lines = _describe_checkpoint(checkpoint)
+
+    summary = "\n".join(lines)
+    logging.info(f"Model information from {str(path)!r}\n\n{summary}")
+
+
+def _describe_checkpoint(checkpoint: Dict[str, Any]) -> List[str]:
+    architecture_name = checkpoint.get("architecture_name")
+    model_ckpt_version = checkpoint.get("model_ckpt_version")
+    trainer_ckpt_version = checkpoint.get("trainer_ckpt_version")
+    epoch = checkpoint.get("epoch")
+    best_epoch = checkpoint.get("best_epoch")
+    best_metric = checkpoint.get("best_metric")
+
+    model = model_from_checkpoint(checkpoint, context="export")
+
+    lines = ["file type: checkpoint"]
+    lines.append(f"architecture: {architecture_name}")
+
+    if model_ckpt_version is not None:
+        lines.append(f"model checkpoint version: {model_ckpt_version}")
+    if trainer_ckpt_version is not None:
+        lines.append(f"trainer checkpoint version: {trainer_ckpt_version}")
+    if epoch is not None:
+        lines.append(f"epoch: {epoch}")
+    if best_epoch is not None:
+        lines.append(f"best epoch: {best_epoch}")
+    if best_metric is not None:
+        lines.append(f"best validation metric: {best_metric}")
+
+    lines += _describe_metadata(model.metadata)
+
+    dataset_info = model.dataset_info
+    lines.append("")
+    lines.append(f"length unit: {dataset_info.length_unit or '(unknown)'}")
+    lines.append(f"atomic types: {_format_atomic_types(dataset_info.atomic_types)}")
+
+    lines.append("")
+    lines.append("targets:")
+    for name, target_info in dataset_info.targets.items():
+        lines += _describe_target(name, target_info)
+
+    auxiliary_outputs = sorted(
+        set(model.supported_outputs()) - set(dataset_info.targets)
+    )
+    if auxiliary_outputs:
+        lines.append("")
+        lines.append("auxiliary outputs:")
+        for name in auxiliary_outputs:
+            lines.append(f"  - {name}")
+
+    return lines
+
+
+def _describe_exported_model(model: AtomisticModel) -> List[str]:
+    capabilities = model.capabilities()
+
+    lines = ["file type: exported model"]
+    lines += _describe_metadata(model.metadata())
+
+    lines.append("")
+    lines.append(f"length unit: {capabilities.length_unit or '(unknown)'}")
+    lines.append(f"atomic types: {_format_atomic_types(capabilities.atomic_types)}")
+
+    interaction_range = f"interaction range: {capabilities.interaction_range}"
+    if capabilities.length_unit:
+        interaction_range += f" {capabilities.length_unit}"
+    lines.append(interaction_range)
+
+    lines.append(f"dtype: {capabilities.dtype}")
+    lines.append(f"supported devices: {', '.join(capabilities.supported_devices)}")
+
+    lines.append("")
+    lines.append("outputs:")
+    for name, output in capabilities.outputs.items():
+        lines += _describe_output(name, output)
+
+    return lines
+
+
+def _describe_metadata(metadata: ModelMetadata) -> List[str]:
+    lines = []
+
+    if metadata.name:
+        lines.append(f"  name: {metadata.name}")
+    if metadata.description:
+        lines.append(f"  description: {metadata.description}")
+    if metadata.authors:
+        lines.append("  authors: " + ", ".join(metadata.authors))
+
+    references = []
+    for section, section_references in metadata.references.items():
+        for reference in section_references:
+            references.append(f"    - ({section}) {reference}")
+    if references:
+        lines.append("  references:")
+        lines += references
+
+    if lines:
+        return ["", "metadata:"] + lines
+    else:
+        return []
+
+
+def _describe_target(name: str, target_info: TargetInfo) -> List[str]:
+    lines = [f"  {name}:"]
+
+    if target_info.quantity:
+        lines.append(f"    quantity: {target_info.quantity}")
+
+    lines.append(f"    unit: {target_info.unit or '(none)'}")
+    lines.append(f"    type: {_target_type(target_info)}")
+    lines.append(f"    sampled per: {target_info.sample_kind}")
+
+    if target_info.gradients:
+        lines.append("    gradients: " + ", ".join(target_info.gradients))
+
+    if target_info.description:
+        lines.append(f"    description: {target_info.description}")
+
+    return lines
+
+
+def _describe_output(name: str, output: ModelOutput) -> List[str]:
+    lines = [f"  {name}:"]
+
+    if output.quantity:
+        lines.append(f"    quantity: {output.quantity}")
+
+    lines.append(f"    unit: {output.unit or '(none)'}")
+    lines.append(f"    sampled per: {output.sample_kind}")
+
+    if output.explicit_gradients:
+        lines.append("    explicit gradients: " + ", ".join(output.explicit_gradients))
+
+    if output.description:
+        lines.append(f"    description: {output.description}")
+
+    return lines
+
+
+def _target_type(target_info: TargetInfo) -> str:
+    if target_info.is_scalar:
+        return "scalar"
+    elif target_info.is_cartesian:
+        return "cartesian"
+    elif target_info.is_spherical:
+        return "spherical"
+    elif target_info.is_atomic_basis:
+        return "atomic basis"
+    else:
+        return "unknown"
+
+
+def _format_atomic_types(atomic_types: List[int]) -> str:
+    entries = []
+    for atomic_type in atomic_types:
+        if 0 < atomic_type < len(chemical_symbols):
+            entries.append(f"{chemical_symbols[atomic_type]} ({atomic_type})")
+        else:
+            entries.append(str(atomic_type))
+
+    return ", ".join(entries)

src/metatrain/share/metatrain-completion.bash

@@ -8,6 +8,7 @@ _metatrain()
   local yaml='!*@(.yml|.yaml)'
   local ckpt='!*.ckpt'
   local pt='!*.pt'
+  local model='!*@(.ckpt|.pt)'
   local architecture_names=$(python -c "
 from metatrain.utils.architectures import find_all_architectures
 print(' '.join(find_all_architectures()))
@@ -57,6 +58,27 @@ print(' '.join(find_all_architectures()))
       COMPREPLY=( $(compgen -W "${opts}" -- "${cur_word}") )
       return 0
       ;;
+    show)
+      case "${prev_word}" in
+        -h|--help)
+          COMPREPLY=( )
+          return 0
+          ;;
+        -e|--extensions-dir)
+          COMPREPLY=( $(compgen -d -- "${cur_word}") )
+          return 0
+          ;;
+        *)
+          if [[ $COMP_CWORD -eq 2 ]]; then
+            COMPREPLY=( $(compgen -f -X "$model" -- "${cur_word}") )
+            return 0
+          fi
+          ;;
+      esac
+      local opts="-h --help -e --extensions-dir"
+      COMPREPLY=( $(compgen -W "${opts}" -- "${cur_word}") )
+      return 0
+      ;;
     eval)
       case "${prev_word}" in
         -h|--help|-o|--output|-b|--batch-size|--check-consistency)
@@ -85,7 +107,7 @@ print(' '.join(find_all_architectures()))
   esac
 
   # Complete the basic metatrain commands.
-  local opts="eval export train -h --help --debug --version"
+  local opts="eval export show train -h --help --debug --version"
   COMPREPLY=( $(compgen -W "${opts}" -- "${cur_word}") )
   return 0
 }