From 76299d5645f30ee57d2398020431b25f2c08adac Mon Sep 17 00:00:00 2001 From: Simon Meierhans Date: Tue, 1 Sep 2026 00:09:18 -0700 Subject: [PATCH] Move timedelta extraction into normalization stage. PiperOrigin-RevId: 974312015 --- dgf/src/api/transform.py | 3 +- dgf/src/transform/BUILD | 1 - dgf/src/transform/normalize.py | 102 ++++++++++++ dgf/src/transform/normalize_test.py | 140 ++++++++++++++++ dgf/src/transform/timeseries.py | 164 +------------------ dgf/src/transform/timeseries_test.py | 231 +-------------------------- 6 files changed, 245 insertions(+), 396 deletions(-) diff --git a/dgf/src/api/transform.py b/dgf/src/api/transform.py index 398906e..5a0d7c8 100644 --- a/dgf/src/api/transform.py +++ b/dgf/src/api/transform.py @@ -28,6 +28,7 @@ from dgf.src.transform.normalize import IdentityNormalizer from dgf.src.transform.normalize import SoftQuantileNormalizer from dgf.src.transform.normalize import SinusoidTimedeltaNormalizer +from dgf.src.transform.normalize import TimedeltaNormalizer from dgf.src.transform.extract import filter_graph from dgf.src.transform.extract import drop_edge_features @@ -53,6 +54,4 @@ 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 TimestampFeatureExtractor -from dgf.src.transform.timeseries import TimestampFeatureExtractorConfig diff --git a/dgf/src/transform/BUILD b/dgf/src/transform/BUILD index 7ac246f..4dd4f9f 100644 --- a/dgf/src/transform/BUILD +++ b/dgf/src/transform/BUILD @@ -140,7 +140,6 @@ py_library( "//dgf/src/data:in_memory_graph", "//dgf/src/data:schema", "//dgf/src/io:feature_format", - "//dgf/src/util:temporal", # numpy dep, ], ) diff --git a/dgf/src/transform/normalize.py b/dgf/src/transform/normalize.py index c18e9b6..8e160f5 100644 --- a/dgf/src/transform/normalize.py +++ b/dgf/src/transform/normalize.py @@ -560,6 +560,108 @@ def normalize_tensorflow(self, value: tf.Tensor) -> Dict[str, tf.Tensor]: return {self.output_feature_name: emb} +@normalizer_registry.register +@dataclasses_json.dataclass_json +@dataclasses.dataclass(kw_only=True) +class TimedeltaNormalizer(AbstractFeatureNormalizer): + """Normalizes a TIMESTAMP feature into a TIMEDELTA feature relative to seed timestamps. + + Subtracts the timestamp value from the node's seed timestamp: + delta = seed_timestamp - t_i + """ + + input_schema: schema_lib.FeatureSchema + output_feature_name: str + type: str = dataclasses.field(default="TimedeltaNormalizer", init=False) + + @classmethod + def create( + cls, + feature_name: str, + input_schema: schema_lib.FeatureSchema, + ) -> "TimedeltaNormalizer": + if input_schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP: + raise ValueError( + f"Feature '{feature_name}' has semantic '{input_schema.semantic}'," + " but TimedeltaNormalizer only supports TIMESTAMP features." + ) + + if not input_schema.is_static_shape(): + raise ValueError( + "TimedeltaNormalizer requires fixed-length feature tensors," + f" but feature '{feature_name}' has a dynamic shape" + f" ({input_schema.shape}). Please run padding first." + ) + + return TimedeltaNormalizer( + input_feature=feature_name, + input_schema=input_schema, + output_feature_name=f"{feature_name}_seed_delta", + ) + + def output_schema(self) -> schema_lib.FeatureSetSchema: + ts_group = self.input_schema.group or ( + self.input_feature if self.input_schema.is_timeseries else None + ) + return { + self.output_feature_name: schema_lib.FeatureSchema( + format=self.input_schema.format, + semantic=schema_lib.FeatureSemantic.TIMEDELTA, + shape=self.input_schema.shape, + is_timeseries=self.input_schema.is_timeseries, + group=ts_group, + ) + } + + def normalize_numpy( + self, + value: np.ndarray, + seed_timestamps: Optional[np.ndarray] = None, + ) -> Dict[str, np.ndarray]: + assert seed_timestamps is not None, ( + "seed_timestamps must be provided to normalize timestamp feature" + f" '{self.input_feature}'." + ) + + assert value.dtype != np.object_, ( + "TimedeltaNormalizer requires fixed-length feature tensors," + f" but feature '{self.input_feature}' is a variable-length object" + " array. Please run padding first." + ) + + seed_arr = seed_timestamps + if value.ndim > 1: + seed_arr = seed_arr.reshape(seed_arr.shape + (1,) * (value.ndim - 1)) + + deltas = seed_arr - value + return {self.output_feature_name: deltas} + + def normalize_tensorflow( + self, + value: tf.Tensor, + seed_timestamps: Optional[tf.Tensor] = None, + ) -> Dict[str, tf.Tensor]: + assert seed_timestamps is not None, ( + "seed_timestamps must be provided to normalize timestamp feature" + f" '{self.input_feature}'." + ) + + assert value.shape.rank is not None, ( + "TimedeltaNormalizer requires fixed-length feature tensors," + f" but feature '{self.input_feature}' has unknown rank." + " Please run padding first." + ) + + seed_tensor = tf.cast(seed_timestamps, value.dtype) + if value.shape.rank > 1: + seed_tensor = tf.reshape( + seed_tensor, (-1,) + (1,) * (value.shape.rank - 1) + ) + + deltas = tf.subtract(seed_tensor, value) + return {self.output_feature_name: deltas} + + @dataclasses_json.dataclass_json @dataclasses.dataclass class AutoNormalizeConfig: diff --git a/dgf/src/transform/normalize_test.py b/dgf/src/transform/normalize_test.py index 451ec1d..53ed1c9 100644 --- a/dgf/src/transform/normalize_test.py +++ b/dgf/src/transform/normalize_test.py @@ -899,5 +899,145 @@ def test_auto_normalize_mask(self): ) +class TimedeltaNormalizerTest(parameterized.TestCase): + + def test_output_schema(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(3,), + is_timeseries=True, + group="sensor", + ) + normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema) + expected_schema = { + "time_seed_delta": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMEDELTA, + shape=(3,), + is_timeseries=True, + group="sensor", + ) + } + self.assertEqual(normalizer.output_schema(), expected_schema) + + def test_timedelta_normalizer_numpy_1d(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(), + ) + normalizer = normalize_lib.TimedeltaNormalizer.create("created_at", schema) + raw_val = np.array([100, 300], dtype=np.int64) + seed_timestamps = np.array([500, 500], dtype=np.int64) + out = normalizer.normalize_numpy(raw_val, seed_timestamps=seed_timestamps) + self.assertIn("created_at_seed_delta", out) + np.testing.assert_array_equal( + out["created_at_seed_delta"], np.array([400, 200], dtype=np.int64) + ) + + def test_timedelta_normalizer_numpy_timeseries(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(3,), + is_timeseries=True, + group="time", + ) + normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema) + raw_val = np.array([[100, 200, 300], [400, 450, 500]], dtype=np.int64) + seed_timestamps = np.array([500, 600], dtype=np.int64) + out = normalizer.normalize_numpy(raw_val, seed_timestamps=seed_timestamps) + expected = np.array( + [[400, 300, 200], [200, 150, 100]], dtype=np.int64 + ) + np.testing.assert_array_equal(out["time_seed_delta"], expected) + + def test_timedelta_normalizer_tensorflow(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + is_timeseries=True, + ) + normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema) + raw_np = np.array([[100, 200], [300, 400]], dtype=np.int64) + seeds_np = np.array([500, 1000], dtype=np.int64) + raw_tf = tf.constant(raw_np) + seeds_tf = tf.constant(seeds_np) + + out_tf = normalizer.normalize_tensorflow(raw_tf, seed_timestamps=seeds_tf) + out_np = normalizer.normalize_numpy(raw_np, seed_timestamps=seeds_np) + expected = np.array([[400, 300], [700, 600]], dtype=np.int64) + + np.testing.assert_array_equal(out_np["time_seed_delta"], expected) + np.testing.assert_array_equal( + out_tf["time_seed_delta"].numpy(), out_np["time_seed_delta"] + ) + + def test_timedelta_normalizer_missing_seed_timestamps_asserts(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(), + ) + normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema) + raw_val = np.array([100], dtype=np.int64) + with self.assertRaises(AssertionError): + normalizer.normalize_numpy(raw_val, seed_timestamps=None) + + with self.assertRaises(AssertionError): + normalizer.normalize_tensorflow(tf.constant([100]), seed_timestamps=None) + + def test_timedelta_normalizer_object_array_raises(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + ) + normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema) + with self.assertRaisesRegex(AssertionError, "requires fixed-length"): + normalizer.normalize_numpy( + np.array([np.array([100])], dtype=object), + seed_timestamps=np.array([500]), + ) + + def test_timedelta_normalizer_tensorflow_unknown_rank_asserts(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(2,), + ) + normalizer = normalize_lib.TimedeltaNormalizer.create("time", schema) + + @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.int64)]) + def normalize_fn(val): + return normalizer.normalize_tensorflow( + val, seed_timestamps=tf.constant([500], dtype=tf.int64) + ) + + with self.assertRaisesRegex(AssertionError, "unknown rank"): + normalize_fn(tf.constant([100, 200], dtype=tf.int64)) + + def test_timedelta_normalizer_invalid_semantic_raises(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + shape=(), + ) + with self.assertRaisesRegex(ValueError, "only supports TIMESTAMP"): + normalize_lib.TimedeltaNormalizer.create("num", schema) + + def test_timedelta_normalizer_dynamic_shape_raises(self): + schema = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + shape=(None,), + ) + with self.assertRaisesRegex(ValueError, "requires fixed-length"): + normalize_lib.TimedeltaNormalizer.create("time", schema) + + if __name__ == "__main__": absltest.main() + diff --git a/dgf/src/transform/timeseries.py b/dgf/src/transform/timeseries.py index 2724e56..2f8f163 100644 --- a/dgf/src/transform/timeseries.py +++ b/dgf/src/transform/timeseries.py @@ -21,12 +21,11 @@ # pytype: disable=module-attr import dataclasses import enum -from typing import Any, Optional, Tuple +from typing import 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.util import temporal as temporal_util import numpy as np @@ -58,16 +57,6 @@ class CalendarFeatureExtractorConfig: features: Tuple[CalendarFeature, ...] = _SUPPORTED_CALENDAR_FEATURES -@dataclasses_json.dataclass_json -@dataclasses.dataclass -class TimestampFeatureExtractorConfig: - """Configuration for extracting time delta features. - - Attributes: - fill_value: Value used for masked time steps and missing boundary deltas. - """ - - fill_value: Any = 0 def _compute_calendar_feature( @@ -230,154 +219,3 @@ def __call__( node_sets=new_node_sets, edge_sets=new_edge_sets ) - -def _compute_seed_deltas( - raw_val: np.ndarray, - mask: Optional[np.ndarray], - seed_timestamp: int, - fill_value: Any, -) -> np.ndarray: - """Computes seed_timestamp - t_i.""" - deltas = seed_timestamp - raw_val - if mask is not None: - mask_for_where = temporal_util.expand_mask_dims(mask, raw_val) - deltas = np.where(mask_for_where, deltas, fill_value) - return deltas - - -class TimestampFeatureExtractor: - """Extracts time delta features from timestamp features.""" - - def __init__( - self, - schema: schema_lib.GraphSchema, - config: Optional[TimestampFeatureExtractorConfig] = None, - ): - self.config = config or TimestampFeatureExtractorConfig() - self.schema = schema - - def _compute_feature_set_timestamp_schema( - self, - schemas: schema_lib.FeatureSetSchema, - ) -> schema_lib.FeatureSetSchema: - """Computes schema for a feature set after extracting time delta features.""" - new_schemas = dict(schemas) - for fname, schema in schemas.items(): - if schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP: - continue - if schema.group is not None: - ts_group = schema.group - # If the timestamp features is a creation time timeseries, we need to - # infer a group name to link it to a potential future mask feature. - elif schema.is_timeseries and schema.is_creation_time: - ts_group = fname - else: - ts_group = None - - out_fname = f"{fname}_seed_delta" - new_schemas[out_fname] = schema_lib.FeatureSchema( - format=schema.format, - semantic=schema_lib.FeatureSemantic.TIMEDELTA, - shape=schema.shape, - is_timeseries=schema.is_timeseries, - group=ts_group, - ) - return new_schemas - - def _extract_feature_set_timestamp_features( - self, - values: in_memory_graph.Features, - schemas: schema_lib.FeatureSetSchema, - seed_timestamp: int, - ) -> in_memory_graph.Features: - """Extracts time delta features for a single feature set.""" - new_values: in_memory_graph.Features = {} - assert seed_timestamp is not None, ( - "seed_timestamp must be provided to extract seed deltas." - ) - - for fname, schema in schemas.items(): - raw_val = values[fname] - new_values[fname] = raw_val - - if schema.semantic != schema_lib.FeatureSemantic.TIMESTAMP: - continue - - 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 pad timeseries features" - " (e.g. via pad_timeseries_graph) first." - ) - - mask = None - ts_group = schema.group or ( - fname if schema.is_timeseries and schema.is_creation_time else None - ) - if ts_group is not None: - mask_name = temporal_util.get_mask_feature_name(fname, schemas) - if mask_name is not None and mask_name in values: - mask = values[mask_name] - - out_fname = f"{fname}_seed_delta" - new_values[out_fname] = _compute_seed_deltas( - raw_val, mask, seed_timestamp, self.config.fill_value - ) - - 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(): - new_ns_schemas[ns_name] = schema_lib.NodeSchema( - features=self._compute_feature_set_timestamp_schema( - ns_schema.features - ) - ) - - new_es_schemas = {} - for es_name, es_schema in self.schema.edge_sets.items(): - new_es_schemas[es_name] = schema_lib.EdgeSchema( - source=es_schema.source, - target=es_schema.target, - features=self._compute_feature_set_timestamp_schema( - es_schema.features - ), - ) - - return schema_lib.GraphSchema( - node_sets=new_ns_schemas, edge_sets=new_es_schemas - ) - - def __call__( - self, graph: in_memory_graph.InMemoryGraph, seed_timestamp: int - ) -> in_memory_graph.InMemoryGraph: - """Extracts timedelta features from timestamps relative to seed_timestamp.""" - new_node_sets = {} - for ns_name, ns_schema in self.schema.node_sets.items(): - ns_val = graph.node_sets[ns_name] - new_vals = self._extract_feature_set_timestamp_features( - values=ns_val.features, - schemas=ns_schema.features, - seed_timestamp=seed_timestamp, - ) - 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] - new_vals = self._extract_feature_set_timestamp_features( - values=es_val.features, - schemas=es_schema.features, - seed_timestamp=seed_timestamp, - ) - 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 - ) diff --git a/dgf/src/transform/timeseries_test.py b/dgf/src/transform/timeseries_test.py index 2f5bf53..10448cf 100644 --- a/dgf/src/transform/timeseries_test.py +++ b/dgf/src/transform/timeseries_test.py @@ -306,236 +306,7 @@ def test_extract_calendar_features_parent_timestamp(self): hw_sch.features["master_time_hour"].group, "master_time" ) - def test_extract_timestamp_features(self): - padded_graph, padded_schema = _make_graph_and_schema( - values={ - "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,), - ), - }, - ) - - ts_extractor = timeseries.TimestampFeatureExtractor( - padded_schema, config=timeseries.TimestampFeatureExtractorConfig() - ) - delta_graph = ts_extractor(padded_graph, seed_timestamp=500) - delta_schema = ts_extractor.output_schema() - - 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] - # 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_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, - group="time", - ), - "time_mask": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.BOOL, - semantic=schema_lib.FeatureSemantic.MASK, - shape=(4,), - is_timeseries=True, - group="time", - ), - "time_seed_delta": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.INTEGER_64, - semantic=schema_lib.FeatureSemantic.TIMEDELTA, - shape=(4,), - is_timeseries=True, - group="time", - ), - } - test_util.assert_are_equal(self, hw_sch.features, expected_schemas) - - def test_extract_timestamp_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=(), - ) - } - graph = in_memory_graph.InMemoryGraph( - node_sets={ - "hardware": in_memory_graph.InMemoryNodeSet( - num_nodes=2, features=values - ) - }, - edge_sets={}, - ) - schema = schema_lib.GraphSchema( - node_sets={"hardware": schema_lib.NodeSchema(features=schemas)}, - edge_sets={}, - ) - extractor = timeseries.TimestampFeatureExtractor(schema) - new_graph = extractor(graph, seed_timestamp=500) - new_schema = extractor.output_schema() - np.testing.assert_array_equal( - new_graph.node_sets["hardware"].features["created_at_seed_delta"], - [435, -1679999515], - ) - self.assertFalse( - new_schema.node_sets["hardware"] - .features["created_at_seed_delta"] - .is_timeseries - ) - - def test_extract_timestamp_features_parent_timestamp(self): - schemas = { - "event_time": _ts_schema( - sem=schema_lib.FeatureSemantic.TIMESTAMP, - group="master_time", - shape=(2,), - ), - "master_time": _ts_schema( - sem=schema_lib.FeatureSemantic.TIMESTAMP, - group="master_time", - shape=(2,), - ), - } - schema = schema_lib.GraphSchema( - node_sets={"hardware": schema_lib.NodeSchema(features=schemas)}, - edge_sets={}, - ) - extractor = timeseries.TimestampFeatureExtractor(schema) - new_schema = extractor.output_schema() - hw_sch = new_schema.node_sets["hardware"] - self.assertEqual( - hw_sch.features["event_time_seed_delta"].group, "master_time" - ) - self.assertEqual( - hw_sch.features["master_time_seed_delta"].group, "master_time" - ) - self.assertEqual( - hw_sch.features["event_time_seed_delta"].semantic, - schema_lib.FeatureSemantic.TIMEDELTA, - ) - - def test_extract_timestamp_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, - ) - }, - ) - extractor = timeseries.TimestampFeatureExtractor( - schema, config=timeseries.TimestampFeatureExtractorConfig() - ) - with self.assertRaisesRegex( - AssertionError, - "TimestampFeatureExtractor requires fixed-length timestamp tensors", - ): - extractor(graph, seed_timestamp=500) - - def test_extract_timestamp_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, - group="time", - 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,), - ) - - extractor = timeseries.TimestampFeatureExtractor( - schema, config=timeseries.TimestampFeatureExtractorConfig() - ) - delta_graph = extractor(graph, seed_timestamp=500) - delta_schema = extractor.output_schema() - es_val = delta_graph.edge_sets["ts_edges"] - es_sch = delta_schema.edge_sets["ts_edges"] - - self.assertIn("time_seed_delta", es_val.features) - self.assertEqual(es_sch.features["time_seed_delta"].group, "time") - self.assertEqual( - es_sch.features["time_seed_delta"].semantic, - schema_lib.FeatureSemantic.TIMEDELTA, - ) - - def test_extract_timestamp_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, - ) - } - graph = in_memory_graph.InMemoryGraph( - node_sets={ - "hardware": in_memory_graph.InMemoryNodeSet( - num_nodes=1, features=values - ) - }, - edge_sets={}, - ) - schema = schema_lib.GraphSchema( - node_sets={"hardware": schema_lib.NodeSchema(features=schemas)}, - edge_sets={}, - ) - extractor = timeseries.TimestampFeatureExtractor(schema) - new_graph = extractor(graph, seed_timestamp=500) - new_schema = extractor.output_schema() - self.assertEqual(new_graph.node_sets["hardware"].features, values) - self.assertEqual(new_schema.node_sets["hardware"].features, schemas) - - @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) - if __name__ == "__main__": absltest.main() +