Skip to content

Add temporal sampling to link prediction. #160

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add temporal sampling to link prediction. #160

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions dgf/src/api/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,7 @@ py_library(
"//dgf/src/sampling:beam_semi_distributed_sampler_v2",
"//dgf/src/sampling:config",
"//dgf/src/sampling:in_memory_sampler",
"//dgf/src/sampling:temporal",
"//dgf/src/sampling/gcp:spanner_graph_sampler",
],
)
Expand Down
6 changes: 4 additions & 2 deletions dgf/src/learning/ten_lines/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -271,9 +273,10 @@ py_library(
"//dgf/src/io:jax",
"//dgf/src/learning/jax:common",
"//dgf/src/sampling:config",
"//dgf/src/sampling:in_memory_sampler",
"//dgf/src/sampling:temporal",
"//dgf/src/transform:normalize",
"//dgf/src/util:log",
"//dgf/src/util:temporal",
"//dgf/src/util:util_py",
# jax dep,
# numpy dep,
Expand Down Expand Up @@ -348,7 +351,6 @@ py_test(
"//dgf/src/data:schema",
"//dgf/src/io:tf_graph_sample",
"//dgf/src/sampling:config",
"//dgf/src/transform:timeseries",
"//dgf/src/util:gen_test_graph",
"//dgf/src/validate:in_memory_graph",
# numpy dep,
Expand Down
91 changes: 3 additions & 88 deletions dgf/src/learning/ten_lines/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,10 +118,6 @@ 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.
"""

graph: Graph
Expand All @@ -144,12 +140,6 @@ 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
] = None

num_seed_nodes: Optional[int] = dataclasses.field(init=False)
batch_iterator: BatchSampleGeneratorIteratorFn = dataclasses.field(init=False)
Expand All @@ -163,15 +153,6 @@ 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)
_has_per_sample_transforms: bool = dataclasses.field(
init=False, default=False
)

def __post_init__(self):
if self.seed_node_idxs is not None:
Expand Down Expand Up @@ -231,60 +212,14 @@ def __post_init__(self):
len(self.seed_node_idxs) if self.seed_node_idxs is not None else None
)

# 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:
if self.temporal and self.format != GraphFormat.IN_MEMORY_GRAPH:
raise ValueError(
"Temporal sampling (`temporal=True`) and per-sample transformations"
" (`timeseries_pad_and_cap`, `timedelta_extraction`) are"
"Temporal sampling (`temporal=True`) is"
" only supported for GraphFormat.IN_MEMORY_GRAPH, but got format:"
f" {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
)
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 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
)
)
current_schema = self._timestamp_extractor.output_schema()

self._output_schema = current_schema

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`."
)
self._output_schema = self.schema

self.batch_iterator, self.single_iterator = self.iterator_builder()

Expand All @@ -296,12 +231,6 @@ def set_sampler_returns_node_idxs_only(
self, sampler_returns_node_idxs_only: bool
):
"""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`."
)
self.sampler_returns_node_idxs_only = sampler_returns_node_idxs_only
if self.in_memory_sampler is not None:
self.in_memory_sampler.set_return_options(
Expand Down Expand Up @@ -380,20 +309,6 @@ def batch_generator():
else:
graph_samples = self.in_memory_sampler.sample(node_idxs)

# 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

try:
yield merge_lib.merge_graphs(
graph_samples, merge_schema, padding=self.padding
Expand Down
Loading