Fix schedule length for small number of adaptation steps

Author: rlouf

aehmc/window_adaptation.py

@@ -252,7 +252,7 @@ def build_schedule(
 
     # Give up on mass matrix adaptation when the number of warmup steps is too small.
     if num_steps < 20:
-        schedule += [(0, False)] * (num_steps - 1)
+        schedule += [(0, False)] * num_steps
     else:
         # When the number of warmup steps is smaller that the sum of the provided (or default)
         # window sizes we need to resize the different windows.

tests/test_adaptation.py

@@ -6,7 +6,7 @@
 @pytest.mark.parametrize(
     "num_steps, expected_schedule",
     [
-        (19, [(0, False)] * 18),  # no mass matrix adaptation
+        (19, [(0, False)] * 19),  # no mass matrix adaptation
         (
             100,
             [(0, False)] * 15 + [(1, False)] * 74 + [(1, True)] + [(0, False)] * 10,
@@ -24,4 +24,5 @@
 )
 def test_adaptation_schedule(num_steps, expected_schedule):
     adaptation_schedule = window_adaptation.build_schedule(num_steps)
+    assert num_steps == len(adaptation_schedule)
     assert adaptation_schedule == expected_schedule