diff --git a/dgf/src/api/transform.py b/dgf/src/api/transform.py index db92eb6..e078ed9 100644 --- a/dgf/src/api/transform.py +++ b/dgf/src/api/transform.py @@ -52,4 +52,6 @@ from dgf.src.transform.timeseries import pad_and_cap_timeseries_features from dgf.src.transform.timeseries import PadAndCapTimeseriesConfig +from dgf.src.transform.timeseries import extract_calendar_features +from dgf.src.transform.timeseries import CalendarFeatureConfig diff --git a/dgf/src/transform/BUILD b/dgf/src/transform/BUILD index ed8de60..cd914b3 100644 --- a/dgf/src/transform/BUILD +++ b/dgf/src/transform/BUILD @@ -255,6 +255,7 @@ py_test( # absl/testing:absltest dep, "//dgf/src/data:in_memory_graph", "//dgf/src/data:schema", + "//dgf/src/util:test_util", # numpy dep, ], ) diff --git a/dgf/src/transform/timeseries.py b/dgf/src/transform/timeseries.py index 1f695d9..bb7bf44 100644 --- a/dgf/src/transform/timeseries.py +++ b/dgf/src/transform/timeseries.py @@ -16,6 +16,7 @@ # pytype: disable=module-attr import dataclasses +import enum from typing import Any, List, Optional, Tuple import dataclasses_json @@ -112,15 +113,70 @@ def _pad_and_cap_single_feature( return padded_matrix, mask_matrix +class CalendarFeature(str, enum.Enum): + """Supported calendar features to extract from timestamps.""" + + SECOND = "second" + MINUTE = "minute" + HOUR = "hour" + DAY_OF_WEEK = "day_of_week" + MONTH = "month" + YEAR = "year" + + +_SUPPORTED_CALENDAR_FEATURES = tuple(CalendarFeature) + + +@dataclasses_json.dataclass_json +@dataclasses.dataclass +class CalendarFeatureConfig: + """Configuration for extracting calendar features from timestamps. + + Attributes: + features: Tuple of calendar feature enums to extract. Supported values: + CalendarFeature.SECOND, CalendarFeature.MINUTE, CalendarFeature.HOUR, + CalendarFeature.DAY_OF_WEEK, CalendarFeature.MONTH, CalendarFeature.YEAR. + """ + + features: Tuple[CalendarFeature, ...] = _SUPPORTED_CALENDAR_FEATURES + + +def _compute_calendar_feature( + ts_array: np.ndarray, feature: CalendarFeature +) -> np.ndarray: + """Computes a single vectorized calendar feature from an int64 timestamp array.""" + + if feature == CalendarFeature.SECOND: + return (ts_array % 60).astype(np.float32) + if feature == CalendarFeature.MINUTE: + return ((ts_array // 60) % 60).astype(np.float32) + if feature == CalendarFeature.HOUR: + return ((ts_array // 3600) % 24).astype(np.float32) + if feature == CalendarFeature.DAY_OF_WEEK: + return (((ts_array // 86400) + 3) % 7).astype(np.float32) + + dt = ts_array.astype("datetime64[s]") + + if feature == CalendarFeature.MONTH: + return (dt.astype("datetime64[M]").astype(int) % 12 + 1).astype(np.float32) + if feature == CalendarFeature.YEAR: + return (dt.astype("datetime64[Y]").astype(int) + 1970).astype(np.float32) + + raise ValueError( + f"Unsupported calendar feature: '{feature}'. Supported features:" + f" {[f.value for f in _SUPPORTED_CALENDAR_FEATURES]}" + ) + + def _process_feature_set( - features: in_memory_graph.Features, - feature_schemas: schema_lib.FeatureSetSchema, + values: in_memory_graph.Features, + schemas: schema_lib.FeatureSetSchema, ts_specs: List[temporal_util.TimeseriesGroupSpec], config: PadAndCapTimeseriesConfig, ) -> Tuple[in_memory_graph.Features, schema_lib.FeatureSetSchema]: """Extracts fixed-dimension sequence features for a feature set.""" - new_features: in_memory_graph.Features = {} - new_feat_schemas: schema_lib.FeatureSetSchema = {} + new_values: in_memory_graph.Features = {} + new_schemas: schema_lib.FeatureSetSchema = {} seq_len = config.sequence_length # Map timeseries feature names to their associated timestamp feature name. @@ -129,46 +185,42 @@ def _process_feature_set( for fname in group.feature_names: ts_features[fname] = group.timestamp_feature_name - for fname, fschema in feature_schemas.items(): + for fname, schema in schemas.items(): if fname not in ts_features: - new_features[fname] = features[fname] - new_feat_schemas[fname] = fschema + new_values[fname] = values[fname] + new_schemas[fname] = schema if not ts_features: - return new_features, new_feat_schemas + return new_values, new_schemas for fname in ts_features: - fschema = feature_schemas[fname] + schema = schemas[fname] - dtype = feature_format.FEATURE_FORMAT_TO_NP_DTYPE[fschema.format] - feat_shape = fschema.shape[1:] if fschema.shape is not None else () + dtype = feature_format.FEATURE_FORMAT_TO_NP_DTYPE[schema.format] + feat_shape = temporal_util.get_timeseries_step_shape(schema) padded_matrix, mask_matrix = _pad_and_cap_single_feature( - raw_series=features[fname], + raw_series=values[fname], seq_len=seq_len, feat_shape=feat_shape, padding_value=config.padding_value, dtype=dtype, - is_static_shape=fschema.is_static_shape(), + is_static_shape=schema.is_static_shape(), ) - new_features[fname] = padded_matrix - new_feat_schemas[fname] = dataclasses.replace( - fschema, - shape=(seq_len,) + feat_shape, - is_timeseries=True, - ) + new_values[fname] = padded_matrix + new_schemas[fname] = temporal_util.with_sequence_length(schema, seq_len) - new_features[f"{fname}_mask"] = mask_matrix - new_feat_schemas[f"{fname}_mask"] = schema_lib.FeatureSchema( + new_values[f"{fname}_mask"] = mask_matrix + new_schemas[f"{fname}_mask"] = schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.BOOL, semantic=schema_lib.FeatureSemantic.NUMERICAL, shape=(seq_len,) + feat_shape, - is_timeseries=fschema.is_timeseries, - timestamps=fschema.timestamps, + is_timeseries=schema.is_timeseries, + timestamps=schema.timestamps, ) - return new_features, new_feat_schemas + return new_values, new_schemas def pad_and_cap_timeseries_features( @@ -221,14 +273,14 @@ def pad_and_cap_timeseries_features( new_ns_schemas[ns_name] = ns_schema continue - new_feats, new_schemas = _process_feature_set( - features=ns_val.features, - feature_schemas=ns_schema.features, + new_vals, new_schemas = _process_feature_set( + values=ns_val.features, + schemas=ns_schema.features, ts_specs=ts_specs, config=config, ) new_node_sets[ns_name] = in_memory_graph.InMemoryNodeSet( - num_nodes=ns_val.num_nodes, features=new_feats + num_nodes=ns_val.num_nodes, features=new_vals ) new_ns_schemas[ns_name] = schema_lib.NodeSchema(features=new_schemas) @@ -243,14 +295,148 @@ def pad_and_cap_timeseries_features( new_es_schemas[es_name] = es_schema continue - new_feats, new_schemas = _process_feature_set( - features=es_val.features, - feature_schemas=es_schema.features, + new_vals, new_schemas = _process_feature_set( + values=es_val.features, + schemas=es_schema.features, ts_specs=ts_specs, config=config, ) new_edge_sets[es_name] = in_memory_graph.InMemoryEdgeSet( - adjacency=es_val.adjacency, features=new_feats + 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 + ), + ) + + +def _extract_feature_set_calendar_features( + values: in_memory_graph.Features, + schemas: schema_lib.FeatureSetSchema, + config: CalendarFeatureConfig, +) -> Tuple[in_memory_graph.Features, schema_lib.FeatureSetSchema]: + """Extracts calendar 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(): + if fname not in values: + raise KeyError( + f"Feature '{fname}' defined in schema but missing from feature dict." + ) + 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_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." + ) + + # Case 1: For non-timeseries timestamps, derived calendar features do not + # reference a timestamp sequence. + if not schema.is_timeseries: + calendar_timestamps = None + # Case 2: For timeseries features that reference a timestamp sequence. + elif schema.timestamps is not None: + 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: + calendar_timestamps = fname + + for cal_feat in config.features: + out_fname = f"{fname}_{cal_feat.value}" + cal_arr = _compute_calendar_feature(raw_val, cal_feat) + new_values[out_fname] = cal_arr + new_schemas[out_fname] = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + shape=schema.shape, + is_timeseries=schema.is_timeseries, + timestamps=calendar_timestamps, + ) + + return new_values, new_schemas + + +def extract_calendar_features( + graph: in_memory_graph.InMemoryGraph, + schema: schema_lib.GraphSchema, + config: Optional[CalendarFeatureConfig] = None, +) -> Tuple[in_memory_graph.InMemoryGraph, schema_lib.GraphSchema]: + """Extracts calendar features (e.g. hour, day_of_week) from timestamp features. + + Requires fixed-length timestamp tensors (e.g., produced after running + `pad_and_cap_timeseries_features`). + + Usage example: + + ```python + graph, schema = dgf.transform.pad_and_cap_timeseries_features( + graph, schema, cap_config + ) + graph, schema = dgf.transform.extract_calendar_features(graph, schema) + ``` + + Args: + graph: The input in-memory graph. + schema: The graph schema containing timestamp features. + config: Optional `CalendarFeatureConfig`. + + Returns: + Tuple `(new_graph, new_schema)` containing original and extracted calendar + features. + """ + + # Default to extracting all calendar features + if config is None: + config = CalendarFeatureConfig() + + 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_calendar_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_calendar_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, diff --git a/dgf/src/transform/timeseries_test.py b/dgf/src/transform/timeseries_test.py index 5e2c939..c93a4da 100644 --- a/dgf/src/transform/timeseries_test.py +++ b/dgf/src/transform/timeseries_test.py @@ -14,43 +14,48 @@ """Tests for padding and capping timeseries sequence features.""" -from typing import Any from absl.testing import absltest from dgf.src.data import in_memory_graph from dgf.src.data import schema as schema_lib from dgf.src.transform import timeseries +from dgf.src.util import test_util import numpy as np def _make_graph_and_schema( - features: dict[str, Any], - feature_schemas: dict[str, schema_lib.FeatureSchema], + values: dict[str, np.ndarray], + schemas: dict[str, schema_lib.FeatureSchema], node_set_name: str = "hardware", num_nodes: int = 1, + edge_values: dict[str, np.ndarray] | None = None, + edge_schemas: dict[str, schema_lib.FeatureSchema] | None = None, + edge_set_name: str = "edges", ) -> tuple[in_memory_graph.InMemoryGraph, schema_lib.GraphSchema]: - np_features = {} - for k, v in features.items(): - if isinstance(v, np.ndarray): - np_features[k] = v - else: - try: - np_features[k] = np.asarray(v) - except ValueError: - np_features[k] = np.asarray(v, dtype=np.object_) + edge_sets = {} + edge_set_schemas = {} + if edge_values is not None and edge_schemas is not None: + edge_sets[edge_set_name] = in_memory_graph.InMemoryEdgeSet( + adjacency=np.array([[0], [0]]), + features=dict(edge_values), + ) + edge_set_schemas[edge_set_name] = schema_lib.EdgeSchema( + source=node_set_name, + target=node_set_name, + features=edge_schemas, + ) + return ( in_memory_graph.InMemoryGraph( node_sets={ node_set_name: in_memory_graph.InMemoryNodeSet( - num_nodes=num_nodes, features=np_features + num_nodes=num_nodes, features=dict(values) ) }, - edge_sets={}, + edge_sets=edge_sets, ), schema_lib.GraphSchema( - node_sets={ - node_set_name: schema_lib.NodeSchema(features=feature_schemas) - }, - edge_sets={}, + node_sets={node_set_name: schema_lib.NodeSchema(features=schemas)}, + edge_sets=edge_set_schemas, ), ) @@ -59,7 +64,7 @@ def _ts_schema( fmt: schema_lib.FeatureFormat = schema_lib.FeatureFormat.FLOAT_32, sem: schema_lib.FeatureSemantic = schema_lib.FeatureSemantic.NUMERICAL, timestamps: str | None = None, - shape: schema_lib.Shape = None, + shape: schema_lib.Shape = (None,), ) -> schema_lib.FeatureSchema: return schema_lib.FeatureSchema( format=fmt, @@ -73,17 +78,23 @@ def _ts_schema( class TimeseriesTest(absltest.TestCase): def test_capping_and_padding(self): - # 'time' is np.ndarray, 'signal' is raw Python list (tests both input types) + # Both 'time' and 'signal' are variable-length sequence object arrays. graph, schema = _make_graph_and_schema( - features={ + values={ "time": np.array( [np.array([10, 20, 30, 40, 50]), np.array([5, 15])], dtype=np.object_, ), - "signal": [[1.0, 2.0, 3.0, 4.0, 5.0], [0.5, 1.5]], + "signal": np.array( + [ + np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32), + np.array([0.5, 1.5], dtype=np.float32), + ], + dtype=np.object_, + ), "id": np.array([101, 102]), }, - feature_schemas={ + schemas={ "time": _ts_schema( fmt=schema_lib.FeatureFormat.INTEGER_64, sem=schema_lib.FeatureSemantic.TIMESTAMP, @@ -104,18 +115,16 @@ def test_capping_and_padding(self): hw_val = new_graph.node_sets["hardware"] hw_sch = new_schema.node_sets["hardware"] - # Node 0 (5 steps > K=3): capped to last 3 steps [30, 40, 50] - np.testing.assert_array_equal(hw_val.features["time"][0], [30, 40, 50]) - np.testing.assert_array_equal(hw_val.features["signal"][0], [3.0, 4.0, 5.0]) - np.testing.assert_array_equal(hw_val.features["time_mask"][0], [1, 1, 1]) - - # Node 1 (2 steps < K=3): left-padded with 0 to [0, 5, 15] - np.testing.assert_array_equal(hw_val.features["time"][1], [0, 5, 15]) - np.testing.assert_array_equal(hw_val.features["signal"][1], [0.0, 0.5, 1.5]) - np.testing.assert_array_equal(hw_val.features["time_mask"][1], [0, 1, 1]) - - # Static non-timeseries feature preserved unchanged - np.testing.assert_array_equal(hw_val.features["id"], [101, 102]) + expected_features = { + "time": np.array([[30, 40, 50], [0, 5, 15]], dtype=np.int64), + "signal": np.array( + [[3.0, 4.0, 5.0], [0.0, 0.5, 1.5]], dtype=np.float32 + ), + "signal_mask": np.array([[True, True, True], [False, True, True]]), + "time_mask": np.array([[True, True, True], [False, True, True]]), + "id": np.array([101, 102]), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) self.assertEqual(hw_sch.features["time"].shape, (3,)) self.assertTrue(hw_sch.features["time"].is_timeseries) @@ -193,13 +202,13 @@ def test_edge_sets_and_non_timeseries(self): def test_multidimensional_sequence(self): graph, schema = _make_graph_and_schema( - features={ + values={ "emb": np.array( [np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)], dtype=np.object_, ) }, - feature_schemas={ + schemas={ "emb": _ts_schema( sem=schema_lib.FeatureSemantic.EMBEDDING, shape=(None, 2), @@ -214,21 +223,20 @@ def test_multidimensional_sequence(self): hw_val = new_graph.node_sets["hardware"] hw_sch = new_schema.node_sets["hardware"] - expected = np.array( - [[[0.0, 0.0], [1.0, 2.0], [3.0, 4.0]]], dtype=np.float32 - ) - np.testing.assert_array_equal(hw_val.features["emb"], expected) - expected_mask = np.array( - [[[0.0, 0.0], [1.0, 1.0], [1.0, 1.0]]], dtype=np.float32 - ) - np.testing.assert_array_equal(hw_val.features["emb_mask"], expected_mask) + expected_features = { + "emb": np.array( + [[[0.0, 0.0], [1.0, 2.0], [3.0, 4.0]]], dtype=np.float32 + ), + "emb_mask": np.array([[[False, False], [True, True], [True, True]]]), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) self.assertEqual(hw_sch.features["emb"].shape, (3, 2)) self.assertEqual(hw_sch.features["emb_mask"].shape, (3, 2)) def test_empty_sequence(self): graph, schema = _make_graph_and_schema( - features={"time": np.array([np.array([])], dtype=np.object_)}, - feature_schemas={ + values={"time": np.array([np.array([])], dtype=np.object_)}, + schemas={ "time": _ts_schema( fmt=schema_lib.FeatureFormat.INTEGER_64, sem=schema_lib.FeatureSemantic.TIMESTAMP, @@ -246,8 +254,8 @@ def test_empty_sequence(self): def test_no_timeseries_in_graph(self): graph, schema = _make_graph_and_schema( - features={"age": np.array([30])}, - feature_schemas={ + values={"age": np.array([30])}, + schemas={ "age": schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.INTEGER_64, semantic=schema_lib.FeatureSemantic.NUMERICAL, @@ -313,14 +321,37 @@ def test_missing_entity_set_raises(self): timeseries.pad_and_cap_timeseries_features( graph, schema, timeseries.PadAndCapTimeseriesConfig(sequence_length=3) ) + with self.assertRaises(KeyError): + timeseries.extract_calendar_features(graph, schema) + + def test_missing_feature_raises(self): + graph, schema = _make_graph_and_schema( + values={}, + schemas={ + "absent_feature": _ts_schema( + fmt=schema_lib.FeatureFormat.INTEGER_64, + sem=schema_lib.FeatureSemantic.TIMESTAMP, + ) + }, + ) + with self.assertRaises(KeyError): + timeseries.pad_and_cap_timeseries_features( + graph, + schema, + timeseries.PadAndCapTimeseriesConfig(sequence_length=3), + ) + with self.assertRaises(KeyError): + timeseries.extract_calendar_features(graph, schema) def test_custom_padding_value(self): graph, schema = _make_graph_and_schema( - features={ + values={ "time": np.array([np.array([10])], dtype=np.object_), - "signal": [[2.0]], + "signal": np.array( + [np.array([2.0], dtype=np.float32)], dtype=np.object_ + ), }, - feature_schemas={ + schemas={ "time": _ts_schema( fmt=schema_lib.FeatureFormat.INTEGER_64, sem=schema_lib.FeatureSemantic.TIMESTAMP, @@ -337,22 +368,24 @@ def test_custom_padding_value(self): ), ) hw_val = new_graph.node_sets["hardware"] - np.testing.assert_array_equal(hw_val.features["time"][0], [-1, -1, 10]) - np.testing.assert_array_equal( - hw_val.features["signal"][0], [-1.0, -1.0, 2.0] - ) - np.testing.assert_array_equal(hw_val.features["time_mask"][0], [0, 0, 1]) + expected_features = { + "time": np.array([[-1, -1, 10]], dtype=np.int64), + "signal": np.array([[-1.0, -1.0, 2.0]], dtype=np.float32), + "signal_mask": np.array([[False, False, True]]), + "time_mask": np.array([[False, False, True]]), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) def test_fixed_shape_vectorized_path_capping(self): # Dense 2D array where all 2 nodes have fixed length T=5 >= K=3 graph, schema = _make_graph_and_schema( - features={ + values={ "time": np.array( [[10, 20, 30, 40, 50], [100, 200, 300, 400, 500]], dtype=np.int64, ), }, - feature_schemas={ + schemas={ "time": _ts_schema( fmt=schema_lib.FeatureFormat.INTEGER_64, sem=schema_lib.FeatureSemantic.TIMESTAMP, @@ -366,18 +399,17 @@ def test_fixed_shape_vectorized_path_capping(self): timeseries.PadAndCapTimeseriesConfig(sequence_length=3), ) hw_val = new_graph.node_sets["hardware"] - np.testing.assert_array_equal( - hw_val.features["time"], [[30, 40, 50], [300, 400, 500]] - ) - np.testing.assert_array_equal( - hw_val.features["time_mask"], [[1, 1, 1], [1, 1, 1]] - ) + expected_features = { + "time": np.array([[30, 40, 50], [300, 400, 500]], dtype=np.int64), + "time_mask": np.array([[True, True, True], [True, True, True]]), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) def test_fixed_shape_vectorized_path_padding(self): # Dense 3D array where all 2 nodes have fixed length T=2 < K=3 and feature # dim 2 graph, schema = _make_graph_and_schema( - features={ + values={ "emb": np.array( [ [[1.0, 1.1], [2.0, 2.2]], @@ -386,7 +418,7 @@ def test_fixed_shape_vectorized_path_padding(self): dtype=np.float32, ), }, - feature_schemas={ + schemas={ "emb": _ts_schema( sem=schema_lib.FeatureSemantic.EMBEDDING, shape=(None, 2), @@ -402,22 +434,235 @@ def test_fixed_shape_vectorized_path_padding(self): ), ) hw_val = new_graph.node_sets["hardware"] - expected_emb = np.array( - [ - [[-1.0, -1.0], [1.0, 1.1], [2.0, 2.2]], - [[-1.0, -1.0], [10.0, 10.1], [20.0, 20.2]], - ], - dtype=np.float32, - ) - expected_mask = np.array( - [ - [[0.0, 0.0], [1.0, 1.0], [1.0, 1.0]], - [[0.0, 0.0], [1.0, 1.0], [1.0, 1.0]], - ], - dtype=np.float32, - ) - np.testing.assert_array_equal(hw_val.features["emb"], expected_emb) - np.testing.assert_array_equal(hw_val.features["emb_mask"], expected_mask) + expected_features = { + "emb": np.array( + [ + [[-1.0, -1.0], [1.0, 1.1], [2.0, 2.2]], + [[-1.0, -1.0], [10.0, 10.1], [20.0, 20.2]], + ], + dtype=np.float32, + ), + "emb_mask": np.array([ + [[False, False], [True, True], [True, True]], + [[False, False], [True, True], [True, True]], + ]), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) + + def test_compute_calendar_feature(self): + ts = np.array([65, 3665, 1680000015], dtype=np.int64) + computed = { + feat: timeseries._compute_calendar_feature(ts, feat) + for feat in timeseries._SUPPORTED_CALENDAR_FEATURES + } + expected = { + timeseries.CalendarFeature.SECOND: np.array( + [5.0, 5.0, 15.0], dtype=np.float32 + ), + timeseries.CalendarFeature.MINUTE: np.array( + [1.0, 1.0, 40.0], dtype=np.float32 + ), + timeseries.CalendarFeature.HOUR: np.array( + [0.0, 1.0, 10.0], dtype=np.float32 + ), + timeseries.CalendarFeature.DAY_OF_WEEK: np.array( + [3.0, 3.0, 1.0], dtype=np.float32 + ), + timeseries.CalendarFeature.MONTH: np.array( + [1.0, 1.0, 3.0], dtype=np.float32 + ), + timeseries.CalendarFeature.YEAR: np.array( + [1970.0, 1970.0, 2023.0], dtype=np.float32 + ), + } + test_util.assert_are_equal(self, computed, expected) + + def test_extract_calendar_features(self): + graph, schema = _make_graph_and_schema( + values={ + "time": np.array( + [np.array([65, 1680000015], 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=2), + ) + cal_graph, cal_schema = timeseries.extract_calendar_features( + padded_graph, padded_schema + ) + + hw_val = cal_graph.node_sets["hardware"] + hw_sch = cal_schema.node_sets["hardware"] + + for cal_k in ( + "second", + "minute", + "hour", + "day_of_week", + "month", + "year", + ): + feature_name = f"time_{cal_k}" + self.assertIn(feature_name, hw_val.features) + fschema = hw_sch.features[feature_name] + self.assertEqual( + fschema.semantic, + schema_lib.FeatureSemantic.NUMERICAL, + ) + self.assertEqual(fschema.shape, (2,)) + self.assertTrue(fschema.is_timeseries) + self.assertEqual(fschema.timestamps, "time") + + # 65 -> 1970-01-01 00:01:05 UTC (Thursday=3) + expected_features = { + "time": hw_val.features["time"], + "time_mask": np.array([[True, True]]), + "time_second": np.array([[5.0, 15.0]], dtype=np.float32), + "time_minute": np.array([[1.0, 40.0]], dtype=np.float32), + "time_hour": np.array([[0.0, 10.0]], dtype=np.float32), + "time_day_of_week": np.array([[3.0, 1.0]], dtype=np.float32), + "time_month": np.array([[1.0, 3.0]], dtype=np.float32), + "time_year": np.array([[1970.0, 2023.0]], dtype=np.float32), + } + test_util.assert_are_equal(self, hw_val.features, expected_features) + + def test_extract_calendar_features_requires_fixed_length(self): + graph, schema = _make_graph_and_schema( + values={ + "time": np.array( + [np.array([65, 1680000015], 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_calendar_features requires fixed-length timestamp tensors", + ): + timeseries.extract_calendar_features(graph, schema) + + def test_extract_calendar_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", + edge_values={"weight": np.array([0.5], dtype=np.float32)}, + edge_schemas={ + "weight": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + ) + }, + edge_set_name="static_edges", + ) + cal_graph, cal_schema = timeseries.extract_calendar_features(graph, schema) + self.assertIn("static_nodes", cal_graph.node_sets) + self.assertIn("static_edges", cal_graph.edge_sets) + self.assertEqual( + cal_schema.node_sets["static_nodes"], schema.node_sets["static_nodes"] + ) + self.assertEqual( + cal_schema.edge_sets["static_edges"], schema.edge_sets["static_edges"] + ) + + def test_extract_calendar_features_edge_sets(self): + graph, schema = _make_graph_and_schema( + values={}, + schemas={}, + node_set_name="nodes", + edge_values={"time": np.array([[65, 1680000015]], 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", + ) + cal_graph, cal_schema = timeseries.extract_calendar_features(graph, schema) + self.assertIn("time_hour", cal_graph.edge_sets["ts_edges"].features) + self.assertIn("time_second", cal_graph.edge_sets["ts_edges"].features) + np.testing.assert_array_equal( + cal_graph.edge_sets["ts_edges"].features["time_second"][0], [5, 15] + ) + self.assertIn("time_hour", cal_schema.edge_sets["ts_edges"].features) + + def test_extract_calendar_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=(), + ) + }, + ) + cal_graph, cal_schema = timeseries.extract_calendar_features(graph, schema) + hw_val = cal_graph.node_sets["hardware"] + hw_sch = cal_schema.node_sets["hardware"] + + self.assertIn("created_at_hour", hw_val.features) + fschema = hw_sch.features["created_at_hour"] + self.assertFalse(fschema.is_timeseries) + self.assertIsNone(fschema.timestamps) + np.testing.assert_array_equal( + hw_val.features["created_at_hour"], [0.0, 10.0] + ) + + def test_extract_calendar_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,), + ), + }, + ) + _, cal_schema = timeseries.extract_calendar_features(graph, schema) + hw_sch = cal_schema.node_sets["hardware"] + # Calendar feature derived from event_time inherits timestamps="master_time" + self.assertEqual( + hw_sch.features["event_time_hour"].timestamps, "master_time" + ) + # Calendar feature derived from master_time references master_time + self.assertEqual( + hw_sch.features["master_time_hour"].timestamps, "master_time" + ) if __name__ == "__main__": diff --git a/dgf/src/util/temporal.py b/dgf/src/util/temporal.py index 8b25fb9..ee98d5c 100644 --- a/dgf/src/util/temporal.py +++ b/dgf/src/util/temporal.py @@ -16,7 +16,7 @@ import collections import dataclasses -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple from dgf.src.data import schema as schema_lib @@ -88,3 +88,30 @@ def extract_timeseries_schema_cache( edge_sets=edge_sets_cache, has_timeseries=has_ts, ) + + +def get_timeseries_step_shape( + fschema: schema_lib.FeatureSchema, +) -> Tuple[Optional[int], ...]: + """Returns per-step feature dimension, excluding leading sequence length shape[0].""" + if not fschema.is_timeseries: + raise ValueError("Feature schema must be a timeseries feature.") + elif fschema.shape is None or not fschema.shape: + raise ValueError( + "Timeseries feature schema must have at least 1 dimension (sequence" + f" length at shape[0]), but got shape={fschema.shape}." + ) + return fschema.shape[1:] + + +def with_sequence_length( + fschema: schema_lib.FeatureSchema, + seq_len: int, +) -> schema_lib.FeatureSchema: + """Returns a copy of the schema with shape[0] set to seq_len.""" + step_shape = get_timeseries_step_shape(fschema) + return dataclasses.replace( + fschema, + shape=(seq_len,) + step_shape, + is_timeseries=True, + ) diff --git a/dgf/src/util/temporal_test.py b/dgf/src/util/temporal_test.py index c03d29c..fc5945f 100644 --- a/dgf/src/util/temporal_test.py +++ b/dgf/src/util/temporal_test.py @@ -123,6 +123,59 @@ def test_extract_timeseries_schema_cache_grouped(self): self.assertEqual(e1_group.timestamp_feature_name, "ts") self.assertCountEqual(e1_group.feature_names, ["ts", "weight_series"]) + def test_feature_schema_helpers(self): + scalar_ts = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=(None,), + is_timeseries=True, + ) + self.assertEqual(temporal.get_timeseries_step_shape(scalar_ts), ()) + self.assertEqual(temporal.with_sequence_length(scalar_ts, 30).shape, (30,)) + self.assertTrue(temporal.with_sequence_length(scalar_ts, 30).is_timeseries) + + vector_ts = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=(None, 8), + is_timeseries=True, + ) + self.assertEqual(temporal.get_timeseries_step_shape(vector_ts), (8,)) + self.assertEqual( + temporal.with_sequence_length(vector_ts, 30).shape, (30, 8) + ) + + unknown_ts = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=None, + is_timeseries=True, + ) + with self.assertRaisesRegex( + ValueError, + r"Timeseries feature schema must have at least 1 dimension \(sequence" + r" length at shape\[0\]\), but got shape=None\.", + ): + temporal.get_timeseries_step_shape(unknown_ts) + + empty_tuple_ts = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=(), + is_timeseries=True, + ) + with self.assertRaisesRegex( + ValueError, + r"Timeseries feature schema must have at least 1 dimension \(sequence" + r" length at shape\[0\]\), but got shape=\(\)\.", + ): + temporal.get_timeseries_step_shape(empty_tuple_ts) + + non_ts = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + shape=(10,), + is_timeseries=False, + ) + with self.assertRaisesRegex( + ValueError, r"Feature schema must be a timeseries feature\." + ): + temporal.get_timeseries_step_shape(non_ts) if __name__ == "__main__": absltest.main()