Skip to content

Remove TimeseriesPadAndCap. #171

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

Merged
copybara-service merged 1 commit into from
Aug 31, 2026
+29 −834
Merged

Remove TimeseriesPadAndCap. #171

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
2 changes: 0 additions & 2 deletions dgf/src/api/transform.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,6 @@
from dgf.src.transform.timeseries import CalendarFeature
from dgf.src.transform.timeseries import CalendarFeatureExtractor
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 TimestampFeatureExtractor
from dgf.src.transform.timeseries import TimestampFeatureExtractorConfig

299 changes: 5 additions & 294 deletions dgf/src/transform/timeseries.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

"""Padding and capping for timeseries sequence features in graphs."""
"""Temporal and timeseries feature extraction for graphs."""

# TODO(simonmeierhans): If a timeseries feature has no group defined, and there
# is a group with the same name as the feature it currently implicitly joins
Expand All @@ -21,101 +21,15 @@
# pytype: disable=module-attr
import dataclasses
import enum
from typing import Any, List, Optional, Tuple
from typing import Any, Optional, Tuple

import dataclasses_json
from dgf.src.data import in_memory_graph
from dgf.src.data import schema as schema_lib
from dgf.src.io import feature_format
from dgf.src.util import temporal as temporal_util
import numpy as np


@dataclasses_json.dataclass_json
@dataclasses.dataclass
class PadAndCapTimeseriesConfig:
"""Configuration for padding and capping timeseries features.

Attributes:
sequence_length: Fixed target sequence dimension K. Sequences longer than K
are capped to the most recent K steps (`[-K:]`). Shorter sequences are
left-padded to length K.
padding_value: Scalar value used for left-padding shorter sequences.
"""

sequence_length: int = 30
padding_value: Any = 0


def _pad_and_cap_single_feature(
raw_series: np.ndarray,
seq_len: int,
feat_shape: Tuple[int, ...],
padding_value: Any,
dtype: Any,
is_static_shape: bool,
) -> Tuple[np.ndarray, np.ndarray]:
"""Pads and caps a single sequence feature into (padded_matrix, mask_matrix)."""
num_entities = raw_series.shape[0]
if dtype == np.bytes_ and raw_series.dtype.kind in ("S", "a"):
dtype = raw_series.dtype
if num_entities == 0:
return (
np.empty((0, seq_len) + feat_shape, dtype=dtype),
np.empty((0, seq_len), dtype=np.bool_),
)

# Fast vectorized path when all entities share a fixed sequence length.
if is_static_shape and raw_series.ndim >= 2:
num_steps = raw_series.shape[1]
if num_steps >= seq_len:
padded_matrix = raw_series[:, -seq_len:].astype(dtype, copy=True)
mask_matrix = np.ones((num_entities, seq_len), dtype=np.bool_)
return padded_matrix, mask_matrix

pad_width = [(0, 0), (seq_len - num_steps, 0)] + [(0, 0)] * len(feat_shape)
padded_matrix = np.pad(
raw_series.astype(dtype, copy=False),
pad_width=pad_width,
mode="constant",
constant_values=padding_value,
)
mask_width = [(0, 0), (seq_len - num_steps, 0)]
mask_matrix = np.pad(
np.ones((num_entities, num_steps), dtype=np.bool_),
pad_width=mask_width,
mode="constant",
constant_values=False,
)
return padded_matrix, mask_matrix

padded_matrix = np.full(
(num_entities, seq_len) + feat_shape,
fill_value=padding_value,
dtype=dtype,
)
# Binary mask matrix matching sequence length shape (num_entities, seq_len)
# where True indicates valid observed time steps and False indicates
# left-padded steps.
mask_matrix = np.zeros((num_entities, seq_len), dtype=np.bool_)

# TODO(mesimon): Move into C++ for performance.
for idx in range(num_entities):
raw_arr = raw_series[idx]
if not isinstance(raw_arr, np.ndarray):
raw_arr = np.asarray(raw_arr)

num_steps = len(raw_arr)
if num_steps >= seq_len:
padded_matrix[idx] = raw_arr[-seq_len:]
mask_matrix[idx] = True
elif num_steps > 0:
padded_matrix[idx, -num_steps:] = raw_arr
mask_matrix[idx, -num_steps:] = True

