#!/usr/bin/env python3
"""Split multi-image .con files into per-image structures.
.. versionadded:: 0.0.2
"""
# /// script
# requires-python = ">=3.11"
# dependencies = [
# "click",
# "ase",
# "rich",
# "readcon>=0.7.0",
# "rgpycrumbs>=1.10.2",
# ]
# ///
import logging
import sys
from enum import Enum
from pathlib import Path
import click
from ase.io import read as aseread
from rich.console import Console
from rich.logging import RichHandler
from rgpycrumbs.geom.api.alignment import IRAConfig, align_structure_robust
[docs]
_CON_SUFFIXES = {".con", ".convel"}
[docs]
def _is_con_path(path: Path) -> bool:
return path.suffix.lower() in _CON_SUFFIXES
[docs]
def _read_trajectory_frames(path: Path) -> tuple[list, list | None]:
"""Load trajectory frames; CON/convel go through readcon to keep metadata.
Returns ``(ase_atoms_list, con_frames_or_none)``. ``con_frames`` is populated
only for ``.con``/``.convel`` inputs so writers can preserve per-frame energy.
"""
if _is_con_path(path):
from readcon import read_con
con_frames = list(read_con(str(path)))
atoms_list = [frame.to_ase() for frame in con_frames]
for atoms, frame in zip(atoms_list, con_frames, strict=True):
if frame.energy is not None:
atoms.info["energy"] = float(frame.energy)
return atoms_list, con_frames
return list(aseread(path, index=":")), None
[docs]
def _energy_from_atoms(atoms) -> float | None:
energy = atoms.info.get("energy") if getattr(atoms, "info", None) else None
if energy is None and getattr(atoms, "calc", None) is not None:
try:
energy = atoms.get_potential_energy()
except Exception:
energy = None
if energy is None:
return None
try:
return float(energy)
except (TypeError, ValueError):
return None
[docs]
def _apply_frame_energy(frame, energy: float | None):
"""Set per-frame total energy on a ``ConFrame`` copy-friendly handle.
Prefers ``set_energy`` / ``set_scalar_metadata`` (readcon>=0.13). Older
wheels that expose only a read-only ``energy`` property are reconstructed
with string metadata values accepted by the constructor.
"""
if energy is None:
return frame
value = float(energy)
setter = getattr(frame, "set_energy", None)
if callable(setter):
setter(value)
return frame
scalar = getattr(frame, "set_scalar_metadata", None)
if callable(scalar):
scalar("energy", value)
return frame
from readcon import ConFrame
metadata = frame.metadata
if hasattr(metadata, "items"):
metadata_dict = {str(k): v for k, v in metadata.items()}
else:
metadata_dict = dict(metadata or {})
metadata_dict.pop("energy", None)
clone = ConFrame(
frame.cell,
frame.angles,
frame.atoms,
frame.prebox_header,
frame.postbox_header,
metadata_dict,
)
clone_setter = getattr(clone, "set_energy", None)
if callable(clone_setter):
clone_setter(value)
return clone
metadata_dict["energy"] = str(value)
return ConFrame(
frame.cell,
frame.angles,
frame.atoms,
frame.prebox_header,
frame.postbox_header,
metadata_dict,
)
[docs]
def _write_con_atoms(dest: Path, atoms, energy: float | None = None) -> None:
"""Write one ASE structure as a CON file via readcon."""
from readcon import ConFrame, write_con
if energy is None:
energy = _energy_from_atoms(atoms)
frame = _apply_frame_energy(ConFrame.from_ase(atoms), energy)
write_con(str(dest), [frame])
[docs]
def _write_con_frame(dest: Path, frame, atoms=None, energy: float | None = None) -> None:
"""Write a readcon ``ConFrame``, optionally replacing positions from *atoms*."""
from readcon import ConFrame, write_con
if atoms is not None:
if energy is None:
energy = (
frame.energy if frame.energy is not None else _energy_from_atoms(atoms)
)
updated = _apply_frame_energy(ConFrame.from_ase(atoms), energy)
write_con(str(dest), [updated])
return
if energy is not None:
frame = _apply_frame_energy(frame, energy)
write_con(str(dest), [frame])
# Optional IRA import logic
try:
from rgpycrumbs._aux import _import_from_parent_env
[docs]
ira_mod = _import_from_parent_env("ira_mod")
except ImportError:
ira_mod = None
logging.basicConfig(
level=logging.INFO,
format="%(message)s",
datefmt="[%X]",
handlers=[
RichHandler(
console=CONSOLE,
rich_tracebacks=True,
markup=True,
show_path=False,
show_level=True,
show_time=True,
)
],
)
[docs]
class AlignMode(Enum):
"""Defines structural alignment strategies."""
[docs]
ENDPOINTS = "endpoints"
[docs]
class SplitMode(Enum):
"""Defines trajectory validation strictness."""
[docs]
NEB = "neb" # Strict: must be a clean multiple of images_per_path
[docs]
FLEX = "flex" # Flexible: allows partial paths or simple slicing
[docs]
def align_path(frames, mode: AlignMode, iraconf: IRAConfig):
"""Applies the selected alignment strategy to the image sequence."""
if mode == AlignMode.NONE or len(frames) < 2:
return frames
ref = frames[0]
if mode == AlignMode.ALL:
logging.info("Aligning [bold]all[/bold] images to reactant reference.")
return [ref.copy()] + [
align_structure_robust(
ref, f.copy(), IRAConfig(iraconf.use_ira, iraconf.kmax)
).atoms
for f in frames[1:]
]
if mode == AlignMode.ENDPOINTS:
logging.info("Aligning [bold]endpoints[/bold] (product to reactant) only.")
aligned_product = align_structure_robust(ref, frames[-1].copy(), iraconf).atoms
# Intermediate frames remain unchanged in this specific mode logic,
# Usually, endpoint alignment implies ensuring the BCs match.
new_frames = [f.copy() for f in frames]
new_frames[-1] = aligned_product
return new_frames
return frames
@click.command()
@click.argument(
"neb_trajectory_file",
type=click.Path(exists=True, dir_okay=False, readable=True, path_type=Path),
)
@click.option(
"--mode",
type=click.Choice([m.value for m in SplitMode]),
default=SplitMode.FLEX.value,
help="Validation mode: 'neb' (strict multiples) or 'normal' (flexible).",
)
@click.option(
"--output-dir",
type=click.Path(file_okay=False, writable=True, path_type=Path),
default=None,
help="Directory to save output files. Defaults to the input filename stem.",
)
@click.option(
"--images-per-path",
type=int,
required=True,
help="Number of images in a single NEB path (e.g., 7). [REQUIRED]",
)
@click.option(
"--path-index",
type=int,
default=-1,
show_default=True,
help="Index of the NEB path to extract (0-based). Use -1 for the last path.",
)
@click.option(
"--center/--no-center",
default=False,
help="Center the atomic coordinates around the origin.",
)
@click.option(
"--box-diagonal",
nargs=3,
type=(float, float, float),
default=(25.0, 25.0, 25.0),
show_default=True,
help="Override the unit cell dimensions (Å) during processing.",
)
@click.option(
"--align-type",
type=click.Choice([m.value for m in AlignMode]),
default=AlignMode.NONE.value,
help="Alignment: 'all' (every image), 'endpoints' (reactant/product), or 'none'.",
)
@click.option(
"--use-ira",
is_flag=True,
help="Enable Iterative Reordering and Alignment (requires ira_mod).",
)
@click.option(
"--ira-kmax",
type=float,
default=1.8,
help="kmax factor for the IRA matching algorithm.",
)
@click.option(
"--path-list-filename",
default="ipath.dat",
help="Name of the file listing the generated .con absolute paths.",
)
[docs]
def con_splitter(
neb_trajectory_file: Path,
mode: str,
output_dir: Path | None,
images_per_path: int,
path_index: int,
center: bool,
box_diagonal: tuple[float, float, float],
align_type: str,
use_ira: bool,
ira_kmax: float,
path_list_filename: str,
):
"""
Splits a multi-step trajectory file (.traj, .con, etc.) into
individual .con files for a *single* specified path.
This script reads a trajectory file, which may contain multiple NEB
optimization steps (paths), and extracts only the frames corresponding
to a single specified path.
It writes each frame of that path into a separate .con file
(e.g., ipath_000.con, ipath_001.con, ...).
It also generates a text file (default: 'ipath.dat') that lists the
absolute paths of all created .con files.
This utility extracts specific optimization steps and applies physical
chemistry refinements such as centering, cell overrides, and structural
alignment (RMSD minimization).
"""
if output_dir is None:
output_dir = Path(neb_trajectory_file.stem)
output_dir.mkdir(parents=True, exist_ok=True)
CONSOLE.rule(f"[bold green]Processing {neb_trajectory_file.name}[/bold green]")
if images_per_path <= 0:
logging.critical("--images-per-path must be a positive integer.")
sys.exit(1)
try:
all_frames, all_con_frames = _read_trajectory_frames(Path(neb_trajectory_file))
if not all_frames:
logging.error("No frames found in input file.")
sys.exit(1)
except Exception as e:
logging.critical(f"Failed to read trajectory: {e}")
sys.exit(1)
total_frames = len(all_frames)
num_paths = total_frames // images_per_path
remainder = total_frames % images_per_path
# Validation Logic based on Mode
if mode == SplitMode.NEB.value and remainder != 0:
logging.warning(
f"Trajectory has {total_frames} frames,"
f" which is not a multiple of {images_per_path}. "
f"This often indicates an interrupted NEB calculation."
)
if total_frames < images_per_path:
logging.critical(
f"Total frames ({total_frames})"
f" is less than images per path ({images_per_path})."
)
sys.exit(1)
target_idx = num_paths - 1 if path_index == -1 else path_index
if not (0 <= target_idx < num_paths):
logging.critical(
f"Path index {target_idx} is out of bounds (0 to {num_paths - 1})."
)
sys.exit(1)
start, end = target_idx * images_per_path, (target_idx + 1) * images_per_path
frames = all_frames[start:end]
con_frames = all_con_frames[start:end] if all_con_frames is not None else None
logging.info(f"Extracted [cyan]Path {target_idx}[/cyan] with {len(frames)} images.")
geometry_mutated = False
if center:
logging.info("Centering structures...")
if box_diagonal:
logging.info("Overriding box...")
if center and len(frames) > 0:
geometry_mutated = True
ref_atoms = frames[0].copy()
ref_center = ref_atoms.get_center_of_mass()
box_center = [d / 2.0 for d in box_diagonal]
shift = box_center - ref_center
for atoms in frames:
atoms.set_cell(box_diagonal)
atoms.translate(shift)
align_strategy = AlignMode(align_type)
if align_strategy != AlignMode.NONE:
geometry_mutated = True
frames = align_path(
frames, align_strategy, IRAConfig(enabled=use_ira, kmax=ira_kmax)
)
created_paths = []
for i, atoms in enumerate(frames):
name = f"ipath_{i:03d}.con"
dest = output_dir / name
energy = _energy_from_atoms(atoms)
if con_frames is not None and not geometry_mutated:
_write_con_frame(dest, con_frames[i], energy=energy)
elif con_frames is not None:
if energy is None and con_frames[i].energy is not None:
energy = float(con_frames[i].energy)
_write_con_frame(dest, con_frames[i], atoms=atoms, energy=energy)
else:
_write_con_atoms(dest, atoms, energy=energy)
created_paths.append(str(dest.resolve()))
logging.info(f" - Saved [green]{name}[/green]")
with open(output_dir / path_list_filename, "w") as f:
f.write("\n".join(created_paths) + "\n")
logging.info(f"Path list saved to [magenta]{path_list_filename}[/magenta]")
CONSOLE.rule("[bold green]Complete[/bold green]")
if __name__ == "__main__":
con_splitter()
