from __future__ import annotations
from pathlib import Path
import shutil, json
from typing import Any, TYPE_CHECKING
import numpy as np
if TYPE_CHECKING:
from ensemble_analyzer.conformer.conformer import Conformer
from ensemble_analyzer.protocol.protocol import Protocol
[docs]
def mkdir(directory: str) -> bool:
"""
Create a directory, ensuring parent directories exist.
Args:
directory (str): Path of the folder to be created.
Returns:
bool: True if creation was successful (or directory already exists).
"""
# os.makedirs(directory, exist_ok=True)
Path(directory).mkdir(parents=True, exist_ok=True)
return True
[docs]
def move_files(conf: Conformer, protocol: Protocol, label: str) -> None:
"""
Move calculation output files to the conformer's specific folder.
Identifies files generated by the calculator (based on label) and moves
them to `conf_X/protocol_Y/`.
Args:
conf (Conformer): The conformer instance associated with the files.
protocol (Protocol): The protocol instance defining the step.
label (str): The calculator label used as a prefix/suffix for files.
Returns:
None
"""
cwd = Path.cwd()
files = [f for f in cwd.iterdir() if f.name.startswith(label)]
dest_folder = cwd / conf.folder / f"protocol_{protocol.number}"
mkdir(str(dest_folder))
for file in files:
dst = dest_folder / f"{conf.number}_p{protocol.number}_{file.name}"
shutil.move(str(file), str(dst))
[docs]
def tail(file_path: str, num_lines:int=100) -> str:
"""
Read the last N lines of a file.
Useful for extracting summary information from large log files.
Args:
file_path (str): Path to the file.
num_lines (int, optional): Number of lines to read from the end. Defaults to 100.
Returns:
str: The tail of the file content as a single string.
"""
with Path(file_path).open() as f:
fl = f.readlines()
return "".join(fl[-num_lines:])
[docs]
class SerialiseEncoder(json.JSONEncoder):
"""
Custom JSON Encoder for Ensemble Analyzer objects.
Handles serialization of:
- NumPy arrays (converted to lists).
- NumPy numeric types (converted to Python scalars).
- NaN / Inf floats (converted to None).
- Complex numbers (converted to [real, imag]).
- Custom objects with '__dict__' (via obj.__dict__).
"""
[docs]
def default(self, obj) -> Any:
"""
Override default serialization method.
Args:
obj (Any): The object to serialize.
Returns:
Any: JSON-serializable representation of the object.
"""
if isinstance(obj, np.ndarray):
return obj.tolist()
if isinstance(obj, np.floating):
return None if np.isnan(obj) or np.isinf(obj) else float(obj)
if isinstance(obj, np.integer):
return int(obj)
if isinstance(obj, np.bool_):
return bool(obj)
if hasattr(obj, "__dict__"):
return obj.__dict__
# Let the base class raise TypeError for anything else
return super().default(obj)