diff --git a/dgf/src/learning/ten_lines/BUILD b/dgf/src/learning/ten_lines/BUILD index f8d2a0d..8c361e3 100644 --- a/dgf/src/learning/ten_lines/BUILD +++ b/dgf/src/learning/ten_lines/BUILD @@ -171,6 +171,7 @@ py_library( "//dgf/src/transform:normalize", "//dgf/src/util:filesystem", "//dgf/src/util:log", + "//dgf/src/util:temporal", "//dgf/src/util:util_py", # jax dep, # jaxtyping dep, @@ -249,6 +250,7 @@ py_library( "//dgf/src/transform:merge", "//dgf/src/transform:normalize", "//dgf/src/util:log", + "//dgf/src/util:temporal", "//dgf/src/util:util_py", # jax dep, # numpy dep, diff --git a/dgf/src/learning/ten_lines/link_prediction_dataset.py b/dgf/src/learning/ten_lines/link_prediction_dataset.py index 509b03c..2a53fdc 100644 --- a/dgf/src/learning/ten_lines/link_prediction_dataset.py +++ b/dgf/src/learning/ten_lines/link_prediction_dataset.py @@ -39,6 +39,7 @@ from dgf.src.transform import merge as merge_lib from dgf.src.transform import normalize as normalize_lib 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 @@ -108,6 +109,7 @@ class NodeIdsBatch: pos_trg_node_idxs: np.ndarray neg_trg_node_idxs: np.ndarray edge_idxs: np.ndarray + seed_timestamps: Optional[np.ndarray] = None @dataclasses.dataclass(kw_only=True) @@ -198,6 +200,10 @@ class GNNLinkDatasetPreparator: ) 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 + ) # The is_prepared data computed by the `prepare()` method. live: Optional[LiveData] = dataclasses.field(init=False, default=None) @@ -211,6 +217,22 @@ def __post_init__(self): " `mask_target_edgeset` is True, the entire target edgeset is removed" " from the sampling schema, making `mask_seed_edge` redundant." ) + if self.cache_normalized_features and ( + self.temporal_sampling + or temporal_util.schema_has_dynamic_timeseries_features(self.schema) + ): + raise ValueError( + "`cache_normalized_features=True` is not currently supported with" + " `temporal_sampling=True` or dynamic timeseries features." + ) + if self.temporal_sampling: + if self.target_edgeset not in self.edgeset_timestamp_features: + raise ValueError( + f"The target edgeset '{self.target_edgeset}' must have a creation" + " time feature when temporal_sampling=True. Set" + " is_creation_time=True on the creation time feature in" + f" schema.edge_sets['{self.target_edgeset}']." + ) def get_live(self) -> LiveData: if self.live is None: @@ -344,10 +366,11 @@ def _sampling_schema(self) -> schema_lib.GraphSchema: return self.schema def _edgeset_to_mask(self) -> Optional[str]: - if self.mask_seed_edge: - return self.target_edgeset - else: - return None + return ( + self.target_edgeset + if self.mask_seed_edge and not self.temporal_sampling + else None + ) def _get_merge_schema( self, base_schema: schema_lib.GraphSchema @@ -401,11 +424,18 @@ def _in_memory_generate_node_ids( shuffle=self.shuffle, ): nei = nei_generator.generate(batch_edge_idxs) + seed_timestamps = None + if self.temporal_sampling: + ts_feature = self.edgeset_timestamp_features[self.target_edgeset] + seed_timestamps = target_edgeset_data.features[ts_feature][ + batch_edge_idxs + ] yield NodeIdsBatch( pos_src_node_idxs=nei.pos_src_node_idxs, pos_trg_node_idxs=nei.pos_trg_node_idxs, neg_trg_node_idxs=nei.neg_trg_node_idxs, edge_idxs=batch_edge_idxs, + seed_timestamps=seed_timestamps, ) def _prepare_on_in_memory_graph(self): @@ -433,24 +463,43 @@ def _prepare_on_in_memory_graph(self): sampling_schema = self._sampling_schema() edgeset_to_mask = self._edgeset_to_mask() + source_plan = ( + self.source_sampling_plan + if self.source_sampling_plan is not None + else sampling_config_lib.simple_sampling_config_to_sampling_plan( + self.sampling_config, self.schema + ) + if isinstance( + self.sampling_config, sampling_config_lib.SimpleSamplingConfig + ) + else self.sampling_config + ) + target_plan_config = dataclasses.replace( + self.sampling_config, + seed_nodeset=self.schema.edge_sets[self.target_edgeset].target, + ) + target_plan = ( + self.target_sampling_plan + if self.target_sampling_plan is not None + else sampling_config_lib.simple_sampling_config_to_sampling_plan( + target_plan_config, self.schema + ) + if isinstance( + target_plan_config, sampling_config_lib.SimpleSamplingConfig + ) + else target_plan_config + ) + source_sampler = in_memory_sampler_lib.create_sampler( graph=self.graph, - plan=self.source_sampling_plan - if self.source_sampling_plan is not None - else self.sampling_config, + plan=source_plan, schema=sampling_schema, batch_size=self.batch_size, edgeset_to_mask=edgeset_to_mask, ) - target_plan = dataclasses.replace( - self.sampling_config, - seed_nodeset=self.schema.edge_sets[self.target_edgeset].target, - ) target_sampler = in_memory_sampler_lib.create_sampler( graph=self.graph, - plan=self.target_sampling_plan - if self.target_sampling_plan is not None - else target_plan, + plan=target_plan, schema=sampling_schema, batch_size=self.batch_size * self.num_negative_nodes, edgeset_to_mask=edgeset_to_mask, @@ -464,11 +513,15 @@ def gen_raw_source_samples() -> Iterator[in_memory_graph_lib.InMemoryGraph]: """Generate graph samples for the source nodeset.""" while True: for batch_seed in self._in_memory_generate_node_ids(nei_generator): + masked_edge_idxs = ( + batch_seed.edge_idxs + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None + ) samples = source_sampler.sample( batch_seed.pos_src_node_idxs, - masked_edge_idxs=batch_seed.edge_idxs - if self.mask_seed_edge - else None, + masked_edge_idxs=masked_edge_idxs, + seed_timestamps=batch_seed.seed_timestamps, ) for sample in samples: yield sample @@ -477,20 +530,31 @@ def gen_raw_target_samples() -> Iterator[in_memory_graph_lib.InMemoryGraph]: """Generate graph samples for the target nodeset.""" while True: for batch_seed in self._in_memory_generate_node_ids(nei_generator): + masked_edge_idxs = ( + batch_seed.edge_idxs + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None + ) # We generate both positive and negative samples. pos_trg_samples = target_sampler.sample( batch_seed.pos_trg_node_idxs, - masked_edge_idxs=batch_seed.edge_idxs - if self.mask_seed_edge - else None, + masked_edge_idxs=masked_edge_idxs, + seed_timestamps=batch_seed.seed_timestamps, + ) + neg_seed_timestamps = ( + np.repeat(batch_seed.seed_timestamps, self.num_negative_nodes) + if batch_seed.seed_timestamps is not None + else None + ) + neg_masked_edge_idxs = ( + np.repeat(batch_seed.edge_idxs, self.num_negative_nodes, axis=0) + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None ) neg_trg_samples = target_sampler.sample( batch_seed.neg_trg_node_idxs.flatten(), - masked_edge_idxs=np.repeat( - batch_seed.edge_idxs, self.num_negative_nodes, axis=0 - ) - if self.mask_seed_edge - else None, + masked_edge_idxs=neg_masked_edge_idxs, + seed_timestamps=neg_seed_timestamps, ) for sample in pos_trg_samples: yield sample @@ -550,11 +614,15 @@ def gen_normalized_merged_positive_source_samples() -> ( ): while True: for batch_seed in self._in_memory_generate_node_ids(nei_generator): + masked_edge_idxs = ( + batch_seed.edge_idxs + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None + ) samples = source_sampler.sample( batch_seed.pos_src_node_idxs, - masked_edge_idxs=batch_seed.edge_idxs - if self.mask_seed_edge - else None, + masked_edge_idxs=masked_edge_idxs, + seed_timestamps=batch_seed.seed_timestamps, ) merged_samples, _ = merge_lib.merge_graphs( samples, sampling_schema, padding=None, sentinel_offset=True @@ -566,11 +634,15 @@ def gen_normalized_merged_positive_target_samples() -> ( ): while True: for batch_seed in self._in_memory_generate_node_ids(nei_generator): + masked_edge_idxs = ( + batch_seed.edge_idxs + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None + ) samples = target_sampler.sample( batch_seed.pos_trg_node_idxs, - masked_edge_idxs=batch_seed.edge_idxs - if self.mask_seed_edge - else None, + masked_edge_idxs=masked_edge_idxs, + seed_timestamps=batch_seed.seed_timestamps, ) merged_samples, _ = merge_lib.merge_graphs( samples, sampling_schema, padding=None, sentinel_offset=True @@ -582,13 +654,20 @@ def gen_normalized_merged_negative_target_samples() -> ( ): while True: for batch_seed in self._in_memory_generate_node_ids(nei_generator): + neg_seed_timestamps = ( + np.repeat(batch_seed.seed_timestamps, self.num_negative_nodes) + if batch_seed.seed_timestamps is not None + else None + ) + neg_masked_edge_idxs = ( + np.repeat(batch_seed.edge_idxs, self.num_negative_nodes, axis=0) + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None + ) samples = target_sampler.sample( batch_seed.neg_trg_node_idxs.flatten(), - masked_edge_idxs=np.repeat( - batch_seed.edge_idxs, self.num_negative_nodes, axis=0 - ) - if self.mask_seed_edge - else None, + masked_edge_idxs=neg_masked_edge_idxs, + seed_timestamps=neg_seed_timestamps, ) merged_samples, _ = merge_lib.merge_graphs( samples, sampling_schema, padding=None, sentinel_offset=True @@ -662,18 +741,8 @@ def gen_normalized_merged_negative_target_samples() -> ( positive_source_padding=positive_source_padding, positive_target_padding=positive_target_padding, negative_target_padding=negative_target_padding, - source_sampling_plan=sampling_config_lib.simple_sampling_config_to_sampling_plan( - self.sampling_config, self.schema - ) - if isinstance( - self.sampling_config, sampling_config_lib.SimpleSamplingConfig - ) - else self.sampling_config, - target_sampling_plan=sampling_config_lib.simple_sampling_config_to_sampling_plan( - target_plan, self.schema - ) - if isinstance(target_plan, sampling_config_lib.SimpleSamplingConfig) - else target_plan, + source_sampling_plan=source_plan, + target_sampling_plan=target_plan, source_sampler=source_sampler, target_sampler=target_sampler, num_edges_in_seed_edgeset=len(self.seed_edge_idxs) @@ -737,10 +806,16 @@ def _sample_and_merge( target. - A tuple of merge offsets dictionaries for each merged graph. """ + masked_edge_idxs = ( + batch_seed.edge_idxs + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None + ) # Positive source pos_src_samples = live.source_sampler.sample( batch_seed.pos_src_node_idxs, - masked_edge_idxs=batch_seed.edge_idxs if self.mask_seed_edge else None, + masked_edge_idxs=masked_edge_idxs, + seed_timestamps=batch_seed.seed_timestamps, ) pos_src_merged, pos_src_offsets = merge_lib.merge_graphs( pos_src_samples, @@ -752,7 +827,8 @@ def _sample_and_merge( # Positive target pos_trg_samples = live.target_sampler.sample( batch_seed.pos_trg_node_idxs, - masked_edge_idxs=batch_seed.edge_idxs if self.mask_seed_edge else None, + masked_edge_idxs=masked_edge_idxs, + seed_timestamps=batch_seed.seed_timestamps, ) pos_trg_merged, pos_trg_offsets = merge_lib.merge_graphs( pos_trg_samples, @@ -762,13 +838,20 @@ def _sample_and_merge( ) # Negative target + neg_seed_timestamps = ( + np.repeat(batch_seed.seed_timestamps, self.num_negative_nodes) + if batch_seed.seed_timestamps is not None + else None + ) + neg_masked_edge_idxs = ( + np.repeat(batch_seed.edge_idxs, self.num_negative_nodes, axis=0) + if self.mask_seed_edge and batch_seed.seed_timestamps is None + else None + ) neg_trg_samples = live.target_sampler.sample( batch_seed.neg_trg_node_idxs.flatten(), - masked_edge_idxs=np.repeat( - batch_seed.edge_idxs, self.num_negative_nodes, axis=0 - ) - if self.mask_seed_edge - else None, + masked_edge_idxs=neg_masked_edge_idxs, + seed_timestamps=neg_seed_timestamps, ) neg_trg_merged, neg_trg_offsets = merge_lib.merge_graphs( neg_trg_samples, diff --git a/dgf/src/learning/ten_lines/link_prediction_dataset_test.py b/dgf/src/learning/ten_lines/link_prediction_dataset_test.py index 56e46dd..edad747 100644 --- a/dgf/src/learning/ten_lines/link_prediction_dataset_test.py +++ b/dgf/src/learning/ten_lines/link_prediction_dataset_test.py @@ -522,6 +522,85 @@ def test_caching(self, cache_features, cache_device): num_batches += 1 self.assertEqual(num_batches, 2) + 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, + ) + preparator = link_prediction_dataset.GNNLinkDatasetPreparator( + graph=graph, + schema=schema, + sampling_config=sampling_config, + batch_size=2, + drop_remainder=False, + shuffle=True, + target_edgeset="e1", + num_negative_nodes=2, + seed_edge_idxs=None, + temporal_sampling=True, + edgeset_timestamp_features={"e1": "timestamp"}, + edge_neighbor_generator=edge_neighbor_generator_lib.RandomEdgeNeighborGeneratorConfig(), + ) + self.assertFalse(preparator.is_prepared()) + preparator.prepare() + self.assertTrue(preparator.is_prepared()) + + num_batches = 0 + live = preparator.get_live() + for sample in preparator.generate(): + in_memory_graph_validate_lib.validate_graph( + sample.positive_source_graph, + live.source_normalizer.output_schema(), + raise_on_warning=False, + ) + in_memory_graph_validate_lib.validate_graph( + sample.positive_target_graph, + live.target_normalizer.output_schema(), + raise_on_warning=False, + ) + in_memory_graph_validate_lib.validate_graph( + sample.negative_target_graph, + live.target_normalizer.output_schema(), + raise_on_warning=False, + ) + num_batches += 1 + self.assertEqual(num_batches, 2) + + def test_temporal_cache_normalized_features_raises(self): + graph, schema = gen_test_graph.generate_temporal_in_memory_graph(False) + sampling_config = sampling_config_lib.SimpleSamplingConfig( + seed_nodeset="n1", + num_hops=1, + hop_width=2, + reverse=True, + temporal_sampling=True, + ) + with self.assertRaisesRegex( + ValueError, + "`cache_normalized_features=True` is not currently supported", + ): + link_prediction_dataset.GNNLinkDatasetPreparator( + graph=graph, + schema=schema, + sampling_config=sampling_config, + batch_size=2, + drop_remainder=False, + shuffle=True, + target_edgeset="e1", + num_negative_nodes=2, + seed_edge_idxs=None, + temporal_sampling=True, + cache_normalized_features=True, + edgeset_timestamp_features={"e1": "timestamp"}, + edge_neighbor_generator=( + edge_neighbor_generator_lib.RandomEdgeNeighborGeneratorConfig() + ), + ) + if __name__ == "__main__": absltest.main() diff --git a/dgf/src/learning/ten_lines/link_prediction_model.py b/dgf/src/learning/ten_lines/link_prediction_model.py index b48ae32..6ceb7d3 100644 --- a/dgf/src/learning/ten_lines/link_prediction_model.py +++ b/dgf/src/learning/ten_lines/link_prediction_model.py @@ -18,7 +18,7 @@ import dataclasses import os -from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple +from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple, Union import dataclasses_json from dgf.src.data import in_memory_graph @@ -102,6 +102,13 @@ class ModelData: source_sampling_plan: sampling_config_lib.SamplingPlan target_sampling_plan: sampling_config_lib.SamplingPlan training_stats: TrainingStats + temporal_sampling: bool = False + nodeset_timestamp_features: Dict[str, str] = dataclasses.field( + default_factory=dict + ) + edgeset_timestamp_features: Dict[str, str] = dataclasses.field( + default_factory=dict + ) edge_neighbor_generator: ( edge_neighbor_generator_lib.EdgeNeighborGeneratorConfig ) = edge_neighbor_generator_lib.registry.field( @@ -370,6 +377,7 @@ def predict( source_node_idxs: common.SeedNodeIdxs, target_node_idxs: common.SeedNodeIdxs, *, + edge_timestamps: Optional[np.ndarray] = None, all_combinations: bool = False, verbose: int = 2, ) -> np.ndarray: @@ -413,6 +421,7 @@ def predict( graph: The graph containing nodes and features. source_node_idxs: A list of node indices for edge sources. target_node_idxs: A list of node indices for edge targets. + edge_timestamps: Optional timestamps for the evaluated edges. all_combinations: If False (default), `len(target_node_idxs)` must be divisible by `len(source_node_idxs)`. If `m = len(target_node_idxs) // len(source_node_idxs)`, edge `source_node_idxs[i // m] -> @@ -440,7 +449,8 @@ def predict( graph, source_node_idxs, target_node_idxs, - all_combinations, + all_combinations=all_combinations, + edge_timestamps=edge_timestamps, verbose=verbose, ): prediction_list.append(batch.predictions) @@ -487,6 +497,7 @@ def _build_samplers( "edgeset_to_mask": ( self._data.task.target_edgeset if self._data.hparams.message_passing_on_target_edgeset + and not self._data.temporal_sampling else None ), } @@ -504,6 +515,7 @@ def predict_batch( source_node_idxs: common.SeedNodeIdxs, target_node_idxs: common.SeedNodeIdxs, all_combinations: bool = False, + edge_timestamps: Optional[np.ndarray] = None, verbose: int = 2, source_sampler: Optional[in_memory_sampler_lib.Sampler] = None, target_sampler: Optional[in_memory_sampler_lib.Sampler] = None, @@ -514,7 +526,10 @@ def predict_batch( np_source_node_idxs = np.asarray(source_node_idxs) np_target_node_idxs = np.asarray(target_node_idxs) - mask_edges = self._data.hparams.message_passing_on_target_edgeset + mask_edges = ( + self._data.hparams.message_passing_on_target_edgeset + and not self._data.temporal_sampling + ) edge_lookup = None target_edgeset = self._data.task.target_edgeset if mask_edges: @@ -531,11 +546,30 @@ def predict_batch( flat_sources = grid[0].flatten() flat_targets = grid[1].flatten() num_examples = len(flat_sources) + if edge_timestamps is not None: + assert edge_timestamps.shape == ( + len(np_source_node_idxs), + len(np_target_node_idxs), + ) + flat_timestamps = edge_timestamps.flatten() + else: + flat_timestamps = None else: num_examples = len(np_target_node_idxs) m = num_examples // len(np_source_node_idxs) flat_sources = np.repeat(np_source_node_idxs, m) flat_targets = np_target_node_idxs + if edge_timestamps is not None: + assert edge_timestamps.shape == (num_examples,) + flat_timestamps = edge_timestamps + else: + flat_timestamps = None + + if self._data.temporal_sampling and edge_timestamps is None: + raise ValueError( + "`edge_timestamps` must be provided in `predict()` / `predict_batch()` " + "when `temporal_sampling=True`." + ) source_nodeset = self._data.schema.edge_sets[ self._data.task.target_edgeset @@ -613,7 +647,20 @@ def merge_and_predict( batch_src = flat_sources[batch_indices] batch_trg = flat_targets[batch_indices] - if mask_edges: + batch_timestamps = ( + flat_timestamps[batch_indices] + if flat_timestamps is not None + else None + ) + + if self._data.temporal_sampling: + source_samples = source_sampler.sample( + batch_src, seed_timestamps=batch_timestamps + ) + target_samples = target_sampler.sample( + batch_trg, seed_timestamps=batch_timestamps + ) + elif mask_edges: assert edge_lookup is not None queries = np.stack([batch_src, batch_trg], axis=0) masked_edge_idxs = edge_lookup.query_array(queries) @@ -641,6 +688,7 @@ def predict_embedding( node_idxs: common.SeedNodeIdxs, encoder: Literal["source", "target"], *, + node_timestamps: Optional[np.ndarray] = None, verbose: int = 2, ) -> np.ndarray: """Predicts node embeddings for source or target sides. @@ -677,6 +725,7 @@ def predict_embedding( node_idxs: A list of node indices for which to predict the embedding. encoder: Whether to predict the embedding using the "source" or "target" encoder of the model. Must be one of {"source", "target"}. + node_timestamps: Optional timestamps for the evaluated nodes. verbose: Verbosity level. Returns: @@ -715,6 +764,22 @@ def predict_embedding( else: raise ValueError(f"Invalid encoder: {encoder}") + flat_node_timestamps = None + if node_timestamps is not None: + assert node_timestamps.shape == (len(np_node_idxs),) + flat_node_timestamps = node_timestamps + elif self._data.temporal_sampling: + if nodeset not in self._data.nodeset_timestamp_features: + raise ValueError( + "Cannot perform temporal sampling during embedding prediction" + f" because node set '{nodeset}' has no creation timestamp feature." + " Please provide `node_timestamps` in `predict_embedding()`." + ) + ts_feature = self._data.nodeset_timestamp_features[nodeset] + flat_node_timestamps = graph.node_sets[nodeset].features[ts_feature][ + np_node_idxs + ] + generator = util.batch_indices_generator( np.arange(len(np_node_idxs)), batch_size=self._data.hparams.batch_size, @@ -750,7 +815,12 @@ def merge_and_predict_emb(sub_samples: List[in_memory_graph.InMemoryGraph]): embeddings_list = [] for batch_indices in generator: nodes = np_node_idxs[batch_indices] - samples = sampler.sample(nodes) + seed_timestamps = ( + flat_node_timestamps[batch_indices] + if flat_node_timestamps is not None + else None + ) + samples = sampler.sample(nodes, seed_timestamps=seed_timestamps) for emb in self.execute_with_split_on_error( merge_and_predict_emb, samples @@ -1215,10 +1285,19 @@ def evaluate( examples_per_seed_edge = 1 + num_negative_nodes carry_over_probs = np.array([]) + full_timestamps = None + if self._data.temporal_sampling: + ts_feature = self._data.edgeset_timestamp_features[target_edgeset] + seed_timestamps = graph.edge_sets[target_edgeset].features[ts_feature][ + seed_edge_idxs + ] + full_timestamps = np.repeat(seed_timestamps, examples_per_seed_edge) + for batch_pred in self.predict_batch( graph, full_src.tolist(), full_trg.tolist(), + edge_timestamps=full_timestamps, verbose=verbose, source_sampler=source_sampler, target_sampler=target_sampler, diff --git a/dgf/src/learning/ten_lines/link_prediction_test.py b/dgf/src/learning/ten_lines/link_prediction_test.py index d09435f..c7a78db 100644 --- a/dgf/src/learning/ten_lines/link_prediction_test.py +++ b/dgf/src/learning/ten_lines/link_prediction_test.py @@ -67,6 +67,7 @@ def gen_toy_graph( num_nodes_b: int = 500, num_categorical_values: int = 200, random_seed: int = 42, + has_timestamp_feature: bool = False, ) -> Tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema]: """Generates a toy dataset for link prediction testing. @@ -83,6 +84,7 @@ def gen_toy_graph( num_nodes_b: Number of nodes in node set "B". num_categorical_values: Number of categories for feature "f2". random_seed: Seed for reproducibility of features and edge selection. + has_timestamp_feature: If True, adds timestamp features to nodes and edges. Returns: A tuple containing the generated InMemoryGraph and its GraphSchema. @@ -90,48 +92,64 @@ def gen_toy_graph( rng = np.random.default_rng(random_seed) + node_a_features_schema = { + "#id": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.BYTES, + semantic=schema_lib.FeatureSemantic.PRIMARY_ID, + ), + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.EMBEDDING, + ), + "f2": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_32, + semantic=schema_lib.FeatureSemantic.CATEGORICAL, + num_categorical_values=num_categorical_values, + ), + } + node_b_features_schema = { + "#id": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.BYTES, + semantic=schema_lib.FeatureSemantic.PRIMARY_ID, + ), + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.FLOAT_32, + semantic=schema_lib.FeatureSemantic.EMBEDDING, + ), + "f2": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_32, + semantic=schema_lib.FeatureSemantic.CATEGORICAL, + num_categorical_values=num_categorical_values, + ), + } + edge_features_schema = {} + if has_timestamp_feature: + node_a_features_schema["timestamp"] = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_creation_time=True, + ) + node_b_features_schema["timestamp"] = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_creation_time=True, + ) + edge_features_schema["timestamp"] = schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64, + semantic=schema_lib.FeatureSemantic.TIMESTAMP, + is_creation_time=True, + ) + schema = schema_lib.GraphSchema( node_sets={ - "A": schema_lib.NodeSchema( - features={ - "#id": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.BYTES, - semantic=schema_lib.FeatureSemantic.PRIMARY_ID, - ), - "f1": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.FLOAT_32, - semantic=schema_lib.FeatureSemantic.EMBEDDING, - ), - "f2": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.INTEGER_32, - semantic=schema_lib.FeatureSemantic.CATEGORICAL, - num_categorical_values=num_categorical_values, - ), - } - ), - "B": schema_lib.NodeSchema( - features={ - "#id": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.BYTES, - semantic=schema_lib.FeatureSemantic.PRIMARY_ID, - ), - "f1": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.FLOAT_32, - semantic=schema_lib.FeatureSemantic.EMBEDDING, - ), - "f2": schema_lib.FeatureSchema( - format=schema_lib.FeatureFormat.INTEGER_32, - semantic=schema_lib.FeatureSemantic.CATEGORICAL, - num_categorical_values=num_categorical_values, - ), - } - ), + "A": schema_lib.NodeSchema(features=node_a_features_schema), + "B": schema_lib.NodeSchema(features=node_b_features_schema), }, edge_sets={ "A_to_B": schema_lib.EdgeSchema( source="A", target="B", - features={}, + features=edge_features_schema, ), }, ) @@ -178,33 +196,43 @@ def gen_toy_graph( np.mean(degrees), ) + node_a_features = { + "f1": f1_a, + "f2": f2_a, + "#id": np.array([f"A_{i}".encode() for i in range(num_nodes_a)]), + } + node_b_features = { + "f1": f1_b, + "f2": f2_b, + "#id": np.array([f"B_{i}".encode() for i in range(num_nodes_b)]), + } + edge_features = {} + if has_timestamp_feature: + node_a_features["timestamp"] = rng.integers( + 0, 1000, size=num_nodes_a, dtype=np.int64 + ) + node_b_features["timestamp"] = rng.integers( + 0, 1000, size=num_nodes_b, dtype=np.int64 + ) + edge_features["timestamp"] = rng.integers( + 0, 1000, size=len(src_idxs), dtype=np.int64 + ) + graph = in_memory_graph_lib.InMemoryGraph( node_sets={ "A": in_memory_graph_lib.InMemoryNodeSet( - features={ - "f1": f1_a, - "f2": f2_a, - "#id": np.array( - [f"A_{i}".encode() for i in range(num_nodes_a)] - ), - }, + features=node_a_features, num_nodes=num_nodes_a, ), "B": in_memory_graph_lib.InMemoryNodeSet( - features={ - "f1": f1_b, - "f2": f2_b, - "#id": np.array( - [f"B_{i}".encode() for i in range(num_nodes_b)] - ), - }, + features=node_b_features, num_nodes=num_nodes_b, ), }, edge_sets={ "A_to_B": in_memory_graph_lib.InMemoryEdgeSet( adjacency=np.stack([src_idxs, tgt_idxs], axis=0), - features={}, + features=edge_features, ), }, ) @@ -909,5 +937,154 @@ def test_evaluate(self): # Note: The random negative sampling make the evaluation non deterministic. +class LinkPredictionRealLookingTemporal(parameterized.TestCase): + + @classmethod + def setUpClass(cls): + super().setUpClass() + + cls.graph, cls.schema = gen_toy_graph(has_timestamp_feature=True) + + def train_model() -> link_prediction_model.LinkPredictionModel: + return link_prediction_train.train_link_model( + graph=cls.graph, + schema=cls.schema, + target_edgeset="A_to_B", + time_aware=True, + **RAPID_TRAINING_KWARGS, + ) + + cls.model = train_model() + + def test_evaluate(self): + evaluation = self.model.evaluate(self.graph) + logging.info("evaluation:\n%s", evaluation) + assert evaluation.num_examples is not None + self.assertGreater(evaluation.num_examples, 0) + + def test_predict(self): + src_nodes = self.graph.edge_sets["A_to_B"].adjacency[0][:2] + trg_nodes = self.graph.edge_sets["A_to_B"].adjacency[1][:2] + edge_timestamps = self.graph.edge_sets["A_to_B"].features["timestamp"][:2] + predictions = self.model.predict( + graph=self.graph, + source_node_idxs=src_nodes, + target_node_idxs=trg_nodes, + edge_timestamps=edge_timestamps, + ) + self.assertEqual(predictions.shape, (2,)) + self.assertTrue(np.all(predictions >= 0.0) and np.all(predictions <= 1.0)) + + @parameterized.named_parameters( + ( + "2d_matrix_combinations", + np.array([[100, 200, 300], [400, 500, 600]]), + True, + (2, 3), + ), + ("1d_vector_pairwise", np.array([300, 400]), False, (2,)), + ) + def test_predict_with_edge_timestamps( + self, edge_timestamps, all_combinations, expected_shape + ): + predictions = self.model.predict( + graph=self.graph, + source_node_idxs=[0, 1], + target_node_idxs=[0, 1, 2] if all_combinations else [0, 1], + edge_timestamps=edge_timestamps, + all_combinations=all_combinations, + ) + self.assertEqual(predictions.shape, expected_shape) + self.assertTrue(np.all(predictions >= 0.0) and np.all(predictions <= 1.0)) + + def test_predict_embedding(self): + emb_dim = self.model.data().hparams.node_embedding_dim + emb = self.model.predict_embedding( + self.graph, + node_idxs=[0, 1], + encoder="source", + node_timestamps=np.array([100, 200]), + ) + self.assertEqual(emb.shape, (2, emb_dim)) + + def test_save_and_load(self): + with tempfile.TemporaryDirectory() as tmpdir: + self.model.save(tmpdir) + restored_model = common_lib.load_model(tmpdir) + + assert isinstance(restored_model, link_prediction_model.LinkPredictionModel) + self.assertTrue(restored_model.data().temporal_sampling) + self.assertTrue( + restored_model.data().source_sampling_plan.temporal_sampling + ) + self.assertTrue( + restored_model.data().target_sampling_plan.temporal_sampling + ) + test_util.assert_are_equal(self, self.model.data(), restored_model.data()) + + src_nodes = self.graph.edge_sets["A_to_B"].adjacency[0][:2] + trg_nodes = self.graph.edge_sets["A_to_B"].adjacency[1][:2] + edge_timestamps = self.graph.edge_sets["A_to_B"].features["timestamp"][:2] + predictions = self.model.predict( + graph=self.graph, + source_node_idxs=src_nodes, + target_node_idxs=trg_nodes, + edge_timestamps=edge_timestamps, + ) + restored_predictions = restored_model.predict( + graph=self.graph, + source_node_idxs=src_nodes, + target_node_idxs=trg_nodes, + edge_timestamps=edge_timestamps, + ) + self.assertEqual(predictions.shape, restored_predictions.shape) + self.assertTrue( + np.all(restored_predictions >= 0.0) + and np.all(restored_predictions <= 1.0) + ) + + def test_errors_and_validation(self): + with self.assertRaisesRegex( + ValueError, + "`edge_timestamps` must be provided", + ): + self.model.predict( + graph=self.graph, + source_node_idxs=[0], + target_node_idxs=[0], + ) + + with self.assertRaises(AssertionError): + self.model.predict( + graph=self.graph, + source_node_idxs=[0, 1], + target_node_idxs=[0, 1, 2], + edge_timestamps=np.array([100, 200, 300, 400]), + all_combinations=True, + ) + + with self.assertRaises(AssertionError): + self.model.predict_embedding( + self.graph, + node_idxs=[0, 1], + encoder="source", + node_timestamps=np.array([100, 200, 300]), + ) + + graph, schema = gen_toy_graph(has_timestamp_feature=True) + del schema.edge_sets["A_to_B"].features["timestamp"] + with self.assertRaisesRegex( + ValueError, + "The target edgeset 'A_to_B' must have a creation time feature", + ): + link_prediction_train.train_link_model( + graph=graph, + schema=schema, + target_edgeset="A_to_B", + time_aware=True, + **RAPID_TRAINING_KWARGS, + ) + + if __name__ == "__main__": absltest.main() diff --git a/dgf/src/learning/ten_lines/link_prediction_train.py b/dgf/src/learning/ten_lines/link_prediction_train.py index f4cc345..728d256 100644 --- a/dgf/src/learning/ten_lines/link_prediction_train.py +++ b/dgf/src/learning/ten_lines/link_prediction_train.py @@ -39,6 +39,7 @@ from dgf.src.transform import normalize as normalize_lib from dgf.src.util import filesystem as fs 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 @@ -168,6 +169,8 @@ def prepare_datasets( cache_normalized_features_device: Literal["host", "device"], source_sampling_plan: Optional[sampling_config_lib.SamplingPlan], target_sampling_plan: Optional[sampling_config_lib.SamplingPlan], + temporal_sampling: bool, + edgeset_timestamp_features: Dict[str, str], ) -> Tuple[ link_prediction_dataset.GNNLinkDatasetPreparator, Optional[link_prediction_dataset.GNNLinkDatasetPreparator], @@ -197,6 +200,7 @@ def prepare_datasets( num_hops=hparams.num_sampling_hops, hop_width=hparams.sampling_width, reverse=True, + temporal_sampling=temporal_sampling, ) # TODO(gbm): Parametrize. @@ -235,6 +239,8 @@ def prepare_datasets( "edge_neighbor_generator": edge_neighbor_generator, "cache_normalized_features": cache_normalized_features, "cache_normalized_features_device": cache_normalized_features_device, + "temporal_sampling": temporal_sampling, + "edgeset_timestamp_features": edgeset_timestamp_features, } train_dataset = link_prediction_dataset.GNNLinkDatasetPreparator( @@ -309,6 +315,7 @@ def train_link_model( node_embedding_dim: int = 128, learning_rate: float = 1e-3, cache_valid_dataset: bool = True, + time_aware: bool = False, num_negative_nodes: int = 8, message_passing_on_target_edgeset: bool = True, negative_edges: Literal["random", "random-walk"] = "random", @@ -370,6 +377,9 @@ def train_link_model( node_embedding_dim: Node embedding dimension. learning_rate: Learning rate. cache_valid_dataset: If True, the validation dataset is cached in memory. + time_aware: Enables temporal-aware training. If `False` (default), no + temporal sampling is applied. If `True`, timestamp features are inferred + from the schema (via features marked as creation timestamps). num_negative_nodes: Number of negative target nodes to sample for each edge. message_passing_on_target_edgeset: If True, message passing is allowed to use edges from the `target_edgeset`. Otherwise, these edges are excluded @@ -413,7 +423,6 @@ def train_link_model( A LinkPredictionModel instance. """ - # TODO(gbm): Add support for temporal aware sampling. # TODO(gbm): Add support for decomposable and non-decomposable decoders. # TODO(gbm): Add support for other type of graph inputs. @@ -437,6 +446,20 @@ def train_link_model( " than one edgeset." ) + if time_aware: + nodeset_ts_features = temporal_util.nodeset_timestamp_features(schema) + edgeset_ts_features = temporal_util.edgeset_timestamp_features(schema) + if target_edgeset not in edgeset_ts_features: + raise ValueError( + f"The target edgeset '{target_edgeset}' must have a creation time" + " feature. Set is_creation_time=True on the creation time feature" + f" (e.g. `schema.edge_sets['{target_edgeset}'].features[].is_creation_time = True`)." + ) + else: + nodeset_ts_features = {} + edgeset_ts_features = {} + if verbose >= 2: log.info( "Graph input schema:\n%s", @@ -486,6 +509,8 @@ def train_link_model( cache_normalized_features_device=cache_normalized_features_device, source_sampling_plan=source_sampling_plan, target_sampling_plan=target_sampling_plan, + temporal_sampling=time_aware, + edgeset_timestamp_features=edgeset_ts_features, ) source_normalized_schema = ( train_dataset.get_live().source_normalizer.output_schema() @@ -772,6 +797,9 @@ def valid_step(params, opt_state, batch): target_feature_stats=train_dataset.get_live().target_feature_stats, source_sampling_plan=train_dataset.get_live().source_sampling_plan, target_sampling_plan=train_dataset.get_live().target_sampling_plan, + temporal_sampling=time_aware, + nodeset_timestamp_features=nodeset_ts_features, + edgeset_timestamp_features=edgeset_ts_features, training_stats=TrainingStats( num_train_seed_edges=train_dataset.num_edge_in_seed_edgeset(), num_valid_seed_edges=valid_dataset.num_edge_in_seed_edgeset()