instance splitter fix

pts/transform/sampler.py

@@ -101,6 +101,7 @@ class ExpectedNumInstanceSampler(InstanceSampler):
         self.lookup = np.arange(2 ** 13)
 
     def __call__(self, ts: np.ndarray, a: int, b: int) -> np.ndarray:
+        assert a <= b, "First index must be less than or equal to the last index."
         while ts.shape[-1] >= len(self.lookup):
             self.lookup = np.arange(2 * len(self.lookup))
 

pts/transform/split.py

@@ -146,28 +146,35 @@ class InstanceSplitter(FlatMapTransformation):
 
         len_target = target.shape[-1]
 
+        minimum_length = (
+            self.future_length
+            if self.pick_incomplete
+            else self.past_length + self.future_length
+        )
+
         if is_train:
-            if len_target < self.future_length:
-                # We currently cannot handle time series that are shorter than
-                # the prediction length during training, so we just skip these.
-                # If we want to include them we would need to pad and to mask
-                # the loss.
-                sampling_indices: Union[np.ndarray, List[int]] = []
-            else:
-                if self.pick_incomplete:
-                    sampling_indices = self.train_sampler(
-                        target, 0, len_target - self.future_length
-                    )
-                else:
-                    sampling_indices = self.train_sampler(
-                        target, self.past_length, len_target - self.future_length,
-                    )
+            sampling_bounds = (
+                (0, len_target - self.future_length)
+                if self.pick_incomplete
+                else (self.past_length, len_target - self.future_length)
+            )
+
+            # We currently cannot handle time series that are
+            # too short during training, so we just skip these.
+            # If we want to include them we would need to pad and to
+            # mask the loss.
+            sampled_indices = (
+                np.array([], dtype=int)
+                if len_target < minimum_length
+                else self.train_sampler(target, *sampling_bounds)
+            )
         else:
-            sampling_indices = [len_target]
-        for i in sampling_indices:
+            assert self.pick_incomplete or len_target >= self.past_length
+            sampled_indices = np.array([len_target], dtype=int)
+        for i in sampled_indices:
             pad_length = max(self.past_length - i, 0)
             if not self.pick_incomplete:
-                assert pad_length == 0
+                assert pad_length == 0, f"pad_length should be zero, got {pad_length}"
             d = data.copy()
             for ts_field in slice_cols:
                 if i > self.past_length:

test/test_transform.py

@@ -1,12 +1,26 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
+# Standard library imports
 from typing import Tuple
 
-import pytest
-
+# Third-party imports
 import numpy as np
 import pandas as pd
-
 import torch
+import pytest
 
