"""IDEA 3: does the PIMT transformer itself exploit the physics channel to predict the time-evolved (pyramid) descriptor profile? Idea 2 showed (with GBM/linear probes on hand-built features) that the physics trajectory DOUBLES explained variance on the drydown profile. This evaluator tests the same claim end-to-end in the trained model: for each held-out formula, run the objective head with the REAL physics tensor vs a ZEROED physics tensor (physics-off ablation at inference), and score the predicted 3-tier pyramid against the stored pyramid target -- per tier (top/mid/base) and overall. If physics carries non-redundant signal for the time-evolved profile, physics-on should beat physics-off, with the largest gain on the BASE tier (drydown). Metric: per-descriptor binary cross-entropy (the training objective) and mean cosine similarity between predicted and target tier profiles. Lower BCE / higher cosine with physics-on than physics-off = physics helps. """ from __future__ import annotations import argparse import json from pathlib import Path import numpy as np import torch from pino.heads import PIMTHeads from pino.pimt_model import ( FragranceTrajectoryDataset, PhysicsInformedMixtureTransformer, objective_targets_to_pyramid, ) from pino.upload_data import create_molecule_disjoint_split_from_records def load_model(ckpt_path: Path, objective_dim: int): sd = torch.load(ckpt_path, map_location="cpu", weights_only=False) state = sd.get("model_state_dict", sd.get("model", sd)) heads_state = sd.get("heads_state_dict", sd.get("heads", {})) state_dim = int(state["gating.physics_scaler"].shape[0]) embedding_dim = int(state["input_proj.weight"].shape[1]) hidden_dim = int(state["input_proj.weight"].shape[0]) num_layers = sum(1 for k in state if k.endswith("self_attn.in_proj_weight")) model = PhysicsInformedMixtureTransformer( embedding_dim=embedding_dim, state_dim=state_dim, hidden_dim=hidden_dim, num_heads=4, num_layers=max(num_layers, 1), ) model.load_state_dict(state, strict=False) heads = PIMTHeads(hidden_dim=hidden_dim, objective_dim=objective_dim) if heads_state: heads.load_state_dict(heads_state, strict=False) model.eval(); heads.eval() return model, heads, {"embedding_dim": embedding_dim, "state_dim": state_dim, "hidden_dim": hidden_dim, "num_layers": num_layers} def tier_cosine(pred, targ): # pred/targ: (3, 138); cosine per tier out = [] for t in range(3): a = pred[t]; b = targ[t] na = np.linalg.norm(a); nb = np.linalg.norm(b) out.append(float(np.dot(a, b) / (na * nb)) if na > 0 and nb > 0 else 0.0) return out # [top, mid, base] def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--checkpoint", required=True) ap.add_argument("--data", required=True) ap.add_argument("--structural-source", default="morgan") ap.add_argument("--train-ratio", type=float, default=0.85) ap.add_argument("--seed", type=int, default=42) ap.add_argument("--split", choices=["molecule", "formula"], default="formula") ap.add_argument("--arm-label", default=None) ap.add_argument("--use-gamma", action="store_true") ap.add_argument("--physics-mode", choices=["full", "off", "gamma_only", "headspace_only"], default="full", help="Match the arm's training conditioning: gamma_only zeroes channels 0-1 of the 'on' tensor, headspace_only zeroes channel 2, off zeroes all. Ensures each factorial arm is evaluated with the same masked physics it was trained on.") ap.add_argument("--max-eval", type=int, default=400) ap.add_argument("--output", required=True) args = ap.parse_args() objective_dim = 575 if args.structural_source == "pom_alltags" else 138 engine_source = "openpom_256" if args.structural_source == "pom_alltags" else args.structural_source model, heads, dims = load_model(Path(args.checkpoint), objective_dim) state_dim = dims["state_dim"] with open(args.data) as f: all_records = [json.loads(line) for line in f if line.strip()] if args.split == "molecule": split = create_molecule_disjoint_split_from_records(all_records, train_ratio=args.train_ratio, seed=args.seed) val_records = split["validation"] else: rng = np.random.default_rng(args.seed) n = len(all_records) n_val = max(1, int(round(n * (1.0 - args.train_ratio)))) perm = rng.permutation(n) val_records = [all_records[i] for i in sorted(set(perm[:n_val].tolist()))] # restrict to blends with a real trajectory + pyramid target blends = [r for r in val_records if not r.get("is_control") and len(r.get("formula", [])) >= 3 and r.get("objective_targets") and len(r.get("trajectory", [])) > 1] blends = blends[: args.max_eval] ds = FragranceTrajectoryDataset(records=blends, use_embedding_fallback=True, structural_source=engine_source, objective_dim=objective_dim, use_gamma=args.use_gamma) device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device); heads.to(device) bce_on, bce_off = [], [] cos_on = np.zeros(3); cos_off = np.zeros(3) used = 0 bce = torch.nn.functional.binary_cross_entropy with torch.no_grad(): for i in range(len(ds)): item = ds[i] tokens = item["tokens"].unsqueeze(0).to(device) physics = item["physics"].unsqueeze(0).to(device) if physics.size(-1) != state_dim: if physics.size(-1) < state_dim: pad = torch.zeros(1, physics.size(1), physics.size(2), state_dim - physics.size(-1), device=device) physics = torch.cat([physics, pad], dim=-1) else: physics = physics[..., :state_dim] # mask: all real tokens present in a single (unpadded) item S = tokens.size(1) mask = torch.zeros(1, S, dtype=torch.bool, device=device) # target_obj is already the (3,138) pyramid for this dataset target = item["target_obj"].clamp(0.0, 1.0).to(device) # (3,138) # physics ON -- with the arm's training conditioning (factorial modes # mask channels so eval matches what the model saw during training) if args.physics_mode == "off": physics = torch.zeros_like(physics) elif args.physics_mode == "gamma_only": physics = physics.clone(); physics[..., 0] = 0.0; physics[..., 1] = 0.0 elif args.physics_mode == "headspace_only": physics = physics.clone(); physics[..., 2] = 0.0 lat_on = model(tokens, physics) pyr_on = heads(lat_on, physics, mask)["objective"][0] # (3,138) # physics OFF (zeroed) phys_zero = torch.zeros_like(physics) lat_off = model(tokens, phys_zero) pyr_off = heads(lat_off, phys_zero, mask)["objective"][0] bce_on.append(float(bce(pyr_on, target).cpu())) bce_off.append(float(bce(pyr_off, target).cpu())) c_on = tier_cosine(pyr_on.cpu().numpy(), target.cpu().numpy()) c_off = tier_cosine(pyr_off.cpu().numpy(), target.cpu().numpy()) cos_on += np.array(c_on); cos_off += np.array(c_off) used += 1 cos_on /= max(used, 1); cos_off /= max(used, 1) result = { "arm": args.arm_label or Path(args.checkpoint).stem, "use_gamma": bool(args.use_gamma), "physics_mode": args.physics_mode, "seed": args.seed, "model_dims": dims, "task": "3-tier descriptor pyramid (top/mid/base); physics-on vs physics-off inference ablation", "n_eval": used, "bce_physics_on": float(np.mean(bce_on)), "bce_physics_off": float(np.mean(bce_off)), "bce_delta_off_minus_on": float(np.mean(bce_off) - np.mean(bce_on)), "cosine_top_on": float(cos_on[0]), "cosine_mid_on": float(cos_on[1]), "cosine_base_on": float(cos_on[2]), "cosine_top_off": float(cos_off[0]), "cosine_mid_off": float(cos_off[1]), "cosine_base_off": float(cos_off[2]), "cosine_base_delta_on_minus_off": float(cos_on[2] - cos_off[2]), "cosine_mean_on": float(cos_on.mean()), "cosine_mean_off": float(cos_off.mean()), "verdict": ("physics helps (on beats off)" if np.mean(bce_on) < np.mean(bce_off) else "physics does not help"), } out = Path(args.output); out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps(result, indent=2)) print(json.dumps({k: result[k] for k in ["arm", "n_eval", "bce_physics_on", "bce_physics_off", "bce_delta_off_minus_on", "cosine_base_on", "cosine_base_off", "cosine_base_delta_on_minus_off", "verdict"]}, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())