From a566da25c9efaea583414b5de337c4ad5357a670 Mon Sep 17 00:00:00 2001 From: Simon Meierhans Date: Wed, 8 Jul 2026 09:06:32 -0700 Subject: [PATCH] Add time delta feature extractor. PiperOrigin-RevId: 944531942 --- dgf/src/api/transform.py | 2 + dgf/src/transform/timeseries.py | 173 ++++++++++++ dgf/src/transform/timeseries_test.py | 381 +++++++++++++++++++++++++++ 3 files changed, 556 insertions(+) diff --git a/dgf/src/api/transform.py b/dgf/src/api/transform.py index e078ed9..cd12da3 100644 --- a/dgf/src/api/transform.py +++ b/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 diff --git a/dgf/src/transform/timeseries.py b/dgf/src/transform/timeseries.py index 8acb370..b82a414 100644 --- a/dgf/src/transform/timeseries.py +++ b/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 + ), + ) diff --git a/dgf/src/transform/timeseries_test.py b/dgf/src/transform/timeseries_test.py index c93a4da..7e01758 100644 --- a/dgf/src/transform/timeseries_test.py +++ b/dgf/src/transform/timeseries_test.py @@ -664,6 +664,387 @@ def test_extract_calendar_features_parent_timestamp(self): hw_sch.features["master_time_hour"].timestamps, "master_time" ) + def test_extract_time_delta_features(self): + graph, schema = _make_graph_and_schema( + values={ + "time": np.array( + [np.array([100, 250, 300], dtype=np.int64)], dtype=np.object_ + ) + }, + schemas={ + "time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + ) + }, + ) + padded_graph, padded_schema = timeseries.pad_and_cap_timeseries_features( + graph, + schema, + timeseries.PadAndCapTimeseriesConfig(sequence_length=4), + ) + delta_graph, delta_schema = timeseries.extract_time_delta_features( + padded_graph, + padded_schema, + timeseries.TimeDeltaFeatureConfig(seed_timestamp=500), + ) + + hw_val = delta_graph.node_sets["hardware"] + hw_sch = delta_schema.node_sets["hardware"] + + # Padded sequence: [0, 100, 250, 300] with mask [0, 1, 1, 1] + # Consecutive delta: [0, 0, 150, 50] + # Seed delta (seed=500): [0, 400, 250, 200] + expected_features = { + "time": np.array([[0, 100, 250, 300]], dtype=np.int64), + "time_mask": np.array([[False, True, True, True]]), + "time_consecutive_delta": np.array([[0, 0, 150, 50]], dtype=np.int64), + "time_seed_delta": np.array([[0, 400, 250, 200]], dtype=np.int64), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) + + expected_schemas = { + "time": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(4,), + is_timeseries=True, + timestamps=None, + ), + "time_mask": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.BOOL, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + shape=(4,), + is_timeseries=True, + timestamps=None, + ), + "time_consecutive_delta": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + shape=(4,), + is_timeseries=True, + timestamps="time", + ), + "time_seed_delta": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(4,), + is_timeseries=True, + timestamps="time", + ), + } + test_util.assert_are_equal(self, hw_sch.features, expected_schemas) + + def test_extract_time_delta_features_static_timestamp(self): + # Non-timeseries timestamp feature. + graph, schema = _make_graph_and_schema( + values={"created_at": np.array([65, 1680000015], dtype=np.int64)}, + schemas={ + "created_at": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_timeseries=False, + shape=(), + ) + }, + ) + delta_graph, delta_schema = timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + hw_val = delta_graph.node_sets["hardware"] + hw_sch = delta_schema.node_sets["hardware"] + + # Delta is not required to be positive. + expected_features = { + "created_at": np.array([65, 1680000015], dtype=np.int64), + "created_at_seed_delta": np.array([435, -1679999515], dtype=np.int64), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) + + expected_schemas = { + "created_at": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(), + is_timeseries=False, + timestamps=None, + ), + "created_at_seed_delta": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(), + is_timeseries=False, + timestamps=None, + ), + } + test_util.assert_are_equal(self, hw_sch.features, expected_schemas) + + def test_extract_time_delta_features_parent_timestamp(self): + # Timestamp feature that references other timestamp feature. + graph, schema = _make_graph_and_schema( + values={ + "event_time": np.array([[65, 3665]], dtype=np.int64), + "master_time": np.array([[65, 3665]], dtype=np.int64), + }, + schemas={ + "event_time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + timestamps="master_time", + shape=(2,), + ), + "master_time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + ), + }, + ) + _, delta_schema = timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + hw_sch = delta_schema.node_sets["hardware"] + # Time delta feature derived from event_time inherits + # timestamps="master_time" + self.assertEqual( + hw_sch.features["event_time_consecutive_delta"].timestamps, + "master_time", + ) + # Time delta feature derived from master_time references master_time + self.assertEqual( + hw_sch.features["master_time_consecutive_delta"].timestamps, + "master_time", + ) + + def test_extract_time_delta_features_requires_fixed_length(self): + graph, schema = _make_graph_and_schema( + values={ + "time": np.array( + [np.array([100, 250], dtype=np.int64)], dtype=np.object_ + ) + }, + schemas={ + "time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + ) + }, + ) + with self.assertRaisesRegex( + ValueError, + "extract_time_delta_features requires fixed-length timestamp tensors", + ): + timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + + def test_extract_time_delta_features_edge_sets(self): + graph, schema = _make_graph_and_schema( + values={}, + schemas={}, + node_set_name="nodes", + edge_values={"time": np.array([[100, 250]], dtype=np.int64)}, + edge_schemas={ + "time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + ) + }, + edge_set_name="ts_edges", + ) + # Give it a pre-existing mask like pad_and_cap would + graph.edge_sets["ts_edges"].features["time_mask"] = np.array( + [[1, 1]], dtype=np.bool_ + ) + schema.edge_sets["ts_edges"].features["time_mask"] = _ts_schema( + fmt=schema_lib.FeatureFormat.BOOL, + sem=schema_lib.FeatureSemantic.NUMERICAL, + shape=(2,), + ) + + delta_graph, delta_schema = timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + es_val = delta_graph.edge_sets["ts_edges"] + es_sch = delta_schema.edge_sets["ts_edges"] + + expected_features = { + "time": np.array([[100, 250]], dtype=np.int64), + "time_mask": np.array([[True, True]]), + "time_consecutive_delta": np.array([[0, 150]], dtype=np.int64), + "time_seed_delta": np.array([[400, 250]], dtype=np.int64), + } + test_util.assert_are_equal(self, es_val.features, expected_features) + + expected_schemas = { + "time": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + is_timeseries=True, + timestamps=None, + ), + "time_mask": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.BOOL, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + shape=(2,), + is_timeseries=True, + timestamps=None, + ), + "time_consecutive_delta": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + shape=(2,), + is_timeseries=True, + timestamps="time", + ), + "time_seed_delta": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + is_timeseries=True, + timestamps="time", + ), + } + test_util.assert_are_equal(self, es_sch.features, expected_schemas) + + def test_extract_time_delta_features_non_timeseries(self): + graph, schema = _make_graph_and_schema( + values={"x": np.array([1.0], dtype=np.float32)}, + schemas={ + "x": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + ) + }, + node_set_name="static_nodes", + ) + delta_graph, delta_schema = timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + self.assertIn("static_nodes", delta_graph.node_sets) + self.assertEqual( + delta_schema.node_sets["static_nodes"], schema.node_sets["static_nodes"] + ) + + def test_extract_time_delta_features_missing_mask(self): + graph, schema = _make_graph_and_schema( + values={"time": np.array([[0, 100, 250]], dtype=np.int64)}, + schemas={ + "time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(3,), + ) + }, + ) + # Purposely omit 'time_mask' from values to hit the fallback + delta_graph, _ = timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + hw_val = delta_graph.node_sets["hardware"] + # The implicit mask should be [0, 1, 1] since raw_val != 0 + # Seed delta (seed=500) = (500 - time) * mask + np.testing.assert_array_equal( + hw_val.features["time_seed_delta"][0], [0, 400, 250] + ) + + def test_extract_time_delta_features_custom_fill_value(self): + graph, schema = _make_graph_and_schema( + values={ + "time": np.array([[0, 100, 250]], dtype=np.int64), + "time_mask": np.array([[False, True, True]]), + }, + schemas={ + "time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(3,), + ), + "time_mask": _ts_schema( + fmt=schema_lib.FeatureFormat.BOOL, + sem=schema_lib.FeatureSemantic.NUMERICAL, + shape=(3,), + ), + }, + ) + delta_graph, _ = timeseries.extract_time_delta_features( + graph, + schema, + timeseries.TimeDeltaFeatureConfig(seed_timestamp=500, fill_value=-999), + ) + hw_val = delta_graph.node_sets["hardware"] + expected_features = { + "time": np.array([[0, 100, 250]], dtype=np.int64), + "time_mask": np.array([[False, True, True]]), + "time_consecutive_delta": np.array([[-999, -999, 150]], dtype=np.int64), + "time_seed_delta": np.array([[-999, 400, 250]], dtype=np.int64), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) + + def test_extract_time_delta_features_invalid_semantic(self): + graph, schema = _make_graph_and_schema( + values={ + "not_a_ts": np.array([[1, 2]], dtype=np.int64), + }, + schemas={ + "not_a_ts": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.NUMERICAL, + shape=(2,), + ), + }, + ) + delta_graph, _ = timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + self.assertNotIn( + "not_a_ts_consecutive_delta", delta_graph.node_sets["hardware"].features + ) + + def test_extract_time_delta_features_missing_feature_raises(self): + graph, schema = _make_graph_and_schema( + values={}, + schemas={ + "missing_time": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + ), + }, + ) + with self.assertRaises(KeyError): + timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + + def test_extract_time_delta_features_non_timeseries_edge_sets(self): + graph, schema = _make_graph_and_schema( + values={"x": np.array([1.0], dtype=np.float32)}, + schemas={ + "x": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + ) + }, + edge_values={"y": np.array([2.0], dtype=np.float32)}, + edge_schemas={ + "y": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + ) + }, + edge_set_name="static_edges", + ) + delta_graph, delta_schema = timeseries.extract_time_delta_features( + graph, schema, timeseries.TimeDeltaFeatureConfig(seed_timestamp=500) + ) + self.assertIn("static_edges", delta_graph.edge_sets) + self.assertEqual( + delta_schema.edge_sets["static_edges"], schema.edge_sets["static_edges"] + ) + if __name__ == "__main__": absltest.main()