+# First-party imports
 from pts.dataset import (
     ProcessStartField,
     FieldName,
@@ -59,7 +73,7 @@ def test_align_timestamp():
 @pytest.mark.parametrize("is_train", TEST_VALUES["is_train"])
 @pytest.mark.parametrize("target", TEST_VALUES["target"])
 @pytest.mark.parametrize("start", TEST_VALUES["start"])
-def test_AddTimeFeatures(start, target, is_train):
+def test_AddTimeFeatures(start, target, is_train: bool):
     pred_length = 13
     t = transform.AddTimeFeatures(
         start_field=FieldName.START,
@@ -69,7 +83,6 @@ def test_AddTimeFeatures(start, target, is_train):
         time_features=[time_feature.DayOfWeek(), time_feature.DayOfMonth()],
     )
 
-
     data = {"start": start, "target": target}
     res = t.map_transform(data, is_train=is_train)
     mat = res["myout"]
@@ -83,7 +96,7 @@ def test_AddTimeFeatures(start, target, is_train):
 @pytest.mark.parametrize("is_train", TEST_VALUES["is_train"])
 @pytest.mark.parametrize("target", TEST_VALUES["target"])
 @pytest.mark.parametrize("start", TEST_VALUES["start"])
-def test_AddTimeFeatures_empty_time_features(start, target, is_train):
+def test_AddTimeFeatures_empty_time_features(start, target, is_train: bool):
     pred_length = 13
     t = transform.AddTimeFeatures(
         start_field=FieldName.START,
@@ -101,7 +114,7 @@ def test_AddTimeFeatures_empty_time_features(start, target, is_train):
 @pytest.mark.parametrize("is_train", TEST_VALUES["is_train"])
 @pytest.mark.parametrize("target", TEST_VALUES["target"])
 @pytest.mark.parametrize("start", TEST_VALUES["start"])
-def test_AddAgeFeatures(start, target, is_train):
+def test_AddAgeFeatures(start, target, is_train: bool):
     pred_length = 13
     t = transform.AddAgeFeature(
         pred_length=pred_length,
@@ -110,7 +123,6 @@ def test_AddAgeFeatures(start, target, is_train):
         log_scale=True,
     )
 
-
     data = {"start": start, "target": target}
     out = t.map_transform(data, is_train=is_train)
     expected_length = len(target) + (0 if is_train else pred_length)
@@ -121,10 +133,11 @@ def test_AddAgeFeatures(start, target, is_train):
     )
 
 
+@pytest.mark.parametrize("pick_incomplete", TEST_VALUES["allow_target_padding"])
 @pytest.mark.parametrize("is_train", TEST_VALUES["is_train"])
 @pytest.mark.parametrize("target", TEST_VALUES["target"])
 @pytest.mark.parametrize("start", TEST_VALUES["start"])
-def test_InstanceSplitter(start, target, is_train):
+def test_InstanceSplitter(start, target, is_train: bool, pick_incomplete: bool):
     train_length = 100
     pred_length = 13
     t = transform.InstanceSplitter(
@@ -136,10 +149,9 @@ def test_InstanceSplitter(start, target, is_train):
         past_length=train_length,
         future_length=pred_length,
         time_series_fields=["some_time_feature"],
-        pick_incomplete=True,
+        pick_incomplete=pick_incomplete,
     )
 
-
     other_feat = np.arange(len(target) + 100)
     data = {
         "start": start,
@@ -148,10 +160,17 @@ def test_InstanceSplitter(start, target, is_train):
         "some_other_col": "ABC",
     }
 
-    out = list(t.flatmap_transform(data, is_train=is_train))
+    if not is_train and not pick_incomplete and len(target) < train_length:
+        with pytest.raises(AssertionError):
+            out = list(t.flatmap_transform(data, is_train=is_train))
+        return
+    else:
+        out = list(t.flatmap_transform(data, is_train=is_train))
 
     if is_train:
-        assert len(out) == max(0, len(target) - pred_length + 1)
+        assert len(out) == max(
+            0, len(target) - pred_length + 1 - (0 if pick_incomplete else train_length),
+        )
     else:
         assert len(out) == 1
 
@@ -183,7 +202,11 @@ def test_InstanceSplitter(start, target, is_train):
 )
 @pytest.mark.parametrize("allow_target_padding", TEST_VALUES["allow_target_padding"])
 def test_CanonicalInstanceSplitter(
-    start, target, is_train, use_prediction_features, allow_target_padding
+    start,
+    target,
+    is_train: bool,
+    use_prediction_features: bool,
+    allow_target_padding: bool,
 ):
     train_length = 100
     pred_length = 13
@@ -285,10 +308,10 @@ def test_Transformation():
 def test_multi_dim_transformation(is_train):
     train_length = 10
 
-    first_dim = np.arange(1, 11, 1).tolist()
+    first_dim: list = list(np.arange(1, 11, 1))
     first_dim[-1] = "NaN"
 
-    second_dim = np.arange(11, 21, 1).tolist()
+    second_dim: list = list(np.arange(11, 21, 1))
     second_dim[0] = "NaN"
 
     ds = ListDataset(
@@ -782,3 +805,4 @@ def assert_padded_array(
         f"Sampled and reference arrays do not match. '"
         f"Got {sampled_no_padding} but should be {reference_no_padding}."
     )
+