import contextlib
import copy
import io
from pathlib import Path
from types import SimpleNamespace
from typing import List, Optional
import numpy as np
from pycurves_lib.io.base_reference import BaseReferenceLibrary
from pycurves_lib.io.curves_mol_loader import MolecularLoader
from pycurves_lib.core.curves_dataclasses import (
BaseGeometryConstants,
BaseLocator,
CurvesContext,
MolecularStructure,
)
from pycurves_lib.io.curves_config_loader import ConfigLoader
from pycurves_lib.core.curves_analyzer import BackboneAnalyzer, HelicalOptimizer
from pycurves_lib.core.curves_calculator import HelicalCalculator
from pycurves_lib.core.parameter_conventions import build_axis_reference_frames
from pycurves_lib.core.axis_support import axis_support_weights
from pycurves_lib.io.curves_output import CurvesOutputFormatter
from pycurves_lib.topology.topology_inferrer import RobustTopologyInferrer
from pycurves_lib.topology.base_annotations import annotate_context
class CurvesWrapper:
"""High-level pyCurves runner.
It can run from an existing Curves `.inp` file or infer one from a PDB.
"""
def __init__(
self,
pdbfile: Optional[str] = None,
inpfile: Optional[str] = None,
output_dir: str = ".",
continuous_strands: bool = False,
altloc: Optional[str] = None,
frame_convention: str = "standard",
axis_convention: str = "global",
axis_weighting: Optional[bool] = None,
fit_override: Optional[bool] = None,
grv_override: Optional[bool] = None,
mini_override: Optional[bool] = None,
comb_override: Optional[bool] = None,
ends_override: Optional[bool] = None,
auto_generate_inp: bool = True,
duplex_only: bool = False,
):
if pdbfile is None and inpfile is None:
raise ValueError("Provide at least a PDB file or an input file.")
self.pdbfile = pdbfile
self.inpfile = inpfile
self.output_dir = output_dir
self.continuous_strands = continuous_strands
self.duplex_only = duplex_only
self.altloc = MolecularLoader.normalize_altloc(altloc)
self.frame_convention, self.axis_convention = self.normalize_conventions(frame_convention, axis_convention)
self.axis_weighting = axis_weighting
self.fit_override = fit_override
self.grv_override = grv_override
self.mini_override = mini_override
self.comb_override = comb_override
self.ends_override = ends_override
self.generated_inpfiles: List[str] = []
self.cfg = None
self.ctx = None
self.bak = None
self.opt = None
self.calc = None
self.analysis_log = ""
self._parsed_config_cache = None
self._parsed_config_cache_key = None
self._reference_library_cache = {}
if self.inpfile is None and auto_generate_inp:
self.generated_inpfiles = self.generate_inp(
pdbfile=self.pdbfile,
output_dir=output_dir,
continuous_strands=continuous_strands,
duplex_only=duplex_only,
)
if not self.generated_inpfiles:
raise ValueError(f"Could not infer a Curves input file from {self.pdbfile!r}.")
self.inpfile = self.generated_inpfiles[0]
@classmethod
def from_file(
cls,
path: str,
output_dir: str = ".",
continuous_strands: bool = False,
altloc: Optional[str] = None,
frame_convention: str = "standard",
axis_convention: str = "global",
duplex_only: bool = False,
):
suffix = Path(path).suffix.lower()
if suffix == ".inp":
pdbfile = cls._pdbfile_from_inp(path)
return cls(
pdbfile=pdbfile,
inpfile=path,
output_dir=output_dir,
continuous_strands=continuous_strands,
duplex_only=duplex_only,
altloc=altloc,
frame_convention=frame_convention,
axis_convention=axis_convention,
)
return cls(
pdbfile=path,
output_dir=output_dir,
continuous_strands=continuous_strands,
duplex_only=duplex_only,
altloc=altloc,
frame_convention=frame_convention,
axis_convention=axis_convention,
)
[docs]
def analyze(
self,
inpfile: Optional[str] = None,
pdbfile: Optional[str] = None,
mini: Optional[bool] = None,
verbose: bool = False,
continuous_strands: Optional[bool] = None,
altloc: Optional[str] = None,
frame_convention: Optional[str] = None,
axis_convention: Optional[str] = None,
duplex_only: Optional[bool] = None,
):
if inpfile is not None:
self.inpfile = inpfile
if pdbfile is not None:
self.pdbfile = pdbfile
if continuous_strands is not None:
self.continuous_strands = continuous_strands
if duplex_only is not None:
self.duplex_only = duplex_only
if altloc is not None:
self.altloc = MolecularLoader.normalize_altloc(altloc)
next_frame_convention = self.frame_convention if frame_convention is None else frame_convention
next_axis_convention = self.axis_convention if axis_convention is None else axis_convention
self.frame_convention, self.axis_convention = self.normalize_conventions(
next_frame_convention,
next_axis_convention,
)
if self.inpfile is None:
self.generated_inpfiles = self.generate_inp(
pdbfile=self.pdbfile,
output_dir=self.output_dir,
continuous_strands=self.continuous_strands,
duplex_only=self.duplex_only,
)
self.inpfile = self.generated_inpfiles[0]
mini_override = self.mini_override if self.mini_override is not None else mini
self.cfg = self._load_config(mini_override=mini_override)
if self.pdbfile is None:
self.pdbfile = self._pdbfile_from_inp(self.inpfile)
if self.pdbfile is None:
raise ValueError("No PDB file is available. Pass pdbfile=... or include file=... in the .inp.")
self.ctx = CurvesContext(self.cfg)
MolecularLoader.load(self.pdbfile, self.ctx, altloc=self.altloc)
return self._analyze_loaded_context(verbose=verbose)
[docs]
def analyze_molecule(
self,
molecule: MolecularStructure,
inpfile: Optional[str] = None,
mini: Optional[bool] = None,
verbose: bool = False,
prev_opt_helical: Optional[np.ndarray] = None,
axis_sign_reference: Optional[np.ndarray] = None,
):
"""Run pyCurves on an already populated MolecularStructure.
This is intended for trajectory workflows where atom metadata and
topology stay fixed while coordinates change frame by frame.
"""
if inpfile is not None:
self.inpfile = inpfile
if self.inpfile is None:
raise ValueError("analyze_molecule requires an existing Curves .inp file.")
mini_override = self.mini_override if self.mini_override is not None else mini
self.cfg = self._load_config(mini_override=mini_override)
self.ctx = CurvesContext(self.cfg)
self.ctx.molecule = molecule
return self._analyze_loaded_context(
verbose=verbose,
prev_opt_helical=prev_opt_helical,
axis_sign_reference=axis_sign_reference,
)
def _analyze_loaded_context(
self,
verbose: bool = False,
prev_opt_helical: Optional[np.ndarray] = None,
axis_sign_reference: Optional[np.ndarray] = None,
):
log_parts = []
local_axis = str(getattr(self.ctx.cfg, "axis_convention", self.axis_convention)).lower() == "local"
if local_axis:
# The local axis is derived from standard base-pair frames and
# Curves+ smooth.f-style construction, not the global minimizer.
self.ctx.cfg.mini = False
mini = bool(self.ctx.cfg.mini)
self._validate_supported_legacy_options(mini=mini)
reference_library = self._reference_library(getattr(self.ctx.cfg, "frame_convention", self.frame_convention))
self.ctx.base_reference_library = reference_library
locator = BaseLocator(BaseGeometryConstants(), reference_library=reference_library)
log_parts.append(self._capture_call(lambda: locator.locate_all(self.ctx), echo=verbose))
annotate_context(self.ctx)
self.ctx.axis_support_weights = axis_support_weights(
np.swapaxes(self.ctx.params.frames, 0, 1),
self.ctx.li,
enabled=bool(
self.ctx.cfg.axis_weighting
and self.ctx.cfg.comb
and self.ctx.nst > 1
),
)
build_axis_reference_frames(self.ctx)
self.bak = BackboneAnalyzer()
log_parts.append(self._capture_call(lambda: self.bak.analyze(self.ctx), echo=verbose))
if axis_sign_reference is not None:
self.ctx.axis_direction_sign_reference = np.asarray(axis_sign_reference, dtype=int).copy()
if local_axis or self.ctx.cfg.zaxe:
self.opt = HelicalOptimizer(self.ctx)
else:
from pycurves_lib.core.curves_optimizer_jax import HelicalOptimizerJAX
self.opt = HelicalOptimizerJAX(self.ctx)
if prev_opt_helical is not None:
# HelicalOptimizer.prepare() seeds XYTP values during construction.
# Apply trajectory warm starts after that initialization boundary.
self.ctx.params.helical = prev_opt_helical.copy()
log_parts.append(self._capture_call(lambda: self.opt.print_fortran_setup_report(self.ctx), echo=verbose))
log_parts.append(self._capture_call(lambda: self.opt.run(mini=mini), echo=verbose))
if mini:
log_parts.append(self._capture_call(lambda: self.opt.print_final_report(), echo=verbose))
self.analysis_log = "".join(log_parts)
self.calc = HelicalCalculator(self.ctx, self.opt)
self.calc.calculate_all()
return self
def _validate_supported_legacy_options(self, mini: bool = True) -> None:
"""Reject legacy modes that would currently produce misleading results."""
cfg = self.ctx.cfg
if mini and cfg.line:
raise NotImplementedError(
"line=.t. is not implemented in the current JAX optimizer. "
"Curves 5.3 treats it as a distinct straight-axis minimization mode."
)
if mini and cfg.dinu:
raise NotImplementedError(
"dinu=.t. is not implemented in the current JAX optimizer. "
"Curves 5.3 changes the objective from adjacent-base to dinucleotide terms."
)
def _load_config(self, mini_override: Optional[bool] = None) -> dict:
"""Return a fresh runtime config from a cached parse of the .inp file."""
if self.inpfile is None:
raise ValueError("No Curves input file is available.")
overrides = self._config_overrides()
if mini_override is not None:
overrides["mini"] = bool(mini_override)
cache_key = (
str(Path(self.inpfile).resolve()),
tuple(sorted(overrides.items())),
)
if cache_key != self._parsed_config_cache_key:
self._parsed_config_cache = ConfigLoader.parse_inp(
self.inpfile,
config_overrides=overrides,
)
self._parsed_config_cache_key = cache_key
return copy.deepcopy(self._parsed_config_cache)
def _reference_library(self, convention: str):
convention = self._normalize_frame_convention(convention)
if convention not in self._reference_library_cache:
self._reference_library_cache[convention] = BaseReferenceLibrary.load(convention)
return self._reference_library_cache[convention]
def _config_overrides(self) -> dict:
return {
"fit": self.fit_override,
"grv": self.grv_override,
"mini": self.mini_override,
"comb": self.comb_override,
"ends": self.ends_override,
"frame_convention": self.frame_convention,
"axis_convention": self.axis_convention,
"axis_weighting": self.axis_weighting,
}
[docs]
def run(self, output: bool = True, mini: Optional[bool] = None, verbose: bool = False, output_format: str = "curves", visualization: bool = False):
self.analyze(mini=mini, verbose=verbose)
if output:
self.output(fmt=output_format, visualization=visualization)
return self
[docs]
def output(self, fmt: str = "curves", file=None, visualization: bool = False) -> str:
if self.ctx is None or self.opt is None or self.calc is None:
raise RuntimeError("Call analyze() before output().")
text = CurvesOutputFormatter(self, visualization=visualization).render(fmt)
if file is None:
print(text, end="")
else:
Path(file).write_text(text, encoding="utf-8")
return text
[docs]
def write_output(self, path: str, fmt: str = "curves", visualization: bool = False) -> str:
return self.output(fmt=fmt, file=path, visualization=visualization)
[docs]
def getFeatures(self):
if self.calc is None:
raise RuntimeError("Call analyze() before getFeatures().")
return {
"helical": self.ctx.params.helical,
"inter_base": self.ctx.params.inter_base,
"global_axis_steps": self.calc.global_axis_steps,
"vkin": self.calc.vkin, # Backward-compatible Fortran name.
"local_inter_base": self.calc.local_inter_base,
"local_inter_base_pair": self.calc.local_inter_base_pair,
"bend": self.calc.bend,
"annotations": CurvesOutputFormatter(self)._annotations(),
}
@staticmethod
def _capture_call(func, echo: bool = False) -> str:
buf = io.StringIO()
with contextlib.redirect_stdout(buf):
func()
text = buf.getvalue()
if echo and text:
print(text, end="")
return text
def generate_inp(
self,
pdbfile: Optional[str] = None,
output_dir: Optional[str] = None,
prefix: Optional[str] = None,
continuous_strands: bool = False,
altloc: Optional[str] = None,
duplex_only: Optional[bool] = None,
) -> List[str]:
pdbfile = pdbfile or self.pdbfile
if pdbfile is None:
raise ValueError("generate_inp requires a PDB file.")
output_dir = output_dir or self.output_dir
mol_holder = SimpleNamespace(molecule=MolecularStructure())
selected_altloc = self.altloc if altloc is None else MolecularLoader.normalize_altloc(altloc)
selected_duplex_only = self.duplex_only if duplex_only is None else duplex_only
MolecularLoader.load(pdbfile, mol_holder, altloc=selected_altloc)
inferrer = RobustTopologyInferrer(mol_holder.molecule, pdbfile=Path(pdbfile).name)
stem = prefix or f"{Path(pdbfile).stem}_auto"
return inferrer.write_inp_files(
output_dir=output_dir,
prefix=stem,
continuous_strands=continuous_strands,
duplex_only=selected_duplex_only,
fit_override=getattr(self, "fit_override", None),
grv_override=getattr(self, "grv_override", None),
comb_override=getattr(self, "comb_override", None),
ends_override=getattr(self, "ends_override", None),
)
@staticmethod
def _normalize_frame_convention(value: str) -> str:
normalized = str(value or "standard").strip().lower().replace("-", "_")
if normalized in {"legacy"}:
return "legacy"
if normalized in {"standard", "curves_plus", "curves+", "curvesplus", "x3dna", "3dna"}:
return "standard"
raise ValueError(f"Unknown frame convention {value!r}; use 'legacy' or 'standard'.")
@staticmethod
def _normalize_axis_convention(value: str) -> str:
normalized = str(value or "global").strip().lower().replace("-", "_")
if normalized in {"global", "legacy", "pycurves"}:
return "global"
if normalized in {"local", "curves_plus", "curves+", "curvesplus", "canal"}:
return "local"
raise ValueError(f"Unknown axis convention {value!r}; use 'global' or 'local'.")
[docs]
@classmethod
def normalize_conventions(cls, frame_convention: str, axis_convention: str):
frame = cls._normalize_frame_convention(frame_convention)
axis = cls._normalize_axis_convention(axis_convention)
if axis == "local":
# Local/Curves+ axis construction is defined on standard base-pair
# reference frames. Treat it as a complete mode request.
frame = "standard"
return frame, axis
@staticmethod
def _pdbfile_from_inp(inpfile: str) -> Optional[str]:
inp_path = Path(inpfile)
text = inp_path.read_text(encoding="utf-8", errors="ignore")
lowered = text.lower()
key = "file="
pos = lowered.find(key)
if pos < 0:
return None
rest = text[pos + len(key):]
value = rest.split(",", 1)[0].split()[0].strip().strip("'\"")
if not value:
return None
pdb_path = Path(value)
if pdb_path.is_absolute():
return str(pdb_path)
sibling = inp_path.parent / pdb_path
if sibling.exists():
return str(sibling)
return value