Unit tests for Entropic Time Scheduler

Signed-off-by: btrentini <brunoxtrentini@gmail.com>

Author: btrentini

sub-packages/bionemo-moco/tests/bionemo/moco/schedules/test_inference_schedules.py

@@ -160,3 +160,109 @@ def test_uniform_dt_padding_dilation(timesteps, device, direction, padding, dila
         assert schedule[0] > schedule[-1]
         for i in range(padding):
             assert schedule[-1 * (i + 1)] == 0
+
+@pytest.mark.parametrize("timesteps", [10, 20])
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("direction", [TimeDirection.UNIFIED, TimeDirection.DIFFUSION])
+def test_entropic_schedule(timesteps, device, direction):
+    """Test the EntropicInferenceSchedule for correctness.
+    Using a tractable predictor function to ensure the scheduler 
+    produces a non-uniform schedule with the correct properties (shape, device, direction, bounds)."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    # Dummy dim for the scheduler
+    dim = 2
+    # A simple time-dependent predictor. Divergence is D*(2t-1) 
+    # non-uniform entropy profile.
+    predictor = lambda t, x: (2 * t - 1) * x
+    x_0_sampler = lambda bs: torch.randn(bs, dim, device=device)
+    x_1_sampler = lambda bs: torch.randn(bs, dim, device=device)
+
+    scheduler = EntropicInferenceSchedule(
+        predictor=predictor,
+        x_0_sampler=x_0_sampler,
+        x_1_sampler=x_1_sampler,
+        nsteps=timesteps,
+        n_approx_entropy_points=25,  # Fewer points for faster testing
+        batch_size=32,
+        direction=direction,
+        device=device,
+    )
+
+    schedule = scheduler.generate_schedule()
+
+    assert schedule.shape == (timesteps,)
+    assert schedule.device.type == device
+
+    # Check that values are within the correct [0, 1] bounds
+    assert torch.all(schedule >= 0) and torch.all(schedule <= 1)
+
+    # Check for correct ordering based on direction
+    if direction == TimeDirection.UNIFIED:
+        assert schedule[0] < schedule[-1]
+        assert torch.all(torch.diff(schedule) >= 0)  # Increase 0 to 1
+    else:
+        assert schedule[0] > schedule[-1]
+        assert torch.all(torch.diff(schedule) <= 0)  # Decrease 1 to 0
+
+    # Check that the schedule is non-uniform, confirming the entropic logic is active
+    # Round to avoid float precision issues making all diffs unique
+    diffs = torch.diff(torch.abs(schedule)).round(decimals=5)
+    # Expect more than one unique step size, unlike a linear schedule
+    if timesteps > 5:
+        assert len(torch.unique(diffs)) > 1
+
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_entropic_schedule_reproducibility(device):
+    """Checks that the the EntropicInferenceSchedule produce reproducible results when 
+    a torch.Generator with a fixed seed is provided."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    timesteps = 10
+    dim = 2
+    predictor = lambda t, x: (2 * t - 1) * x
+    x_0_sampler = lambda bs: torch.randn(bs, dim, device=device)
+    x_1_sampler = lambda bs: torch.randn(bs, dim, device=device)
+
+    gen1 = torch.Generator(device=device).manual_seed(42)
+    scheduler1 = EntropicInferenceSchedule(
+        predictor=predictor,
+        x_0_sampler=x_0_sampler,
+        x_1_sampler=x_1_sampler,
+        nsteps=timesteps,
+        device=device,
+        generator=gen1,
+    )
+    schedule1 = scheduler1.generate_schedule()
+
+    # Run again with the same seed...
+    gen2 = torch.Generator(device=device).manual_seed(42)
+    scheduler2 = EntropicInferenceSchedule(
+        predictor=predictor,
+        x_0_sampler=x_0_sampler,
+        x_1_sampler=x_1_sampler,
+        nsteps=timesteps,
+        device=device,
+        generator=gen2,
+    )
+    schedule2 = scheduler2.generate_schedule()
+
+    # Compare again with another seed
+    gen3 = torch.Generator(device=device).manual_seed(99)
+    scheduler3 = EntropicInferenceSchedule(
+        predictor=predictor,
+        x_0_sampler=x_0_sampler,
+        x_1_sampler=x_1_sampler,
+        nsteps=timesteps,
+        device=device,
+        generator=gen3,
+    )
+    schedule3 = scheduler3.generate_schedule()
+
+    # Schedules from identical seeds should be identical
+    assert torch.allclose(schedule1, schedule2)
+
+    # Schedule from a different seed should be different
+    assert not torch.allclose(schedule1, schedule3)