Add calendar features extractor.

dgf/src/api/transform.py

@@ -52,4 +52,6 @@
 
 from dgf.src.transform.timeseries import pad_and_cap_timeseries_features
 from dgf.src.transform.timeseries import PadAndCapTimeseriesConfig
+from dgf.src.transform.timeseries import extract_calendar_features
+from dgf.src.transform.timeseries import CalendarFeatureConfig
 

dgf/src/transform/timeseries.py

@@ -111,6 +111,52 @@ def _pad_and_cap_single_feature(
 
   return padded_matrix, mask_matrix
 
+_SUPPORTED_CALENDAR_FEATURES = (
+    "second",
+    "minute",
+    "hour",
+    "day_of_week",
+    "month",
+    "year",
+)
+
+
+@dataclasses_json.dataclass_json
+@dataclasses.dataclass
+class CalendarFeatureConfig:
+  """Configuration for extracting calendar features from timestamps.
+
+  Attributes:
+    features: Tuple of calendar feature names to extract. Supported values:
+      "second", "minute", "hour", "day_of_week", "month", "year".
+  """
+
+  features: Tuple[str, ...] = _SUPPORTED_CALENDAR_FEATURES
+
+
+def _compute_calendar_feature(ts_array: np.ndarray, feature: str) -> np.ndarray:
+  """Computes a single vectorized calendar feature from an int64 timestamp array."""
+  if feature == "second":
+    return (ts_array % 60).astype(np.float32)
+  if feature == "minute":
+    return ((ts_array // 60) % 60).astype(np.float32)
+  if feature == "hour":
+    return ((ts_array // 3600) % 24).astype(np.float32)
+  if feature == "day_of_week":
+    return (((ts_array // 86400) + 3) % 7).astype(np.float32)
+
+  dt = ts_array.astype("datetime64[s]")
+
+  if feature == "month":
+    return (dt.astype("datetime64[M]").astype(int) % 12 + 1).astype(np.float32)
+  if feature == "year":
+    return (dt.astype("datetime64[Y]").astype(int) + 1970).astype(np.float32)
+
+  raise ValueError(
+      f"Unsupported calendar feature: '{feature}'. Supported features:"
+      f" {_SUPPORTED_CALENDAR_FEATURES}"
+  )
+
 
 def _process_feature_set(
     features: in_memory_graph.Features,
@@ -141,7 +187,7 @@ def _process_feature_set(
     fschema = feature_schemas[fname]
 
     dtype = feature_format.FEATURE_FORMAT_TO_NP_DTYPE[fschema.format]
-    feat_shape = fschema.shape[1:] if fschema.shape is not None else ()
+    feat_shape = temporal_util.get_timeseries_step_shape(fschema)
 
     padded_matrix, mask_matrix = _pad_and_cap_single_feature(
         raw_series=features[fname],
@@ -153,10 +199,8 @@ def _process_feature_set(
     )
 
     new_features[fname] = padded_matrix
-    new_feat_schemas[fname] = dataclasses.replace(
-        fschema,
-        shape=(seq_len,) + feat_shape,
-        is_timeseries=True,
+    new_feat_schemas[fname] = temporal_util.with_sequence_length(
+        fschema, seq_len
     )
 
     new_features[f"{fname}_mask"] = mask_matrix
@@ -266,3 +310,159 @@ def pad_and_cap_timeseries_features(
           node_sets=new_ns_schemas, edge_sets=new_es_schemas
       ),
   )
+
+
+def _extract_entity_set_calendar_features(
+    features: in_memory_graph.Features,
+    feature_schemas: schema_lib.FeatureSetSchema,
+    ts_specs: List[temporal_util.TimeseriesGroupSpec],
+    config: CalendarFeatureConfig,
+) -> Tuple[in_memory_graph.Features, schema_lib.FeatureSetSchema]:
+  """Extracts calendar features from timestamp features of a single entity set."""
+  new_features: in_memory_graph.Features = dict(features)
+  new_feat_schemas: schema_lib.FeatureSetSchema = dict(feature_schemas)
+
+  timestamp_fnames = {
+      spec.timestamp_feature_name
+      for spec in ts_specs
+      if spec.timestamp_feature_name is not None
+  }
+
+  for fname in timestamp_fnames:
+    if fname not in features:
+      continue
+    fschema = feature_schemas[fname]
+    if fschema.semantic != schema_lib.FeatureSemantic.TIMESTAMP:
+      continue
+
+    raw_val = features[fname]
+    if isinstance(raw_val, np.ndarray) and raw_val.dtype == np.object_:
+      raise ValueError(
+          "extract_calendar_features requires fixed-length timestamp tensors,"
+          f" but feature '{fname}' is a variable-length object array."
+          " Please run pad_and_cap_timeseries_features first."
+      )
+
+    if not fschema.is_timeseries:
+      cal_timestamps = None
+    elif fschema.timestamps is not None:
+      cal_timestamps = fschema.timestamps
+    else:
+      cal_timestamps = fname
+
+    for cal_name in config.features:
+      out_fname = f"{fname}_{cal_name}"
+      cal_arr = _compute_calendar_feature(raw_val, cal_name)
+      new_features[out_fname] = cal_arr
+      new_feat_schemas[out_fname] = schema_lib.FeatureSchema(
+          format=schema_lib.FeatureFormat.FLOAT_32,
+          semantic=schema_lib.FeatureSemantic.NUMERICAL,
+          shape=fschema.shape,
+          is_timeseries=fschema.is_timeseries,
+          timestamps=cal_timestamps,
+      )
+
+  return new_features, new_feat_schemas
+
+
+def extract_calendar_features(
+    graph: in_memory_graph.InMemoryGraph,
+    schema: schema_lib.GraphSchema,
+    config: Optional[CalendarFeatureConfig] = None,
+    schema_cache: Optional[temporal_util.TimeseriesSchemaCache] = None,
+) -> Tuple[in_memory_graph.InMemoryGraph, schema_lib.GraphSchema]:
+  """Extracts calendar features (e.g. hour, day_of_week) from timestamp features.
+
+  Requires fixed-length timestamp tensors (e.g., produced after running
+  `pad_and_cap_timeseries_features`).
+
+  Usage example:
+
+  ```python
+  graph, schema = dgf.transform.pad_and_cap_timeseries_features(
+      graph, schema, cap_config
+  )
+  graph, schema = dgf.transform.extract_calendar_features(graph, schema)
+  ```
+
+  Args:
+    graph: The input in-memory graph.
+    schema: The graph schema containing timestamp features.
+    config: Optional `CalendarFeatureConfig`.
+    schema_cache: Optional pre-computed `TimeseriesSchemaCache` for reuse.
+
+  Returns:
+    Tuple `(new_graph, new_schema)` containing original and extracted calendar
+    features.
+  """
+  if config is None:
+    config = CalendarFeatureConfig()
+
+  for f in config.features:
+    if f not in _SUPPORTED_CALENDAR_FEATURES:
+      raise ValueError(
+          f"Unsupported calendar feature: '{f}'. Supported features:"
+          f" {_SUPPORTED_CALENDAR_FEATURES}"
+      )
+
+  # TODO(mesimon): Pre-compute and reuse schema_cache across dataset examples
+  # rather than recomputing cache for every graph sample.
+  if schema_cache is None:
+    schema_cache = temporal_util.extract_timeseries_schema_cache(schema)
+
+  new_node_sets = {}
+  new_ns_schemas = {}
+
+  for ns_name, ns_schema in schema.node_sets.items():
+    ns_val = graph.node_sets[ns_name]
+    ts_specs = schema_cache.node_sets.get(ns_name, [])
+    if not ts_specs:
+      new_node_sets[ns_name] = ns_val
+      new_ns_schemas[ns_name] = ns_schema
+      continue
+
+    new_feats, new_schemas = _extract_entity_set_calendar_features(
+        features=ns_val.features,
+        feature_schemas=ns_schema.features,
+        ts_specs=ts_specs,
+        config=config,
+    )
+    new_node_sets[ns_name] = in_memory_graph.InMemoryNodeSet(
+        num_nodes=ns_val.num_nodes, features=new_feats
+    )
+    new_ns_schemas[ns_name] = schema_lib.NodeSchema(features=new_schemas)
+
+  new_edge_sets = {}
+  new_es_schemas = {}
+
+  for es_name, es_schema in schema.edge_sets.items():
+    es_val = graph.edge_sets[es_name]
+    ts_specs = schema_cache.edge_sets.get(es_name, [])
+    if not ts_specs:
+      new_edge_sets[es_name] = es_val
+      new_es_schemas[es_name] = es_schema
+      continue
+
+    new_feats, new_schemas = _extract_entity_set_calendar_features(
+        features=es_val.features,
+        feature_schemas=es_schema.features,
+        ts_specs=ts_specs,
+        config=config,
+    )
+    new_edge_sets[es_name] = in_memory_graph.InMemoryEdgeSet(
+        adjacency=es_val.adjacency, features=new_feats
+    )
+    new_es_schemas[es_name] = schema_lib.EdgeSchema(
+        source=es_schema.source,
+        target=es_schema.target,
+        features=new_schemas,
+    )
+
+  return (
+      in_memory_graph.InMemoryGraph(
+          node_sets=new_node_sets, edge_sets=new_edge_sets
+      ),
+      schema_lib.GraphSchema(
+          node_sets=new_ns_schemas, edge_sets=new_es_schemas
+      ),
+  )