Add post_pos as a model target via travel_dir decomposition

step_length already encodes |post_pos - pre_pos| by definition, so a raw
post_pos target would duplicate that magnitude and could drift inconsistent
with step_length during sampling. Instead add travel_dir, a unit vector
(local frame) giving only the direction of pre_pos->post_pos; post_pos is
reconstructed at inference as pre_pos + step_length * travel_dir, keeping
the two self-consistent. Target grows from 6D to 9D.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

tests/test_flow.py

@@ -10,7 +10,7 @@ def _small_model():
 
 
 def _batch(B=8):
-    x1 = torch.randn(B, 6)
+    x1 = torch.randn(B, 9)
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     return x1, cond_cont, cond_cat
@@ -40,7 +40,7 @@ def test_sample_flow_shape():
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     out = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
-    assert out.shape == (B, 6)
+    assert out.shape == (B, 9)
 
 
 def test_ddpm_loss_nonneg():
@@ -56,4 +56,4 @@ def test_sample_ddim_shape():
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     out = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
-    assert out.shape == (B, 6)
+    assert out.shape == (B, 9)

tests/test_network.py

@@ -18,7 +18,7 @@ def test_sinusoidal_embedding_batch_1():
 def test_denoising_mlp_output_shape():
     B = 8
     model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
-    x_t = torch.randn(B, 6)
+    x_t = torch.randn(B, 9)
     t = torch.rand(B)
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.stack([
@@ -26,13 +26,13 @@ def test_denoising_mlp_output_shape():
         torch.randint(0, 3, (B,)),
     ], dim=1)
     out = model(x_t, t, cond_cont, cond_cat)
-    assert out.shape == (B, 6)
+    assert out.shape == (B, 9)
 
 
 def test_denoising_mlp_gradients_flow():
     B = 4
     model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
-    x_t = torch.randn(B, 6)
+    x_t = torch.randn(B, 9)
     t = torch.rand(B)
     cond_cont = torch.randn(B, 9)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)

tests/test_transforms.py

@@ -5,6 +5,8 @@ from giant.data.transforms import (
     local_frame_rotation,
     log_transform,
     Normalizer,
+    reconstruct_post_pos,
+    travel_direction,
 )
 
 
@@ -50,6 +52,58 @@ def test_local_frame_rotation_preserves_norm():
     np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)
 
 
+def test_travel_direction_is_unit_norm():
+    rng = np.random.default_rng(5)
+    N = 50
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
+    result = travel_direction(pre_pos, post_pos)
+    np.testing.assert_allclose(np.linalg.norm(result, axis=1), 1.0, atol=1e-5)
+
+
+def test_travel_direction_matches_normalized_displacement():
+    rng = np.random.default_rng(6)
+    N = 50
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    disp = rng.standard_normal((N, 3)).astype(np.float32)
+    post_pos = pre_pos + disp
+    expected = disp / np.linalg.norm(disp, axis=1, keepdims=True)
+    np.testing.assert_allclose(travel_direction(pre_pos, post_pos), expected, atol=1e-5)
+
+
+def test_reconstruct_post_pos_straight_line():
+    """When post_pos = pre_pos + L * pre_dir, travel_dir equals pre_dir, so its
+    local-frame encoding is ẑ — reconstruction must recover post_pos exactly."""
+    rng = np.random.default_rng(7)
+    N = 20
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
+    step_length = rng.uniform(0.1, 5.0, size=N).astype(np.float32)
+    post_pos = pre_pos + step_length[:, None] * pre_dir
+
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
+    np.testing.assert_allclose(travel_dir_local, np.tile([0, 0, 1], (N, 1)), atol=1e-4)
+
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
+    np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
+
+
+def test_reconstruct_post_pos_general_roundtrip():
+    """Full encode (build_features-style) -> decode (cli.py predict-style) path."""
+    rng = np.random.default_rng(8)
+    N = 100
+    pre_pos = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
+    post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
+    step_length = np.linalg.norm(post_pos - pre_pos, axis=1).astype(np.float32)
+
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
+    np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
+
+
 def test_normalizer_roundtrip():
     rng = np.random.default_rng(3)
     X = rng.standard_normal((200, 9)).astype(np.float32)