from __future__ import annotations
import argparse
import json
import os
import sys
import tempfile
from collections import deque
from concurrent.futures import ProcessPoolExecutor
from multiprocessing import get_context
from pathlib import Path
from types import SimpleNamespace
from typing import Callable, Dict, List, Optional, Sequence, Tuple
import numpy as np
from tqdm import tqdm
try:
import msgspec
except ImportError: # pragma: no cover - exercised by minimal installations
msgspec = None
from pycurves_lib.cli.pycurves_cli_options import add_pycurves_topology_options
from pycurves_lib.io.curves_output import _to_jsonable
from pycurves_lib.md.trajectory_loader import TrajectoryLoader
from pycurves_lib.md.trajectory_statistics import (
BatchSummaryAccumulator,
)
from pycurves_lib.md.batch_curvesplus import BatchCurvesPlusMDAnalyzer
from pycurves_md import MDTrajectoryAnalyzer, make_frame_selector
_WORKER_ANALYZER: Optional[BatchCurvesPlusMDAnalyzer] = None
def _encode_json(value) -> bytes:
if msgspec is not None:
return msgspec.json.encode(value)
return json.dumps(
value,
ensure_ascii=False,
separators=(",", ":"),
allow_nan=False,
).encode("utf-8")
def _flush_batch(
analyzer: BatchCurvesPlusMDAnalyzer,
coordinates: Sequence[np.ndarray] | np.ndarray,
frame_indices: Sequence[int],
times: Sequence[Optional[float]],
mode: str,
frame_payloads: List[Dict],
summary_accumulator: Optional[BatchSummaryAccumulator] = None,
frame_sink: Optional[Callable[[List[Dict]], None]] = None,
) -> int:
if len(coordinates) == 0:
return 0
batch_coordinates = np.asarray(coordinates, dtype=float)
if mode == "summary" and summary_accumulator is not None:
return analyzer.accumulate_batch_summary(batch_coordinates, frame_indices, times, summary_accumulator)
batch_frames, _ = analyzer.analyze_batch(
batch_coordinates,
frame_indices,
times,
accumulator=summary_accumulator if mode == "both" else None,
collect_table_records=False,
)
if mode in {"per-frame", "both"}:
if frame_sink is None:
frame_payloads.extend(batch_frames)
else:
frame_sink(batch_frames)
return len(coordinates)
def _initialize_worker(analyzer: BatchCurvesPlusMDAnalyzer, numba_threads: int) -> None:
global _WORKER_ANALYZER
_WORKER_ANALYZER = analyzer
if analyzer.include_grooves:
try:
import numba
numba.set_num_threads(numba_threads)
except ImportError:
pass
def _process_batch_worker(
coordinates: np.ndarray,
frame_indices: Sequence[int],
times: Sequence[Optional[float]],
mode: str,
) -> Tuple[int, List[Dict], Optional[BatchSummaryAccumulator]]:
if _WORKER_ANALYZER is None: # pragma: no cover - process-pool guard
raise RuntimeError("Batch worker was not initialized.")
frame_payloads: List[Dict] = []
accumulator = BatchSummaryAccumulator() if mode in {"summary", "both"} else None
processed = _flush_batch(
_WORKER_ANALYZER,
coordinates,
frame_indices,
times,
mode,
frame_payloads,
accumulator,
)
return processed, frame_payloads, accumulator
def _write_csv_payload(payload: Dict, prefix: str) -> None:
prefix_path = Path(prefix)
prefix_path.parent.mkdir(parents=True, exist_ok=True)
if "summary" in payload:
for name, rows in payload["summary"].items():
MDTrajectoryAnalyzer._write_rows_csv(f"{prefix}_{name}_summary.csv", rows)
if "frames" in payload:
per_table: Dict[str, List[Dict]] = {}
for frame in payload["frames"]:
for name, rows in frame["dataframes"].items():
if not isinstance(rows, list):
continue
for row in rows:
out = dict(row)
out["frame"] = frame["frame"]
out["time"] = frame["time"]
per_table.setdefault(name, []).append(out)
for name, rows in per_table.items():
MDTrajectoryAnalyzer._write_rows_csv(f"{prefix}_{name}_frames.csv", rows)
def run_batch(args, frame_sink: Optional[Callable[[List[Dict]], None]] = None) -> Dict:
if args.axis_convention.lower().replace("-", "_") not in {"local", "curvesplus", "curves_plus", "curves+", "canal"}:
raise SystemExit("pycurves-md-batch currently supports only --axis-convention local.")
if args.frame_convention.lower().replace("-", "_") not in {"standard", "curvesplus", "curves_plus", "curves+", "x3dna", "3dna"}:
raise SystemExit("pycurves-md-batch currently supports only --frame-convention standard.")
if args.batch_size <= 0:
raise SystemExit("--batch-size must be positive.")
workers = getattr(args, "workers", 1)
if workers <= 0:
raise SystemExit("--workers must be positive.")
frame_selector, selection = make_frame_selector(args.frames, args.start, args.stop, args.step)
reference_topology = MDTrajectoryAnalyzer._reference_topology(args.topology, args.trajectory, args.output_dir)
analyzer = BatchCurvesPlusMDAnalyzer(
topology_file=reference_topology,
inpfile=args.inp,
output_dir=args.output_dir,
continuous_strands=args.continuous_strands,
duplex_only=getattr(args, "duplex_only", False),
fit_override=args.fit,
comb_override=args.comb,
ends_override=args.ends,
include_grooves=args.grooves,
include_curvesplus_axis_steps=args.curvesplus_axis_steps,
include_fit_quality=args.fit_quality,
axis_weighting=args.axis_weighting,
altloc=args.altloc,
)
frame_payloads: List[Dict] = []
summary_accumulator = BatchSummaryAccumulator() if args.mode in {"summary", "both"} else None
processed = 0
iterator = TrajectoryLoader.iter_batches(
args.topology, args.trajectory, frame_selector, args.batch_size
)
if workers == 1:
for batch in tqdm(iterator, desc="Processing frame batches", unit="batch"):
processed += _flush_batch(
analyzer,
batch.coordinates,
batch.indices,
batch.times,
args.mode,
frame_payloads,
summary_accumulator,
frame_sink,
)
else:
threads_per_worker = max(1, (os.cpu_count() or workers) // workers)
pending = deque()
def collect_result(result) -> None:
nonlocal processed
count, frames, partial_summary = result
processed += count
if frames:
if frame_sink is None:
frame_payloads.extend(frames)
else:
frame_sink(frames)
if partial_summary is not None and summary_accumulator is not None:
summary_accumulator.merge(partial_summary)
with ProcessPoolExecutor(
max_workers=workers,
mp_context=get_context("spawn"),
initializer=_initialize_worker,
initargs=(analyzer, threads_per_worker),
) as executor, tqdm(desc="Processing frame batches", unit="batch") as progress:
for batch in iterator:
pending.append(executor.submit(
_process_batch_worker,
batch.coordinates,
batch.indices,
batch.times,
args.mode,
))
if len(pending) >= workers:
collect_result(pending.popleft().result())
progress.update()
while pending:
collect_result(pending.popleft().result())
progress.update()
if processed == 0:
raise SystemExit("No trajectory frames matched the requested frame selection.")
payload = {
"program": "pyCurves",
"format": "pycurves-trajectory-batch-curvesplus-v1",
"inputs": {
"topology_file": args.topology,
"reference_topology_file": reference_topology,
"trajectory_file": args.trajectory,
"inpfile": analyzer.inpfile,
"generated_inpfiles": analyzer.generated_inpfiles,
},
"analysis_options": {
"batch_size": args.batch_size,
"workers": workers,
"continuous_strands": args.continuous_strands,
"duplex_only": getattr(args, "duplex_only", False),
"altloc": args.altloc or "first",
"fit": True if args.fit is None else args.fit,
"comb": True if args.comb is None else args.comb,
"ends": False if args.ends is None else args.ends,
"mini": False,
"axis_convention": "local",
"axis_weighting": bool(analyzer.ctx.cfg.axis_weighting),
"grooves": analyzer.include_grooves,
"curvesplus_axis_steps": args.curvesplus_axis_steps,
"fit_quality": args.fit_quality,
},
"selection": {
**selection,
"processed_frames": processed,
"engine": "batch_curvesplus",
},
"frame_convention": {
"name": "standard",
"compatible_with": ["Curves+", "3DNA", "x3dna"],
"axis_convention": "local",
},
}
if args.mode in {"per-frame", "both"} and frame_sink is None:
payload["frames"] = frame_payloads
if args.mode in {"summary", "both"}:
payload["summary"] = summary_accumulator.to_summary() if summary_accumulator is not None else {}
return payload
def _write_streamed_json(args) -> Dict:
"""Write frame batches as they finish, then atomically publish the JSON file."""
output_path = Path(args.output_file)
output_path.parent.mkdir(parents=True, exist_ok=True)
temporary_path: Optional[Path] = None
try:
with tempfile.NamedTemporaryFile(
mode="wb",
dir=output_path.parent,
prefix=f".{output_path.name}.",
suffix=".tmp",
delete=False,
) as stream:
temporary_path = Path(stream.name)
stream.write(b"{")
first_member = True
frame_sink = None
if args.mode in {"per-frame", "both"}:
stream.write(b'"frames":[')
first_frame = True
def write_frames(frames: List[Dict]) -> None:
nonlocal first_frame
if not frames:
return
if not first_frame:
stream.write(b",")
stream.write(_encode_json(frames)[1:-1])
first_frame = False
frame_sink = write_frames
first_member = False
payload = run_batch(args, frame_sink=frame_sink)
if frame_sink is not None:
stream.write(b"]")
for key, value in payload.items():
if not first_member:
stream.write(b",")
stream.write(_encode_json(key))
stream.write(b":")
stream.write(_encode_json(_to_jsonable(value)))
first_member = False
stream.write(b"}\n")
temporary_path.replace(output_path)
return payload
except BaseException:
if temporary_path is not None:
temporary_path.unlink(missing_ok=True)
raise
[docs]
def analyze_trajectory_batch(
topology_file: str,
trajectory_file: Optional[str] = None,
inpfile: Optional[str] = None,
output_dir: str = ".",
frames: Optional[str] = None,
start: Optional[int] = None,
stop: Optional[int] = None,
step: int = 1,
batch_size: int = 256,
mode: str = "per-frame",
continuous_strands: bool = False,
altloc: Optional[str] = None,
fit: Optional[bool] = None,
grooves: Optional[bool] = None,
comb: Optional[bool] = None,
ends: Optional[bool] = None,
frame_convention: str = "standard",
axis_convention: str = "local",
curvesplus_axis_steps: bool = False,
fit_quality: bool = False,
axis_weighting: Optional[bool] = None,
workers: int = 1,
duplex_only: bool = False,
) -> Dict:
"""Run the vectorized local-axis/standard-frame MD path from Python.
This is the notebook-friendly equivalent of ``pycurves-md-batch``. It is
intended for canonical, standard-frame analyses supported by the batch
engine.
"""
if mode not in {"per-frame", "summary", "both"}:
raise ValueError("mode must be one of: per-frame, summary, both")
args = SimpleNamespace(
topology=topology_file,
trajectory=trajectory_file,
inp=inpfile,
output_dir=output_dir,
output_file=None,
mode=mode,
format="json",
frames=frames,
start=start,
stop=stop,
step=step,
batch_size=batch_size,
workers=workers,
continuous_strands=continuous_strands,
duplex_only=duplex_only,
fit=fit,
grooves=grooves,
comb=comb,
ends=ends,
frame_convention=frame_convention,
axis_convention=axis_convention,
curvesplus_axis_steps=curvesplus_axis_steps,
fit_quality=fit_quality,
axis_weighting=axis_weighting,
altloc=altloc,
)
try:
return run_batch(args)
except SystemExit as exc:
raise ValueError(str(exc)) from exc
def main() -> None:
from pycurves_lib import __version__
parser = argparse.ArgumentParser(
description=(
"Vectorized pyCurves MD runner for standard-frame "
"local-axis (Curves+) analyses."
)
)
parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
parser.add_argument("topology", help="Topology/reference structure file, usually PDB or CIF.")
parser.add_argument("trajectory", nargs="?", help="Trajectory file. Omit for multi-model PDB input.")
parser.add_argument("--inp", help="Existing Curves .inp file. If omitted, inferred from the topology.")
parser.add_argument("--output-dir", default=".", help="Directory for auto-generated .inp files.")
parser.add_argument("--output-file", default="pycurves_trajectory_batch.json", help="JSON file or CSV prefix.")
parser.add_argument("--mode", choices=["per-frame", "summary", "both"], default="summary")
parser.add_argument("--format", choices=["json", "csv"], default="json")
parser.add_argument("--frames", help="Frame list/ranges, e.g. '0,10,20:100:5'. Overrides start/stop/step.")
parser.add_argument("--start", type=int, help="First frame index to include.")
parser.add_argument("--stop", type=int, help="Stop before this frame index.")
parser.add_argument("--step", type=int, default=1, help="Frame stride.")
parser.add_argument("--batch-size", type=int, default=256, help="Number of selected frames processed per vectorized batch.")
parser.add_argument(
"--workers",
type=int,
default=1,
help="Batch worker processes (default: 1). Use 2-4 for long trajectories and benchmark on your system.",
)
add_pycurves_topology_options(parser)
parser.add_argument("--fit", action=argparse.BooleanOptionalAction, default=None, help="Override least-squares base fitting; batch mode currently requires true.")
parser.add_argument("--comb", action=argparse.BooleanOptionalAction, default=None, help="Override combined strand analysis; batch mode currently requires true.")
parser.add_argument("--ends", action=argparse.BooleanOptionalAction, default=None, help="Override terminal virtual end levels; batch mode currently requires false.")
parser.add_argument("--grooves", action=argparse.BooleanOptionalAction, default=None, help="Override groove analysis; defaults to the .inp grv setting.")
parser.add_argument("--frame-convention", default="standard", help="Currently only standard is supported.")
parser.add_argument(
"--axis-convention",
default="local",
help="Currently only local is supported.",
)
parser.add_argument(
"--axis-weighting",
action=argparse.BooleanOptionalAction,
default=None,
help="Weight smooth-axis construction by fitted-pair geometry (default: disabled).",
)
parser.add_argument("--curvesplus-axis-steps", action="store_true", help="Include the Curves+ smooth-axis inter-base-pair step table.")
parser.add_argument("--fit-quality", action="store_true", help="Include vectorized base-fitting RMSD diagnostics.")
args = parser.parse_args()
try:
if args.format == "json":
payload = _write_streamed_json(args)
else:
payload = run_batch(args)
_write_csv_payload(payload, args.output_file.removesuffix(".csv"))
except (ValueError, ImportError, NotImplementedError) as exc:
raise SystemExit(str(exc)) from exc
generated = payload["inputs"].get("generated_inpfiles") or []
if generated:
print("Generated input file(s):", file=sys.stderr)
for path in generated:
print(f" {path}", file=sys.stderr)
print(f"Processed {payload['selection']['processed_frames']} frame(s).", file=sys.stderr)
if __name__ == "__main__":
main()