Spaces:
Paused
Paused
| """Fixed-stratum generation and plotting for v6 protocol validation.""" | |
| from __future__ import annotations | |
| import hashlib | |
| from dataclasses import fields, replace | |
| from pathlib import Path | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import torch | |
| from pff_pk.tasks.batch import FlowBatch | |
| from pff_pk.tasks.batching import inverse_concentration | |
| CONTEXT_COLOR = "#ff7f0e" | |
| GENERATED_COLOR = "#1f77b4" | |
| TRUTH_COLOR = "#202020" | |
| PROTOCOL_CATEGORIES = { | |
| "identity": ( | |
| "reference_identity", | |
| "global_identity", | |
| "add_on_identity", | |
| ), | |
| "dose_scaling": ( | |
| "reference_to_lower_scale", | |
| "reference_to_higher_scale", | |
| "lower_scale_to_reference", | |
| "higher_scale_to_reference", | |
| "global_scale_cross", | |
| ), | |
| "add_on_dose": ( | |
| "reference_to_early_add_on", | |
| "reference_to_late_add_on", | |
| "add_on_to_reference", | |
| "add_on_cross", | |
| ), | |
| "cross_intervention": ( | |
| "global_to_add_on", | |
| "add_on_to_global", | |
| ), | |
| } | |
| ACQUISITIONS = ("unscheduled", "exact", "pseudo_scheduled") | |
| def select_representatives(bank: dict) -> list[tuple[str, int, int]]: | |
| """Select one deterministic bank row for every protocol/acquisition stratum.""" | |
| candidates: dict[str, list[tuple[int, int]]] = {} | |
| for record_index, record in enumerate(bank["batches"]): | |
| candidates.setdefault(record["stratum"], []).extend( | |
| (record_index, row) for row in range(len(record["study_id"])) | |
| ) | |
| expected = { | |
| f"{case}__{acquisition}" | |
| for cases in PROTOCOL_CATEGORIES.values() | |
| for case in cases | |
| for acquisition in ACQUISITIONS | |
| } | |
| if set(candidates) != expected: | |
| raise ValueError("validation bank does not contain the expected 42 v6 strata") | |
| selected = [] | |
| for stratum in sorted(candidates): | |
| digest = hashlib.blake2b(stratum.encode(), digest_size=8).digest() | |
| choices = candidates[stratum] | |
| record_index, row = choices[int.from_bytes(digest, "little") % len(choices)] | |
| selected.append((stratum, record_index, row)) | |
| return selected | |
| def protocol_seed(stratum: str) -> int: | |
| """Return the fixed source-process seed for a protocol stratum.""" | |
| digest = hashlib.blake2b(stratum.encode(), digest_size=4).digest() | |
| return 2026 + int.from_bytes(digest, "little") | |
| def _select(batch: FlowBatch, row: int, device: torch.device) -> FlowBatch: | |
| index = torch.tensor([row], dtype=torch.long) | |
| values = {} | |
| for field in fields(batch): | |
| value = getattr(batch, field.name) | |
| if value is None: | |
| values[field.name] = None | |
| elif isinstance(value, torch.Tensor): | |
| values[field.name] = value.index_select(0, index).to(device) | |
| else: | |
| values[field.name] = [value[row]] | |
| return replace(batch, **values) | |
| def _event_arrays(batch: FlowBatch) -> tuple[np.ndarray, np.ndarray]: | |
| assert batch.context_protocol_events is not None | |
| assert batch.context_protocol_event_mask is not None | |
| assert batch.target_protocol_events is not None | |
| assert batch.target_protocol_event_mask is not None | |
| context_mask = batch.context_protocol_event_mask[0, 0].cpu().numpy() | |
| target_mask = batch.target_protocol_event_mask[0].cpu().numpy() | |
| context = batch.context_protocol_events[0, 0].cpu().numpy()[context_mask] | |
| target = batch.target_protocol_events[0].cpu().numpy()[target_mask] | |
| return context, target | |
| def generate_protocol_draws( | |
| model, cpu_batch: FlowBatch, row: int, samples: int, seed: int | |
| ) -> dict: | |
| """Generate persistent conditional draws for one fixed bank row.""" | |
| device = next(model.parameters()).device | |
| batch = _select(cpu_batch, row, device) | |
| repeated_values = {} | |
| for field in fields(batch): | |
| value = getattr(batch, field.name) | |
| if value is None: | |
| repeated_values[field.name] = None | |
| elif isinstance(value, torch.Tensor): | |
| repeated_values[field.name] = value.expand(samples, *value.shape[1:]).clone() | |
| else: | |
| repeated_values[field.name] = value * samples | |
| repeated = replace(batch, **repeated_values) | |
| cuda_devices = [device.index] if device.type == "cuda" else [] | |
| with torch.random.fork_rng(devices=cuda_devices), torch.inference_mode(): | |
| torch.manual_seed(seed) | |
| generated = inverse_concentration(model.sample(repeated), repeated)[..., 0] | |
| context = inverse_concentration(batch.context_value, batch)[0, ..., 0] | |
| target = inverse_concentration(batch.target_value, batch)[0, ..., 0] | |
| context_events, target_events = _event_arrays(batch) | |
| return { | |
| "study_id": batch.study_id[0], | |
| "context_time": batch.context_time[0, ..., 0].cpu(), | |
| "context_value": context.cpu(), | |
| "context_mask": batch.context_mask[0].cpu(), | |
| "target_time": batch.target_time[0, :, 0].cpu(), | |
| "target_value": target.cpu(), | |
| "target_mask": batch.target_mask[0].cpu(), | |
| "generated": generated.cpu(), | |
| "context_events": torch.from_numpy(context_events), | |
| "target_events": torch.from_numpy(target_events), | |
| } | |
| def _schedule(events: torch.Tensor) -> str: | |
| labels = [] | |
| for time, amount, duration, oral in events.tolist(): | |
| route = "PO" if oral >= 0.5 else "IV" | |
| label = f"{route} {amount:.3g}× @ {time:.3g}" | |
| if duration > 0: | |
| label += f" over {duration:.3g}" | |
| labels.append(label) | |
| return " + ".join(labels) if labels else "no dose" | |
| def _human_case(value: str) -> str: | |
| return value.replace("_", " ").replace("to", "→") | |
| def plot_protocol_category( | |
| records: dict[str, dict], | |
| category: str, | |
| cases: tuple[str, ...], | |
| output: Path, | |
| step: int, | |
| ) -> tuple[Path, Path]: | |
| """Render one protocol category over the three acquisition designs.""" | |
| figure, axes = plt.subplots( | |
| len(cases), | |
| len(ACQUISITIONS), | |
| figsize=(15.6, 3.25 * len(cases) + 1.4), | |
| squeeze=False, | |
| constrained_layout=True, | |
| ) | |
| legend_handles = None | |
| for row, case in enumerate(cases): | |
| for column, acquisition in enumerate(ACQUISITIONS): | |
| axis = axes[row, column] | |
| record = records[f"{case}__{acquisition}"] | |
| context_time = record["context_time"].numpy() | |
| context_value = record["context_value"].numpy() | |
| context_mask = record["context_mask"].numpy() | |
| for times, values, mask in zip( | |
| context_time, context_value, context_mask, strict=True | |
| ): | |
| axis.plot( | |
| times[mask], | |
| values[mask], | |
| color=CONTEXT_COLOR, | |
| linewidth=0.5, | |
| alpha=0.28, | |
| zorder=1, | |
| ) | |
| target_mask = record["target_mask"].numpy() | |
| target_time = record["target_time"].numpy()[target_mask] | |
| for generated in record["generated"].numpy(): | |
| axis.plot( | |
| target_time, | |
| generated[target_mask], | |
| color=GENERATED_COLOR, | |
| linewidth=1.0, | |
| alpha=0.72, | |
| zorder=3, | |
| ) | |
| truth = axis.plot( | |
| target_time, | |
| record["target_value"].numpy()[target_mask], | |
| color=TRUTH_COLOR, | |
| linestyle="--", | |
| linewidth=1.5, | |
| zorder=4, | |
| label="held-out truth", | |
| )[0] | |
| context_line = axis.plot([], [], color=CONTEXT_COLOR, label="context cohort")[0] | |
| generated_line = axis.plot([], [], color=GENERATED_COLOR, label="generated draws")[0] | |
| legend_handles = (context_line, generated_line, truth) | |
| axis.text( | |
| 0.02, | |
| 0.97, | |
| f"study {record['study_id'][-8:]}\n" | |
| f"context: {_schedule(record['context_events'])}\n" | |
| f"target: {_schedule(record['target_events'])}", | |
| transform=axis.transAxes, | |
| fontsize=7.2, | |
| ha="left", | |
| va="top", | |
| bbox={"facecolor": "white", "alpha": 0.78, "edgecolor": "none", "pad": 2}, | |
| zorder=10, | |
| ) | |
| if row == 0: | |
| axis.set_title( | |
| acquisition.replace("_", " "), | |
| fontsize=10.5, | |
| fontweight="bold", | |
| pad=8, | |
| ) | |
| if column == 0: | |
| axis.set_ylabel(f"{_human_case(case)}\nC / Cref", fontsize=8) | |
| if row == len(cases) - 1: | |
| axis.set_xlabel("dimensionless time tau", fontsize=8) | |
| axis.set_yscale("log") | |
| axis.tick_params(labelsize=7) | |
| axis.grid(alpha=0.13, linewidth=0.4) | |
| figure.suptitle( | |
| f"v6 validation: {category.replace('_', ' ')} | checkpoint step {step}\n" | |
| "one held-out family per protocol/acquisition stratum; five conditional draws", | |
| fontsize=13, | |
| ) | |
| if legend_handles is not None: | |
| figure.legend( | |
| legend_handles, | |
| [handle.get_label() for handle in legend_handles], | |
| loc="outside lower center", | |
| ncol=3, | |
| frameon=False, | |
| ) | |
| png = output / f"protocol_validation_{category}_step{step:06d}.png" | |
| pdf = png.with_suffix(".pdf") | |
| figure.savefig(png, dpi=220, facecolor="white") | |
| figure.savefig(pdf, facecolor="white") | |
| plt.close(figure) | |
| return png, pdf | |