From 436ce932998620063a89c2400db617d98f181068 Mon Sep 17 00:00:00 2001 From: Simon Meierhans Date: Wed, 5 Aug 2026 00:35:30 -0700 Subject: [PATCH] Integrate temporal graphs into node prediction 10 lines of code dataset generation. PiperOrigin-RevId: 959479016 --- dgf/src/api/transform.py | 5 +- dgf/src/learning/ten_lines/BUILD | 14 + dgf/src/learning/ten_lines/dataset.py | 212 ++++++---- dgf/src/learning/ten_lines/dataset_test.py | 28 +- .../ten_lines/node_prediction_dataset.py | 228 ++++++++--- .../ten_lines/node_prediction_dataset_test.py | 365 ++++++++++++++++-- .../ten_lines/node_prediction_model.py | 100 ++++- .../ten_lines/node_prediction_test.py | 281 ++++++++++++-- .../ten_lines/node_prediction_train.py | 37 +- dgf/src/transform/BUILD | 1 + dgf/src/transform/temporal.py | 57 ++- dgf/src/transform/temporal_test.py | 151 +++++++- dgf/src/transform/timeseries.py | 124 +++++- dgf/src/transform/timeseries_test.py | 104 ++++- dgf/src/util/gen_test_graph.py | 103 +++++ 15 files changed, 1520 insertions(+), 290 deletions(-) diff --git a/dgf/src/api/transform.py b/dgf/src/api/transform.py index 90987e8..92ccc82 100644 --- a/dgf/src/api/transform.py +++ b/dgf/src/api/transform.py @@ -29,7 +29,7 @@ from dgf.src.transform.normalize import SoftQuantileNormalizer from dgf.src.transform.normalize import SinusoidTimedeltaNormalizer -from dgf.src.transform.extract import filter_schema +from dgf.src.transform.extract import filter_schema as filter_schema_by_features from dgf.src.transform.extract import filter_graph from dgf.src.transform.extract import drop_edge_features @@ -56,6 +56,7 @@ 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 PerSampleTransform +from dgf.src.transform.timeseries import PerSampleTransformConfig from dgf.src.transform.timeseries import TimestampFeatureExtractor from dgf.src.transform.timeseries import TimestampFeatureExtractorConfig - diff --git a/dgf/src/learning/ten_lines/BUILD b/dgf/src/learning/ten_lines/BUILD index dcbca37..559dfaa 100644 --- a/dgf/src/learning/ten_lines/BUILD +++ b/dgf/src/learning/ten_lines/BUILD @@ -68,6 +68,8 @@ py_library( "//dgf/src/sampling:in_memory_sampler", "//dgf/src/transform:merge", "//dgf/src/transform:normalize", + "//dgf/src/transform:temporal", + "//dgf/src/transform:timeseries", "//dgf/src/util:filesystem", "//dgf/src/util:log", "//dgf/src/util:temporal", @@ -273,7 +275,10 @@ py_library( "//dgf/src/sampling:config", "//dgf/src/sampling:in_memory_sampler", "//dgf/src/transform:normalize", + "//dgf/src/transform:temporal", + "//dgf/src/transform:timeseries", "//dgf/src/util:log", + "//dgf/src/util:temporal", "//dgf/src/util:util_py", # jax dep, # numpy dep, @@ -317,8 +322,11 @@ py_test( "//dgf/src/io:tf_graph_sample", "//dgf/src/learning/jax/layers:standard", "//dgf/src/sampling:in_memory_sampler", + "//dgf/src/transform:temporal", + "//dgf/src/transform:timeseries", "//dgf/src/util:filesystem", "//dgf/src/util:gen_test_graph", + "//dgf/src/util:temporal", "//dgf/src/util:test_util", # jax dep, # numpy dep, @@ -422,11 +430,17 @@ py_test( ":node_prediction_dataset", # absl/testing:absltest dep, # absl/testing:parameterized dep, + "//dgf/src/data:in_memory_graph", + "//dgf/src/data:jax_in_memory_graph", + "//dgf/src/data:schema", "//dgf/src/io:jax", "//dgf/src/io:tf_graph_sample", "//dgf/src/sampling:config", + "//dgf/src/transform:normalize", "//dgf/src/util:gen_test_graph", + "//dgf/src/util:temporal", "//dgf/src/validate:in_memory_graph", + # jax dep, # numpy dep, ], ) diff --git a/dgf/src/learning/ten_lines/dataset.py b/dgf/src/learning/ten_lines/dataset.py index 38a45d2..14199ec 100644 --- a/dgf/src/learning/ten_lines/dataset.py +++ b/dgf/src/learning/ten_lines/dataset.py @@ -35,6 +35,9 @@ str, ] +PerSampleTransformConfig = timeseries_transform.PerSampleTransformConfig +PerSampleTransform = timeseries_transform.PerSampleTransform + class GraphFormat(enum.Enum): """The format of the input graph. @@ -69,6 +72,37 @@ class GraphFormat(enum.Enum): ] +# TODO(simonmeierhans): This should be removed once we support edge features in +# sample generation. It is currently necessary because propagating +# timestamps to edges is necessary for temporal sampling, but edge features are +# not supported. +def sanitize_schema_for_sample_generator( + schema: schema_lib.GraphSchema, +) -> schema_lib.GraphSchema: + """Returns a schema copy with empty feature schemas on edge sets. + + In-memory sampled subgraphs for node prediction do not populate edge feature + arrays. Clearing feature schemas on edge sets prevents + TimestampFeatureExtractor from attempting to access unpopulated edge features + during sample generation. + + Args: + schema: The graph schema to sanitize. + + Returns: + A sanitized schema with empty feature schemas on edge sets. + """ + if not any(bool(es.features) for es in schema.edge_sets.values()): + return schema + return dataclasses.replace( + schema, + edge_sets={ + es_name: dataclasses.replace(es, features={}) + for es_name, es in schema.edge_sets.items() + }, + ) + + @dataclasses.dataclass class SampleGeneratorFromAnything: """Converts a user input graph into a generator of batched graph samples. @@ -118,10 +152,9 @@ class SampleGeneratorFromAnything: sampler_returns_node_idxs_only: If `True`, the sampler returns only node indices without feature values. If `False` (default), the sampler returns feature values. - timeseries_pad_and_cap: Configuration for padding and capping timeseries - sequence features per sample before merging. - timedelta_extraction: Configuration for time delta extraction per sample - before merging. + per_sample_transforms: Configuration for per-sample preprocessing transforms + (e.g., timeseries padding/capping and timedelta extraction) before + merging. """ graph: Graph @@ -144,11 +177,8 @@ class SampleGeneratorFromAnything: default_factory=dict ) sampler_returns_node_idxs_only: bool = False - timeseries_pad_and_cap: Optional[ - timeseries_transform.PadAndCapTimeseriesConfig - ] = None - timedelta_extraction: Optional[ - timeseries_transform.TimestampFeatureExtractorConfig + per_sample_transforms: Optional[ + timeseries_transform.PerSampleTransformConfig ] = None num_seed_nodes: Optional[int] = dataclasses.field(init=False) @@ -163,12 +193,9 @@ class SampleGeneratorFromAnything: _seed_timestamps_all: Optional[np.ndarray] = dataclasses.field( init=False, default=None ) - _pad_and_cap_transformer: Optional[ - timeseries_transform.PadAndCapTimeseries - ] = dataclasses.field(init=False, default=None) - _timestamp_extractor: Optional[ - timeseries_transform.TimestampFeatureExtractor - ] = dataclasses.field(init=False, default=None) + _per_sample_transform: timeseries_transform.PerSampleTransform = ( + dataclasses.field(init=False) + ) _has_per_sample_transforms: bool = dataclasses.field( init=False, default=False ) @@ -192,6 +219,9 @@ def __post_init__(self): ) ) + if self.format == GraphFormat.IN_MEMORY_GRAPH: + self.schema = sanitize_schema_for_sample_generator(self.schema) + self._target_nodeset = self.sampling_config.root.nodeset # pyrefly: ignore[missing-attribute] self._ts_feature = temporal_util.creation_time_feature_name( self.schema.node_sets[self._target_nodeset].features @@ -233,61 +263,56 @@ def __post_init__(self): # Initializing per-sample transforms and precomputing the output schema. current_schema = self.schema - self._has_per_sample_transforms = ( - self.timeseries_pad_and_cap is not None - or self.timedelta_extraction is not None - ) - - if ( - self.temporal or self._has_per_sample_transforms - ) and self.format != GraphFormat.IN_MEMORY_GRAPH: - raise ValueError( - "Temporal sampling (`temporal=True`) and per-sample transformations" - " (`timeseries_pad_and_cap`, `timedelta_extraction`) are" - " only supported for GraphFormat.IN_MEMORY_GRAPH, but got format:" - f" {self.format}" + if self.temporal and self.format != GraphFormat.IN_MEMORY_GRAPH: + raise NotImplementedError( + "Temporal sampling (`temporal=True`) is only supported for" + f" GraphFormat.IN_MEMORY_GRAPH, but got format: {self.format}" ) - if ( - self.timeseries_pad_and_cap is not None - and temporal_util.schema_has_timeseries_features(current_schema) - ): - self._pad_and_cap_transformer = timeseries_transform.PadAndCapTimeseries( - current_schema, self.timeseries_pad_and_cap + if self.per_sample_transforms is not None: + if ( + self.per_sample_transforms.timedelta_extraction is not None + and self._ts_feature is None + ): + raise ValueError( + "Timedelta extraction is configured, but no creation timestamp" + f" feature was found in node set '{self._target_nodeset}'." + ) + self._per_sample_transform = self.per_sample_transforms.make( + current_schema + ) + else: + self._per_sample_transform = timeseries_transform.PerSampleTransform( + timeseries_transform.PerSampleTransformConfig(), current_schema ) - current_schema = self._pad_and_cap_transformer.output_schema() + + current_schema = self._per_sample_transform.output_schema() if temporal_util.schema_has_dynamic_timeseries_features(current_schema): raise ValueError( "Dynamic shape timeseries features were detected in the schema;" - " please configure `timeseries_pad_and_cap` to pad/cap the sequences" - " first." - ) - - if self.timedelta_extraction is not None: - if self._ts_feature is None: - raise ValueError( - "Timedelta extraction is configured, but no creation timestamp" - f" feature was found in node set '{self._target_nodeset}'." - ) - self._timestamp_extractor = ( - timeseries_transform.TimestampFeatureExtractor( - current_schema, self.timedelta_extraction - ) + " please configure `timeseries_pad_and_cap` in" + " `per_sample_transforms` to pad/cap the sequences first." ) - current_schema = self._timestamp_extractor.output_schema() self._output_schema = current_schema + self._has_per_sample_transforms = ( + self._per_sample_transform.has_transforms() + ) if self.sampler_returns_node_idxs_only and self._has_per_sample_transforms: raise ValueError( - "Per-sample transformations (`timeseries_pad_and_cap`," - " `timedelta_extraction`) cannot be used when" - " `sampler_returns_node_idxs_only=True`." + "Per-sample transformations (`per_sample_transforms`) cannot be" + " used when `sampler_returns_node_idxs_only=True`." ) self.batch_iterator, self.single_iterator = self.iterator_builder() + @property + def has_per_sample_transforms(self) -> bool: + """Returns True if any per-sample transformations are active.""" + return self._has_per_sample_transforms + def output_schema(self) -> schema_lib.GraphSchema: """Returns the schema of the generated graph samples.""" return self._output_schema @@ -298,9 +323,8 @@ def set_sampler_returns_node_idxs_only( """Changes whether the sampler returns only node indices.""" if sampler_returns_node_idxs_only and self._has_per_sample_transforms: raise ValueError( - "Per-sample transformations (`timeseries_pad_and_cap`," - " `timedelta_extraction`) cannot be used when" - " `sampler_returns_node_idxs_only=True`." + "Per-sample transformations (`per_sample_transforms`) cannot be" + " used when `sampler_returns_node_idxs_only=True`." ) self.sampler_returns_node_idxs_only = sampler_returns_node_idxs_only if self.in_memory_sampler is not None: @@ -346,6 +370,35 @@ def _get_merge_schema(self) -> schema_lib.GraphSchema: return schema_lib.GraphSchema(node_sets=node_sets, edge_sets=edge_sets) + def _extract_seed_timestamp( + self, sample: in_memory_graph.InMemoryGraph + ) -> Optional[int]: + """Extracts the root seed node timestamp from a sample if available.""" + if ( + self.per_sample_transforms is not None + and self.per_sample_transforms.timedelta_extraction is not None + and self._ts_feature is not None + and self._target_nodeset in sample.node_sets + ): + feats = sample.node_sets[self._target_nodeset].features + if self._ts_feature in feats and len(feats[self._ts_feature]) > 0: + return int(feats[self._ts_feature][0]) + return None + + def _transform_sample( + self, + sample: in_memory_graph.InMemoryGraph, + seed_timestamp: Optional[int] = None, + ) -> in_memory_graph.InMemoryGraph: + """Applies per-sample transformations to a sample.""" + if not self._has_per_sample_transforms: + return sample + if seed_timestamp is None: + seed_timestamp = self._extract_seed_timestamp(sample) + return self._per_sample_transform.transform_sample( + sample, seed_timestamp=seed_timestamp + ) + def _generator_from_in_memory_graph( self, ) -> Tuple[BatchSampleGeneratorIteratorFn, SingleSampleGeneratorIteratorFn]: @@ -362,7 +415,7 @@ def _generator_from_in_memory_graph( def batch_generator(): assert self.in_memory_sampler is not None for node_idxs in util.batch_indices_generator( - self.seed_node_idxs # pyrefly: ignore[bad-argumet-type] + self.seed_node_idxs # pyrefly: ignore[bad-argument-type] if self.seed_node_idxs is not None else self.num_seed_nodes, batch_size=self.batch_size, @@ -382,17 +435,9 @@ def batch_generator(): # Check if per-sample transforms are configured for efficiency. if self._has_per_sample_transforms: - transformed_samples = [] - for i, sample in enumerate(graph_samples): - curr_sg = sample - if self._pad_and_cap_transformer is not None: - curr_sg = self._pad_and_cap_transformer(curr_sg) - if self._timestamp_extractor is not None: - assert seed_timestamps is not None - st = int(seed_timestamps[i]) - curr_sg = self._timestamp_extractor(curr_sg, seed_timestamp=st) - transformed_samples.append(curr_sg) - graph_samples = transformed_samples + graph_samples = self._per_sample_transform.transform_sample_list( + graph_samples, seed_timestamps=seed_timestamps + ) try: yield merge_lib.merge_graphs( @@ -413,17 +458,19 @@ def single_generator(): shuffle=self.shuffle, ): node_idx = int(node_idxs[0]) - if self.temporal: - # Get the seed node timestamp - seed_timestamp = None - if self._seed_timestamps_all is not None: - seed_timestamp = int(self._seed_timestamps_all[node_idx]) + seed_timestamp = None + if self._seed_timestamps_all is not None: + seed_timestamp = int(self._seed_timestamps_all[node_idx]) - yield self.in_memory_sampler.sample( + if self.temporal: + sample = self.in_memory_sampler.sample( node_idx, seed_timestamps=seed_timestamp ) else: - yield self.in_memory_sampler.sample(node_idx) + sample = self.in_memory_sampler.sample(node_idx) + + sample = self._transform_sample(sample, seed_timestamp=seed_timestamp) + yield sample return batch_generator, single_generator @@ -432,6 +479,7 @@ def _generator_from_path_tf_sample( ) -> Tuple[BatchSampleGeneratorIteratorFn, SingleSampleGeneratorIteratorFn]: """Creates a SampleGenerator from a path to a bagz file.""" assert isinstance(self.graph, str) + merge_schema = self.output_schema() def batch_generator(): # TODO(gbm): Use pygrain and add shuffling? @@ -442,27 +490,31 @@ def batch_generator(): batch = [] try: for _ in range(self.batch_size): - batch.append(next(it)) + sample = self._transform_sample(next(it)) + batch.append(sample) except StopIteration: if batch and not self.drop_remainder: try: yield merge_lib.merge_graphs( - batch, self.schema, padding=self.padding + batch, merge_schema, padding=self.padding ) except merge_lib.InsufficientPaddingError as e: if not self.skip_overflow_padding_error: raise e return try: - yield merge_lib.merge_graphs(batch, self.schema, padding=self.padding) + yield merge_lib.merge_graphs( + batch, merge_schema, padding=self.padding + ) except merge_lib.InsufficientPaddingError as e: if not self.skip_overflow_padding_error: raise e def single_generator(): - return tf_graph_sample.read_tfgnn_graphs( + for sample in tf_graph_sample.read_tfgnn_graphs( self.graph, self.schema, container_type=container_type - ) + ): + yield self._transform_sample(sample) return batch_generator, single_generator diff --git a/dgf/src/learning/ten_lines/dataset_test.py b/dgf/src/learning/ten_lines/dataset_test.py index ab59bb5..57a8511 100644 --- a/dgf/src/learning/ten_lines/dataset_test.py +++ b/dgf/src/learning/ten_lines/dataset_test.py @@ -278,8 +278,10 @@ def test_per_sample_transformations(self): temporal_sampling=True, ), temporal=True, - timeseries_pad_and_cap=pad_and_cap_config, - timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig(), + per_sample_transforms=timeseries_transform.PerSampleTransformConfig( + timeseries_pad_and_cap=pad_and_cap_config, + timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig(), + ), drop_remainder=False, shuffle=False, ) @@ -406,7 +408,9 @@ def test_sampler_returns_node_idxs_only_with_transforms_raises(self): temporal_sampling=True, ), temporal=True, - timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig(), + per_sample_transforms=timeseries_transform.PerSampleTransformConfig( + timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig() + ), sampler_returns_node_idxs_only=True, drop_remainder=False, shuffle=False, @@ -415,7 +419,7 @@ def test_sampler_returns_node_idxs_only_with_transforms_raises(self): def test_non_in_memory_format_with_transforms_raises(self): _, schema = self._create_temporal_test_graph_and_schema() with self.assertRaisesRegex( - ValueError, + NotImplementedError, "only supported for GraphFormat.IN_MEMORY_GRAPH", ): dataset.SampleGeneratorFromAnything( @@ -446,10 +450,12 @@ def test_timedelta_extraction_with_temporal_false(self): num_hops=1, hop_width=2, ), - timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig( - sequence_length=5 + per_sample_transforms=timeseries_transform.PerSampleTransformConfig( + timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig( + sequence_length=5 + ), + timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig(), ), - timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig(), temporal=False, drop_remainder=False, shuffle=False, @@ -484,7 +490,9 @@ def test_dynamic_set_sampler_returns_node_idxs_only_raises(self): temporal_sampling=True, ), temporal=True, - timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig(), + per_sample_transforms=timeseries_transform.PerSampleTransformConfig( + timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig() + ), drop_remainder=False, shuffle=False, ) @@ -519,7 +527,9 @@ def test_timedelta_extraction_without_pad_and_cap_dynamic_ts_raises(self): temporal_sampling=True, ), temporal=True, - timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig(), + per_sample_transforms=timeseries_transform.PerSampleTransformConfig( + timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig() + ), drop_remainder=False, shuffle=False, ) diff --git a/dgf/src/learning/ten_lines/node_prediction_dataset.py b/dgf/src/learning/ten_lines/node_prediction_dataset.py index 0614e4b..5d9060c 100644 --- a/dgf/src/learning/ten_lines/node_prediction_dataset.py +++ b/dgf/src/learning/ten_lines/node_prediction_dataset.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prepare training datast for the 10-line node prediction model. +"""Prepare training dataset for the 10-line node prediction model. The main utility GNNDatasetPreparator generates normalized, padded and batched graph samples for node prediction. @@ -35,7 +35,10 @@ from dgf.src.sampling import config as sampling_config_lib from dgf.src.sampling import in_memory_sampler as in_memory_sampler_lib from dgf.src.transform import normalize as normalize_lib +from dgf.src.transform import temporal as temporal_transform +from dgf.src.transform import timeseries as timeseries_transform from dgf.src.util import log +from dgf.src.util import temporal as temporal_util from dgf.src.util import util import jax import jax.numpy as jnp @@ -119,22 +122,18 @@ class GNNDatasetPreparator: verbose_preparation: If true, display a progress bar during the preparation stage that appears only the first time `generate` is called. skip_overflow_padding_error: If padding is set, the merging stage can fail - if the "padding" is not large enought. If - skip_overflow_padding_error=True, such batch is skipped. If - skip_overflow_padding_error=False, and error is raised. + if the "padding" is not large enough. If skip_overflow_padding_error=True, + such batches are skipped. If skip_overflow_padding_error=False, an error + is raised. format: The format of the graph. If set to AUTO, the format will be inferred from the graph object. seed_node_idxs: Optional array of node indices to use as seeds for sampling. If None, all nodes in the seed nodeset are used. temporal_sampling: True if the sampling relies on timestamp features to condition the sampling. - edgeset_timestamp_features: Optional dictionary mapping edgeset names to the - feature name containing timestamp information. - nodeset_timestamp_features: Optional dictionary mapping nodeset names to the - feature name containing timestamp information. cache_normalized_features: If True, pre-compute the normalized features during the preparation stage instead of computing them on the fly in the - generator. This optioncan speeds up data generation/training but increases + generator. This option speeds data generation/training but increases memory consumption. This option is only available for in memory graph inputs. cache_normalized_features_device: Specifies the device ("host" for RAM or @@ -162,27 +161,63 @@ class GNNDatasetPreparator: verbose_preparation: bool = True skip_overflow_padding_error: bool = False temporal_sampling: bool = False - edgeset_timestamp_features: Dict[str, str] = dataclasses.field( - default_factory=dict - ) - nodeset_timestamp_features: Dict[str, str] = dataclasses.field( - default_factory=dict - ) cache_normalized_features: bool = True cache_normalized_features_device: Literal["host", "device"] = "device" - # The is_prepared data computed by the `prepare()` method. + # Computed fields + edgeset_timestamp_features: Dict[str, str] = dataclasses.field(init=False) + nodeset_timestamp_features: Dict[str, str] = dataclasses.field(init=False) + per_sample_transforms: Optional[ + timeseries_transform.PerSampleTransformConfig + ] = dataclasses.field(init=False, default=None) live: Optional[LiveData] = dataclasses.field(init=False, default=None) + @property + def target_has_creation_time(self) -> bool: + """Returns True if the target nodeset has a creation time feature.""" + target_nodeset = self.sampling_plan.root.nodeset + return ( + temporal_util.creation_time_feature_name( + self.schema.node_sets[target_nodeset].features + ) + is not None + ) + def __post_init__(self): if self.batch_size <= 0: raise ValueError(f"batch_size must be positive, got {self.batch_size}") + self.nodeset_timestamp_features = temporal_util.nodeset_timestamp_features( + self.schema + ) + self.edgeset_timestamp_features = temporal_util.edgeset_timestamp_features( + self.schema + ) + + def _get_per_sample_transforms( + self, + ) -> Optional[timeseries_transform.PerSampleTransformConfig]: + pad_and_cap_config = ( + timeseries_transform.PadAndCapTimeseriesConfig() + if temporal_util.schema_has_timeseries_features(self.schema) + else None + ) + timestamp_config = ( + timeseries_transform.TimestampFeatureExtractorConfig() + if self.target_has_creation_time + else None + ) + if pad_and_cap_config is not None or timestamp_config is not None: + return timeseries_transform.PerSampleTransformConfig( + timeseries_pad_and_cap=pad_and_cap_config, + timedelta_extraction=timestamp_config, + ) + return None def get_live(self) -> LiveData: if self.live is None: raise RuntimeError( - "The dataset preparator has not been is_prepared yet. Call" - " `prepare()` or `generate()` at least once before calling" + "The dataset preparator has not been prepared yet. Call `prepare()`" + " or `prepare_from_existing_one()` before calling `generate()` or" " `generated_schema()`." ) return self.live @@ -201,16 +236,20 @@ def generated_schema(self) -> schema_lib.GraphSchema: def prepare_from_existing_one(self, other: "GNNDatasetPreparator"): """Pre-compute data (same as "prepare") but with an already computed cache. - Instead of beeing recomputed, the following are grabbed from "other": + Instead of being recomputed, the following are grabbed from "other": - Feature statistics - - Normalier + - Normalizer - Padding - Sampling plan Args: other: Another dataset preparator from which to copy the data. """ + if self.live is not None: + raise ValueError("prepare_from_existing_one() can only be called once.") + other_live = other.get_live() + self.per_sample_transforms = other.per_sample_transforms sample_generator = dataset.SampleGeneratorFromAnything( graph=self.graph, @@ -224,8 +263,9 @@ def prepare_from_existing_one(self, other: "GNNDatasetPreparator"): skip_overflow_padding_error=self.skip_overflow_padding_error, padding=other_live.padding, temporal=self.temporal_sampling, - edgeset_timestamp_features=self.edgeset_timestamp_features, - nodeset_timestamp_features=self.nodeset_timestamp_features, + edgeset_timestamp_features=self.edgeset_timestamp_features or {}, + nodeset_timestamp_features=self.nodeset_timestamp_features or {}, + per_sample_transforms=self.per_sample_transforms, ) self.live = LiveData( @@ -237,6 +277,12 @@ def prepare_from_existing_one(self, other: "GNNDatasetPreparator"): sample_generator=sample_generator, ) + # If the sample generator has per-sample transforms, we cannot cache the + # normalized features. + # TODO(simonmeierhans): Find a more robust way to ensure this. + if sample_generator.has_per_sample_transforms: + self.cache_normalized_features = False + if self.cache_normalized_features: if isinstance(self.graph, in_memory_graph_lib.InMemoryGraph): sample_generator.set_sampler_returns_node_idxs_only(True) @@ -261,10 +307,9 @@ def prepare_from_existing_one(self, other: "GNNDatasetPreparator"): self.cache_normalized_features = False def prepare(self): - """Pre-compute and pre-pare what is necessary for the generation. + """Pre-compute and prepare what is necessary for the generation. - Can only be called once. Called automatically the first time "generate" is - called. + Can only be called once. """ if self.live is not None: raise ValueError("prepare() can only be called once.") @@ -273,8 +318,11 @@ def prepare(self): if self.verbose_preparation: log.info("Create graph sampler") + per_sample_transforms = self._get_per_sample_transforms() + self.per_sample_transforms = per_sample_transforms + # Convert the user input graph into a generator of batched graph samples. - # Note: Those samples are neither normalized not padded. + # Note: Those samples are neither normalized nor padded. sample_generator = dataset.SampleGeneratorFromAnything( graph=self.graph, schema=self.schema, @@ -286,8 +334,9 @@ def prepare(self): format=self.format, skip_overflow_padding_error=self.skip_overflow_padding_error, temporal=self.temporal_sampling, - edgeset_timestamp_features=self.edgeset_timestamp_features, - nodeset_timestamp_features=self.nodeset_timestamp_features, + edgeset_timestamp_features=self.edgeset_timestamp_features or {}, + nodeset_timestamp_features=self.nodeset_timestamp_features or {}, + per_sample_transforms=per_sample_transforms, ) # A generator of non-normalized graph samples. @@ -299,21 +348,23 @@ def gen_raw_samples(): if self.num_samples_for_stats is not None: gen_raw_samples_iter = itertools.islice( gen_raw_samples_iter, - self.num_samples_for_stats // self.batch_size, + max(1, self.num_samples_for_stats // self.batch_size), ) + effective_schema = sample_generator.output_schema() + # Compute the feature statistics (for the feature normalization) if self.verbose_preparation: log.info("Compute feature statistics") feature_stats = ( in_process_feature_statistics_lib.feature_statistics_from_graphs( - gen_raw_samples_iter, self.schema + gen_raw_samples_iter, effective_schema ) ) if self.verbose_preparation: log.info(" %s", feature_stats) normalizer = normalize_lib.auto_normalize( - schema=self.schema, + schema=effective_schema, stats=feature_stats, config=self.auto_normalize_config, ) @@ -328,14 +379,14 @@ def gen_normalized_samples(): if self.num_samples_for_stats is not None: gen_normalized_samples_iter = itertools.islice( gen_normalized_samples_iter, - self.num_samples_for_stats // self.batch_size, + max(1, self.num_samples_for_stats // self.batch_size), ) # Compute the batched graph statistics (for the padding) if self.verbose_preparation: log.info("Compute graph statistics for padding") padding = padding_lib.padding_from_graph_generator( - self.schema, gen_normalized_samples_iter + effective_schema, gen_normalized_samples_iter ) if self.verbose_preparation: log.info( @@ -353,6 +404,9 @@ def gen_normalized_samples(): sample_generator=sample_generator, ) + if sample_generator.has_per_sample_transforms: + self.cache_normalized_features = False + if self.cache_normalized_features: if isinstance(self.graph, in_memory_graph_lib.InMemoryGraph): # Cache the normalized data. @@ -496,6 +550,8 @@ def compute_train_and_valid_node_idxs( train_seed_nodes: Optional[common.SeedNodeIdxs], valid_seed_nodes: Optional[common.SeedNodeIdxs], max_num_valid_examples: Optional[int], + schema: Optional[schema_lib.GraphSchema] = None, + batch_size: int = 1, ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: """Computes the training and validation seed node indices.""" if not isinstance(graph, in_memory_graph_lib.InMemoryGraph) or ( @@ -510,6 +566,9 @@ def compute_train_and_valid_node_idxs( ) return None, None + if schema is not None and target_nodeset not in schema.node_sets: + raise ValueError(f"Target node set '{target_nodeset}' not found in schema.") + num_graph_seed_nodes = graph.node_sets[target_nodeset].num_nodes if valid_graph is None: num_valid_graph_seed_nodes = num_graph_seed_nodes @@ -521,14 +580,14 @@ def compute_train_and_valid_node_idxs( if train_seed_nodes is not None: return np.array(train_seed_nodes), ( - np.array(valid_seed_nodes) if valid_seed_nodes else None + np.array(valid_seed_nodes) if valid_seed_nodes is not None else None ) if valid_seed_nodes is not None: if valid_graph is None: raise ValueError( "`valid_seed_nodes` can only be specified when `train_seed_nodes` is" - " also specified if not validation graph (valid_graph) is provided." + " also specified if no validation graph (valid_graph) is provided." ) return None, np.array(valid_seed_nodes) @@ -541,12 +600,30 @@ def compute_train_and_valid_node_idxs( ) return None, None - train_seed_node_idxs, valid_seed_node_idxs = util.split_train_valid( - num_graph_seed_nodes, - validation_ratio, - random_seed, - max_num_valid_examples=max_num_valid_examples, - ) + ts_feature = None + if schema is not None: + ts_feature = temporal_util.creation_time_feature_name( + schema.node_sets[target_nodeset].features + ) + + if ts_feature is not None: + timestamps = graph.node_sets[target_nodeset].features[ts_feature] + train_seed_node_idxs, valid_seed_node_idxs = ( + util.split_train_valid_temporal( + creation_times=timestamps, + validation_ratio=validation_ratio, + batch_size=batch_size, + max_num_valid_examples=max_num_valid_examples, + ) + ) + else: + train_seed_node_idxs, valid_seed_node_idxs = util.split_train_valid( + num_graph_seed_nodes, + validation_ratio, + random_seed, + batch_size=batch_size, + max_num_valid_examples=max_num_valid_examples, + ) log.info( "Num. training seed nodes: %d, Num. validation seed nodes: %d", len(train_seed_node_idxs), @@ -557,7 +634,7 @@ def compute_train_and_valid_node_idxs( def prepare_datasets( graph: common.Graph, - valid_graph: common.Graph, + valid_graph: Optional[common.Graph], schema: schema_lib.GraphSchema, target_nodeset: str, random_seed: int, @@ -567,16 +644,14 @@ def prepare_datasets( verbose: int, graph_format: Union[dataset.GraphFormat, str], validation_ratio: float, - train_seed_nodes: Optional[common.SeedNodeIdxs], - valid_seed_nodes: Optional[common.SeedNodeIdxs], - temporal_sampling: bool, - nodeset_timestamp_features: dict[str, str], - edgeset_timestamp_features: dict[str, str], - num_valid_steps: Optional[int], - cache_valid_dataset: bool, - cache_normalized_features: bool, - cache_normalized_features_device: Literal["host", "device"], - sampling_plan: Optional[sampling_config_lib.SamplingPlan], + train_seed_nodes: Optional[common.SeedNodeIdxs] = None, + valid_seed_nodes: Optional[common.SeedNodeIdxs] = None, + temporal_sampling: bool = False, + num_valid_steps: Optional[int] = None, + cache_valid_dataset: bool = False, + cache_normalized_features: bool = True, + cache_normalized_features_device: Literal["host", "device"] = "device", + sampling_plan: Optional[sampling_config_lib.SamplingPlan] = None, auto_normalize_config: Optional[normalize_lib.AutoNormalizeConfig] = None, keep_raw_features: Optional[set[str]] = None, ) -> Tuple["GNNDatasetPreparator", Optional["GNNDatasetPreparator"]]: @@ -597,9 +672,43 @@ def prepare_datasets( train_seed_nodes=train_seed_nodes, valid_seed_nodes=valid_seed_nodes, max_num_valid_examples=max_num_valid_examples, + schema=schema, + batch_size=batch_size, + ) + ) + + target_has_creation_time = ( + temporal_util.creation_time_feature_name( + schema.node_sets[target_nodeset].features ) + is not None ) + if target_has_creation_time: + + orig_schema = schema + if isinstance(graph, in_memory_graph_lib.InMemoryGraph): + graph, schema = temporal_transform.propagate_timestamp_to_edges( + graph, schema + ) + else: + log.warning( + "Automatic edge timestamp propagation is only supported for" + " InMemoryGraph inputs. Skipping edge timestamp propagation for" + " training graph." + ) + if valid_graph is not None: + if isinstance(valid_graph, in_memory_graph_lib.InMemoryGraph): + valid_graph, _ = temporal_transform.propagate_timestamp_to_edges( + valid_graph, orig_schema + ) + else: + log.warning( + "Automatic edge timestamp propagation is only supported for" + " InMemoryGraph inputs. Skipping edge timestamp propagation for" + " validation graph." + ) + if sampling_plan is None: sampling_config = sampling_config_lib.SimpleSamplingConfig( seed_nodeset=target_nodeset, @@ -614,11 +723,18 @@ def prepare_datasets( if auto_normalize_config is None: auto_normalize_config = normalize_lib.AutoNormalizeConfig( - keep_raw_features=keep_raw_features or set(), + keep_raw_features=set(keep_raw_features or ()), ignore_features_without_stats=True, + timedelta_sinusoid=target_has_creation_time, + ) + else: + merged_raw_features = set(auto_normalize_config.keep_raw_features) + if keep_raw_features is not None: + merged_raw_features.update(keep_raw_features) + auto_normalize_config = dataclasses.replace( + auto_normalize_config, + keep_raw_features=merged_raw_features, ) - elif keep_raw_features is not None: - auto_normalize_config.keep_raw_features.update(keep_raw_features) common_kwargs = { "format": graph_format, @@ -630,8 +746,6 @@ def prepare_datasets( "auto_normalize_config": auto_normalize_config, "skip_overflow_padding_error": True, "temporal_sampling": temporal_sampling, - "nodeset_timestamp_features": nodeset_timestamp_features, - "edgeset_timestamp_features": edgeset_timestamp_features, "cache_normalized_features": cache_normalized_features, "cache_normalized_features_device": cache_normalized_features_device, } diff --git a/dgf/src/learning/ten_lines/node_prediction_dataset_test.py b/dgf/src/learning/ten_lines/node_prediction_dataset_test.py index fe5f10a..e19ec68 100644 --- a/dgf/src/learning/ten_lines/node_prediction_dataset_test.py +++ b/dgf/src/learning/ten_lines/node_prediction_dataset_test.py @@ -12,17 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Test converting between heterogeneous graph edge types.""" +"""Tests for node prediction dataset preparation.""" import os from absl.testing import absltest from absl.testing import parameterized +from dgf.src.data import in_memory_graph +from dgf.src.data import jax_in_memory_graph +from dgf.src.data import schema as schema_lib from dgf.src.io import jax as jax_io_lib from dgf.src.io import tf_graph_sample from dgf.src.learning.ten_lines import node_prediction_dataset from dgf.src.sampling import config as sampling_config_lib +from dgf.src.transform import normalize as normalize_lib from dgf.src.util import gen_test_graph from dgf.src.validate import in_memory_graph as in_memory_graph_validate_lib +import jax import numpy as np @@ -36,7 +41,6 @@ class GNNDatasetPreparatorTest(parameterized.TestCase): def test_in_memory_graph(self, cache_features, cache_device): graph = gen_test_graph.generate_in_memory_graph(True, False) schema = gen_test_graph.generate_schema(True, False, True, False) - print("schema:\n", schema) sampling_config = sampling_config_lib.SimpleSamplingConfig( seed_nodeset="n1", num_hops=2, @@ -58,8 +62,6 @@ def test_in_memory_graph(self, cache_features, cache_device): ) self.assertFalse(preparator.is_prepared()) preparator.prepare() - self.assertTrue(preparator.is_prepared()) - num_batches = 0 num_graphs = 0 normalized_schema = preparator.get_live().normalizer.output_schema() @@ -91,15 +93,15 @@ def sanitize(g): def test_in_memory_temporal_graph(self): graph, schema = gen_test_graph.generate_temporal_in_memory_graph(False) - sampling_config = sampling_config_lib.SimpleSamplingConfig( - seed_nodeset="n1", - num_hops=2, - hop_width=3, - reverse=True, - temporal_sampling=True, - ) sampling_plan = sampling_config_lib.simple_sampling_config_to_sampling_plan( - sampling_config, schema + sampling_config_lib.SimpleSamplingConfig( + seed_nodeset="n1", + num_hops=2, + hop_width=3, + reverse=True, + temporal_sampling=True, + ), + schema, ) preparator = node_prediction_dataset.GNNDatasetPreparator( graph=graph, @@ -109,25 +111,13 @@ def test_in_memory_temporal_graph(self): drop_remainder=True, shuffle=True, temporal_sampling=True, - nodeset_timestamp_features={"n1": "timestamp"}, - edgeset_timestamp_features={"e1": "timestamp"}, cache_normalized_features_device="host", ) self.assertFalse(preparator.is_prepared()) preparator.prepare() - self.assertTrue(preparator.is_prepared()) - - num_batches = 0 - num_graphs = 0 - normalized_schema = preparator.get_live().normalizer.output_schema() - for graph_sample, merge_offset in preparator.generate(): - in_memory_graph_validate_lib.validate_graph( - graph_sample, normalized_schema, raise_on_warning=False - ) - num_batches += 1 - num_graphs += len(merge_offset["n1"]) - 1 - self.assertEqual(num_batches, 2) - self.assertEqual(num_graphs, 4) + batches = list(preparator.generate()) + self.assertLen(batches, 2) + self.assertEqual(sum(len(offset["n1"]) - 1 for _, offset in batches), 4) def test_tf_gnn_samples(self): tmpdir = self.create_tempdir().full_path @@ -167,8 +157,6 @@ def in_mem_graphs(): ) self.assertFalse(preparator.is_prepared()) preparator.prepare() - self.assertTrue(preparator.is_prepared()) - num_batches = 0 num_graphs = 0 normalized_schema = preparator.get_live().normalizer.output_schema() @@ -181,6 +169,325 @@ def in_mem_graphs(): self.assertEqual(num_batches, 10) self.assertEqual(num_graphs, 20) + def _create_preparator( + self, + graph: in_memory_graph.InMemoryGraph, + schema: schema_lib.GraphSchema, + seed_nodeset: str, + ) -> node_prediction_dataset.GNNDatasetPreparator: + sampling_plan = sampling_config_lib.simple_sampling_config_to_sampling_plan( + sampling_config_lib.SimpleSamplingConfig( + seed_nodeset=seed_nodeset, num_hops=1, hop_width=2, reverse=True + ), + schema, + ) + return node_prediction_dataset.GNNDatasetPreparator( + graph=graph, + schema=schema, + sampling_plan=sampling_plan, + batch_size=2, + shuffle=True, + drop_remainder=True, + ) + + def test_compute_train_and_valid_node_idxs_temporal(self): + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph() + ) + train_idx, valid_idx = ( + node_prediction_dataset.compute_train_and_valid_node_idxs( + graph=graph, + valid_graph=None, + graph_format="IN_MEMORY_GRAPH", + target_nodeset="alerts", + random_seed=42, + validation_ratio=0.33, + train_seed_nodes=None, + valid_seed_nodes=None, + max_num_valid_examples=None, + schema=schema, + batch_size=2, + ) + ) + self.assertIsNotNone(train_idx) + self.assertIsNotNone(valid_idx) + np.testing.assert_array_equal(train_idx, np.array([0, 1, 2, 3])) + np.testing.assert_array_equal(valid_idx, np.array([4, 5])) + + def test_compute_train_and_valid_node_idxs_temporal_unsorted(self): + schema = schema_lib.GraphSchema( + node_sets={ + "alerts": schema_lib.NodeSchema( + features={ + "creation_time": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_creation_time=True, + ), + } + ), + }, + edge_sets={}, + ) + # Timestamps out of order: index 1 is first (100) and index 3 is last (600) + alerts_nodes = in_memory_graph.InMemoryNodeSet( + num_nodes=6, + features={ + "creation_time": np.array( + [500, 100, 300, 600, 200, 400], dtype=np.int64 + ), + }, + ) + graph = in_memory_graph.InMemoryGraph( + node_sets={"alerts": alerts_nodes}, edge_sets={} + ) + train_idx, valid_idx = ( + node_prediction_dataset.compute_train_and_valid_node_idxs( + graph=graph, + valid_graph=None, + graph_format="IN_MEMORY_GRAPH", + target_nodeset="alerts", + random_seed=42, + validation_ratio=0.33, + train_seed_nodes=None, + valid_seed_nodes=None, + max_num_valid_examples=None, + schema=schema, + batch_size=2, + ) + ) + self.assertIsNotNone(train_idx) + self.assertIsNotNone(valid_idx) + # Chronological sorted order of timestamps: + # 1 (100), 4 (200), 2 (300), 5 (400) -> train + # 0 (500), 3 (600) -> valid + np.testing.assert_array_equal(train_idx, np.array([1, 4, 2, 5])) + np.testing.assert_array_equal(valid_idx, np.array([0, 3])) + + def test_compute_node_idxs_invalid_target_nodeset_raises(self): + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph() + ) + with self.assertRaisesRegex( + ValueError, "Target node set 'unknown_nodes' not found in schema." + ): + node_prediction_dataset.compute_train_and_valid_node_idxs( + graph=graph, + valid_graph=None, + graph_format="IN_MEMORY_GRAPH", + target_nodeset="unknown_nodes", + random_seed=42, + validation_ratio=0.33, + train_seed_nodes=None, + valid_seed_nodes=None, + max_num_valid_examples=None, + schema=schema, + batch_size=2, + ) + + def test_compute_train_and_valid_node_idxs_missing_timestamp_falls_back(self): + graph = gen_test_graph.generate_in_memory_graph(True, False) + schema = gen_test_graph.generate_schema(True, False, True, False) + train_idx, valid_idx = ( + node_prediction_dataset.compute_train_and_valid_node_idxs( + graph=graph, + valid_graph=None, + graph_format="IN_MEMORY_GRAPH", + target_nodeset="n1", + random_seed=42, + validation_ratio=0.33, + train_seed_nodes=None, + valid_seed_nodes=None, + max_num_valid_examples=None, + schema=schema, + batch_size=1, + ) + ) + self.assertIsNotNone(train_idx) + self.assertIsNotNone(valid_idx) + + def test_prepare_datasets_temporal_auto_detection(self): + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph() + ) + train_dataset, valid_dataset = node_prediction_dataset.prepare_datasets( + graph=graph, + valid_graph=None, + schema=schema, + target_nodeset="alerts", + random_seed=42, + batch_size=2, + num_sampling_hops=1, + sampling_width=2, + verbose=0, + graph_format="IN_MEMORY_GRAPH", + validation_ratio=0.33, + ) + self.assertIsNotNone(valid_dataset) + assert valid_dataset is not None + train_samples = [sample for sample, _ in train_dataset.generate()] + self.assertLen(train_samples, 2) + for sample in train_samples: + self.assertIn("time_mask", sample.node_sets["hardware"].features) + self.assertIn( + "time_seed_delta_SINUSOID", sample.node_sets["hardware"].features + ) + self.assertIn( + "creation_time_seed_delta_SINUSOID", + sample.node_sets["alerts"].features, + ) + self.assertEqual( + sample.node_sets["hardware"] + .features["signal_SOFT_QUANTILE"] + .shape[1], + 30, + ) + + valid_samples = [sample for sample, _ in valid_dataset.generate()] + self.assertLen(valid_samples, 1) + for sample in valid_samples: + self.assertIn( + "creation_time_seed_delta_SINUSOID", + sample.node_sets["alerts"].features, + ) + + def test_prepare_datasets_temporal_jax_generation(self): + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph() + ) + train_dataset, _ = node_prediction_dataset.prepare_datasets( + graph=graph, + valid_graph=None, + schema=schema, + target_nodeset="alerts", + random_seed=42, + batch_size=2, + num_sampling_hops=1, + sampling_width=2, + verbose=0, + graph_format="IN_MEMORY_GRAPH", + validation_ratio=0.33, + ) + batches = list(train_dataset.generate_jax()) + self.assertLen(batches, 2) + for jax_sample, jax_offsets in batches: + self.assertIsInstance(jax_sample, jax_in_memory_graph.JaxInMemoryGraph) + self.assertIn("time_mask", jax_sample.node_sets["hardware"].features) + self.assertIn( + "time_seed_delta_SINUSOID", + jax_sample.node_sets["hardware"].features, + ) + self.assertIsInstance( + jax_sample.node_sets["alerts"].features[ + "creation_time_seed_delta_SINUSOID" + ], + (jax.Array, np.ndarray), + ) + self.assertIn("alerts", jax_offsets) + self.assertIsInstance(jax_offsets["alerts"], (jax.Array, np.ndarray)) + + def test_edge_timestamp_propagation_during_prepare_datasets(self): + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph() + ) + self.assertEmpty(schema.edge_sets["alert_to_hw"].features) + self.assertEmpty(graph.edge_sets["alert_to_hw"].features) + + train_dataset, _ = node_prediction_dataset.prepare_datasets( + graph=graph, + valid_graph=None, + schema=schema, + target_nodeset="alerts", + random_seed=42, + batch_size=2, + num_sampling_hops=1, + sampling_width=2, + verbose=0, + graph_format="IN_MEMORY_GRAPH", + validation_ratio=0.33, + ) + self.assertIn( + "timestamps", + train_dataset.schema.edge_sets["alert_to_hw"].features, + ) + + def test_prepare_small_num_samples_for_stats(self): + graph = gen_test_graph.generate_in_memory_graph(True, False) + schema = gen_test_graph.generate_schema(True, False, True, False) + sampling_plan = sampling_config_lib.simple_sampling_config_to_sampling_plan( + sampling_config_lib.SimpleSamplingConfig(seed_nodeset="n1"), schema + ) + preparator = node_prediction_dataset.GNNDatasetPreparator( + graph=graph, + schema=schema, + sampling_plan=sampling_plan, + batch_size=2, + drop_remainder=True, + shuffle=True, + num_samples_for_stats=1, # Smaller than batch_size + ) + preparator.prepare() + + def test_get_transform_configs_auto_detection(self): + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph() + ) + transforms = self._create_preparator( + graph, schema, "alerts" + )._get_per_sample_transforms() + self.assertIsNotNone(transforms) + self.assertIsNotNone(transforms.timeseries_pad_and_cap) + self.assertIsNotNone(transforms.timedelta_extraction) + + def test_get_per_sample_transforms_non_temporal_returns_none(self): + graph = gen_test_graph.generate_in_memory_graph(True, False) + schema = gen_test_graph.generate_schema(True, False, True, False) + transforms = self._create_preparator( + graph, schema, "n1" + )._get_per_sample_transforms() + self.assertIsNone(transforms) + + def test_prepare_datasets_does_not_mutate_caller_auto_normalize_config(self): + graph = gen_test_graph.generate_in_memory_graph(True, False) + schema = gen_test_graph.generate_schema(True, False, True, False) + caller_config = normalize_lib.AutoNormalizeConfig( + keep_raw_features=set(["f2"]) + ) + node_prediction_dataset.prepare_datasets( + graph=graph, + valid_graph=None, + schema=schema, + target_nodeset="n1", + random_seed=42, + batch_size=1, + num_sampling_hops=1, + sampling_width=2, + verbose=0, + graph_format="IN_MEMORY_GRAPH", + validation_ratio=0.5, + keep_raw_features=set(["extra_feature"]), + auto_normalize_config=caller_config, + ) + self.assertEqual(caller_config.keep_raw_features, set(["f2"])) + + def test_compute_train_and_valid_node_idxs_numpy_seed_nodes(self): + graph = gen_test_graph.generate_in_memory_graph(True, False) + train_idx, valid_idx = ( + node_prediction_dataset.compute_train_and_valid_node_idxs( + graph=graph, + valid_graph=None, + graph_format="IN_MEMORY_GRAPH", + target_nodeset="n1", + random_seed=42, + validation_ratio=0.5, + train_seed_nodes=np.array([0, 1]), + valid_seed_nodes=np.array([2, 3]), + max_num_valid_examples=None, + ) + ) + np.testing.assert_array_equal(train_idx, np.array([0, 1])) + np.testing.assert_array_equal(valid_idx, np.array([2, 3])) + if __name__ == "__main__": absltest.main() diff --git a/dgf/src/learning/ten_lines/node_prediction_model.py b/dgf/src/learning/ten_lines/node_prediction_model.py index 7d535e4..7126417 100644 --- a/dgf/src/learning/ten_lines/node_prediction_model.py +++ b/dgf/src/learning/ten_lines/node_prediction_model.py @@ -37,6 +37,7 @@ from dgf.src.learning.jax.layers import classification as classification_lib from dgf.src.learning.jax.layers import regression as regression_lib from dgf.src.learning.ten_lines import common +from dgf.src.learning.ten_lines import dataset from dgf.src.learning.ten_lines import evaluation from dgf.src.learning.ten_lines import node_prediction_core_model from dgf.src.learning.ten_lines import report @@ -44,6 +45,8 @@ from dgf.src.sampling import in_memory_sampler as in_memory_sampler_lib from dgf.src.transform import merge as merge_lib from dgf.src.transform import normalize as normalize_lib +from dgf.src.transform import temporal as temporal_transform +from dgf.src.transform import timeseries as timeseries_transform from dgf.src.util import log from dgf.src.util import temporal as temporal_util from dgf.src.util import util @@ -160,6 +163,9 @@ class ModelData: edgeset_timestamp_features: Dict[str, str] = dataclasses.field( default_factory=dict ) + per_sample_transforms: Optional[ + timeseries_transform.PerSampleTransformConfig + ] = None final_evaluation: Optional[evaluation.Evaluation] = None # This field is serialized / deserialized manually. @@ -206,6 +212,30 @@ def name(cls) -> str: def data(self) -> ModelData: return self._data + @property + def target_has_creation_time(self) -> bool: + """Returns True if the target nodeset has a creation timestamp feature.""" + return ( + temporal_util.creation_time_feature_name( + self._data.schema.node_sets[self._data.task.target_nodeset].features + ) + is not None + ) + + def inference_schema( + self, input_features_only: bool = True + ) -> schema_lib.GraphSchema: + """Returns the schema of pre-sampled subgraphs expected during inference.""" + schema = self._data.schema + if input_features_only: + schema = schema_to_input_feature_schema(schema, self._data.task) + current_schema = dataset.sanitize_schema_for_sample_generator(schema) + if self._data.per_sample_transforms is not None: + return self._data.per_sample_transforms.make( + current_schema + ).output_schema() + return current_schema + def _internal_save(self, path: str) -> None: # TODO(gbm): Have the params saving logic in common. checkpointer = ocp.StandardCheckpointer() @@ -364,6 +394,15 @@ def predict_batch( # TODO(gbm): Implement a better fall-back option. batch_size = max(1, self._data.hparams.batch_size // 2) + # TODO(simonmeierhans): Factor out such preparation logic such that it is + # guaranteed to be consistent. + if self.target_has_creation_time and isinstance( + graph, in_memory_graph.InMemoryGraph + ): + graph, _ = temporal_transform.propagate_timestamp_to_edges( + graph, self._data.schema + ) + sampler = in_memory_sampler_lib.create_sampler( graph=graph, plan=self._data.sampling_plan, @@ -391,26 +430,44 @@ def predict_batch( ), ) - for batch_seed_node_idxs in batch_seed_node_idxs_generator: + current_schema = dataset.sanitize_schema_for_sample_generator(schema) + per_sample_transform = None + if self._data.per_sample_transforms is not None: + per_sample_transform = self._data.per_sample_transforms.make( + current_schema + ) + current_schema = per_sample_transform.output_schema() + for batch_seed_node_idxs in batch_seed_node_idxs_generator: # TODO(gbm): The sampler should consume np array directly. + seed_timestamps = None + target_nodeset = self._data.task.target_nodeset + ts_feature = temporal_util.creation_time_feature_name( + self._data.schema.node_sets[target_nodeset].features + ) + if ts_feature is not None and isinstance( + graph, in_memory_graph.InMemoryGraph + ): + timestamps = graph.node_sets[target_nodeset].features[ts_feature] + seed_timestamps = timestamps[batch_seed_node_idxs] + if self._data.temporal_sampling: - target_nodeset = self._data.task.target_nodeset - ts_feature = temporal_util.creation_time_feature_name( - self._data.schema.node_sets[target_nodeset].features - ) - seed_timestamps = None - if ts_feature is not None: - timestamps = graph.node_sets[target_nodeset].features[ts_feature] - seed_timestamps = timestamps[batch_seed_node_idxs] graph_samples = sampler.sample( batch_seed_node_idxs, seed_timestamps=seed_timestamps ) else: graph_samples = sampler.sample(batch_seed_node_idxs) + if ( + per_sample_transform is not None + and per_sample_transform.has_transforms() + ): + graph_samples = per_sample_transform.transform_sample_list( + graph_samples, seed_timestamps=seed_timestamps + ) + yield from self._predict_sub_batch( - live, schema, batch_seed_node_idxs, graph_samples + live, current_schema, batch_seed_node_idxs, graph_samples ) def _predict_sub_batch( @@ -471,7 +528,7 @@ def predict_on_graph_sample_batch( An array of probabilities for each graph sample. """ live = self._get_live() - schema = schema_to_input_feature_schema(self._data.schema, self._data.task) + schema = self.inference_schema() merged_graph, merge_offsets = merge_lib.merge_graphs( graph_samples, @@ -587,11 +644,12 @@ def evaluate_generator( if num_eval_steps is not None: graph_samples = itertools.islice(graph_samples, num_eval_steps) + current_schema = self.inference_schema(input_features_only=False) + batch_size = max(1, self._data.hparams.batch_size // 2) def batch_prediction_generator(): live = self._get_live() - schema = self._data.schema iterator = graph_samples if verbose >= 2: iterator = tqdm.tqdm(iterator, desc="Evaluation", total=num_eval_steps) @@ -601,12 +659,18 @@ def batch_prediction_generator(): batch.append(sample) if len(batch) == batch_size: yield from self._predict_sub_batch( - live, schema, np.zeros(len(batch), dtype=np.int32), batch + live, + current_schema, + np.zeros(len(batch), dtype=np.int32), + batch, ) batch = [] if batch: yield from self._predict_sub_batch( - live, schema, np.zeros(len(batch), dtype=np.int32), batch + live, + current_schema, + np.zeros(len(batch), dtype=np.int32), + batch, ) if verbose >= 1: @@ -759,10 +823,10 @@ def to_tensorflow_function( polymorphic_shapes=[(None, "(b,)")], native_serialization_platforms=["cpu", "cuda"], ) - schema = self.data().schema + current_schema = self.inference_schema() - graph_spec = io_tf_lib.schema_to_spec(schema) - graph_dict_spec = io_tf_lib.schema_to_dict_spec(schema) + graph_spec = io_tf_lib.schema_to_spec(current_schema) + graph_dict_spec = io_tf_lib.schema_to_dict_spec(current_schema) seed_node_idxs_spec = tf.TensorSpec( shape=[None], dtype=tf.int32, name="seed_node_idxs" ) @@ -849,7 +913,7 @@ def __call__( wrapper = wrapper_class( tf_apply_core_model, live.normalizer, - self.data().schema, + current_schema, self.data().padding, ) diff --git a/dgf/src/learning/ten_lines/node_prediction_test.py b/dgf/src/learning/ten_lines/node_prediction_test.py index ef545ef..2749daf 100644 --- a/dgf/src/learning/ten_lines/node_prediction_test.py +++ b/dgf/src/learning/ten_lines/node_prediction_test.py @@ -20,7 +20,7 @@ import tempfile from typing import Tuple import unittest -from unittest import mock + from absl import logging from absl.testing import absltest from absl.testing import parameterized @@ -34,8 +34,11 @@ from dgf.src.learning.ten_lines import node_prediction_model from dgf.src.learning.ten_lines import node_prediction_train as node_prediction_lib from dgf.src.sampling import in_memory_sampler as in_memory_sampler_lib +from dgf.src.transform import temporal as temporal_transform +from dgf.src.transform import timeseries as timeseries_transform from dgf.src.util import filesystem as fs from dgf.src.util import gen_test_graph +from dgf.src.util import temporal as temporal_util from dgf.src.util import test_util import jax import jax.numpy as jnp @@ -320,13 +323,32 @@ def test_save_and_load(self): self.assertTrue(np.allclose(np.sum(restored_predictions, axis=1), 1.0)) # Check that predictions on a pre-sampled graph are exactly equal + graph_for_sampling = self.graph + seed_timestamps = None + if getattr(self.model, "target_has_creation_time", False): + graph_for_sampling, _ = temporal_transform.propagate_timestamp_to_edges( + self.graph, self.schema + ) + target_nodeset = self.model.data().task.target_nodeset + ts_name = temporal_util.creation_time_feature_name( + self.schema.node_sets[target_nodeset].features + ) + if ts_name is not None: + seed_timestamps = graph_for_sampling.node_sets[target_nodeset].features[ + ts_name + ] sampler = in_memory_sampler_lib.create_sampler( - graph=self.graph, + graph=graph_for_sampling, plan=self.model.data().sampling_plan, schema=self.model.data().schema, batch_size=5, ) - sample = sampler.sample(0) + sample = sampler.sample( + 0, + seed_timestamps=int(seed_timestamps[0]) + if seed_timestamps is not None + else None, + ) predictions_on_sample = self.model.predict(graph=sample, seed_node_idxs=[0]) restored_predictions_on_sample = restored_model.predict( graph=sample, seed_node_idxs=[0] @@ -344,24 +366,75 @@ def test_save_and_load(self): ) def test_to_tensorflow_function(self, consume_tf_graph_dict: bool): + graph_for_sampling = self.graph + seed_timestamps = None + if getattr(self.model, "target_has_creation_time", False): + graph_for_sampling, _ = temporal_transform.propagate_timestamp_to_edges( + self.graph, self.schema + ) + target_nodeset = self.model.data().task.target_nodeset + ts_name = temporal_util.creation_time_feature_name( + self.schema.node_sets[target_nodeset].features + ) + if ts_name is not None: + seed_timestamps = graph_for_sampling.node_sets[target_nodeset].features[ + ts_name + ] sampler = in_memory_sampler_lib.create_sampler( - graph=self.graph, + graph=graph_for_sampling, plan=self.model.data().sampling_plan, schema=self.model.data().schema, batch_size=5, ) # Note: We use different samples to test that the traced/frozen model # can run on samples with different sizes. - sample_1 = sampler.sample(0) - sample_2 = sampler.sample(1) - - tf_sample_1 = tf_io.graph_to_tf_graph( - sample_1, schema=self.model.data().schema + sample_1 = sampler.sample( + 0, + seed_timestamps=int(seed_timestamps[0]) + if seed_timestamps is not None + else None, ) - tf_sample_2 = tf_io.graph_to_tf_graph( - sample_2, schema=self.model.data().schema + sample_2 = sampler.sample( + 1, + seed_timestamps=int(seed_timestamps[1]) + if seed_timestamps is not None + else None, ) + if getattr(self.model, "target_has_creation_time", False): + extractor = timeseries_transform.TimestampFeatureExtractor( + self.schema, + timeseries_transform.TimestampFeatureExtractorConfig(), + ) + target_nodeset = self.model.data().task.target_nodeset + ts_name = temporal_util.creation_time_feature_name( + self.schema.node_sets[target_nodeset].features + ) + assert ts_name is not None + st1 = int(sample_1.node_sets[target_nodeset].features[ts_name][0]) + sample_1 = extractor(sample_1, seed_timestamp=st1) + st2 = int(sample_2.node_sets[target_nodeset].features[ts_name][0]) + sample_2 = extractor(sample_2, seed_timestamp=st2) + + if temporal_util.schema_has_timeseries_features(self.schema): + pad_cap = timeseries_transform.PadAndCapTimeseries( + self.schema, + timeseries_transform.PadAndCapTimeseriesConfig(), + ) + sample_1 = pad_cap(sample_1) + sample_2 = pad_cap(sample_2) + + target_column = self.model.data().task.target_column + target_nodeset = self.model.data().task.target_nodeset + if target_column in sample_1.node_sets[target_nodeset].features: + del sample_1.node_sets[target_nodeset].features[target_column] + if target_column in sample_2.node_sets[target_nodeset].features: + del sample_2.node_sets[target_nodeset].features[target_column] + + inference_schema = self.model.inference_schema() + tf_sample_1 = tf_io.graph_to_tf_graph(sample_1, schema=inference_schema) + tf_sample_2 = tf_io.graph_to_tf_graph(sample_2, schema=inference_schema) + if consume_tf_graph_dict: kwargs_call_1 = { **tf_io.tf_graph_to_tf_graph_dict(tf_sample_1), @@ -382,8 +455,12 @@ def test_to_tensorflow_function(self, consume_tf_graph_dict: bool): prediction_sample_1 = tf_predict_fn(**kwargs_call_1) # pyrefly: ignore[not-callable] # Expected prediction - expected_prediction_sample_1 = self.model.predict(sample_1, [0]) - expected_prediction_sample_2 = self.model.predict(sample_2, [0]) + expected_prediction_sample_1 = self.model.predict_on_graph_sample_batch( + [sample_1] + ) + expected_prediction_sample_2 = self.model.predict_on_graph_sample_batch( + [sample_2] + ) np.testing.assert_allclose( prediction_sample_1.numpy(), expected_prediction_sample_1, atol=1e-5 @@ -429,16 +506,25 @@ def test_to_tensorflow_function(self, consume_tf_graph_dict: bool): (f"nodes_client_{h}id", [None]), ("nodes_client_city", [None]), ("nodes_client_age", [None]), - (f"nodes_client_categorical{u}label", [None]), ("nodes_transaction_reserved_size", ()), (f"nodes_transaction_{h}id", [None]), - ("nodes_transaction_date", [None]), ("nodes_transaction_amount", [None]), + ("nodes_transaction_country", [None]), ( f"edges_transation{u}to{u}client_reserved_adjacency", [2, None], ), ] + if getattr(self.model, "target_has_creation_time", False): + expected_keys_and_shapes.extend([ + (f"nodes_client_created{u}at{u}seed{u}delta", [None]), + (f"nodes_transaction_date{u}seed{u}delta", [None]), + ]) + else: + expected_keys_and_shapes.extend([ + (f"nodes_client_created{u}at", [None]), + ("nodes_transaction_date", [None]), + ]) for key, expected_shape in expected_keys_and_shapes: self.assertTrue( @@ -536,8 +622,12 @@ def test_graph_samples(self): tmpdir = self.create_tempdir().full_path path = os.path.join(tmpdir, "samples@5.tfrecord") + schema = copy.deepcopy(self.schema) + schema.node_sets["client"].features["created_at"].is_creation_time = False + schema.node_sets["transaction"].features["date"].is_creation_time = False + subgraph = synthetic_lib.generate_synthetic_graph( - self.schema, + schema, synthetic_lib.SyntheticGraphConfig(num_nodes=5, num_edges=5), ) @@ -548,14 +638,14 @@ def in_mem_graphs(): tf_graph_sample_lib.write_tfgnn_graphs( in_mem_graphs(), path, - schema=self.schema, + schema=schema, container_type="TF_RECORD", ) model = node_prediction_lib.train_node_model( graph=path, valid_graph=path, - schema=self.schema, + schema=schema, target_nodeset="client", target_column="categorical_label", **RAPID_TRAINING_KWARGS, @@ -563,6 +653,22 @@ def in_mem_graphs(): self.assertIsNone(model.data().training_stats.num_train_seed_nodes) self.assertIsNone(model.data().training_stats.num_valid_seed_nodes) + def test_graph_samples_temporal_raises_not_implemented(self): + tmpdir = self.create_tempdir().full_path + path = os.path.join(tmpdir, "samples@5.tfrecord") + with self.assertRaisesRegex( + NotImplementedError, + "Temporal sampling.*is only supported for.*IN_MEMORY_GRAPH", + ): + node_prediction_lib.train_node_model( + graph=path, + valid_graph=path, + schema=self.schema, + target_nodeset="client", + target_column="categorical_label", + **RAPID_TRAINING_KWARGS, + ) + def test_predict_batch_insufficient_padding(self): """Tests that predict_batch handles InsufficientPaddingError by splitting.""" @@ -777,7 +883,6 @@ def train_model() -> node_prediction_lib.NodePredictionModel: schema=cls.schema, target_nodeset="client", target_column="categorical_label", - time_aware=True, **RAPID_TRAINING_KWARGS, ) @@ -862,24 +967,6 @@ def test_predict(self): class NodePredictionTemporalValidationTest(absltest.TestCase): - def test_temporal_training_no_target_nodeset_timestamp_raises(self): - graph, schema = gen_test_graph.generate_temporal_in_memory_graph( - include_e2=False - ) - schema.node_sets["n1"].features["timestamp"].is_creation_time = False - - with self.assertRaisesRegex( - ValueError, - "The target nodeset 'n1' must have a creation time feature", - ): - node_prediction_lib.train_node_model( - graph=graph, - schema=schema, - target_nodeset="n1", - target_column="timestamp", - time_aware=True, - ) - def test_temporal_training_no_edgeset_timestamp_plan_is_temporal(self): graph, schema = _gen_graph_real_looking(has_timestamp_feature=True) for es_schema in schema.edge_sets.values(): @@ -891,7 +978,6 @@ def test_temporal_training_no_edgeset_timestamp_plan_is_temporal(self): schema=schema, target_nodeset="client", target_column="categorical_label", - time_aware=True, num_train_steps=2, valid_every_n_steps=1, batch_size=1, @@ -903,7 +989,124 @@ def test_temporal_training_no_edgeset_timestamp_plan_is_temporal(self): ) self.assertTrue(model.data().temporal_sampling) self.assertTrue(model.data().sampling_plan.temporal_sampling) - self.assertEmpty(model.data().sampling_plan.edgeset_timestamp_features) + self.assertEqual( + model.data().sampling_plan.edgeset_timestamp_features, + {"transation_to_client": "timestamps"}, + ) + + +class NodePredictionPerSampleTransformTest(parameterized.TestCase): + + def test_train_temporal_timeseries_auto_infers_per_sample_transforms(self): + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph( + num_alerts=20, num_hardware=5 + ) + ) + + model = node_prediction_lib.train_node_model( + graph=graph, + schema=schema, + target_nodeset="alerts", + target_column="label", + num_train_steps=4, + valid_every_n_steps=2, + batch_size=2, + sampling_width=2, + num_sampling_hops=1, + node_embedding_dim=8, + num_layers=1, + verbose=0, + ) + + # Check that ModelData contains the auto-inferred config + per_sample_transforms = model.data().per_sample_transforms + self.assertIsNotNone(per_sample_transforms) + self.assertIsNotNone(per_sample_transforms.timeseries_pad_and_cap) + self.assertIsNotNone(per_sample_transforms.timedelta_extraction) + self.assertEqual( + per_sample_transforms.timeseries_pad_and_cap.sequence_length, 30 + ) + self.assertEqual( + per_sample_transforms.timeseries_pad_and_cap.padding_value, 0 + ) + self.assertEqual( + per_sample_transforms.timedelta_extraction.fill_value, 0 + ) + + # Verify inference_schema matches default sequence_length=30 + inf_schema = model.inference_schema() + self.assertEqual( + inf_schema.node_sets["hardware"].features["signal"].shape, (30,) + ) + self.assertEqual( + inf_schema.node_sets["hardware"].features["time_mask"].shape, (30,) + ) + self.assertEqual( + inf_schema.node_sets["hardware"].features["time_seed_delta"].shape, + (30,), + ) + self.assertIn( + "creation_time_seed_delta", + inf_schema.node_sets["alerts"].features, + ) + self.assertNotIn("label", inf_schema.node_sets["alerts"].features) + self.assertTrue(bool(model.data().edgeset_timestamp_features)) + self.assertTrue(bool(model.data().nodeset_timestamp_features)) + + # Run predictions and verify they succeed + predictions = model.predict(graph, seed_node_idxs=[0, 1, 2], verbose=0) + self.assertEqual(predictions.shape, (3,)) + + # Save and restore model and verify per_sample_transforms is preserved + with tempfile.TemporaryDirectory() as tmpdir: + model.save(tmpdir) + restored_model = common_lib.load_model(tmpdir) + + assert isinstance(restored_model, node_prediction_lib.NodePredictionModel) + self.assertEqual( + restored_model.data().per_sample_transforms, + per_sample_transforms, + ) + restored_predictions = restored_model.predict( + graph, seed_node_idxs=[0, 1, 2], verbose=0 + ) + np.testing.assert_allclose(predictions, restored_predictions, atol=1e-4) + + def test_model_data_serialization_with_per_sample_transforms(self): + config = timeseries_transform.PerSampleTransformConfig( + timeseries_pad_and_cap=timeseries_transform.PadAndCapTimeseriesConfig( + sequence_length=12, padding_value=0 + ), + timedelta_extraction=timeseries_transform.TimestampFeatureExtractorConfig(), + ) + graph, schema = ( + gen_test_graph.generate_temporal_timeseries_in_memory_graph( + num_alerts=20, num_hardware=5 + ) + ) + model = node_prediction_lib.train_node_model( + graph=graph, + schema=schema, + target_nodeset="alerts", + target_column="label", + num_train_steps=2, + valid_every_n_steps=1, + batch_size=2, + sampling_width=1, + num_sampling_hops=1, + node_embedding_dim=4, + num_layers=1, + verbose=0, + ) + model.data().per_sample_transforms = config + model_data = model.data() + json_str = model_data.to_json() # pyrefly: ignore[missing-attribute] + restored_data = node_prediction_model.ModelData.from_json(json_str) # pyrefly: ignore[missing-attribute] + self.assertEqual( + restored_data.per_sample_transforms, + config, + ) if __name__ == "__main__": diff --git a/dgf/src/learning/ten_lines/node_prediction_train.py b/dgf/src/learning/ten_lines/node_prediction_train.py index 477204e..ecb6310 100644 --- a/dgf/src/learning/ten_lines/node_prediction_train.py +++ b/dgf/src/learning/ten_lines/node_prediction_train.py @@ -163,7 +163,6 @@ def train_node_model( node_embedding_dim: int = 128, learning_rate: float = 1e-3, cache_valid_dataset: bool = True, - time_aware: bool = False, message_pooling: str = "sum", experimental_preprocess_core_model_config: Optional[ Callable[[CoreModelConfig], CoreModelConfig] @@ -225,9 +224,6 @@ def train_node_model( cache_valid_dataset: Whether to cache the validation dataset in memory. This can speed up validation if sample generation and normalization are time-consuming, but it will increase memory usage. - time_aware: Enables temporal-aware training. If `False` (default), no - temporal masking is applied. If `True`, timestamp features are inferred - from the schema (via features marked as creation timestamps). message_pooling: The pooling method to use for aggregating messages. experimental_preprocess_core_model_config: Advanced option. An optional callable to modify the `CoreModelConfig` before it is used to build the @@ -288,20 +284,12 @@ def train_node_model( f" {list(schema.node_sets.keys())}" ) - # Note: Maybe one day, the timestamp features will be used for something else. - if time_aware: - nodeset_ts_features = temporal_util.nodeset_timestamp_features(schema) - edgeset_ts_features = temporal_util.edgeset_timestamp_features(schema) - if target_nodeset not in nodeset_ts_features: - raise ValueError( - f"The target nodeset '{target_nodeset}' must have a creation time" - " feature. Set is_creation_time=True on the creation time feature" - " (e.g. `schema.node_sets[].features[].is_creation_time = True`)." + target_has_creation_time = ( + temporal_util.creation_time_feature_name( + schema.node_sets[target_nodeset].features ) - else: - nodeset_ts_features = {} - edgeset_ts_features = {} + is not None + ) if verbose >= 2: log.info( @@ -354,9 +342,7 @@ def train_node_model( validation_ratio=validation_ratio, train_seed_nodes=train_seed_nodes, valid_seed_nodes=valid_seed_nodes, - temporal_sampling=time_aware, - nodeset_timestamp_features=nodeset_ts_features, - edgeset_timestamp_features=edgeset_ts_features, + temporal_sampling=target_has_creation_time, num_valid_steps=num_valid_steps, cache_valid_dataset=cache_valid_dataset, cache_normalized_features=cache_normalized_features, @@ -675,9 +661,12 @@ def valid_step(params, opt_state, batch: Batch): padding=dataset_preparator.padding, sampling_plan=dataset_preparator.sampling_plan, feature_stats=dataset_preparator.feature_stats, - temporal_sampling=time_aware, - nodeset_timestamp_features=nodeset_ts_features, - edgeset_timestamp_features=edgeset_ts_features, + temporal_sampling=dataset_preparator.sampling_plan.temporal_sampling, + nodeset_timestamp_features=temporal_util.nodeset_timestamp_features( + schema + ), + edgeset_timestamp_features=dataset_preparator.sampling_plan.edgeset_timestamp_features, + per_sample_transforms=train_dataset.per_sample_transforms, training_stats=TrainingStats( num_train_seed_nodes=train_dataset.num_nodes_in_seed_nodeset(), num_valid_seed_nodes=valid_dataset.num_nodes_in_seed_nodeset() @@ -697,7 +686,7 @@ def valid_step(params, opt_state, batch: Batch): if verbose >= 1: log.info("Final model evaluation") model.data().final_evaluation = model.evaluate_generator( - valid_dataset.get_live().sample_generator.single_iterator() + valid_dataset.get_live().sample_generator.single_iterator() # pyrefly: ignore[missing-attribute] ) model.metadata.captured_logs = captured_logs diff --git a/dgf/src/transform/BUILD b/dgf/src/transform/BUILD index 7d51812..d2f7842 100644 --- a/dgf/src/transform/BUILD +++ b/dgf/src/transform/BUILD @@ -112,6 +112,7 @@ py_library( deps = [ "//dgf/src/data:in_memory_graph", "//dgf/src/data:schema", + "//dgf/src/util:log", "//dgf/src/util:temporal", # numpy dep, ], diff --git a/dgf/src/transform/temporal.py b/dgf/src/transform/temporal.py index 1d2d868..32c8684 100644 --- a/dgf/src/transform/temporal.py +++ b/dgf/src/transform/temporal.py @@ -17,6 +17,7 @@ from typing import List, Optional from dgf.src.data import in_memory_graph from dgf.src.data import schema as schema_lib +from dgf.src.util import log from dgf.src.util import temporal as temporal_util import numpy as np @@ -55,8 +56,17 @@ def propagate_timestamp_to_edges( """ new_edge_sets = dict(graph.edge_sets) new_edge_set_schemas = dict(schema.edge_sets) + if target_edgesets is not None: + unknown_edgesets = set(target_edgesets) - set(schema.edge_sets) + if unknown_edgesets: + raise ValueError( + f"Unknown target edgesets: {sorted(unknown_edgesets)}. Valid" + f" edgesets in schema are: {sorted(schema.edge_sets)}." + ) - def get_node_ts(nodeset_name: str): + def get_node_ts( + nodeset_name: str, + ) -> tuple[Optional[np.ndarray], Optional[schema_lib.FeatureFormat]]: feat_name = temporal_util.creation_time_feature_name( schema.node_sets[nodeset_name].features ) @@ -74,38 +84,63 @@ def get_node_ts(nodeset_name: str): if target_edgesets is not None and edgeset_name not in target_edgesets: continue - if target_feature in edgeset_schema.features: + ts_name = temporal_util.creation_time_feature_name(edgeset_schema.features) + if ( + ts_name is not None + and ts_name in graph.edge_sets[edgeset_name].features + ): + log.info( + "Skipping edge set '%s' because it already has a creation time" + " feature '%s'.", + edgeset_name, + ts_name, + ) + continue + feat_name = ts_name if ts_name is not None else target_feature + if ts_name is None and feat_name in edgeset_schema.features: raise ValueError( - f"Target feature '{target_feature}' already exists in edgeset" - f" '{edgeset_name}'." + f"Feature '{feat_name}' already exists in edgeset '{edgeset_name}'" + " and is not marked as a creation time feature." ) src_ts, src_format = get_node_ts(edgeset_schema.source) tgt_ts, tgt_format = get_node_ts(edgeset_schema.target) if src_ts is None and tgt_ts is None: - raise ValueError( - f"Neither source nodeset '{edgeset_schema.source}' nor target nodeset" - f" '{edgeset_schema.target}' has timestamps for edgeset" - f" '{edgeset_name}'." + if target_edgesets is not None: + raise ValueError( + f"Neither source nodeset '{edgeset_schema.source}' nor target" + f" nodeset '{edgeset_schema.target}' has timestamps for edgeset" + f" '{edgeset_name}'." + ) + log.info( + "Skipping edge set '%s' because neither source nodeset '%s' nor" + " target nodeset '%s' has timestamps.", + edgeset_name, + edgeset_schema.source, + edgeset_schema.target, ) + continue edgeset_value = graph.edge_sets[edgeset_name] src_indices, tgt_indices = edgeset_value.adjacency if src_ts is not None and tgt_ts is not None: + assert src_format is not None edge_ts = np.maximum(src_ts[src_indices], tgt_ts[tgt_indices]) ts_format = src_format elif src_ts is not None: + assert src_format is not None edge_ts = src_ts[src_indices] ts_format = src_format else: - edge_ts = tgt_ts[tgt_indices] # pyrefly: ignore[unsupported-operation] + assert tgt_ts is not None and tgt_format is not None + edge_ts = tgt_ts[tgt_indices] ts_format = tgt_format new_edge_sets[edgeset_name] = in_memory_graph.InMemoryEdgeSet( adjacency=edgeset_value.adjacency, - features={**edgeset_value.features, target_feature: edge_ts}, + features={**edgeset_value.features, feat_name: edge_ts}, ) new_edge_set_schemas[edgeset_name] = schema_lib.EdgeSchema( @@ -113,7 +148,7 @@ def get_node_ts(nodeset_name: str): target=edgeset_schema.target, features={ **edgeset_schema.features, - target_feature: schema_lib.FeatureSchema( + feat_name: schema_lib.FeatureSchema( format=ts_format, semantic=schema_lib.FeatureSemantic.TIMESTAMP, is_creation_time=True, diff --git a/dgf/src/transform/temporal_test.py b/dgf/src/transform/temporal_test.py index 4546d93..b113696 100644 --- a/dgf/src/transform/temporal_test.py +++ b/dgf/src/transform/temporal_test.py @@ -28,10 +28,12 @@ def test_propagate_timestamp_to_edges_basic(self): graph = in_memory_graph.InMemoryGraph( node_sets={ "n1": in_memory_graph.InMemoryNodeSet( - num_nodes=2, features={"timestamps": np.array([10, 20], dtype=np.int64)} + num_nodes=2, + features={"timestamps": np.array([10, 20], dtype=np.int64)}, ), "n2": in_memory_graph.InMemoryNodeSet( - num_nodes=2, features={"timestamps": np.array([5, 25], dtype=np.int64)} + num_nodes=2, + features={"timestamps": np.array([5, 25], dtype=np.int64)}, ), }, edge_sets={ @@ -84,7 +86,8 @@ def test_propagate_timestamp_to_edges_single_node_timestamp(self): graph = in_memory_graph.InMemoryGraph( node_sets={ "n1": in_memory_graph.InMemoryNodeSet( - num_nodes=2, features={"timestamps": np.array([10, 20], dtype=np.int64)} + num_nodes=2, + features={"timestamps": np.array([10, 20], dtype=np.int64)}, ), "n2": in_memory_graph.InMemoryNodeSet(num_nodes=2, features={}), }, @@ -140,13 +143,138 @@ def test_propagate_timestamp_to_edges_fail_no_timestamps(self): ) with self.assertRaisesRegex(ValueError, "Neither source nodeset"): + temporal.propagate_timestamp_to_edges( + graph, schema, target_edgesets=["e1"] + ) + + def test_propagate_timestamp_to_edges_unknown_edgeset_raises(self): + graph = in_memory_graph.InMemoryGraph( + node_sets={ + "n1": in_memory_graph.InMemoryNodeSet( + num_nodes=2, + features={"timestamps": np.array([10, 20], dtype=np.int64)}, + ), + }, + edge_sets={ + "e1": in_memory_graph.InMemoryEdgeSet( + adjacency=np.array([[0, 1], [1, 0]]), + ) + }, + ) + schema = schema_lib.GraphSchema( + node_sets={ + "n1": schema_lib.NodeSchema( + features={ + "timestamps": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_creation_time=True, + ) + } + ), + }, + edge_sets={"e1": schema_lib.EdgeSchema(source="n1", target="n1")}, + ) + with self.assertRaisesRegex(ValueError, "Unknown target edgesets:"): + temporal.propagate_timestamp_to_edges( + graph, schema, target_edgesets=["non_existent_edgeset"] + ) + + def test_propagate_timestamp_to_edges_skips_static_auxiliary_edges(self): + graph = in_memory_graph.InMemoryGraph( + node_sets={ + "temporal_node": in_memory_graph.InMemoryNodeSet( + num_nodes=2, + features={"timestamps": np.array([10, 20], dtype=np.int64)}, + ), + "static_node": in_memory_graph.InMemoryNodeSet( + num_nodes=2, + features={}, + ), + }, + edge_sets={ + "temporal_edge": in_memory_graph.InMemoryEdgeSet( + adjacency=np.array([[0, 1], [1, 0]]), + ), + "static_edge": in_memory_graph.InMemoryEdgeSet( + adjacency=np.array([[0, 1], [1, 0]]), + ), + }, + ) + schema = schema_lib.GraphSchema( + node_sets={ + "temporal_node": schema_lib.NodeSchema( + features={ + "timestamps": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_creation_time=True, + ) + } + ), + "static_node": schema_lib.NodeSchema(features={}), + }, + edge_sets={ + "temporal_edge": schema_lib.EdgeSchema( + source="temporal_node", target="temporal_node" + ), + "static_edge": schema_lib.EdgeSchema( + source="static_node", target="static_node" + ), + }, + ) + _, new_schema = temporal.propagate_timestamp_to_edges(graph, schema) + self.assertIn("timestamps", new_schema.edge_sets["temporal_edge"].features) + self.assertNotIn("timestamps", new_schema.edge_sets["static_edge"].features) + + def test_propagate_timestamp_to_edges_collision_raises(self): + graph = in_memory_graph.InMemoryGraph( + node_sets={ + "n1": in_memory_graph.InMemoryNodeSet( + num_nodes=2, + features={"timestamps": np.array([10, 20], dtype=np.int64)}, + ), + }, + edge_sets={ + "e1": in_memory_graph.InMemoryEdgeSet( + adjacency=np.array([[0], [1]]), + features={"timestamps": np.array([1.5], dtype=np.float32)}, + ) + }, + ) + schema = schema_lib.GraphSchema( + node_sets={ + "n1": schema_lib.NodeSchema( + features={ + "timestamps": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + is_creation_time=True, + ) + } + ), + }, + edge_sets={ + "e1": schema_lib.EdgeSchema( + source="n1", + target="n1", + features={ + "timestamps": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + is_creation_time=False, + ) + }, + ) + }, + ) + with self.assertRaisesRegex( + ValueError, "already exists in edgeset 'e1' and is not marked" + ): temporal.propagate_timestamp_to_edges(graph, schema) def test_propagate_timestamp_to_edges_custom_names(self): graph = in_memory_graph.InMemoryGraph( node_sets={ "n1": in_memory_graph.InMemoryNodeSet( - num_nodes=2, features={"time": np.array([10, 20], dtype=np.int64)} + num_nodes=2, + features={"time": np.array([10, 20], dtype=np.int64)}, ), "n2": in_memory_graph.InMemoryNodeSet( num_nodes=2, features={"ts": np.array([5, 25], dtype=np.int64)} @@ -198,16 +326,17 @@ def test_propagate_timestamp_to_edges_custom_names(self): new_schema.edge_sets["e1"].features["edge_ts"].is_creation_time ) - def test_propagate_timestamp_to_edges_fail_existing_feature(self): + def test_propagate_timestamp_to_edges_skips_existing_creation_time(self): graph = in_memory_graph.InMemoryGraph( node_sets={ "n1": in_memory_graph.InMemoryNodeSet( - num_nodes=2, features={"timestamps": np.array([10, 20], dtype=np.int64)} + num_nodes=2, + features={"timestamps": np.array([10, 20], dtype=np.int64)}, ), }, edge_sets={ "e1": in_memory_graph.InMemoryEdgeSet( - adjacency=np.array([[0, 1]]), + adjacency=np.array([[0], [1]]), features={"timestamps": np.array([1], dtype=np.int64)}, ) }, @@ -219,6 +348,7 @@ def test_propagate_timestamp_to_edges_fail_existing_feature(self): features={ "timestamps": schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.INTEGER_64, + is_creation_time=True, ) } ), @@ -230,14 +360,17 @@ def test_propagate_timestamp_to_edges_fail_existing_feature(self): features={ "timestamps": schema_lib.FeatureSchema( format=schema_lib.FeatureFormat.INTEGER_64, + is_creation_time=True, ) }, ) }, ) - with self.assertRaisesRegex(ValueError, "already exists in edgeset"): - temporal.propagate_timestamp_to_edges(graph, schema) + new_graph, _ = temporal.propagate_timestamp_to_edges(graph, schema) + np.testing.assert_array_equal( + new_graph.edge_sets["e1"].features["timestamps"], np.array([1]) + ) if __name__ == "__main__": diff --git a/dgf/src/transform/timeseries.py b/dgf/src/transform/timeseries.py index bf6fb5d..778b3e9 100644 --- a/dgf/src/transform/timeseries.py +++ b/dgf/src/transform/timeseries.py @@ -581,9 +581,9 @@ def _extract_feature_set_timestamp_features( ) -> 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." - ) + 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] @@ -670,3 +670,121 @@ def __call__( return in_memory_graph.InMemoryGraph( node_sets=new_node_sets, edge_sets=new_edge_sets ) + + +@dataclasses_json.dataclass_json +@dataclasses.dataclass +class PerSampleTransformConfig: + """Configuration for per-sample preprocessing transforms. + + Attributes: + timeseries_pad_and_cap: Optional configuration for padding and capping + timeseries features. + timedelta_extraction: Optional configuration for extracting relative + timestamp features (seed deltas). + """ + + timeseries_pad_and_cap: Optional[PadAndCapTimeseriesConfig] = None + timedelta_extraction: Optional[TimestampFeatureExtractorConfig] = None + + def make(self, schema: schema_lib.GraphSchema) -> "PerSampleTransform": + return PerSampleTransform(self, schema) + + +class PerSampleTransform: + """Executes per-sample transforms on InMemoryGraph samples. + + Attributes: + config: The `PerSampleTransformConfig` specifying active transforms. + schema: The input `GraphSchema`. + """ + + def __init__( + self, + config: PerSampleTransformConfig, + schema: schema_lib.GraphSchema, + ): + self.config = config + self.schema = schema + + current_schema = schema + self._pad_and_cap_transformer: Optional[PadAndCapTimeseries] = None + self._timestamp_extractor: Optional[TimestampFeatureExtractor] = None + + if ( + config.timeseries_pad_and_cap is not None + and temporal_util.schema_has_timeseries_features(current_schema) + ): + self._pad_and_cap_transformer = PadAndCapTimeseries( + current_schema, config.timeseries_pad_and_cap + ) + current_schema = self._pad_and_cap_transformer.output_schema() + + if temporal_util.schema_has_dynamic_timeseries_features(current_schema): + raise ValueError( + "Dynamic shape timeseries features were detected in the schema;" + " please configure `timeseries_pad_and_cap` to pad/cap the sequences" + " first." + ) + + if config.timedelta_extraction is not None: + self._timestamp_extractor = TimestampFeatureExtractor( + current_schema, config.timedelta_extraction + ) + current_schema = self._timestamp_extractor.output_schema() + + self._output_schema = current_schema + + def output_schema(self) -> schema_lib.GraphSchema: + """Returns the transformed GraphSchema.""" + return self._output_schema + + def has_transforms(self) -> bool: + """Returns True if any per-sample transformer is active.""" + return ( + self._pad_and_cap_transformer is not None + or self._timestamp_extractor is not None + ) + + def transform_sample( + self, + sample: in_memory_graph.InMemoryGraph, + seed_timestamp: Optional[int] = None, + ) -> in_memory_graph.InMemoryGraph: + """Transforms a single InMemoryGraph sample.""" + curr_sg = sample + if self._pad_and_cap_transformer is not None: + curr_sg = self._pad_and_cap_transformer(curr_sg) + if self._timestamp_extractor is not None: + assert ( + seed_timestamp is not None + ), "seed_timestamp must be provided to extract seed deltas." + curr_sg = self._timestamp_extractor( + curr_sg, seed_timestamp=seed_timestamp + ) + return curr_sg + + def transform_sample_list( + self, + samples: List[in_memory_graph.InMemoryGraph], + seed_timestamps: Optional[np.ndarray] = None, + ) -> List[in_memory_graph.InMemoryGraph]: + """Transforms a list of InMemoryGraph samples.""" + if not self.has_transforms(): + return samples + transformed_samples = [] + for i, sample in enumerate(samples): + st = int(seed_timestamps[i]) if seed_timestamps is not None else None + transformed_samples.append( + self.transform_sample(sample, seed_timestamp=st) + ) + return transformed_samples + + def __call__( + self, + sample: in_memory_graph.InMemoryGraph, + seed_timestamp: Optional[int] = None, + ) -> in_memory_graph.InMemoryGraph: + """Transforms a single InMemoryGraph sample.""" + return self.transform_sample(sample, seed_timestamp=seed_timestamp) + diff --git a/dgf/src/transform/timeseries_test.py b/dgf/src/transform/timeseries_test.py index c5e24df..c19ba59 100644 --- a/dgf/src/transform/timeseries_test.py +++ b/dgf/src/transform/timeseries_test.py @@ -332,9 +332,7 @@ def test_clashing_mask_name_raises(self): schema, timeseries.PadAndCapTimeseriesConfig(sequence_length=3), ) - with self.assertRaisesRegex( - ValueError, "clashes with an existing feature" - ): + 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): @@ -814,12 +812,8 @@ def test_extract_calendar_features_parent_timestamp(self): extractor = timeseries.CalendarFeatureExtractor(schema) cal_schema = extractor.output_schema() hw_sch = cal_schema.node_sets["hardware"] - self.assertEqual( - hw_sch.features["event_time_hour"].group, "master_time" - ) - self.assertEqual( - hw_sch.features["master_time_hour"].group, "master_time" - ) + self.assertEqual(hw_sch.features["event_time_hour"].group, "master_time") + self.assertEqual(hw_sch.features["master_time_hour"].group, "master_time") def test_extract_timestamp_features(self): graph, schema = _make_graph_and_schema( @@ -1050,6 +1044,98 @@ def test_compute_seed_deltas(self, mask, fill_value, expected): deltas = timeseries._compute_seed_deltas(raw_val, mask, 500, fill_value) np.testing.assert_array_equal(deltas, expected) + def test_per_sample_transform_config_serialization(self): + config = timeseries.PerSampleTransformConfig( + timeseries_pad_and_cap=timeseries.PadAndCapTimeseriesConfig( + sequence_length=15, padding_value=-1 + ), + timedelta_extraction=timeseries.TimestampFeatureExtractorConfig( + fill_value=-999 + ), + ) + json_str = config.to_json() # pyrefly: ignore[missing-attribute] + restored = timeseries.PerSampleTransformConfig.from_json(json_str) # pyrefly: ignore[missing-attribute] + self.assertEqual(config, restored) + self.assertEqual(restored.timeseries_pad_and_cap.sequence_length, 15) + self.assertEqual(restored.timeseries_pad_and_cap.padding_value, -1) + self.assertEqual(restored.timedelta_extraction.fill_value, -999) + + def test_per_sample_transform_end_to_end(self): + graph, schema = _make_graph_and_schema( + values={ + "time": np.array( + [np.array([100, 250, 300], dtype=np.int64)], dtype=np.object_ + ), + "signal": np.array( + [np.array([1.0, 2.0, 3.0], dtype=np.float32)], dtype=np.object_ + ), + }, + 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"), + }, + ) + config = timeseries.PerSampleTransformConfig( + timeseries_pad_and_cap=timeseries.PadAndCapTimeseriesConfig( + sequence_length=4 + ), + timedelta_extraction=timeseries.TimestampFeatureExtractorConfig(), + ) + transform = config.make(schema) + self.assertTrue(transform.has_transforms()) + + out_schema = transform.output_schema() + self.assertIn("time_mask", out_schema.node_sets["hardware"].features) + self.assertIn("time_seed_delta", out_schema.node_sets["hardware"].features) + + # Test transform single sample + transformed_single = transform.transform_sample(graph, seed_timestamp=500) + hw_single = transformed_single.node_sets["hardware"] + np.testing.assert_array_equal( + hw_single.features["time"], [[0, 100, 250, 300]] + ) + np.testing.assert_array_equal( + hw_single.features["time_mask"], [[False, True, True, True]] + ) + np.testing.assert_array_equal( + hw_single.features["time_seed_delta"], [[0, 400, 250, 200]] + ) + + # Test transform sample list + transformed_list = transform.transform_sample_list( + [graph], seed_timestamps=np.array([500]) + ) + self.assertLen(transformed_list, 1) + hw_list = transformed_list[0].node_sets["hardware"] + np.testing.assert_array_equal( + hw_list.features["time_seed_delta"], [[0, 400, 250, 200]] + ) + + def test_per_sample_transform_no_active_transforms(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, + ) + }, + ) + config = timeseries.PerSampleTransformConfig() + transform = config.make(schema) + self.assertFalse(transform.has_transforms()) + self.assertEqual(transform.output_schema(), schema) + + transformed = transform.transform_sample(graph) + self.assertIs(transformed, graph) + transformed_list = transform.transform_sample_list([graph]) + self.assertIs(transformed_list[0], graph) + if __name__ == "__main__": absltest.main() diff --git a/dgf/src/util/gen_test_graph.py b/dgf/src/util/gen_test_graph.py index 826612e..8c08230 100644 --- a/dgf/src/util/gen_test_graph.py +++ b/dgf/src/util/gen_test_graph.py @@ -2141,6 +2141,109 @@ def generate_temporal_in_memory_graph( return graph, schema +def generate_temporal_timeseries_in_memory_graph( + num_alerts: int = 6, + num_hardware: int = 2, +) -> Tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema]: + """Generates an InMemoryGraph with temporal and timeseries features for testing.""" + schema = schema_lib.GraphSchema( + node_sets={ + "alerts": schema_lib.NodeSchema( + features={ + "creation_time": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_creation_time=True, + ), + "label": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + ), + } + ), + "hardware": schema_lib.NodeSchema( + features={ + "time": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_timeseries=True, + is_creation_time=True, + shape=(None,), + ), + "signal": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.NUMERICAL, + is_timeseries=True, + shape=(None,), + ), + } + ), + }, + edge_sets={ + "alert_to_hw": schema_lib.EdgeSchema( + source="alerts", target="hardware", features={} + ), + "hw_to_alert": schema_lib.EdgeSchema( + source="hardware", target="alerts", features={} + ), + }, + ) + alerts_nodes = in_memory_graph_lib.InMemoryNodeSet( + num_nodes=num_alerts, + features={ + "creation_time": np.arange( + 100, 100 + 100 * num_alerts, 100, dtype=np.int64 + ), + "label": np.arange(1.0, num_alerts + 1.0, dtype=np.float32), + }, + ) + base_times = [ + np.array([50, 80, 120], dtype=np.int64), + np.array([150, 250, 450], dtype=np.int64), + np.array([70, 90, 110, 130], dtype=np.int64), + np.array([60], dtype=np.int64), + np.array([100, 200, 250], dtype=np.int64), + ] + base_signals = [ + np.array([1.5, 2.5, 3.5], dtype=np.float32), + np.array([4.5, 5.5, 6.5], dtype=np.float32), + np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), + np.array([5.0], dtype=np.float32), + np.array([6.0, 7.0, 8.0], dtype=np.float32), + ] + hw_times = [base_times[i % len(base_times)] for i in range(num_hardware)] + hw_signals = [ + base_signals[i % len(base_signals)] for i in range(num_hardware) + ] + + hardware_nodes = in_memory_graph_lib.InMemoryNodeSet( + num_nodes=num_hardware, + features={ + "time": np.array(hw_times, dtype=object), + "signal": np.array(hw_signals, dtype=object), + }, + ) + alert_to_hw = in_memory_graph_lib.InMemoryEdgeSet( + adjacency=np.array( + [np.arange(num_alerts), np.arange(num_alerts) % num_hardware], + dtype=np.int64, + ), + features={}, + ) + hw_to_alert = in_memory_graph_lib.InMemoryEdgeSet( + adjacency=np.array( + [np.arange(num_alerts) % num_hardware, np.arange(num_alerts)], + dtype=np.int64, + ), + features={}, + ) + graph = in_memory_graph_lib.InMemoryGraph( + node_sets={"alerts": alerts_nodes, "hardware": hardware_nodes}, + edge_sets={"alert_to_hw": alert_to_hw, "hw_to_alert": hw_to_alert}, + ) + return graph, schema + + def generate_recommender_like_in_memory_graph() -> ( Tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema] ):