Elgato_dark/chronos-forecasting
1
0
Fork 0
mirror of https://github.com/amazon-science/chronos-forecasting synced 2026-05-24 10:08:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Implement method stubs

Browse source
This commit is contained in:
Oleksandr Shchur 2026-05-13 17:59:51 +00:00
parent a8837cfd88
commit f826a74b4d
1 changed files with 398 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

414
src/chronos/chronos2/preprocess.py
View file

@ -44,10 +44,32 @@ def from_tensor(
-------
list[PreparedInput], one per series
"""
...
if isinstance(data, np.ndarray):
data = torch.from_numpy(data)
if data.ndim != 3:
raise ValueError(
f"Expected 3-d tensor with shape (n_series, n_variates, context_length), got shape {tuple(data.shape)}"
)
data = data.to(dtype=torch.float32)
n_targets = data.shape[1]
results: list[PreparedInput] = []
for i in range(data.shape[0]):
future_cov = torch.full((n_targets, prediction_length), fill_value=torch.nan)
results.append(
PreparedInput(
context=data[i].clone(),
future_covariates=future_cov,
n_targets=n_targets,
n_covariates=0,
n_future_covariates=0,
)
)
return results
def from_tensor_list(
def from_list_of_tensors(
data: "list[torch.Tensor | np.ndarray]",
prediction_length: int,
) -> list[PreparedInput]:
@ -67,7 +89,27 @@ def from_tensor_list(
-------
list[PreparedInput], one per input tensor
"""
...
results: list[PreparedInput] = []
for idx, item in enumerate(data):
if isinstance(item, np.ndarray):
item = torch.from_numpy(item)
if item.ndim > 2:
raise ValueError(
f"Each element should be 1-d or 2-d, found shape {tuple(item.shape)} at index {idx}"
)
context = item.view(-1, item.shape[-1]).to(dtype=torch.float32)
n_targets = context.shape[0]
future_cov = torch.full((n_targets, prediction_length), fill_value=torch.nan)
results.append(
PreparedInput(
context=context,
future_covariates=future_cov,
n_targets=n_targets,
n_covariates=0,
n_future_covariates=0,
)
)
return results
def from_dataframe(
@ -75,6 +117,7 @@ def from_dataframe(
target_columns: list[str],
prediction_length: int,
future_df: "pd.DataFrame | None" = None,
known_covariate_columns: list[str] | None = None,
id_column: str = "item_id",
timestamp_column: str = "timestamp",
use_target_encoding: bool = True,
@ -99,7 +142,12 @@ def from_dataframe(
prediction_length
Number of future time steps
future_df
Optional DataFrame with future covariate values (same id_column, timestamp_column)
Optional DataFrame with future covariate values (same id_column, timestamp_column).
Mutually exclusive with known_covariate_columns.
known_covariate_columns
Optional list of column names that are known-future covariates. Use when future values
are not available (e.g., during training). Future values will be NaN-filled.
Mutually exclusive with future_df.
id_column
Column name for series ID
timestamp_column
@ -114,7 +162,66 @@ def from_dataframe(
-------
list[PreparedInput], one per unique item_id (in original order)
"""
...
if future_df is not None and known_covariate_columns is not None:
raise ValueError("Cannot provide both future_df and known_covariate_columns")
if validate_inputs:
_validate_dataframe(
df=df,
future_df=future_df,
target_columns=target_columns,
prediction_length=prediction_length,
id_column=id_column,
timestamp_column=timestamp_column,
)
import pandas.api.types as ptypes
covariate_columns = [
c for c in df.columns if c not in {id_column, timestamp_column} and c not in target_columns
]
# Determine which covariates are known-future
known_future_columns: set[str] = set()
if future_df is not None:
known_future_columns = {c for c in covariate_columns if c in future_df.columns}
elif known_covariate_columns is not None:
known_future_columns = {c for c in covariate_columns if c in known_covariate_columns}
# Extract target: (n_targets, total_rows)
target = df[target_columns].to_numpy(dtype=np.float32, na_value=np.nan).T
# Extract past covariates
past_covariates: dict[str, np.ndarray] = {}
for col in covariate_columns:
if ptypes.is_numeric_dtype(df[col]):
past_covariates[col] = df[col].to_numpy(dtype=np.float32, na_value=np.nan)
else:
past_covariates[col] = df[col].to_numpy(dtype=object)
# Extract future covariate values: key present = known-future, value = data or None
future_covariates: dict[str, np.ndarray | None] = {}
if future_df is not None:
for col in known_future_columns:
if ptypes.is_numeric_dtype(future_df[col]):
future_covariates[col] = future_df[col].to_numpy(dtype=np.float32, na_value=np.nan)
else:
future_covariates[col] = future_df[col].to_numpy(dtype=object)
else:
for col in known_future_columns:
future_covariates[col] = None
# Compute series lengths
series_lengths = df.groupby(id_column, sort=False).size().tolist()
return _build_prepared_inputs(
target=target,
past_covariates=past_covariates,
future_covariates=future_covariates,
series_lengths=series_lengths,
prediction_length=prediction_length,
use_target_encoding=use_target_encoding,
)
def from_list_of_dicts(
@ -154,13 +261,63 @@ def from_list_of_dicts(
-------
list[PreparedInput], one per dict
"""
...
if validate_inputs:
_validate_list_of_dicts(data=data, prediction_length=prediction_length)
if len(data) == 0:
return []
# Determine covariate structure from first dict
first_past_covariates = data[0].get("past_covariates", {})
first_future_covariates = data[0].get("future_covariates", {})
past_covariate_keys = sorted(first_past_covariates.keys())
known_future_columns = set(first_future_covariates.keys())
# Stack targets: (n_targets, total_context_rows)
target_arrays = []
series_lengths = []
for d in data:
t = np.asarray(d["target"], dtype=np.float32)
if t.ndim == 1:
t = t.reshape(1, -1)
target_arrays.append(t)
series_lengths.append(t.shape[-1])
target = np.concatenate(target_arrays, axis=1)
# Stack past covariates: {name: (total_context_rows,)}
past_covariates: dict[str, np.ndarray] = {}
for key in past_covariate_keys:
arrays = [np.asarray(d.get("past_covariates", {})[key]) for d in data]
stacked = np.concatenate(arrays)
if np.issubdtype(stacked.dtype, np.number):
past_covariates[key] = stacked.astype(np.float32)
else:
past_covariates[key] = stacked.astype(object)
# Stack future covariates: {name: array or None}
future_covariates: dict[str, np.ndarray | None] = {}
for key in known_future_columns:
arrays = [np.asarray(d.get("future_covariates", {})[key]) for d in data]
stacked = np.concatenate(arrays)
if np.issubdtype(stacked.dtype, np.number):
future_covariates[key] = stacked.astype(np.float32)
else:
future_covariates[key] = stacked.astype(object)
return _build_prepared_inputs(
target=target,
past_covariates=past_covariates,
future_covariates=future_covariates,
series_lengths=series_lengths,
prediction_length=prediction_length,
use_target_encoding=use_target_encoding,
)
def _build_prepared_inputs(
target: np.ndarray,
past_covariates: dict[str, np.ndarray],
future_covariates: dict[str, np.ndarray],
future_covariates: dict[str, np.ndarray | None],
series_lengths: list[int],
prediction_length: int,
use_target_encoding: bool,
@ -171,8 +328,10 @@ def _build_prepared_inputs(
Assumptions
-----------
- Arrays are stacked in item order (item 0's rows first, then item 1's, etc.)
- future_covariates keys are a subset of past_covariates keys
- Categorical columns have object dtype; numeric columns have float32 dtype
- future_covariates keys are a subset of past_covariates keys
- Key present in future_covariates = known-future covariate
- Value is the actual future data (shape: n_series * prediction_length) or None if unavailable
Parameters
----------
@ -182,8 +341,8 @@ def _build_prepared_inputs(
{name: values} for all covariates (past-only and known-future)
Each array shape: (total_context_rows,)
future_covariates
{name: values} for known-future covariates only
Each array shape: (n_series * prediction_length,)
{name: values_or_None} for known-future covariates.
Each array shape: (n_series * prediction_length,), or None if values unavailable.
series_lengths
Context length of each series (sum = total_context_rows)
prediction_length
@ -195,7 +354,126 @@ def _build_prepared_inputs(
-------
list[PreparedInput], one per series
"""
...
n_series = len(series_lengths)
n_targets = target.shape[0]
n_covariates = len(past_covariates)
n_future_covariates = len(future_covariates)
# Build item ID codes for target encoding
id_codes = np.repeat(np.arange(n_series), series_lengths)
future_id_codes = np.repeat(np.arange(n_series), prediction_length)
# Encode covariates
encoded_past_covariates: list[np.ndarray] = []
encoded_future_covariates: list[np.ndarray] = []
for key, values in past_covariates.items():
is_known_future = key in future_covariates
future_values = future_covariates.get(key)
if values.dtype == object:
# Categorical: ordinal encode first
all_past_values = values.astype(str)
categories = np.unique(all_past_values[all_past_values != "nan"])
cat_to_code = {cat: i for i, cat in enumerate(categories)}
n_categories = len(categories)
# NaN in past gets its own code
nan_code = n_categories
n_categories_with_nan = n_categories + 1
past_codes = np.array([cat_to_code.get(v, nan_code) for v in all_past_values], dtype=np.intp)
future_codes = None
if future_values is not None:
all_future_values = future_values.astype(str)
future_codes = np.array(
[cat_to_code.get(v, nan_code) for v in all_future_values], dtype=np.intp
)
if use_target_encoding and n_targets == 1:
encoded_past, encoded_future = _target_encode(
id_codes=id_codes,
cat_codes=past_codes,
target=target[0],
n_items=n_series,
n_categories=n_categories_with_nan,
future_id_codes=future_id_codes if future_codes is not None else None,
future_cat_codes=future_codes,
)
encoded_past_covariates.append(encoded_past)
if is_known_future:
encoded_future_covariates.append(
encoded_future if encoded_future is not None
else np.full(n_series * prediction_length, np.nan, dtype=np.float32)
)
else:
encoded_past_covariates.append(past_codes.astype(np.float32))
if is_known_future:
encoded_future_covariates.append(
future_codes.astype(np.float32) if future_codes is not None
else np.full(n_series * prediction_length, np.nan, dtype=np.float32)
)
else:
encoded_past_covariates.append(values)
if is_known_future:
encoded_future_covariates.append(
future_values if future_values is not None
else np.full(n_series * prediction_length, np.nan, dtype=np.float32)
)
if not is_known_future:
encoded_future_covariates.append(
np.full(n_series * prediction_length, np.nan, dtype=np.float32)
)
# Split into per-series PreparedInputs
past_splits = np.cumsum(series_lengths[:-1]).tolist() if n_series > 1 else []
future_splits = (
list(range(prediction_length, n_series * prediction_length, prediction_length))
if n_series > 1
else []
)
results: list[PreparedInput] = []
for i in range(n_series):
# Target slice
p_start = sum(series_lengths[:i])
p_end = p_start + series_lengths[i]
target_i = target[:, p_start:p_end]
# Past covariates slice
if encoded_past_covariates:
past_cov_i = np.stack([arr[p_start:p_end] for arr in encoded_past_covariates])
else:
past_cov_i = np.zeros((0, series_lengths[i]), dtype=np.float32)
# Future covariates slice
f_start = i * prediction_length
f_end = f_start + prediction_length
if encoded_future_covariates:
future_cov_i = np.stack([arr[f_start:f_end] for arr in encoded_future_covariates])
else:
future_cov_i = np.zeros((0, prediction_length), dtype=np.float32)
# Build context: targets then covariates
context = np.concatenate([target_i, past_cov_i], axis=0)
# Build future_covariates: NaN padding for targets, then covariate futures
target_padding = np.full((n_targets, prediction_length), np.nan, dtype=np.float32)
future_full = np.concatenate([target_padding, future_cov_i], axis=0)
results.append(
PreparedInput(
context=torch.from_numpy(context).to(dtype=torch.float32),
future_covariates=torch.from_numpy(future_full).to(dtype=torch.float32),
n_targets=n_targets,
n_covariates=n_covariates,
n_future_covariates=n_future_covariates,
)
)
return results
def _validate_dataframe(
@ -213,13 +491,44 @@ def _validate_dataframe(
- Required columns exist
- Target columns are numeric
- All series have >= 3 points
- Consistent frequency across series
- future_df has same item_ids and exactly prediction_length rows per series
"""
...
required = {id_column, timestamp_column} | set(target_columns)
missing = required - set(df.columns)
if missing:
raise ValueError(f"DataFrame is missing required columns: {missing}")
for col in target_columns:
if not np.issubdtype(df[col].dtype, np.number):
raise ValueError(f"Target column '{col}' must be numeric, got dtype {df[col].dtype}")
series_sizes = df.groupby(id_column, sort=False).size()
short_series = series_sizes[series_sizes < 3]
if len(short_series) > 0:
raise ValueError(
f"All series must have >= 3 points. Found {len(short_series)} series with fewer."
)
if future_df is not None:
future_missing = {id_column} - set(future_df.columns)
if future_missing:
raise ValueError(f"future_df is missing required columns: {future_missing}")
past_ids = df[id_column].unique()
future_ids = future_df[id_column].unique()
if not np.array_equal(np.sort(past_ids), np.sort(future_ids)):
raise ValueError("future_df must have the same item IDs as df")
future_sizes = future_df.groupby(id_column, sort=False).size()
wrong_length = future_sizes[future_sizes != prediction_length]
if len(wrong_length) > 0:
raise ValueError(
f"future_df must have exactly {prediction_length} rows per item. "
f"Found {len(wrong_length)} items with wrong length."
)
def _validate_dict_list(
def _validate_list_of_dicts(
data: list[dict],
prediction_length: int,
) -> None:
@ -227,13 +536,86 @@ def _validate_dict_list(
Validate list[dict] structure. Raises ValueError on failure.
Checks:
- All dicts have same keys
- All dicts have "target" key
- All targets have same n_targets
- All past_covariates have same column names
- All future_covariates have same column names and are subset of past_covariates
- future_covariates have length == prediction_length
- past_covariates have length == target length
"""
...
if len(data) == 0:
return
allowed_keys = {"target", "past_covariates", "future_covariates"}
first_past_keys = sorted(data[0].get("past_covariates", {}).keys())
first_future_keys = sorted(data[0].get("future_covariates", {}).keys())
first_target = np.asarray(data[0]["target"])
first_n_targets = 1 if first_target.ndim == 1 else first_target.shape[0]
if not set(first_future_keys).issubset(set(first_past_keys)):
raise ValueError(
f"future_covariates keys must be a subset of past_covariates keys. "
f"Got past={first_past_keys}, future={first_future_keys}"
)
for idx, d in enumerate(data):
keys = set(d.keys())
if not keys.issubset(allowed_keys):
raise ValueError(
f"Invalid keys at index {idx}. Allowed: {allowed_keys}, found: {keys}"
)
if "target" not in keys:
raise ValueError(f"Element at index {idx} is missing required key 'target'")
target = np.asarray(d["target"])
if target.ndim > 2:
raise ValueError(
f"Target must be 1-d or 2-d, found shape {tuple(target.shape)} at index {idx}"
)
n_targets = 1 if target.ndim == 1 else target.shape[0]
if n_targets != first_n_targets:
raise ValueError(
f"All targets must have same n_targets. Expected {first_n_targets}, "
f"got {n_targets} at index {idx}"
)
history_length = target.shape[-1]
past_covariates = d.get("past_covariates", {})
if not isinstance(past_covariates, dict):
raise ValueError(
f"past_covariates must be a dict at index {idx}, got {type(past_covariates)}"
)
if sorted(past_covariates.keys()) != first_past_keys:
raise ValueError(
f"All past_covariates must have same keys. Expected {first_past_keys}, "
f"got {sorted(past_covariates.keys())} at index {idx}"
)
for key, val in past_covariates.items():
val = np.asarray(val)
if val.ndim != 1 or len(val) != history_length:
raise ValueError(
f"past_covariates['{key}'] must be 1-d with length {history_length}, "
f"got shape {tuple(val.shape)} at index {idx}"
)
future_covariates = d.get("future_covariates", {})
if not isinstance(future_covariates, dict):
raise ValueError(
f"future_covariates must be a dict at index {idx}, got {type(future_covariates)}"
)
if sorted(future_covariates.keys()) != first_future_keys:
raise ValueError(
f"All future_covariates must have same keys. Expected {first_future_keys}, "
f"got {sorted(future_covariates.keys())} at index {idx}"
)
for key, val in future_covariates.items():
val = np.asarray(val)
if val.ndim != 1 or len(val) != prediction_length:
raise ValueError(
f"future_covariates['{key}'] must be 1-d with length {prediction_length}, "
f"got shape {tuple(val.shape)} at index {idx}"
)
def _target_encode(