return padded_matrix, mask_matrix


class CalendarFeature(str, enum.Enum):
"""Supported calendar features to extract from timestamps."""

Expand Down Expand Up @@ -183,208 +97,6 @@ def _compute_calendar_feature(
)


class PadAndCapTimeseries:
"""Pads and caps timeseries sequence features into fixed-dimension tensors.

Transforms variable-length sequence features (`is_timeseries=True`) in the
graph into fixed-length matrices of shape `(num_entities, sequence_length) +
step_shape`. Sequences longer than `sequence_length` are capped from the right
(keeping the most recent steps), while shorter sequences are left-padded.

For each sequence group, a corresponding boolean mask feature (`{group}_mask`)
is generated with semantic `MASK`, where `True` indicates observed time steps
and `False` indicates padded steps.

Attributes:
schema: The input `GraphSchema`.
config: Configuration specifying sequence length and padding value.
schema_cache: Pre-computed `TimeseriesSchemaCache` for fast grouping.
"""

def __init__(
self,
schema: schema_lib.GraphSchema,
config: Optional[PadAndCapTimeseriesConfig] = None,
schema_cache: Optional[temporal_util.TimeseriesSchemaCache] = None,
):
self.config = config or PadAndCapTimeseriesConfig()
self.schema = schema
if schema_cache is None:
schema_cache = temporal_util.extract_timeseries_schema_cache(schema)
self.schema_cache = schema_cache

def _compute_feature_set_pad_and_cap_schema(
self,
schemas: schema_lib.FeatureSetSchema,
ts_specs: List[temporal_util.TimeseriesGroupSpec],
) -> schema_lib.FeatureSetSchema:
"""Computes schema for a feature set after padding/capping."""
new_schemas = dict(schemas)
seq_len = self.config.sequence_length

ts_features = {}
for group in ts_specs:
for feature_name in group.feature_names:
ts_features[feature_name] = group.timestamp_feature_name

for feature_name in ts_features:
feature_schema = schemas[feature_name]
ts_group = feature_schema.group or feature_name
new_schemas[feature_name] = dataclasses.replace(
temporal_util.with_sequence_length(feature_schema, seq_len),
group=ts_group,
)
if feature_schema.semantic == schema_lib.FeatureSemantic.MASK:
continue

mask_name = temporal_util.get_mask_feature_name(feature_name, schemas)
if mask_name is None:
mask_name = f"{ts_group}_mask"
if mask_name in schemas:
raise ValueError(
f"Cannot generate mask for sequence group '{ts_group}'. The"
f" fallback mask name '{mask_name}' clashes with an existing"
" feature in the schema that is not a valid mask. Please"
" explicitly define a mask feature for this group or rename the"
" clashing feature."
)

if mask_name not in new_schemas:
new_schemas[mask_name] = schema_lib.FeatureSchema(
format=schema_lib.FeatureFormat.BOOL,
semantic=schema_lib.FeatureSemantic.MASK,
shape=(seq_len,),
is_timeseries=feature_schema.is_timeseries,
group=ts_group,
)

return new_schemas

def _pad_and_cap_feature_set(
self,
values: in_memory_graph.Features,
schemas: schema_lib.FeatureSetSchema,
ts_specs: List[temporal_util.TimeseriesGroupSpec],
) -> in_memory_graph.Features:
"""Extracts fixed-dimension sequence features for a feature set."""
new_values: in_memory_graph.Features = {}
seq_len = self.config.sequence_length

# Map timeseries feature names to their associated timestamp feature name.
ts_features = {}
for group in ts_specs:
for feature_name in group.feature_names:
ts_features[feature_name] = group.timestamp_feature_name

# Copy over non-timeseries features.
for feature_name in schemas:
if feature_name not in ts_features:
new_values[feature_name] = values[feature_name]

for feature_name in ts_features:
feature_schema = schemas[feature_name]

dtype = feature_format.FEATURE_FORMAT_TO_NP_DTYPE[feature_schema.format]
feat_shape = temporal_util.get_timeseries_step_shape(feature_schema)

padded_matrix, mask_matrix = _pad_and_cap_single_feature(
raw_series=values[feature_name],
seq_len=seq_len,
feat_shape=feat_shape,
padding_value=self.config.padding_value,
dtype=dtype,
is_static_shape=feature_schema.is_static_shape(),
)

