Set timeseries padding parameter to max_timeseries_len

dgf/src/analyse/BUILD

@@ -57,6 +57,7 @@ py_library(
         "//dgf/src/data:in_memory_graph",
         "//dgf/src/data:padding",
         "//dgf/src/data:schema",
+        # numpy dep,
     ],
 )
 
@@ -183,9 +184,12 @@ py_test(
     deps = [
         ":padding",
         # absl/testing:absltest dep,
+        "//dgf/src/data:in_memory_graph",
         "//dgf/src/data:padding",
+        "//dgf/src/data:schema",
         "//dgf/src/util:gen_test_graph",
         "//dgf/src/util:test_util",
+        # numpy dep,
     ],
 )
 

dgf/src/analyse/padding.py

@@ -13,19 +13,35 @@
 # limitations under the License.
 
 import math
-from typing import Iterator, Optional
+from typing import Dict, Iterator, Optional
 from dgf.src.data import in_memory_graph as in_memory_graph_lib
 from dgf.src.data import padding as padding_lib
 from dgf.src.data import schema as schema_lib
 
 # TODO(gbm): Implement Beam version.
 
 
+def _feature_padding_from_schema(
+    features_schema: schema_lib.FeatureSetSchema,
+    max_timeseries_len: Optional[int] = None,
+) -> Dict[str, padding_lib.FeaturePadding]:
+  """Creates feature padding for timeseries features from a feature schema."""
+  features = {}
+  if max_timeseries_len is not None:
+    for feat_name, feat_schema in features_schema.items():
+      if feat_schema.is_timeseries:
+        features[feat_name] = padding_lib.FeaturePadding(
+            max_timeseries_len=max_timeseries_len
+        )
+  return features
+
+
 def padding_from_graph_generator(
     schema: schema_lib.GraphSchema,
     graphs: Iterator[in_memory_graph_lib.InMemoryGraph],
     relative_margin: float = 0.1,
     absolute_margin: int = 1,
+    max_timeseries_len: Optional[int] = None,
 ) -> padding_lib.Padding:
   """Creates a padding configuration from a set of in-memory graphs.
 
@@ -55,6 +71,8 @@ def padding_from_graph_generator(
     graphs: An iterator over in-memory heterogeneous graphs.
     relative_margin: A relative margin.
     absolute_margin: An absolute margin.
+    max_timeseries_len: Optional maximum sequence length for timeseries
+      features.
 
   Returns:
     Padding configuration compatible with the graphs.
@@ -86,16 +104,24 @@ def padding_from_graph_generator(
   padded_node_sets = {}
   for node_set_name, max_n in max_nodes.items():
     padded_n = math.ceil((max_n + absolute_margin) * (1.0 + relative_margin))
+    node_set_schema = schema.node_sets[node_set_name]
     padded_node_sets[node_set_name] = padding_lib.NodeSetPadding(
-        num_nodes=padded_n
+        num_nodes=padded_n,
+        features=_feature_padding_from_schema(
+            node_set_schema.features, max_timeseries_len
+        ),
     )
 
   # Create the padding with a margin.
   padded_edge_sets = {}
   for edge_set_name, max_e in max_edges.items():
     padded_e = math.ceil((max_e + absolute_margin) * (1.0 + relative_margin))
+    edge_set_schema = schema.edge_sets[edge_set_name]
     padded_edge_sets[edge_set_name] = padding_lib.EdgeSetPadding(
-        num_edges=padded_e
+        num_edges=padded_e,
+        features=_feature_padding_from_schema(
+            edge_set_schema.features, max_timeseries_len
+        ),
     )
 
   return padding_lib.Padding(

dgf/src/analyse/padding_test.py

@@ -14,9 +14,12 @@
 
 from absl.testing import absltest
 from dgf.src.analyse import padding as padding_lib
+from dgf.src.data import in_memory_graph as in_memory_graph_lib
 from dgf.src.data import padding as padding_data_lib
+from dgf.src.data import schema as schema_lib
 from dgf.src.util import gen_test_graph
 from dgf.src.util import test_util
+import numpy as np
 
 test_util.disable_diff_truncation()
 
@@ -43,6 +46,110 @@ def test_basic(self):
     )
     test_util.assert_are_equal(self, padding, expected_padding)
 
+  def test_timeseries_padding(self):
+    # 3 features: 1 static timeseries, 2 variable timeseries sharing a group.
+    schema = schema_lib.GraphSchema(
+        node_sets={
+            "n": schema_lib.NodeSchema(
+                features={
+                    "static_ts": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.FLOAT_32,
+                        shape=(5,),
+                        is_timeseries=True,
+                    ),
+                    "var_ts1": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.FLOAT_32,
+                        shape=(None,),
+                        is_timeseries=True,
+                        group="g1",
+                    ),
+                    "var_ts2": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.INTEGER_64,
+                        semantic=schema_lib.FeatureSemantic.TIMESTAMP,
+                        shape=(None,),
+                        is_timeseries=True,
+                        is_creation_time=True,
+                        group="g1",
+                    ),
+                }
+            )
+        },
+        edge_sets={
+            "e": schema_lib.EdgeSchema(
+                source="n",
+                target="n",
+                features={
+                    "edge_ts": schema_lib.FeatureSchema(
+                        format=schema_lib.FeatureFormat.FLOAT_32,
+                        shape=(None,),
+                        is_timeseries=True,
+                    ),
+                },
+            )
+        },
+    )
+    g1 = in_memory_graph_lib.InMemoryGraph(
+        node_sets={
+            "n": in_memory_graph_lib.InMemoryNodeSet(
+                num_nodes=2,
+                features={
+                    "static_ts": np.zeros((2, 5), dtype=np.float32),
+                    "var_ts1": np.array(
+                        [np.array([1.0, 2.0]), np.array([3.0, 4.0, 5.0])],
+                        dtype=object,
+                    ),
+                    "var_ts2": np.array(
+                        [
+                            np.array([10, 20], dtype=np.int64),
+                            np.array([20, 30, 40], dtype=np.int64),
+                        ],
+                        dtype=object,
+                    ),
+                },
+            )
+        },
+        edge_sets={
+            "e": in_memory_graph_lib.InMemoryEdgeSet(
+                adjacency=np.array([[0], [1]], dtype=np.int32),
+                features={
+                    "edge_ts": np.array([np.array([1.0, 2.0])], dtype=object),
+                },
+            )
+        },
+    )
+    padding = padding_lib.padding_from_graph_generator(
+        schema, iter([g1]), max_timeseries_len=4
+    )
+    expected_padding = padding_data_lib.Padding(
+        node_sets={
+            "n": padding_data_lib.NodeSetPadding(
+                num_nodes=4,
+                features={
+                    "static_ts": padding_data_lib.FeaturePadding(
+                        max_timeseries_len=4
+                    ),
+                    "var_ts1": padding_data_lib.FeaturePadding(
+                        max_timeseries_len=4
+                    ),
+                    "var_ts2": padding_data_lib.FeaturePadding(
+                        max_timeseries_len=4
+                    ),
+                },
+            )
+        },
+        edge_sets={
+            "e": padding_data_lib.EdgeSetPadding(
+                num_edges=3,
+                features={
+                    "edge_ts": padding_data_lib.FeaturePadding(
+                        max_timeseries_len=4
+                    ),
+                },
+            )
+        },
+    )
+    test_util.assert_are_equal(self, padding, expected_padding)
+
   def test_print_padding(self):
     padding = padding_data_lib.Padding(
         node_sets={
@@ -64,6 +171,22 @@ def test_print_padding(self):
   e1: 100 edges"""
     self.assertEqual(output, expected_output)
 
+  def test_feature_padding_from_schema(self):
+    features_schema = {
+        "f": schema_lib.FeatureSchema(
+            format=schema_lib.FeatureFormat.FLOAT_32,
+            is_timeseries=True,
+        ),
+    }
+    feature_padding = padding_lib._feature_padding_from_schema(
+        features_schema, max_timeseries_len=5
+    )
+    expected_padding = {
+        "f": padding_data_lib.FeaturePadding(max_timeseries_len=5),
+    }
+    self.assertEqual(feature_padding, expected_padding)
+
 
 if __name__ == "__main__":
   absltest.main()
+

dgf/src/learning/ten_lines/node_prediction_dataset.py

@@ -340,7 +340,11 @@ def gen_normalized_samples():
     if self.verbose_preparation:
       log.info("Compute graph statistics for padding")
     padding = padding_lib.padding_from_graph_generator(
-        self.schema, gen_normalized_samples_iter
+        self.schema,
+        gen_normalized_samples_iter,
+        max_timeseries_len=getattr(
+            sample_generator.sampling_config, "max_timeseries_len", None
+        ),
     )
     if self.verbose_preparation:
       log.info(