Add time delta feature extractor.

dgf/src/api/transform.py

@@ -54,4 +54,6 @@
 from dgf.src.transform.timeseries import PadAndCapTimeseriesConfig
 from dgf.src.transform.timeseries import extract_calendar_features
 from dgf.src.transform.timeseries import CalendarFeatureConfig
+from dgf.src.transform.timeseries import extract_time_delta_features
+from dgf.src.transform.timeseries import TimeDeltaFeatureConfig
 

dgf/src/transform/timeseries.py

@@ -141,6 +141,24 @@ class CalendarFeatureConfig:
   features: Tuple[CalendarFeature, ...] = _SUPPORTED_CALENDAR_FEATURES
 
 
+@dataclasses_json.dataclass_json
+@dataclasses.dataclass
+class TimeDeltaFeatureConfig:
+  """Configuration for extracting time delta features.
+
+  Attributes:
+    seed_timestamp: Target reference timestamp to compute absolute deltas from.
+    extract_consecutive: Whether to extract t_i - t_{i-1} (consecutive gaps).
+    extract_seed_delta: Whether to extract seed_timestamp - t_i (seed deltas).
+    fill_value: Value used for masked time steps and missing boundary deltas.
+  """
+
+  seed_timestamp: int
+  extract_consecutive: bool = True
+  extract_seed_delta: bool = True
+  fill_value: int = 0
+
+
 def _compute_calendar_feature(
     ts_array: np.ndarray, feature: CalendarFeature
 ) -> np.ndarray:
@@ -448,3 +466,158 @@ def extract_calendar_features(
           node_sets=new_ns_schemas, edge_sets=new_es_schemas
       ),
   )
+
+
+def _extract_feature_set_time_delta_features(
+    values: in_memory_graph.Features,
+    schemas: schema_lib.FeatureSetSchema,
+    config: TimeDeltaFeatureConfig,
+) -> Tuple[in_memory_graph.Features, schema_lib.FeatureSetSchema]:
+  """Extracts time delta features from timestamp features of a single feature set."""
+  new_values: in_memory_graph.Features = {}
+  new_schemas: schema_lib.FeatureSetSchema = {}
+
+  for fname, schema in schemas.items():
+    raw_val = values[fname]
+    new_values[fname] = raw_val
+    new_schemas[fname] = schema
+
+    # Skip non-timestamp features.
+    if schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP:
+      continue
+
+    if raw_val.dtype == np.object_:
+      raise ValueError(
+          "extract_time_delta_features requires fixed-length timestamp"
+          f" tensors, but feature '{fname}' is a variable-length object array."
+          " Please run pad_and_cap_timeseries_features first."
+      )
+
+    mask_fname = f"{fname}_mask"
+    if mask_fname in values:
+      mask_matrix = values[mask_fname].astype(np.bool_)
+    else:
+      mask_matrix = (raw_val != 0)
+
+    # Case 1: For non-timeseries timestamps, derived time delta features do not
+    # reference a timestamp sequence.
+    if not schema.is_timeseries:
+      delta_timestamps = None
+    # Case 2: For timeseries features that reference a timestamp sequence.
+    elif schema.timestamps is not None:
+      delta_timestamps = schema.timestamps
+    # Case 3: For timeseries features that do not reference a timestamp
+    # sequence, derived time delta sequences reference the original feature
+    # as their timestamp sequence.
+    else:
+      delta_timestamps = fname
+
+    # Extract consecutive deltas for timeseries. The mask ensures that there
+    # are no boundary effects.
+    if config.extract_consecutive and schema.is_timeseries:
+      pair_mask = mask_matrix[:, :-1] & mask_matrix[:, 1:]
+      valid_diffs = np.where(
+          pair_mask, raw_val[:, 1:] - raw_val[:, :-1], config.fill_value
+      )
+      consecutive_delta = np.hstack([
+          np.full((raw_val.shape[0], 1), config.fill_value, dtype=np.int64),
+          valid_diffs,
+      ])
+      out_fname = f"{fname}_consecutive_delta"
+      new_values[out_fname] = consecutive_delta
+      new_schemas[out_fname] = schema_lib.FeatureSchema(
+          format=schema_lib.FeatureFormat.INTEGER_64,
+          semantic=schema_lib.FeatureSemantic.NUMERICAL,
+          shape=schema.shape,
+          is_timeseries=schema.is_timeseries,
+          timestamps=delta_timestamps,
+      )
+
+    # Extract difference to seed timestamp.
+    if config.extract_seed_delta:
+      seed_delta = np.where(
+          mask_matrix, config.seed_timestamp - raw_val, config.fill_value
+      )
+      out_fname = f"{fname}_seed_delta"
+      new_values[out_fname] = seed_delta
+      new_schemas[out_fname] = schema_lib.FeatureSchema(
+          format=schema_lib.FeatureFormat.INTEGER_64,
+          semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+          shape=schema.shape,
+          is_timeseries=schema.is_timeseries,
+          timestamps=delta_timestamps,
+      )
+
+  return new_values, new_schemas
+
+
+def extract_time_delta_features(
+    graph: in_memory_graph.InMemoryGraph,
+    schema: schema_lib.GraphSchema,
+    config: TimeDeltaFeatureConfig,
+) -> Tuple[in_memory_graph.InMemoryGraph, schema_lib.GraphSchema]:
+  """Extracts time delta features (consecutive gaps and seed deltas) from timestamps.
+
+  Requires fixed-length timestamp tensors (e.g. produced after running
+  `pad_and_cap_timeseries_features`).
+
+  Usage example:
+
+  ```python
+  config = dgf.transform.TimeDeltaFeatureConfig(seed_timestamp=1680000000)
+  graph, schema = dgf.transform.extract_time_delta_features(
+      graph, schema, config
+  )
+  ```
+
+  Args:
+    graph: The input in-memory graph.
+    schema: The graph schema containing timestamp features.
+    config: `TimeDeltaFeatureConfig` specifying seed timestamp and flags.
+
+  Returns:
+    Tuple `(new_graph, new_schema)` containing original and extracted delta
+    features.
+  """
+  new_node_sets = {}
+  new_ns_schemas = {}
+
+  for ns_name, ns_schema in schema.node_sets.items():
+    ns_val = graph.node_sets[ns_name]
+    new_vals, new_schemas = _extract_feature_set_time_delta_features(
+        values=ns_val.features,
+        schemas=ns_schema.features,
+        config=config,
+    )
+    new_node_sets[ns_name] = in_memory_graph.InMemoryNodeSet(
+        num_nodes=ns_val.num_nodes, features=new_vals
+    )
+    new_ns_schemas[ns_name] = schema_lib.NodeSchema(features=new_schemas)
+
+  new_edge_sets = {}
+  new_es_schemas = {}
+
+  for es_name, es_schema in schema.edge_sets.items():
+    es_val = graph.edge_sets[es_name]
+    new_vals, new_schemas = _extract_feature_set_time_delta_features(
+        values=es_val.features,
+        schemas=es_schema.features,
+        config=config,
+    )
+    new_edge_sets[es_name] = in_memory_graph.InMemoryEdgeSet(
+        adjacency=es_val.adjacency, features=new_vals
+    )
+    new_es_schemas[es_name] = schema_lib.EdgeSchema(
+        source=es_schema.source,
+        target=es_schema.target,
+        features=new_schemas,
+    )
+
+  return (
+      in_memory_graph.InMemoryGraph(
+          node_sets=new_node_sets, edge_sets=new_edge_sets
+      ),
+      schema_lib.GraphSchema(
+          node_sets=new_ns_schemas, edge_sets=new_es_schemas
+      ),
+  )