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/BUILD b/dgf/src/transform/BUILD index cd914b3..c29ff68 100644 --- a/dgf/src/transform/BUILD +++ b/dgf/src/transform/BUILD @@ -253,6 +253,7 @@ py_test( deps = [ ":timeseries", # absl/testing:absltest dep, + # absl/testing:parameterized dep, "//dgf/src/data:in_memory_graph", "//dgf/src/data:schema", "//dgf/src/util:test_util", diff --git a/dgf/src/transform/timeseries.py b/dgf/src/transform/timeseries.py index 8acb370..75f536c 100644 --- a/dgf/src/transform/timeseries.py +++ b/dgf/src/transform/timeseries.py @@ -14,6 +14,10 @@ """Padding and capping for timeseries sequence features in graphs.""" +# TODO(simonmeierhans): Split output schema computation (done once at +# preparation) and output value computation (done per node/execution) for +# all functions. + # pytype: disable=module-attr import dataclasses import enum @@ -141,6 +145,22 @@ class CalendarFeatureConfig: features: Tuple[CalendarFeature, ...] = _SUPPORTED_CALENDAR_FEATURES +@dataclasses_json.dataclass_json +@dataclasses.dataclass +class TimeDeltaFeatureConfig: + """Configuration for extracting time delta features. + + Attributes: + 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. + """ + + extract_consecutive: bool = True + extract_seed_delta: bool = True + fill_value: Any = 0 + + def _compute_calendar_feature( ts_array: np.ndarray, feature: CalendarFeature ) -> np.ndarray: @@ -211,6 +231,7 @@ def _process_feature_set( new_values[fname] = padded_matrix new_schemas[fname] = temporal_util.with_sequence_length(schema, seq_len) + # TODO(simonmeierhans): Remove reliance on naming. new_values[f"{fname}_mask"] = mask_matrix new_schemas[f"{fname}_mask"] = schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.BOOL, @@ -338,24 +359,24 @@ def _extract_feature_set_calendar_features( if schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP: continue - if raw_val.dtype == np.object_: + if not schema.is_static_shape() or raw_val.dtype == np.object_: raise ValueError( "extract_calendar_features requires fixed-length timestamp tensors," - f" but feature '{fname}' is a variable-length object array." - " Please run pad_and_cap_timeseries_features first." + f" but feature '{fname}' is a variable-length object array or has a" + f" dynamic shape ({schema.shape}). Please run" + " pad_and_cap_timeseries_features first." ) - # Case 1: For non-timeseries timestamps, derived calendar features do not - # reference a timestamp sequence. - if not schema.is_timeseries: + if schema.is_timeseries is None or not schema.is_timeseries: + # Non-timeseries features do not have associated timestamp sequences. calendar_timestamps = None - # Case 2: For timeseries features that reference a timestamp sequence. elif schema.timestamps is not None: + # This feature belongs to a timeseries group; reference the group's + # master timestamp feature name (e.g., "master_time"). calendar_timestamps = schema.timestamps - # Case 3: For timeseries features that do not reference a timestamp - # sequence, derived calendar sequences reference the original feature - # as their timestamp sequence. else: + # This feature is itself a sequence of timestamps; derived calendar + # features inherit this feature's name as their timestamp sequence. calendar_timestamps = fname for cal_feat in config.features: @@ -448,3 +469,197 @@ def extract_calendar_features( node_sets=new_ns_schemas, edge_sets=new_es_schemas ), ) + + +def _compute_consecutive_deltas( + raw_val: np.ndarray, + mask_matrix: Optional[np.ndarray], + fill_value: Any, +) -> np.ndarray: + """Computes consecutive timestamp differences (t_i - t_{i-1}) with padding.""" + if mask_matrix is not None: + pair_mask = mask_matrix[:, :-1] & mask_matrix[:, 1:] + valid_diffs = np.where( + pair_mask, raw_val[:, 1:] - raw_val[:, :-1], fill_value + ) + else: + valid_diffs = raw_val[:, 1:] - raw_val[:, :-1] + return np.hstack([ + np.full((raw_val.shape[0], 1), fill_value, dtype=np.int64), + valid_diffs, + ]) + + +def _compute_seed_deltas( + raw_val: np.ndarray, + mask_matrix: Optional[np.ndarray], + seed_timestamp: int, + fill_value: Any, +) -> np.ndarray: + """Computes differences to a seed timestamp with masking.""" + if mask_matrix is not None: + return np.where(mask_matrix, seed_timestamp - raw_val, fill_value) + return seed_timestamp - raw_val + + +def _extract_feature_set_time_delta_features( + values: in_memory_graph.Features, + schemas: schema_lib.FeatureSetSchema, + config: TimeDeltaFeatureConfig, + seed_timestamp: Optional[int] = None, +) -> 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 not schema.is_static_shape() or 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" + f" or has a dynamic shape ({schema.shape}). Please run" + " pad_and_cap_timeseries_features first." + ) + + # TODO(simonmeierhans): Remove reliance on naming. + mask_fname = f"{fname}_mask" + if mask_fname in values: + mask_matrix = values[mask_fname].astype(np.bool_) + else: + mask_matrix = None + + if schema.is_timeseries is None or not schema.is_timeseries: + # Non-timeseries features do not have associated timestamp sequences. + delta_timestamps = None + elif schema.timestamps is not None: + # This feature belongs to a timeseries group; reference the group's + # master timestamp feature name (e.g., "master_time"). + delta_timestamps = schema.timestamps + else: + # This feature is itself a sequence of timestamps; derived time delta + # features inherit this feature's name as their timestamp sequence. + 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: + consecutive_delta = _compute_consecutive_deltas( + raw_val, mask_matrix, config.fill_value + ) + 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 = _compute_seed_deltas( + raw_val, mask_matrix, seed_timestamp, 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, + seed_timestamp: Optional[int] = None, +) -> 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(extract_seed_delta=True) + graph, schema = dgf.transform.extract_time_delta_features( + graph, schema, config, seed_timestamp=1680000000 + ) + ``` + + Args: + graph: The input in-memory graph. + schema: The graph schema containing timestamp features. + config: `TimeDeltaFeatureConfig` specifying extraction flags and fill value. + seed_timestamp: Target reference timestamp to compute absolute deltas from. + Required when `config.extract_seed_delta` is True. + + Returns: + Tuple `(new_graph, new_schema)` containing original and extracted delta + features. + """ + if config.extract_seed_delta and seed_timestamp is None: + raise ValueError( + "seed_timestamp must be provided when config.extract_seed_delta is" + " True." + ) + + 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, + seed_timestamp=seed_timestamp, + ) + 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, + seed_timestamp=seed_timestamp, + ) + 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..3e7f734 100644 --- a/dgf/src/transform/timeseries_test.py +++ b/dgf/src/transform/timeseries_test.py @@ -15,6 +15,7 @@ """Tests for padding and capping timeseries sequence features.""" from absl.testing import absltest +from absl.testing import parameterized from dgf.src.data import in_memory_graph from dgf.src.data import schema as schema_lib from dgf.src.transform import timeseries @@ -75,7 +76,7 @@ def _ts_schema( ) -class TimeseriesTest(absltest.TestCase): +class TimeseriesTest(parameterized.TestCase): def test_capping_and_padding(self): # Both 'time' and 'signal' are variable-length sequence object arrays. @@ -664,6 +665,326 @@ 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): + 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=(), + ) + } + new_vals, new_schs = timeseries._extract_feature_set_time_delta_features( + values, schemas, timeseries.TimeDeltaFeatureConfig(), seed_timestamp=500 + ) + np.testing.assert_array_equal( + new_vals["created_at_seed_delta"], [435, -1679999515] + ) + self.assertFalse(new_schs["created_at_seed_delta"].is_timeseries) + + def test_extract_time_delta_features_parent_timestamp(self): + values = { + "event_time": np.array([[65, 3665]], dtype=np.int64), + "master_time": np.array([[65, 3665]], dtype=np.int64), + } + schemas = { + "event_time": _ts_schema( + sem=schema_lib.FeatureSemantic.TIMESTAMP, + timestamps="master_time", + shape=(2,), + ), + "master_time": _ts_schema( + sem=schema_lib.FeatureSemantic.TIMESTAMP, shape=(2,) + ), + } + _, new_schs = timeseries._extract_feature_set_time_delta_features( + values, schemas, timeseries.TimeDeltaFeatureConfig(), seed_timestamp=500 + ) + self.assertEqual( + new_schs["event_time_consecutive_delta"].timestamps, "master_time" + ) + self.assertEqual( + new_schs["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", + ) + 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"] + + self.assertIn("time_consecutive_delta", es_val.features) + self.assertIn("time_seed_delta", es_val.features) + self.assertEqual( + es_sch.features["time_consecutive_delta"].timestamps, "time" + ) + self.assertEqual(es_sch.features["time_seed_delta"].timestamps, "time") + + def test_extract_time_delta_features_non_timeseries(self): + 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, + ) + } + new_vals, new_schs = timeseries._extract_feature_set_time_delta_features( + values, schemas, timeseries.TimeDeltaFeatureConfig(), seed_timestamp=500 + ) + self.assertEqual(new_vals, values) + self.assertEqual(new_schs, schemas) + + @parameterized.parameters( + (np.array([[False, True, True]]), 0, [[0, 0, 150]]), + (None, 0, [[0, 100, 150]]), + (np.array([[False, True, True]]), -999, [[-999, -999, 150]]), + ) + def test_compute_consecutive_deltas(self, mask, fill_value, expected): + raw_val = np.array([[0, 100, 250]], dtype=np.int64) + deltas = timeseries._compute_consecutive_deltas(raw_val, mask, fill_value) + np.testing.assert_array_equal(deltas, expected) + + @parameterized.parameters( + (np.array([[False, True, True]]), 0, [[0, 400, 250]]), + (None, 0, [[500, 400, 250]]), + (np.array([[False, True, True]]), -999, [[-999, 400, 250]]), + ) + def test_compute_seed_deltas(self, mask, fill_value, expected): + raw_val = np.array([[0, 100, 250]], dtype=np.int64) + deltas = timeseries._compute_seed_deltas(raw_val, mask, 500, fill_value) + np.testing.assert_array_equal(deltas, expected) + + 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"] + ) + + def test_extract_time_delta_features_seed_timestamp_validation(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,), + ) + }, + ) + with self.assertRaisesRegex( + ValueError, + "seed_timestamp must be provided when config.extract_seed_delta is" + " True.", + ): + timeseries.extract_time_delta_features( + graph, + schema, + timeseries.TimeDeltaFeatureConfig(extract_seed_delta=True), + ) + + # When extract_seed_delta is False, omitting seed_timestamp is fine. + delta_graph, _ = timeseries.extract_time_delta_features( + graph, + schema, + timeseries.TimeDeltaFeatureConfig(extract_seed_delta=False), + ) + self.assertNotIn( + "time_seed_delta", delta_graph.node_sets["hardware"].features + ) + self.assertIn( + "time_consecutive_delta", delta_graph.node_sets["hardware"].features + ) + if __name__ == "__main__": absltest.main()