| import torch, json, numpy as np |
| from pathlib import Path |
| import sys |
| sys.path.insert(0, 'src') |
|
|
| from pino.embeddings import OlfactoryEmbeddingEngine |
| from pino.heads import PIMTHeads |
| from pino.pimt_model import DEFAULT_EMBEDDING_DIM, PhysicsInformedMixtureTransformer |
| from pino.registry import AromaRegistry |
| from pino.thermo.naturals import NATURAL_PROFILES |
|
|
| |
| device = torch.device('cpu') |
| model = PhysicsInformedMixtureTransformer(embedding_dim=DEFAULT_EMBEDDING_DIM, state_dim=2, hidden_dim=256, num_heads=4, num_layers=4).to(device) |
| model.eval() |
|
|
| |
| with open('data/empirical_dataset_v1.jsonl') as f: |
| for line in f: |
| rec = json.loads(line) |
| if any(str(i.get('cas','')).startswith('NATURAL:') for i in rec.get('formula',[])): |
| sample = rec |
| break |
|
|
| engine = OlfactoryEmbeddingEngine() |
| registry = AromaRegistry() |
|
|
| tokens = [] |
| for item in sample['formula']: |
| cas = str(item['cas']).strip() |
| bare = cas.replace('NATURAL:', '') |
| rec = registry.get(bare) |
| smiles = rec.get('smiles', '') if rec else '' |
| z = engine.get_embedding(smiles, cas=cas) |
| tokens.append(torch.from_numpy(z)) |
| tokens = torch.stack(tokens, dim=0).unsqueeze(0).float().to(device) |
|
|
| physics = np.zeros((len(sample['trajectory']), len(sample['formula']), 2), dtype=np.float32) |
| const_to_token = {} |
| for idx, item in enumerate(sample['formula']): |
| raw_cas = str(item['cas']).strip() |
| bare = raw_cas.replace('NATURAL:', '') |
| profile = NATURAL_PROFILES.get(bare) or NATURAL_PROFILES.get(f'NATURAL:{bare}') |
| if profile: |
| for const_cas in profile['constituents']: |
| const_to_token[const_cas] = idx |
| else: |
| const_to_token[bare] = idx |
| for t, step in enumerate(sample['trajectory']): |
| for const_cas, token_idx in const_to_token.items(): |
| physics[t, token_idx, 0] += step['x_liquid'].get(const_cas, 0.0) |
| physics[t, token_idx, 1] += step['OAV'].get(const_cas, 0.0) |
| physics[:,:,1] = np.log10(np.maximum(physics[:,:,1], 1e-10)) |
| physics_t = torch.from_numpy(physics).unsqueeze(0).float().to(device) |
| mask = torch.zeros(1, tokens.size(1), dtype=torch.bool, device=device) |
|
|
| with torch.no_grad(): |
| gated = model.gating(tokens, physics_t) |
| print('FiLM gated variance across time:', gated.var(dim=1).mean().item()) |
| print('FiLM gated first vs last timestep mean diff:', (gated[0,0] - gated[0,-1]).abs().mean().item()) |
| print('FiLM gated first vs last timestep max diff:', (gated[0,0] - gated[0,-1]).abs().max().item()) |
|
|
| proj = model.input_proj(gated) |
| b, t, s, _ = proj.shape |
| x = proj.reshape(b * t, s, 256) |
| enc_out = model.encoder(x, src_key_padding_mask=mask.unsqueeze(1).expand(-1, t, -1).reshape(b * t, s)) |
| enc_out = enc_out.reshape(b, t, s, 256) |
| print('Encoder output variance across time:', enc_out.var(dim=1).mean().item()) |
| print('Encoder output first vs last timestep mean diff:', (enc_out[0,0] - enc_out[0,-1]).abs().mean().item()) |
| print('Encoder output first vs last timestep max diff:', (enc_out[0,0] - enc_out[0,-1]).abs().max().item()) |
|
|
| |
| physics0 = torch.zeros_like(physics_t) |
| gated0 = model.gating(tokens, physics0) |
| proj0 = model.input_proj(gated0) |
| x0 = proj0.reshape(b * t, s, 256) |
| enc_out0 = model.encoder(x0, src_key_padding_mask=mask.unsqueeze(1).expand(-1, t, -1).reshape(b * t, s)) |
| enc_out0 = enc_out0.reshape(b, t, s, 256) |
| print('Diff encoder vs zero physics (mean):', (enc_out - enc_out0).abs().mean().item()) |
| print('Diff encoder vs zero physics (max):', (enc_out - enc_out0).abs().max().item()) |
|
|
