from __future__ import annotations
import json
from pathlib import Path
from dataclasses import asdict, is_dataclass
import shutil
from typing import TYPE_CHECKING, Any
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)
def _serialise(obj: Any) -> Any:
"""Recursively convert non-JSON-serializable objects to plain Python types."""
if isinstance(obj, np.ndarray):
return _serialise(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 isinstance(obj, dict):
return {str(k) if isinstance(k, (int, np.integer)) else k: _serialise(v) for k, v in obj.items()}
if isinstance(obj, (list, tuple)):
return [_serialise(item) for item in obj]
if is_dataclass(obj) and not isinstance(obj, type):
return _serialise(asdict(obj))
return obj
def _is_matrix_like(val: list) -> bool:
"""True if *val* is a list of lists (2D array) — formatted compactly."""
return bool(val) and isinstance(val[0], list)
def _format_value(val: Any, indent: int = 0, step: int = 4) -> str:
"""Recursively format JSON, keeping matrix rows on single lines."""
prefix = " " * indent
inner = " " * (indent + step)
if isinstance(val, dict):
if not val:
return "{}"
items = []
for k, v in val.items():
items.append(f"{inner}{json.dumps(k)}: {_format_value(v, indent + step, step)}")
return "{\n" + ",\n".join(items) + "\n" + prefix + "}"
if isinstance(val, list):
if not val:
return "[]"
if _is_matrix_like(val):
rows = ",\n".join(inner + json.dumps(row) for row in val)
return "[\n" + rows + "\n" + prefix + "]"
items = []
for item in val:
items.append(inner + _format_value(item, indent + step, step))
return "[\n" + ",\n".join(items) + "\n" + prefix + "]"
return json.dumps(val)
[docs]
def write_json(data: Any, fp, indent: int = 4) -> None:
"""Write JSON with compact matrix formatting.
Matrices (lists of lists) are kept on single inner lines::
[
[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0]
]
instead of expanding every element onto a separate line.
Args:
data: Data to serialize (numpy objects are converted automatically).
fp: File-like object to write to.
indent: Spaces per level (default 4).
"""
fp.write(_format_value(_serialise(data), 0, indent))