"""
ControlStream: top-level parsed representation of a NONMEM control stream.
Parses a .ctl/.mod file into a list of typed record objects and provides
accessor methods for downstream components.
"""
from __future__ import annotations
import os
import warnings
from dataclasses import dataclass, field
from typing import Any, TypeVar
from openpkpd.parser.lexer import RawRecord, split_into_raw_records
from openpkpd.parser.records.abbreviated import AbbreviatedRecord
from openpkpd.parser.records.base import BaseRecord
from openpkpd.parser.records.contr import ContrRecord
from openpkpd.parser.records.covariance import CovarianceRecord
from openpkpd.parser.records.data import DataRecord
from openpkpd.parser.records.des import DESRecord
from openpkpd.parser.records.design import DesignRecord
from openpkpd.parser.records.error import ErrorRecord
from openpkpd.parser.records.estimation import EstimationRecord
from openpkpd.parser.records.input import InputRecord
from openpkpd.parser.records.mixture import MixtureRecord
from openpkpd.parser.records.nonparametric import NonparametricRecord
from openpkpd.parser.records.omega import OmegaRecord
from openpkpd.parser.records.pk import PKRecord
from openpkpd.parser.records.pred import PredRecord
from openpkpd.parser.records.prior import (
OmegaPDRecord,
OmegaPRecord,
PriorRecord,
SigmaPDRecord,
SigmaPRecord,
ThetaPRecord,
ThetaPVRecord,
)
# Import all record types
from openpkpd.parser.records.problem import ProblemRecord
from openpkpd.parser.records.sigma import SigmaRecord
from openpkpd.parser.records.simulation import SimulationRecord
from openpkpd.parser.records.sizes import SizesRecord
from openpkpd.parser.records.subroutines import SubroutinesRecord
from openpkpd.parser.records.table import TableRecord
from openpkpd.parser.records.theta import ThetaRecord
from openpkpd.utils.errors import ParseError, ParseWarning
T = TypeVar("T", bound=BaseRecord)
# Registry: canonical record name → record class
_RECORD_REGISTRY: dict[str, type[BaseRecord]] = {
"PROBLEM": ProblemRecord,
"DATA": DataRecord,
"INPUT": InputRecord,
"SUBROUTINES": SubroutinesRecord,
"THETA": ThetaRecord,
"OMEGA": OmegaRecord,
"SIGMA": SigmaRecord,
"PK": PKRecord,
"DES": DESRecord,
"ERROR": ErrorRecord,
"PRED": PredRecord,
"ESTIMATION": EstimationRecord,
"COVARIANCE": CovarianceRecord,
"TABLE": TableRecord,
"SIMULATION": SimulationRecord,
"PRIOR": PriorRecord,
"THETAP": ThetaPRecord,
"THETAPV": ThetaPVRecord,
"OMEGAP": OmegaPRecord,
"OMEGAPD": OmegaPDRecord,
"SIGMAP": SigmaPRecord,
"SIGMAPD": SigmaPDRecord,
"MIXTURE": MixtureRecord,
"ABBREVIATED": AbbreviatedRecord,
"NONPARAMETRIC": NonparametricRecord,
"SIZES": SizesRecord,
"DESIGN": DesignRecord,
"CONTR": ContrRecord,
}
class UnknownRecord(BaseRecord):
"""Placeholder for unrecognised record types."""
def _parse(self, text: str) -> None:
self.text = text
[docs]
@dataclass
class ControlStream:
"""
Top-level parsed representation of a NONMEM control stream.
"""
source_text: str
records: list[BaseRecord] = field(default_factory=list)
source_path: str | None = None
# ── Accessors ─────────────────────────────────────────────────────────────
[docs]
def get(self, record_type: str) -> BaseRecord | None:
"""Return the first record of the given type, or None."""
upper = record_type.upper()
for rec in self.records:
if rec.record_name.upper() == upper:
return rec
return None
[docs]
def get_all(self, record_type: str) -> list[BaseRecord]:
"""Return all records of the given type."""
upper = record_type.upper()
return [rec for rec in self.records if rec.record_name.upper() == upper]
[docs]
def get_typed(self, record_type: str, cls: type[T]) -> T | None:
"""Type-safe accessor returning a specific record class."""
rec = self.get(record_type)
if rec is None:
return None
if not isinstance(rec, cls):
raise ParseError(
f"Expected {cls.__name__} for ${record_type}, got {type(rec).__name__}"
)
return rec
[docs]
def get_all_typed(self, record_type: str, cls: type[T]) -> list[T]:
"""Type-safe accessor returning a list of specific record class."""
result = []
for rec in self.get_all(record_type):
if not isinstance(rec, cls):
raise ParseError(
f"Expected {cls.__name__} for ${record_type}, got {type(rec).__name__}"
)
result.append(rec)
return result
# Convenience typed accessors
@property
def problem(self) -> ProblemRecord | None:
return self.get_typed("PROBLEM", ProblemRecord)
@property
def data(self) -> DataRecord | None:
return self.get_typed("DATA", DataRecord)
@property
def input(self) -> InputRecord | None:
return self.get_typed("INPUT", InputRecord)
@property
def subroutines(self) -> SubroutinesRecord | None:
return self.get_typed("SUBROUTINES", SubroutinesRecord)
@property
def pk(self) -> PKRecord | None:
return self.get_typed("PK", PKRecord)
@property
def des(self) -> DESRecord | None:
return self.get_typed("DES", DESRecord)
@property
def error(self) -> ErrorRecord | None:
return self.get_typed("ERROR", ErrorRecord)
@property
def pred(self) -> PredRecord | None:
return self.get_typed("PRED", PredRecord)
@property
def theta_records(self) -> list[ThetaRecord]:
return self.get_all_typed("THETA", ThetaRecord)
@property
def omega_records(self) -> list[OmegaRecord]:
return self.get_all_typed("OMEGA", OmegaRecord)
@property
def sigma_records(self) -> list[SigmaRecord]:
return self.get_all_typed("SIGMA", SigmaRecord)
@property
def simulation(self) -> SimulationRecord | None:
return self.get_typed("SIMULATION", SimulationRecord)
@property
def mixture(self) -> MixtureRecord | None:
return self.get_typed("MIXTURE", MixtureRecord)
@property
def prior_record(self) -> PriorRecord | None:
return self.get_typed("PRIOR", PriorRecord)
@property
def thetap_record(self) -> ThetaPRecord | None:
return self.get_typed("THETAP", ThetaPRecord)
@property
def thetapv_record(self) -> ThetaPVRecord | None:
return self.get_typed("THETAPV", ThetaPVRecord)
@property
def omegap_record(self) -> OmegaPRecord | None:
return self.get_typed("OMEGAP", OmegaPRecord)
@property
def omegapd_record(self) -> OmegaPDRecord | None:
return self.get_typed("OMEGAPD", OmegaPDRecord)
@property
def sigmap_record(self) -> SigmaPRecord | None:
return self.get_typed("SIGMAP", SigmaPRecord)
@property
def sigmapd_record(self) -> SigmaPDRecord | None:
return self.get_typed("SIGMAPD", SigmaPDRecord)
@property
def estimation_records(self) -> list[EstimationRecord]:
return self.get_all_typed("ESTIMATION", EstimationRecord)
@property
def covariance(self) -> CovarianceRecord | None:
return self.get_typed("COVARIANCE", CovarianceRecord)
@property
def table_records(self) -> list[TableRecord]:
return self.get_all_typed("TABLE", TableRecord)
# ── Constructors ──────────────────────────────────────────────────────────
[docs]
@classmethod
def from_file(cls, path: str) -> ControlStream:
"""Load and parse a NONMEM control stream from a file."""
if not os.path.exists(path):
raise ParseError(f"Control stream file not found: {path}")
with open(path, encoding="utf-8", errors="replace") as fh:
text = fh.read()
cs = cls.from_string(text)
cs.source_path = os.path.abspath(path)
return cs
[docs]
@classmethod
def from_string(cls, text: str) -> ControlStream:
"""Parse a NONMEM control stream from a string."""
raw_records = split_into_raw_records(text)
records: list[BaseRecord] = []
for raw in raw_records:
rec = _build_record(raw)
records.append(rec)
return cls(source_text=text, records=records)
[docs]
def to_dict(self) -> dict[str, Any]:
"""Serialize to a plain dict for debugging."""
return {
"source_path": self.source_path,
"records": [r.to_dict() for r in self.records],
}
[docs]
def to_string(self) -> str:
"""Serialize the entire control stream to NONMEM .ctl text.
Each record is rendered via its ``to_string()`` method. Structured
records (THETA, OMEGA, SIGMA, ESTIMATION, …) reconstruct from their
parsed attributes so programmatic modifications are preserved. All
other records reproduce their original ``raw_text``.
"""
parts = [rec.to_string() for rec in self.records]
return "\n".join(parts)
[docs]
def write(self, path: str) -> None:
"""Write the control stream to *path*.
Creates parent directories if they do not exist.
"""
os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
with open(path, "w", encoding="utf-8") as fh:
fh.write(self.to_string())
def __repr__(self) -> str:
rec_names = [r.record_name for r in self.records]
return f"ControlStream(records={rec_names})"
def _build_record(raw: RawRecord) -> BaseRecord:
"""Instantiate the appropriate record class from a RawRecord."""
cls = _RECORD_REGISTRY.get(raw.name, UnknownRecord)
try:
rec = cls(raw_text=raw.raw_text, header_line=raw.header_line)
# Patch record_name for UnknownRecord
if isinstance(rec, UnknownRecord):
rec.record_name = raw.name
warnings.warn(
f"Unrecognised record '${raw.name}' — skipped. "
"Check spelling or NONMEM version.",
ParseWarning,
stacklevel=2,
)
return rec
except Exception as exc:
raise ParseError(
f"Error parsing ${raw.name} record",
line=raw.header_line,
context=str(exc),
) from exc