Skills › Scientific Agent Skills › scikit-survival › _common.py
_common.py
463 lines · skills/scikit-survival/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared, standard-library-first helpers for the bundled survival CLIs."""
4from __future__ import annotations
6import csv
7import io
8import json
9import math
10import os
11import stat
12import tempfile
13from pathlib import Path
14from typing import Any, Iterable
17MAX_INPUT_BYTES = 32 * 1024 * 1024
18MAX_REPORT_BYTES = 4 * 1024 * 1024
19MAX_ROWS = 20_000
20MAX_FEATURES = 256
21MAX_TIME_POINTS = 512
22DEFAULT_SEED = 20_260_723
23PINNED_INSTALL = (
24 'uv pip install "scikit-survival==0.28.0" "scikit-learn==1.9.1" '
25 '"numpy==2.5.1" "pandas==3.0.6" "scipy==1.18.1"'
26)
29class CliError(ValueError):
30 """An expected command-line validation error."""
33def bounded_int(minimum: int, maximum: int):
34 """Return an argparse converter for a bounded integer."""
36 def convert(value: str) -> int:
37 try:
38 parsed = int(value)
39 except ValueError as exc:
40 raise CliError(f"expected an integer, got {value!r}") from exc
41 if not minimum <= parsed <= maximum:
42 raise CliError(
43 f"expected an integer from {minimum} through {maximum}, got {parsed}"
44 )
45 return parsed
47 return convert
50def finite_float(value: str) -> float:
51 """Parse a finite floating-point value."""
53 try:
54 parsed = float(value)
55 except ValueError as exc:
56 raise CliError(f"expected a number, got {value!r}") from exc
57 if not math.isfinite(parsed):
58 raise CliError("value must be finite")
59 return parsed
62def positive_float(value: str) -> float:
63 """Parse a finite positive floating-point value."""
65 parsed = finite_float(value)
66 if parsed <= 0:
67 raise CliError("value must be greater than zero")
68 return parsed
71def probability(value: str) -> float:
72 """Parse a probability strictly between zero and one."""
74 parsed = finite_float(value)
75 if not 0 < parsed < 1:
76 raise CliError("value must be strictly between zero and one")
77 return parsed
80def parse_names(value: str | None) -> list[str]:
81 """Parse a comma-separated list of unique, non-empty column names."""
83 if value is None:
84 return []
85 names = [item.strip() for item in value.split(",")]
86 if not names or any(not item for item in names):
87 raise CliError("column lists must contain non-empty comma-separated names")
88 if len(names) != len(set(names)):
89 raise CliError("column lists must not contain duplicates")
90 if len(names) > MAX_FEATURES:
91 raise CliError(f"at most {MAX_FEATURES} columns are allowed")
92 return names
95def parse_floats(value: str | None) -> list[float]:
96 """Parse a comma-separated list of finite floating-point values."""
98 if value is None:
99 return []
100 items = [item.strip() for item in value.split(",")]
101 if not items or any(not item for item in items):
102 raise CliError("expected non-empty comma-separated numbers")
103 values = [finite_float(item) for item in items]
104 if len(values) > MAX_TIME_POINTS:
105 raise CliError(f"at most {MAX_TIME_POINTS} values are allowed")
106 return values
109def checked_input_file(
110 value: str | os.PathLike[str],
111 *,
112 suffixes: Iterable[str],
113 max_bytes: int = MAX_INPUT_BYTES,
114) -> Path:
115 """Return a bounded regular local file, rejecting URLs and symlinks."""
117 raw = os.fspath(value)
118 if "://" in raw:
119 raise CliError("network URLs are not accepted; provide a local file")
120 path = Path(raw).expanduser()
121 if path.is_symlink():
122 raise CliError(f"input must not be a symlink: {path}")
123 try:
124 info = path.stat()
125 except OSError as exc:
126 raise CliError(f"cannot access input file {path}: {exc}") from exc
127 if not stat.S_ISREG(info.st_mode):
128 raise CliError(f"input is not a regular file: {path}")
129 if info.st_size > max_bytes:
130 raise CliError(f"input is {info.st_size} bytes; limit is {max_bytes} bytes")
131 allowed = {suffix.lower() for suffix in suffixes}
132 if path.suffix.lower() not in allowed:
133 raise CliError(f"input suffix must be one of: {', '.join(sorted(allowed))}")
134 return path.resolve()
137def checked_output_file(
138 value: str | os.PathLike[str],
139 *,
140 suffixes: Iterable[str],
141 force: bool = False,
142) -> Path:
143 """Validate an explicit local output without following symlinks."""
145 raw = os.fspath(value)
146 if "://" in raw:
147 raise CliError("network URLs are not accepted as output paths")
148 path = Path(raw).expanduser()
149 if path.name in {"", ".", ".."}:
150 raise CliError("output must name a file")
151 if path.is_symlink():
152 raise CliError(f"output must not be a symlink: {path}")
153 allowed = {suffix.lower() for suffix in suffixes}
154 if path.suffix.lower() not in allowed:
155 raise CliError(f"output suffix must be one of: {', '.join(sorted(allowed))}")
156 parent = path.parent
157 if not parent.exists() or not parent.is_dir() or parent.is_symlink():
158 raise CliError(f"output parent must be an existing regular directory: {parent}")
159 if path.exists():
160 if not path.is_file():
161 raise CliError(f"output exists and is not a regular file: {path}")
162 if not force:
163 raise CliError(f"refusing to overwrite existing output: {path}")
164 return parent.resolve() / path.name
167def atomic_write_bytes(path: Path, payload: bytes, *, force: bool = False) -> None:
168 """Write bytes through a private same-directory temporary file."""
170 destination = checked_output_file(path, suffixes={path.suffix.lower()}, force=force)
171 descriptor, temporary_name = tempfile.mkstemp(
172 prefix=f".{destination.name}.", suffix=".tmp", dir=destination.parent
173 )
174 temporary = Path(temporary_name)
175 try:
176 with os.fdopen(descriptor, "wb") as handle:
177 handle.write(payload)
178 handle.flush()
179 os.fsync(handle.fileno())
180 os.chmod(temporary, 0o600)
181 if destination.exists() and not force:
182 raise CliError(f"refusing to overwrite existing output: {destination}")
183 os.replace(temporary, destination)
184 finally:
185 temporary.unlink(missing_ok=True)
188def _json_bytes(document: Any) -> bytes:
189 """Serialize deterministic strict JSON."""
191 try:
192 payload = (
193 json.dumps(
194 document,
195 indent=2,
196 sort_keys=True,
197 ensure_ascii=False,
198 allow_nan=False,
199 )
200 + "\n"
201 ).encode("utf-8")
202 except (TypeError, ValueError) as exc:
203 raise CliError(f"report is not strict JSON: {exc}") from exc
204 if len(payload) > MAX_REPORT_BYTES:
205 raise CliError(
206 f"report is {len(payload)} bytes; limit is {MAX_REPORT_BYTES} bytes"
207 )
208 return payload
211def emit_json(
212 document: Any,
213 *,
214 output: str | os.PathLike[str] | None = None,
215 force: bool = False,
216) -> None:
217 """Print deterministic JSON or write it atomically."""
219 payload = _json_bytes(document)
220 if output is None:
221 print(payload.decode("utf-8"), end="")
222 return
223 destination = checked_output_file(output, suffixes={".json"}, force=force)
224 atomic_write_bytes(destination, payload, force=force)
227def emit_text(
228 text: str,
229 *,
230 output: str | os.PathLike[str] | None = None,
231 force: bool = False,
232) -> None:
233 """Print text or write it atomically to Markdown."""
235 payload = text.encode("utf-8")
236 if len(payload) > MAX_REPORT_BYTES:
237 raise CliError(
238 f"report is {len(payload)} bytes; limit is {MAX_REPORT_BYTES} bytes"
239 )
240 if output is None:
241 print(text, end="" if text.endswith("\n") else "\n")
242 return
243 destination = checked_output_file(output, suffixes={".md"}, force=force)
244 atomic_write_bytes(destination, payload, force=force)
247def load_json(value: str | os.PathLike[str]) -> Any:
248 """Load bounded strict JSON from a local file."""
250 path = checked_input_file(value, suffixes={".json"})
252 def reject_constant(constant: str) -> None:
253 raise CliError(f"non-standard JSON constant is not allowed: {constant}")
255 try:
256 with path.open("r", encoding="utf-8") as handle:
257 return json.load(handle, parse_constant=reject_constant)
258 except (OSError, json.JSONDecodeError) as exc:
259 raise CliError(f"cannot read valid JSON from {path.name}: {exc}") from exc
262def read_csv(value: str | os.PathLike[str]):
263 """Load a bounded local CSV with pandas, rejecting excessive dimensions."""
265 path = checked_input_file(value, suffixes={".csv"})
266 try:
267 import pandas as pd
268 except ImportError as exc:
269 raise CliError(
270 f"pandas is unavailable; install with `{PINNED_INSTALL}`"
271 ) from exc
272 try:
273 with path.open(newline="", encoding="utf-8-sig") as handle:
274 header = next(csv.reader(handle), [])
275 if not header or any(not name.strip() for name in header):
276 raise CliError("CSV column names must be non-empty")
277 if len(header) != len(set(header)):
278 raise CliError("CSV column names must be unique before parsing")
279 frame = pd.read_csv(path, nrows=MAX_ROWS + 1)
280 except Exception as exc:
281 raise CliError(f"cannot parse CSV {path.name}: {exc}") from exc
282 validate_frame_bounds(frame)
283 return frame, path.name
286def validate_frame_bounds(frame: Any) -> None:
287 """Reject empty or oversized tabular inputs."""
289 rows, columns = frame.shape
290 if rows == 0:
291 raise CliError("input contains no rows")
292 if rows > MAX_ROWS:
293 raise CliError(f"input has {rows} rows; limit is {MAX_ROWS}")
294 if columns > MAX_FEATURES + 2:
295 raise CliError(f"input has {columns} columns; limit is {MAX_FEATURES + 2}")
296 if not frame.columns.is_unique:
297 raise CliError("column names must be unique")
300def normalize_binary_event(values: Any):
301 """Return a strict boolean event vector from bool or 0/1 values."""
303 try:
304 import numpy as np
305 import pandas as pd
306 except ImportError as exc:
307 raise CliError(
308 f"scientific stack unavailable; install with `{PINNED_INSTALL}`"
309 ) from exc
310 series = pd.Series(values)
311 if series.isna().any():
312 raise CliError("event indicator contains missing values")
313 if pd.api.types.is_bool_dtype(series.dtype):
314 result = series.astype(bool).to_numpy()
315 elif pd.api.types.is_numeric_dtype(series.dtype):
316 numeric = pd.to_numeric(series, errors="raise").to_numpy(dtype=float)
317 if not np.isfinite(numeric).all() or not np.isin(numeric, [0.0, 1.0]).all():
318 raise CliError("event indicator must contain only boolean or 0/1 values")
319 result = numeric.astype(bool)
320 else:
321 normalized = series.astype(str).str.strip().str.casefold()
322 mapping = {"true": True, "false": False, "1": True, "0": False}
323 if not normalized.isin(mapping).all():
324 raise CliError(
325 "event indicator strings must be one of true, false, 1, or 0"
326 )
327 result = normalized.map(mapping).to_numpy(dtype=bool)
328 if not result.any():
329 raise CliError("at least one observed event is required")
330 return result
333def normalize_positive_times(values: Any, *, label: str = "time"):
334 """Return a finite, strictly positive float time vector."""
336 try:
337 import numpy as np
338 import pandas as pd
339 except ImportError as exc:
340 raise CliError(
341 f"scientific stack unavailable; install with `{PINNED_INSTALL}`"
342 ) from exc
343 try:
344 times = np.asarray(pd.to_numeric(values, errors="raise"), dtype=float)
345 except (TypeError, ValueError) as exc:
346 raise CliError(f"{label} must contain only numeric values") from exc
347 if times.ndim != 1:
348 raise CliError(f"{label} must be one-dimensional")
349 if not np.isfinite(times).all():
350 raise CliError(f"{label} contains missing or non-finite values")
351 if (times <= 0).any():
352 raise CliError(f"{label} must be strictly positive")
353 return times
356def structured_survival(
357 event: Any,
358 time: Any,
359 *,
360 event_name: str = "event",
361 time_name: str = "time",
362):
363 """Create scikit-survival's two-field structured outcome lazily."""
365 if event_name == time_name:
366 raise CliError("event and time field names must differ")
367 try:
368 from sksurv.util import Surv
369 except ImportError as exc:
370 raise CliError(
371 f"scikit-survival is unavailable; install with `{PINNED_INSTALL}`"
372 ) from exc
373 return Surv.from_arrays(
374 event=normalize_binary_event(event),
375 time=normalize_positive_times(time),
376 name_event=event_name,
377 name_time=time_name,
378 )
381def atomic_save_npy(
382 value: Any, output: str | os.PathLike[str], *, force: bool = False
383) -> None:
384 """Write a NumPy array with pickle disabled."""
386 try:
387 import numpy as np
388 except ImportError as exc:
389 raise CliError(
390 f"NumPy is unavailable; install with `{PINNED_INSTALL}`"
391 ) from exc
392 destination = checked_output_file(output, suffixes={".npy"}, force=force)
393 buffer = io.BytesIO()
394 np.save(buffer, value, allow_pickle=False)
395 atomic_write_bytes(destination, buffer.getvalue(), force=force)
398def atomic_save_npz(
399 arrays: dict[str, Any],
400 output: str | os.PathLike[str],
401 *,
402 force: bool = False,
403) -> None:
404 """Write named NumPy arrays to a compressed archive without object arrays."""
406 try:
407 import numpy as np
408 except ImportError as exc:
409 raise CliError(
410 f"NumPy is unavailable; install with `{PINNED_INSTALL}`"
411 ) from exc
412 if not arrays:
413 raise CliError("at least one array is required")
414 for name, value in arrays.items():
415 array = np.asarray(value)
416 if array.dtype.hasobject:
417 raise CliError(f"array {name!r} has object dtype, which is not allowed")
418 destination = checked_output_file(output, suffixes={".npz"}, force=force)
419 buffer = io.BytesIO()
420 np.savez_compressed(buffer, **arrays)
421 atomic_write_bytes(destination, buffer.getvalue(), force=force)
424def synthetic_survival_frame(
425 *, rows: int = 240, seed: int = DEFAULT_SEED
426) -> tuple[Any, list[str], list[str]]:
427 """Create deterministic, non-clinical right-censored tabular data."""
429 if not 40 <= rows <= MAX_ROWS:
430 raise CliError(f"synthetic rows must be between 40 and {MAX_ROWS}")
431 try:
432 import numpy as np
433 import pandas as pd
434 except ImportError as exc:
435 raise CliError(
436 f"scientific stack unavailable; install with `{PINNED_INSTALL}`"
437 ) from exc
438 rng = np.random.default_rng(seed)
439 x_linear = rng.normal(size=rows)
440 x_noise = rng.normal(size=rows)
441 segment = rng.choice(["alpha", "beta", "gamma"], size=rows, p=[0.4, 0.35, 0.25])
442 segment_effect = np.select(
443 [segment == "beta", segment == "gamma"], [0.35, -0.25], default=0.0
444 )
445 linear_predictor = 0.8 * x_linear + segment_effect
446 event_time = rng.exponential(scale=8.0 * np.exp(-linear_predictor))
447 censor_time = rng.exponential(scale=11.0, size=rows)
448 event = event_time <= censor_time
449 observed_time = np.minimum(event_time, censor_time) + 0.05
450 x_linear = x_linear.astype(float)
451 x_linear[::29] = np.nan
452 segment = segment.astype(object)
453 segment[::37] = np.nan
454 frame = pd.DataFrame(
455 {
456 "event": event,
457 "time": observed_time,
458 "x_linear": x_linear,
459 "x_noise": x_noise,
460 "segment": segment,
461 }
462 )
463 return frame, ["x_linear", "x_noise"], ["segment"]