new_values[feature_name] = padded_matrix
if feature_schema.semantic == schema_lib.FeatureSemantic.MASK:
continue

ts_group = feature_schema.group or feature_name
mask_name = temporal_util.get_mask_feature_name(feature_name, schemas)
if mask_name is None:
mask_name = f"{ts_group}_mask"

# Copy over mask matrix if it doesn't exist yet. Only store once per
# group.
if mask_name not in new_values:
new_values[mask_name] = mask_matrix

return new_values

def output_schema(self) -> schema_lib.GraphSchema:
"""Returns the transformed GraphSchema."""
new_ns_schemas = {}
for ns_name, ns_schema in self.schema.node_sets.items():
ts_specs = self.schema_cache.node_sets[ns_name]
if not ts_specs:
new_ns_schemas[ns_name] = ns_schema
else:
new_ns_schemas[ns_name] = schema_lib.NodeSchema(
features=self._compute_feature_set_pad_and_cap_schema(
ns_schema.features, ts_specs
)
)

new_es_schemas = {}
for es_name, es_schema in self.schema.edge_sets.items():
ts_specs = self.schema_cache.edge_sets[es_name]
if not ts_specs:
new_es_schemas[es_name] = es_schema
else:
new_es_schemas[es_name] = schema_lib.EdgeSchema(
source=es_schema.source,
target=es_schema.target,
features=self._compute_feature_set_pad_and_cap_schema(
es_schema.features, ts_specs
),
)

return schema_lib.GraphSchema(
node_sets=new_ns_schemas, edge_sets=new_es_schemas
)

def __call__(
self, graph: in_memory_graph.InMemoryGraph
) -> in_memory_graph.InMemoryGraph:
"""Transforms timeseries sequence features in the graph."""
new_node_sets = {}
for ns_name, ns_schema in self.schema.node_sets.items():
ns_val = graph.node_sets[ns_name]
ts_specs = self.schema_cache.node_sets[ns_name]
if not ts_specs:
new_node_sets[ns_name] = ns_val
continue
new_vals = self._pad_and_cap_feature_set(
values=ns_val.features,
schemas=ns_schema.features,
ts_specs=ts_specs,
)
new_node_sets[ns_name] = in_memory_graph.InMemoryNodeSet(
num_nodes=ns_val.num_nodes, features=new_vals
)

new_edge_sets = {}
for es_name, es_schema in self.schema.edge_sets.items():
es_val = graph.edge_sets[es_name]
ts_specs = self.schema_cache.edge_sets[es_name]
if not ts_specs:
new_edge_sets[es_name] = es_val
continue
new_vals = self._pad_and_cap_feature_set(
values=es_val.features,
schemas=es_schema.features,
ts_specs=ts_specs,
)
new_edge_sets[es_name] = in_memory_graph.InMemoryEdgeSet(
adjacency=es_val.adjacency, features=new_vals
)

return in_memory_graph.InMemoryGraph(
node_sets=new_node_sets, edge_sets=new_edge_sets
)


class CalendarFeatureExtractor:
"""Extracts calendar features (e.g. hour, day_of_week) from timestamp features.

Expand Down Expand Up @@ -457,8 +169,7 @@ def _extract_feature_set_calendar_features(

assert raw_val.dtype != np.object_, (
"CalendarFeatureExtractor requires fixed-length timestamp tensors,"
f" but feature '{fname}' is a variable-length object array. Please"
" run PadAndCapTimeseries first."
f" but feature '{fname}' is a variable-length object array."
)

for cal_feat in self.config.features:
Expand Down Expand Up @@ -595,8 +306,8 @@ def _extract_feature_set_timestamp_features(
assert schema.is_static_shape() and raw_val.dtype != np.object_, (
"TimestampFeatureExtractor requires fixed-length timestamp tensors,"
f" but feature '{fname}' is a variable-length object array or has"
f" dynamic shape ({schema.shape}). Please run PadAndCapTimeseries"
" first."
f" dynamic shape ({schema.shape}). Please pad timeseries features"
" (e.g. via pad_timeseries_graph) first."
)

mask = None
Expand Down
Loading