Move timedelta extraction into normalization stage.

dgf/src/api/transform.py

@@ -28,6 +28,7 @@
 from dgf.src.transform.normalize import IdentityNormalizer
 from dgf.src.transform.normalize import SoftQuantileNormalizer
 from dgf.src.transform.normalize import SinusoidTimedeltaNormalizer
+from dgf.src.transform.normalize import TimedeltaNormalizer
 
 from dgf.src.transform.extract import filter_graph
 from dgf.src.transform.extract import drop_edge_features
@@ -53,6 +54,4 @@
 from dgf.src.transform.timeseries import CalendarFeature
 from dgf.src.transform.timeseries import CalendarFeatureExtractor
 from dgf.src.transform.timeseries import CalendarFeatureExtractorConfig
-from dgf.src.transform.timeseries import TimestampFeatureExtractor
-from dgf.src.transform.timeseries import TimestampFeatureExtractorConfig
 

dgf/src/transform/BUILD

@@ -140,7 +140,6 @@ py_library(
         "//dgf/src/data:in_memory_graph",
         "//dgf/src/data:schema",
         "//dgf/src/io:feature_format",
-        "//dgf/src/util:temporal",
         # numpy dep,
     ],
 )

dgf/src/transform/normalize.py

@@ -560,6 +560,108 @@ def normalize_tensorflow(self, value: tf.Tensor) -> Dict[str, tf.Tensor]:
     return {self.output_feature_name: emb}
 
 
+@normalizer_registry.register
+@dataclasses_json.dataclass_json
+@dataclasses.dataclass(kw_only=True)
+class TimedeltaNormalizer(AbstractFeatureNormalizer):
+  """Normalizes a TIMESTAMP feature into a TIMEDELTA feature relative to seed timestamps.
+
+  Subtracts the timestamp value from the node's seed timestamp:
+  delta = seed_timestamp - t_i
+  """
+
+  input_schema: schema_lib.FeatureSchema
+  output_feature_name: str
+  type: str = dataclasses.field(default="TimedeltaNormalizer", init=False)
+
+  @classmethod
+  def create(
+      cls,
+      feature_name: str,
+      input_schema: schema_lib.FeatureSchema,
+  ) -> "TimedeltaNormalizer":
+    if input_schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP:
+      raise ValueError(
+          f"Feature '{feature_name}' has semantic '{input_schema.semantic}',"
+          " but TimedeltaNormalizer only supports TIMESTAMP features."
+      )
+
+    if not input_schema.is_static_shape():
+      raise ValueError(
+          "TimedeltaNormalizer requires fixed-length feature tensors,"
+          f" but feature '{feature_name}' has a dynamic shape"
+          f" ({input_schema.shape}). Please run padding first."
+      )
+
+    return TimedeltaNormalizer(
+        input_feature=feature_name,
+        input_schema=input_schema,
+        output_feature_name=f"{feature_name}_seed_delta",
+    )
+
+  def output_schema(self) -> schema_lib.FeatureSetSchema:
+    ts_group = self.input_schema.group or (
+        self.input_feature if self.input_schema.is_timeseries else None
+    )
+    return {
+        self.output_feature_name: schema_lib.FeatureSchema(
+            format=self.input_schema.format,
+            semantic=schema_lib.FeatureSemantic.TIMEDELTA,
+            shape=self.input_schema.shape,
+            is_timeseries=self.input_schema.is_timeseries,
+            group=ts_group,
+        )
+    }
+
+  def normalize_numpy(
+      self,
+      value: np.ndarray,
+      seed_timestamps: Optional[np.ndarray] = None,
+  ) -> Dict[str, np.ndarray]:
+    assert seed_timestamps is not None, (
+        "seed_timestamps must be provided to normalize timestamp feature"
+        f" '{self.input_feature}'."
+    )
+
+    assert value.dtype != np.object_, (
+        "TimedeltaNormalizer requires fixed-length feature tensors,"
+        f" but feature '{self.input_feature}' is a variable-length object"
+        " array. Please run padding first."
+    )
+
+    seed_arr = seed_timestamps
+    if value.ndim > 1:
+      seed_arr = seed_arr.reshape(seed_arr.shape + (1,) * (value.ndim - 1))
+
+    deltas = seed_arr - value
+    return {self.output_feature_name: deltas}
+
+  def normalize_tensorflow(
+      self,
+      value: tf.Tensor,
+      seed_timestamps: Optional[tf.Tensor] = None,
+  ) -> Dict[str, tf.Tensor]:
+    assert seed_timestamps is not None, (
+        "seed_timestamps must be provided to normalize timestamp feature"
+        f" '{self.input_feature}'."
+    )
+
+    assert value.shape.rank is not None, (
+        "TimedeltaNormalizer requires fixed-length feature tensors,"
+        f" but feature '{self.input_feature}' has unknown rank."
+        " Please run padding first."
+    )
+
+    seed_tensor = tf.cast(seed_timestamps, value.dtype)
+    if value.shape.rank > 1:
+      seed_tensor = tf.reshape(
+          seed_tensor, (-1,) + (1,) * (value.shape.rank - 1)
+      )
+
+    deltas = tf.subtract(seed_tensor, value)
+    return {self.output_feature_name: deltas}
+
+
 @dataclasses_json.dataclass_json
 @dataclasses.dataclass
 class AutoNormalizeConfig:

