File size: 1,646 Bytes
4cb605e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | from __future__ import annotations
import torch
from pino.pimt_model import DEFAULT_EMBEDDING_DIM, objective_targets_to_pyramid
from pino.train import pad_trajectory_collate
def _item(target_obj: torch.Tensor, *, n_mol: int = 2, t_steps: int = 49) -> dict:
return {
"tokens": torch.ones(n_mol, DEFAULT_EMBEDDING_DIM),
"physics": torch.ones(t_steps, n_mol, 2),
"target_obj": target_obj,
"target_sub": torch.tensor([0.25, 0.25, 0.25, 0.25, 0.0, 0.5, 0.5]),
}
def test_objective_targets_to_pyramid_keeps_native_pyramid() -> None:
targets = torch.rand(3, 138)
normalized = objective_targets_to_pyramid(targets)
assert normalized.shape == (3, 138)
assert torch.equal(normalized, targets)
def test_pad_trajectory_collate_normalizes_mixed_objective_target_shapes() -> None:
pyramid = torch.zeros(3, 138)
pyramid[0, 1] = 1.0
trajectory = torch.zeros(49, 138)
trajectory[0, 2] = 0.4
trajectory[20, 3] = 0.7
trajectory[-1, 4] = 1.0
batch = pad_trajectory_collate([_item(pyramid), _item(trajectory, n_mol=3)])
assert batch["target_obj"].shape == (2, 3, 138)
assert torch.equal(batch["target_obj"][0], pyramid)
assert batch["target_obj"][1, 0, 2] == 0.4
assert batch["target_obj"][1, 1, 3] == 0.7
assert batch["target_obj"][1, 2, 4] == 1.0
def test_pad_trajectory_collate_normalizes_mixed_objective_target_shapes_when_legacy_first() -> None:
trajectory = torch.zeros(49, 138)
pyramid = torch.zeros(3, 138)
batch = pad_trajectory_collate([_item(trajectory), _item(pyramid)])
assert batch["target_obj"].shape == (2, 3, 138)
|