diff --git a/dgf/src/api/transform.py b/dgf/src/api/transform.py index 5d16b98..398906e 100644 --- a/dgf/src/api/transform.py +++ b/dgf/src/api/transform.py @@ -53,8 +53,6 @@ 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 PadAndCapTimeseries -from dgf.src.transform.timeseries import PadAndCapTimeseriesConfig from dgf.src.transform.timeseries import TimestampFeatureExtractor from dgf.src.transform.timeseries import TimestampFeatureExtractorConfig diff --git a/dgf/src/transform/timeseries.py b/dgf/src/transform/timeseries.py index bf6fb5d..2724e56 100644 --- a/dgf/src/transform/timeseries.py +++ b/dgf/src/transform/timeseries.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Padding and capping for timeseries sequence features in graphs.""" +"""Temporal and timeseries feature extraction for graphs.""" # TODO(simonmeierhans): If a timeseries feature has no group defined, and there # is a group with the same name as the feature it currently implicitly joins @@ -21,101 +21,15 @@ # pytype: disable=module-attr import dataclasses import enum -from typing import Any, List, Optional, Tuple +from typing import Any, Optional, Tuple import dataclasses_json from dgf.src.data import in_memory_graph from dgf.src.data import schema as schema_lib -from dgf.src.io import feature_format from dgf.src.util import temporal as temporal_util import numpy as np -@dataclasses_json.dataclass_json -@dataclasses.dataclass -class PadAndCapTimeseriesConfig: - """Configuration for padding and capping timeseries features. - - Attributes: - sequence_length: Fixed target sequence dimension K. Sequences longer than K - are capped to the most recent K steps (`[-K:]`). Shorter sequences are - left-padded to length K. - padding_value: Scalar value used for left-padding shorter sequences. - """ - - sequence_length: int = 30 - padding_value: Any = 0 - - -def _pad_and_cap_single_feature( - raw_series: np.ndarray, - seq_len: int, - feat_shape: Tuple[int, ...], - padding_value: Any, - dtype: Any, - is_static_shape: bool, -) -> Tuple[np.ndarray, np.ndarray]: - """Pads and caps a single sequence feature into (padded_matrix, mask_matrix).""" - num_entities = raw_series.shape[0] - if dtype == np.bytes_ and raw_series.dtype.kind in ("S", "a"): - dtype = raw_series.dtype - if num_entities == 0: - return ( - np.empty((0, seq_len) + feat_shape, dtype=dtype), - np.empty((0, seq_len), dtype=np.bool_), - ) - - # Fast vectorized path when all entities share a fixed sequence length. - if is_static_shape and raw_series.ndim >= 2: - num_steps = raw_series.shape[1] - if num_steps >= seq_len: - padded_matrix = raw_series[:, -seq_len:].astype(dtype, copy=True) - mask_matrix = np.ones((num_entities, seq_len), dtype=np.bool_) - return padded_matrix, mask_matrix - - pad_width = [(0, 0), (seq_len - num_steps, 0)] + [(0, 0)] * len(feat_shape) - padded_matrix = np.pad( - raw_series.astype(dtype, copy=False), - pad_width=pad_width, - mode="constant", - constant_values=padding_value, - ) - mask_width = [(0, 0), (seq_len - num_steps, 0)] - mask_matrix = np.pad( - np.ones((num_entities, num_steps), dtype=np.bool_), - pad_width=mask_width, - mode="constant", - constant_values=False, - ) - return padded_matrix, mask_matrix - - padded_matrix = np.full( - (num_entities, seq_len) + feat_shape, - fill_value=padding_value, - dtype=dtype, - ) - # Binary mask matrix matching sequence length shape (num_entities, seq_len) - # where True indicates valid observed time steps and False indicates - # left-padded steps. - mask_matrix = np.zeros((num_entities, seq_len), dtype=np.bool_) - - # TODO(mesimon): Move into C++ for performance. - for idx in range(num_entities): - raw_arr = raw_series[idx] - if not isinstance(raw_arr, np.ndarray): - raw_arr = np.asarray(raw_arr) - - num_steps = len(raw_arr) - if num_steps >= seq_len: - padded_matrix[idx] = raw_arr[-seq_len:] - mask_matrix[idx] = True - elif num_steps > 0: - padded_matrix[idx, -num_steps:] = raw_arr - mask_matrix[idx, -num_steps:] = True - - return padded_matrix, mask_matrix - - class CalendarFeature(str, enum.Enum): """Supported calendar features to extract from timestamps.""" @@ -183,208 +97,6 @@ def _compute_calendar_feature( ) -class PadAndCapTimeseries: - """Pads and caps timeseries sequence features into fixed-dimension tensors. - - Transforms variable-length sequence features (`is_timeseries=True`) in the - graph into fixed-length matrices of shape `(num_entities, sequence_length) + - step_shape`. Sequences longer than `sequence_length` are capped from the right - (keeping the most recent steps), while shorter sequences are left-padded. - - For each sequence group, a corresponding boolean mask feature (`{group}_mask`) - is generated with semantic `MASK`, where `True` indicates observed time steps - and `False` indicates padded steps. - - Attributes: - schema: The input `GraphSchema`. - config: Configuration specifying sequence length and padding value. - schema_cache: Pre-computed `TimeseriesSchemaCache` for fast grouping. - """ - - def __init__( - self, - schema: schema_lib.GraphSchema, - config: Optional[PadAndCapTimeseriesConfig] = None, - schema_cache: Optional[temporal_util.TimeseriesSchemaCache] = None, - ): - self.config = config or PadAndCapTimeseriesConfig() - self.schema = schema - if schema_cache is None: - schema_cache = temporal_util.extract_timeseries_schema_cache(schema) - self.schema_cache = schema_cache - - def _compute_feature_set_pad_and_cap_schema( - self, - schemas: schema_lib.FeatureSetSchema, - ts_specs: List[temporal_util.TimeseriesGroupSpec], - ) -> schema_lib.FeatureSetSchema: - """Computes schema for a feature set after padding/capping.""" - new_schemas = dict(schemas) - seq_len = self.config.sequence_length - - ts_features = {} - for group in ts_specs: - for feature_name in group.feature_names: - ts_features[feature_name] = group.timestamp_feature_name - - for feature_name in ts_features: - feature_schema = schemas[feature_name] - ts_group = feature_schema.group or feature_name - new_schemas[feature_name] = dataclasses.replace( - temporal_util.with_sequence_length(feature_schema, seq_len), - group=ts_group, - ) - if feature_schema.semantic == schema_lib.FeatureSemantic.MASK: - continue - - mask_name = temporal_util.get_mask_feature_name(feature_name, schemas) - if mask_name is None: - mask_name = f"{ts_group}_mask" - if mask_name in schemas: - raise ValueError( - f"Cannot generate mask for sequence group '{ts_group}'. The" - f" fallback mask name '{mask_name}' clashes with an existing" - " feature in the schema that is not a valid mask. Please" - " explicitly define a mask feature for this group or rename the" - " clashing feature." - ) - - if mask_name not in new_schemas: - new_schemas[mask_name] = schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.BOOL, - semantic=schema_lib.FeatureSemantic.MASK, - shape=(seq_len,), - is_timeseries=feature_schema.is_timeseries, - group=ts_group, - ) - - return new_schemas - - def _pad_and_cap_feature_set( - self, - values: in_memory_graph.Features, - schemas: schema_lib.FeatureSetSchema, - ts_specs: List[temporal_util.TimeseriesGroupSpec], - ) -> in_memory_graph.Features: - """Extracts fixed-dimension sequence features for a feature set.""" - new_values: in_memory_graph.Features = {} - seq_len = self.config.sequence_length - - # Map timeseries feature names to their associated timestamp feature name. - ts_features = {} - for group in ts_specs: - for feature_name in group.feature_names: - ts_features[feature_name] = group.timestamp_feature_name - - # Copy over non-timeseries features. - for feature_name in schemas: - if feature_name not in ts_features: - new_values[feature_name] = values[feature_name] - - for feature_name in ts_features: - feature_schema = schemas[feature_name] - - dtype = feature_format.FEATURE_FORMAT_TO_NP_DTYPE[feature_schema.format] - feat_shape = temporal_util.get_timeseries_step_shape(feature_schema) - - padded_matrix, mask_matrix = _pad_and_cap_single_feature( - raw_series=values[feature_name], - seq_len=seq_len, - feat_shape=feat_shape, - padding_value=self.config.padding_value, - dtype=dtype, - is_static_shape=feature_schema.is_static_shape(), - ) - - new_values[feature_name] = padded_matrix - if feature_schema.semantic == schema_lib.FeatureSemantic.MASK: - continue - - ts_group = feature_schema.group or feature_name - mask_name = temporal_util.get_mask_feature_name(feature_name, schemas) - if mask_name is None: - mask_name = f"{ts_group}_mask" - - # Copy over mask matrix if it doesn't exist yet. Only store once per - # group. - if mask_name not in new_values: - new_values[mask_name] = mask_matrix - - return new_values - - def output_schema(self) -> schema_lib.GraphSchema: - """Returns the transformed GraphSchema.""" - new_ns_schemas = {} - for ns_name, ns_schema in self.schema.node_sets.items(): - ts_specs = self.schema_cache.node_sets[ns_name] - if not ts_specs: - new_ns_schemas[ns_name] = ns_schema - else: - new_ns_schemas[ns_name] = schema_lib.NodeSchema( - features=self._compute_feature_set_pad_and_cap_schema( - ns_schema.features, ts_specs - ) - ) - - new_es_schemas = {} - for es_name, es_schema in self.schema.edge_sets.items(): - ts_specs = self.schema_cache.edge_sets[es_name] - if not ts_specs: - new_es_schemas[es_name] = es_schema - else: - new_es_schemas[es_name] = schema_lib.EdgeSchema( - source=es_schema.source, - target=es_schema.target, - features=self._compute_feature_set_pad_and_cap_schema( - es_schema.features, ts_specs - ), - ) - - return schema_lib.GraphSchema( - node_sets=new_ns_schemas, edge_sets=new_es_schemas - ) - - def __call__( - self, graph: in_memory_graph.InMemoryGraph - ) -> in_memory_graph.InMemoryGraph: - """Transforms timeseries sequence features in the graph.""" - new_node_sets = {} - for ns_name, ns_schema in self.schema.node_sets.items(): - ns_val = graph.node_sets[ns_name] - ts_specs = self.schema_cache.node_sets[ns_name] - if not ts_specs: - new_node_sets[ns_name] = ns_val - continue - new_vals = self._pad_and_cap_feature_set( - values=ns_val.features, - schemas=ns_schema.features, - ts_specs=ts_specs, - ) - new_node_sets[ns_name] = in_memory_graph.InMemoryNodeSet( - num_nodes=ns_val.num_nodes, features=new_vals - ) - - new_edge_sets = {} - for es_name, es_schema in self.schema.edge_sets.items(): - es_val = graph.edge_sets[es_name] - ts_specs = self.schema_cache.edge_sets[es_name] - if not ts_specs: - new_edge_sets[es_name] = es_val - continue - new_vals = self._pad_and_cap_feature_set( - values=es_val.features, - schemas=es_schema.features, - ts_specs=ts_specs, - ) - new_edge_sets[es_name] = in_memory_graph.InMemoryEdgeSet( - adjacency=es_val.adjacency, features=new_vals - ) - - return in_memory_graph.InMemoryGraph( - node_sets=new_node_sets, edge_sets=new_edge_sets - ) - - class CalendarFeatureExtractor: """Extracts calendar features (e.g. hour, day_of_week) from timestamp features. @@ -457,8 +169,7 @@ def _extract_feature_set_calendar_features( assert raw_val.dtype != np.object_, ( "CalendarFeatureExtractor requires fixed-length timestamp tensors," - f" but feature '{fname}' is a variable-length object array. Please" - " run PadAndCapTimeseries first." + f" but feature '{fname}' is a variable-length object array." ) for cal_feat in self.config.features: @@ -595,8 +306,8 @@ def _extract_feature_set_timestamp_features( assert schema.is_static_shape() and raw_val.dtype != np.object_, ( "TimestampFeatureExtractor requires fixed-length timestamp tensors," f" but feature '{fname}' is a variable-length object array or has" - f" dynamic shape ({schema.shape}). Please run PadAndCapTimeseries" - " first." + f" dynamic shape ({schema.shape}). Please pad timeseries features" + " (e.g. via pad_timeseries_graph) first." ) mask = None diff --git a/dgf/src/transform/timeseries_test.py b/dgf/src/transform/timeseries_test.py index c5e24df..2f5bf53 100644 --- a/dgf/src/transform/timeseries_test.py +++ b/dgf/src/transform/timeseries_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for padding and capping timeseries sequence features.""" +"""Tests for temporal and timeseries feature extractors.""" from absl.testing import absltest from absl.testing import parameterized @@ -80,521 +80,6 @@ def _ts_schema( class TimeseriesTest(parameterized.TestCase): - def test_capping_and_padding(self): - # Both 'time' and 'signal' are variable-length sequence object arrays. - graph, schema = _make_graph_and_schema( - values={ - "time": np.array( - [np.array([10, 20, 30, 40, 50]), np.array([5, 15])], - dtype=np.object_, - ), - "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]), - }, - schemas={ - "time": _ts_schema( - fmt=schema_lib.FeatureFormat.INTEGER_64, - sem=schema_lib.FeatureSemantic.TIMESTAMP, - is_creation_time=True, - group="time", - ), - "signal": _ts_schema(group="time"), - "id": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.INTEGER_64, - semantic=schema_lib.FeatureSemantic.NUMERICAL, - ), - }, - num_nodes=2, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - new_graph = pad_and_cap(graph) - new_schema = pad_and_cap.output_schema() - hw_val = new_graph.node_sets["hardware"] - hw_sch = new_schema.node_sets["hardware"] - - 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 - ), - "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) - self.assertTrue(hw_sch.features["time_mask"].is_timeseries) - self.assertEqual( - hw_sch.features["time_mask"].semantic, schema_lib.FeatureSemantic.MASK - ) - - def test_edge_sets_and_non_timeseries(self): - graph = in_memory_graph.InMemoryGraph( - node_sets={ - "user": in_memory_graph.InMemoryNodeSet( - num_nodes=1, features={"age": np.array([30], dtype=np.int64)} - ), - }, - edge_sets={ - "clicks": in_memory_graph.InMemoryEdgeSet( - adjacency=np.array([[0], [0]], dtype=np.int64), - features={ - "time": np.array([np.array([100, 200])], dtype=np.object_) - }, - ), - "static_edge": in_memory_graph.InMemoryEdgeSet( - adjacency=np.array([[0], [0]], dtype=np.int64), - features={"weight": np.array([1.0], dtype=np.float32)}, - ), - }, - ) - schema = schema_lib.GraphSchema( - node_sets={ - "user": schema_lib.NodeSchema( - features={ - "age": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.INTEGER_64, - semantic=schema_lib.FeatureSemantic.NUMERICAL, - ) - } - ) - }, - edge_sets={ - "clicks": schema_lib.EdgeSchema( - source="user", - target="user", - features={ - "time": _ts_schema( - fmt=schema_lib.FeatureFormat.INTEGER_64, - sem=schema_lib.FeatureSemantic.TIMESTAMP, - group="time", - ) - }, - ), - "static_edge": schema_lib.EdgeSchema( - source="user", - target="user", - features={ - "weight": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.FLOAT_32, - semantic=schema_lib.FeatureSemantic.NUMERICAL, - ) - }, - ), - }, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=2), - ) - new_graph = pad_and_cap(graph) - new_schema = pad_and_cap.output_schema() - - np.testing.assert_array_equal( - new_graph.node_sets["user"].features["age"], [30] - ) - np.testing.assert_array_equal( - new_graph.edge_sets["clicks"].features["time"][0], [100, 200] - ) - self.assertTrue( - new_schema.edge_sets["clicks"].features["time"].is_timeseries - ) - - def test_multidimensional_sequence(self): - graph, schema = _make_graph_and_schema( - values={ - "emb": np.array( - [np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)], - dtype=np.object_, - ) - }, - schemas={ - "emb": _ts_schema( - sem=schema_lib.FeatureSemantic.EMBEDDING, - group="emb", - shape=(None, 2), - ) - }, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - new_graph = pad_and_cap(graph) - new_schema = pad_and_cap.output_schema() - hw_val = new_graph.node_sets["hardware"] - hw_sch = new_schema.node_sets["hardware"] - - 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, 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,)) - self.assertEqual( - hw_sch.features["emb_mask"].semantic, schema_lib.FeatureSemantic.MASK - ) - - def test_custom_mask_name_is_reused(self): - graph, schema = _make_graph_and_schema( - values={ - "emb": np.array( - [np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)], - dtype=np.object_, - ), - "my_custom_mask": np.array( - [np.array([True, True], dtype=bool)], - dtype=np.object_, - ), - }, - schemas={ - "emb": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.FLOAT_32, - semantic=schema_lib.FeatureSemantic.EMBEDDING, - is_timeseries=True, - shape=(None, 2), - group="emb_group", - ), - "my_custom_mask": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.BOOL, - semantic=schema_lib.FeatureSemantic.MASK, - is_timeseries=True, - shape=(None,), - group="emb_group", - ), - }, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - new_graph = pad_and_cap(graph) - new_schema = pad_and_cap.output_schema() - hw_val = new_graph.node_sets["hardware"] - hw_sch = new_schema.node_sets["hardware"] - - expected_features = { - "emb": np.array( - [[[0.0, 0.0], [1.0, 2.0], [3.0, 4.0]]], dtype=np.float32 - ), - "my_custom_mask": np.array([[False, 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["my_custom_mask"].shape, (3,)) - self.assertNotIn("emb_mask", hw_sch.features) - - def test_clashing_mask_name_raises(self): - _, schema = _make_graph_and_schema( - values={ - "emb": np.array( - [np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)], - dtype=np.object_, - ), - "emb_mask": np.array( - [np.array([42.0], dtype=np.float32)], - dtype=np.object_, - ), - }, - schemas={ - "emb": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.FLOAT_32, - semantic=schema_lib.FeatureSemantic.EMBEDDING, - is_timeseries=True, - shape=(None, 2), - group="emb", - ), - # This feature is named "emb_mask", which is the fallback mask name - # for the group "emb", but its semantic is not MASK. - "emb_mask": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.FLOAT_32, - semantic=schema_lib.FeatureSemantic.NUMERICAL, - is_timeseries=True, - shape=(None, 1), - group="emb", - ), - }, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - with self.assertRaisesRegex( - ValueError, "clashes with an existing feature" - ): - pad_and_cap.output_schema() - - def test_pad_and_cap_timeseries_features_auto_assigns_group_when_none(self): - graph, schema = _make_graph_and_schema( - values={ - "signal": np.array( - [np.array([1.0, 2.0], dtype=np.float32)], dtype=object - ) - }, - schemas={ - "signal": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.FLOAT_32, - semantic=schema_lib.FeatureSemantic.NUMERICAL, - is_timeseries=True, - shape=(None,), - group=None, - ), - }, - num_nodes=1, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - new_graph = pad_and_cap(graph) - new_schema = pad_and_cap.output_schema() - hw_sch = new_schema.node_sets["hardware"] - self.assertEqual(hw_sch.features["signal"].group, "signal") - self.assertIn("signal_mask", new_graph.node_sets["hardware"].features) - - def test_empty_sequence(self): - graph, schema = _make_graph_and_schema( - values={"time": np.array([np.array([])], dtype=np.object_)}, - schemas={ - "time": _ts_schema( - fmt=schema_lib.FeatureFormat.INTEGER_64, - sem=schema_lib.FeatureSemantic.TIMESTAMP, - group="time", - ) - }, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - new_graph = pad_and_cap(graph) - hw_val = new_graph.node_sets["hardware"] - np.testing.assert_array_equal(hw_val.features["time"][0], [0, 0, 0]) - np.testing.assert_array_equal(hw_val.features["time_mask"][0], [0, 0, 0]) - - def test_no_timeseries_in_graph(self): - graph, schema = _make_graph_and_schema( - values={"age": np.array([30])}, - schemas={ - "age": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.INTEGER_64, - semantic=schema_lib.FeatureSemantic.NUMERICAL, - ) - }, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, timeseries.PadAndCapTimeseriesConfig() - ) - new_graph = pad_and_cap(graph) - self.assertEqual(new_graph.node_sets["hardware"].features["age"][0], 30) - self.assertNotIn("age_mask", new_graph.node_sets["hardware"].features) - - def test_empty_entity_set(self): - graph = in_memory_graph.InMemoryGraph( - node_sets={ - "user": in_memory_graph.InMemoryNodeSet( - num_nodes=0, features={"time": np.array([], dtype=np.object_)} - ), - }, - edge_sets={}, - ) - schema = schema_lib.GraphSchema( - node_sets={ - "user": schema_lib.NodeSchema( - features={ - "time": _ts_schema( - fmt=schema_lib.FeatureFormat.INTEGER_64, - sem=schema_lib.FeatureSemantic.TIMESTAMP, - group="time", - ) - } - ), - }, - edge_sets={}, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, timeseries.PadAndCapTimeseriesConfig(sequence_length=3) - ) - new_graph = pad_and_cap(graph) - new_schema = pad_and_cap.output_schema() - user_val = new_graph.node_sets["user"] - user_sch = new_schema.node_sets["user"] - self.assertEqual(user_val.num_nodes, 0) - self.assertEqual(user_val.features["time"].shape, (0, 3)) - self.assertEqual(user_val.features["time_mask"].shape, (0, 3)) - self.assertEqual(user_sch.features["time"].shape, (3,)) - self.assertTrue(user_sch.features["time"].is_timeseries) - self.assertTrue(user_sch.features["time_mask"].is_timeseries) - self.assertEqual( - user_sch.features["time_mask"].semantic, schema_lib.FeatureSemantic.MASK - ) - - def test_missing_entity_set_raises(self): - graph = in_memory_graph.InMemoryGraph( - node_sets={}, - edge_sets={}, - ) - schema = schema_lib.GraphSchema( - node_sets={ - "absent": schema_lib.NodeSchema( - features={ - "time": _ts_schema(fmt=schema_lib.FeatureFormat.INTEGER_64) - } - ), - }, - edge_sets={}, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, timeseries.PadAndCapTimeseriesConfig(sequence_length=3) - ) - with self.assertRaises(KeyError): - pad_and_cap(graph) - extractor = timeseries.CalendarFeatureExtractor(schema) - with self.assertRaises(KeyError): - extractor(graph) - - 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, - ) - }, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - with self.assertRaises(KeyError): - pad_and_cap(graph) - extractor = timeseries.CalendarFeatureExtractor(schema) - with self.assertRaises(KeyError): - extractor(graph) - - def test_custom_padding_value(self): - graph, schema = _make_graph_and_schema( - values={ - "time": np.array([np.array([10])], dtype=np.object_), - "signal": np.array( - [np.array([2.0], dtype=np.float32)], dtype=np.object_ - ), - }, - schemas={ - "time": _ts_schema( - fmt=schema_lib.FeatureFormat.INTEGER_64, - sem=schema_lib.FeatureSemantic.TIMESTAMP, - group="time", - ), - "signal": _ts_schema(group="time"), - }, - num_nodes=1, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig( - sequence_length=3, padding_value=-1 - ), - ) - new_graph = pad_and_cap(graph) - hw_val = new_graph.node_sets["hardware"] - expected_features = { - "time": np.array([[-1, -1, 10]], dtype=np.int64), - "signal": np.array([[-1.0, -1.0, 2.0]], dtype=np.float32), - "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( - values={ - "time": np.array( - [[10, 20, 30, 40, 50], [100, 200, 300, 400, 500]], - dtype=np.int64, - ), - }, - schemas={ - "time": _ts_schema( - fmt=schema_lib.FeatureFormat.INTEGER_64, - sem=schema_lib.FeatureSemantic.TIMESTAMP, - group="time", - shape=(5,), - ), - }, - num_nodes=2, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=3), - ) - new_graph = pad_and_cap(graph) - hw_val = new_graph.node_sets["hardware"] - 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( - values={ - "emb": np.array( - [ - [[1.0, 1.1], [2.0, 2.2]], - [[10.0, 10.1], [20.0, 20.2]], - ], - dtype=np.float32, - ), - }, - schemas={ - "emb": _ts_schema( - sem=schema_lib.FeatureSemantic.EMBEDDING, - group="emb", - shape=(2, 2), - ), - }, - num_nodes=2, - ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig( - sequence_length=3, padding_value=-1.0 - ), - ) - new_graph = pad_and_cap(graph) - hw_val = new_graph.node_sets["hardware"] - 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, True, True], - [False, 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 = { @@ -624,26 +109,26 @@ def test_compute_calendar_feature(self): test_util.assert_are_equal(self, computed, expected) def test_extract_calendar_features(self): - graph, schema = _make_graph_and_schema( + padded_graph, padded_schema = _make_graph_and_schema( values={ - "time": np.array( - [np.array([65, 1680000015], dtype=np.int64)], dtype=np.object_ - ) + "time": np.array([[65, 1680000015]], dtype=np.int64), + "time_mask": np.array([[True, True]], dtype=np.bool_), }, schemas={ "time": _ts_schema( fmt=schema_lib.FeatureFormat.INTEGER_64, sem=schema_lib.FeatureSemantic.TIMESTAMP, group="time", - ) + shape=(2,), + ), + "time_mask": _ts_schema( + fmt=schema_lib.FeatureFormat.BOOL, + sem=schema_lib.FeatureSemantic.MASK, + group="time", + shape=(2,), + ), }, ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=2), - ) - padded_graph = pad_and_cap(graph) - padded_schema = pad_and_cap.output_schema() cal_extractor = timeseries.CalendarFeatureExtractor(padded_schema) cal_graph = cal_extractor(padded_graph) @@ -822,25 +307,26 @@ def test_extract_calendar_features_parent_timestamp(self): ) def test_extract_timestamp_features(self): - graph, schema = _make_graph_and_schema( + padded_graph, padded_schema = _make_graph_and_schema( values={ - "time": np.array( - [np.array([100, 250, 300], dtype=np.int64)], dtype=np.object_ - ) + "time": np.array([[0, 100, 250, 300]], dtype=np.int64), + "time_mask": np.array([[False, True, True, True]], dtype=np.bool_), }, schemas={ "time": _ts_schema( fmt=schema_lib.FeatureFormat.INTEGER_64, sem=schema_lib.FeatureSemantic.TIMESTAMP, - ) + group="time", + shape=(4,), + ), + "time_mask": _ts_schema( + fmt=schema_lib.FeatureFormat.BOOL, + sem=schema_lib.FeatureSemantic.MASK, + group="time", + shape=(4,), + ), }, ) - pad_and_cap = timeseries.PadAndCapTimeseries( - schema, - timeseries.PadAndCapTimeseriesConfig(sequence_length=4), - ) - padded_graph = pad_and_cap(graph) - padded_schema = pad_and_cap.output_schema() ts_extractor = timeseries.TimestampFeatureExtractor( padded_schema, config=timeseries.TimestampFeatureExtractorConfig()