dgf/src/transform/normalize_test.py

@@ -899,5 +899,145 @@ def test_auto_normalize_mask(self):
     )
 
 
+class TimedeltaNormalizerTest(parameterized.TestCase):
+
+  def test_output_schema(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(3,),
+        is_timeseries=True,
+        group="sensor",
+    )
+    normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema)
+    expected_schema = {
+        "time_seed_delta": schema_lib.FeatureSchema(
+            format=schema_lib.FeatureFormat.INTEGER_64,
+            semantic=schema_lib.FeatureSemantic.TIMEDELTA,
+            shape=(3,),
+            is_timeseries=True,
+            group="sensor",
+        )
+    }
+    self.assertEqual(normalizer.output_schema(), expected_schema)
+
+  def test_timedelta_normalizer_numpy_1d(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(),
+    )
+    normalizer = normalize_lib.TimedeltaNormalizer.create("created_at", schema)
+    raw_val = np.array([100, 300], dtype=np.int64)
+    seed_timestamps = np.array([500, 500], dtype=np.int64)
+    out = normalizer.normalize_numpy(raw_val, seed_timestamps=seed_timestamps)
+    self.assertIn("created_at_seed_delta", out)
+    np.testing.assert_array_equal(
+        out["created_at_seed_delta"], np.array([400, 200], dtype=np.int64)
+    )
+
+  def test_timedelta_normalizer_numpy_timeseries(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(3,),
+        is_timeseries=True,
+        group="time",
+    )
+    normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema)
+    raw_val = np.array([[100, 200, 300], [400, 450, 500]], dtype=np.int64)
+    seed_timestamps = np.array([500, 600], dtype=np.int64)
+    out = normalizer.normalize_numpy(raw_val, seed_timestamps=seed_timestamps)
+    expected = np.array(
+        [[400, 300, 200], [200, 150, 100]], dtype=np.int64
+    )
+    np.testing.assert_array_equal(out["time_seed_delta"], expected)
+
+  def test_timedelta_normalizer_tensorflow(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(2,),
+        is_timeseries=True,
+    )
+    normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema)
+    raw_np = np.array([[100, 200], [300, 400]], dtype=np.int64)
+    seeds_np = np.array([500, 1000], dtype=np.int64)
+    raw_tf = tf.constant(raw_np)
+    seeds_tf = tf.constant(seeds_np)
+
+    out_tf = normalizer.normalize_tensorflow(raw_tf, seed_timestamps=seeds_tf)
+    out_np = normalizer.normalize_numpy(raw_np, seed_timestamps=seeds_np)
+    expected = np.array([[400, 300], [700, 600]], dtype=np.int64)
+
+    np.testing.assert_array_equal(out_np["time_seed_delta"], expected)
+    np.testing.assert_array_equal(
+        out_tf["time_seed_delta"].numpy(), out_np["time_seed_delta"]
+    )
+
+  def test_timedelta_normalizer_missing_seed_timestamps_asserts(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(),
+    )
+    normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema)
+    raw_val = np.array([100], dtype=np.int64)
+    with self.assertRaises(AssertionError):
+      normalizer.normalize_numpy(raw_val, seed_timestamps=None)
+
+    with self.assertRaises(AssertionError):
+      normalizer.normalize_tensorflow(tf.constant([100]), seed_timestamps=None)
+
+  def test_timedelta_normalizer_object_array_raises(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(2,),
+    )
+    normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema)
+    with self.assertRaisesRegex(AssertionError, "requires fixed-length"):
+      normalizer.normalize_numpy(
+          np.array([np.array([100])], dtype=object),
+          seed_timestamps=np.array([500]),
+      )
+
+  def test_timedelta_normalizer_tensorflow_unknown_rank_asserts(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(2,),
+    )
+    normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema)
+
+    @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.int64)])
+    def normalize_fn(val):
+      return normalizer.normalize_tensorflow(
+          val, seed_timestamps=tf.constant([500], dtype=tf.int64)
+      )
+
+    with self.assertRaisesRegex(AssertionError, "unknown rank"):
+      normalize_fn(tf.constant([100, 200], dtype=tf.int64))
+
+  def test_timedelta_normalizer_invalid_semantic_raises(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.NUMERICAL,
+        shape=(),
+    )
+    with self.assertRaisesRegex(ValueError, "only supports TIMESTAMP"):
+      normalize_lib.TimedeltaNormalizer.create("num", schema)
+
+  def test_timedelta_normalizer_dynamic_shape_raises(self):
+    schema = schema_lib.FeatureSchema(
+        format=schema_lib.FeatureFormat.INTEGER_64,
+        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+        shape=(None,),
+    )
+    with self.assertRaisesRegex(ValueError, "requires fixed-length"):
+      normalize_lib.TimedeltaNormalizer.create("time", schema)
+
+
 if __name__ == "__main__":
   absltest.main()
+