diff --git a/sdm/models/chronos2/__init__.py b/sdm/models/chronos2/__init__.py new file mode 100644 index 000000000..bef2e698b --- /dev/null +++ b/sdm/models/chronos2/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Chronos-2.""" + +from sdm.models.chronos2.model import Chronos2 + +__all__ = [ + "Chronos2", +] diff --git a/sdm/models/chronos2/block.py b/sdm/models/chronos2/block.py new file mode 100644 index 000000000..8e59999bd --- /dev/null +++ b/sdm/models/chronos2/block.py @@ -0,0 +1,74 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: D101, D102 + +from typing import Any + +import torch +from torch import Tensor +from torch.nn import Linear, ReLU, Sequential + +from sdm.nn import RMSNorm, RotaryEmbedding, TransformerBlock + + +class LayerNorm(RMSNorm): + def forward( + self, x: Tensor, rope: RotaryEmbedding | None = None + ) -> Tensor: + # T5 accumulates in float32 and rounds before multiplying half weights. + if rope is not None: + x = rope(x) + variance = x.float().square().mean(-1, keepdim=True) + x = x * (variance + (1e-6 if self.eps is None else self.eps)).rsqrt() + if self.weight.dtype in {torch.float16, torch.bfloat16}: + x = x.to(self.weight.dtype) + return self.weight * x + + +class ResidualBlock(torch.nn.Module): + def __init__( + self, + in_channels: int, + hidden_channels: int, + out_channels: int, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + self.mlp = Sequential( + Linear(in_channels, hidden_channels, **factory_kwargs), + ReLU(), + Linear(hidden_channels, out_channels, **factory_kwargs), + ) + self.res = Linear(in_channels, out_channels, **factory_kwargs) + + def forward(self, x: Tensor) -> Tensor: + return self.mlp(x) + self.res(x) + + +class Chronos2TransformerBlock(TransformerBlock): + def __init__( + self, + channels: int, + num_heads: int, + rope: RotaryEmbedding | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + norm = LayerNorm(channels, eps=1e-6, **factory_kwargs) + super().__init__( + channels=channels, + num_query_heads=num_heads, + query_norm=norm, + key_value_norm=norm, + query_transform=rope, + key_transform=rope, + mlp=None, + scale=1.0, + bias=False, + **factory_kwargs, + ) diff --git a/sdm/models/chronos2/ckpt.py b/sdm/models/chronos2/ckpt.py new file mode 100644 index 000000000..b1177a2ab --- /dev/null +++ b/sdm/models/chronos2/ckpt.py @@ -0,0 +1,53 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import Tensor + + +def remap_ckpt(ckpt: dict[str, Tensor]) -> dict[str, Tensor]: # noqa: D103 + out = {"reg_token": ckpt["shared.weight"][1:2]} + for source, target in ( + ("input_patch_embedding", "patch_embedding"), + ("output_patch_embedding", "head"), + ): + for upstream, local in ( + ("hidden_layer", "mlp.0"), + ("output_layer", "mlp.2"), + ("residual_layer", "res"), + ): + for suffix in ("weight", "bias"): + out[f"{target}.{local}.{suffix}"] = ckpt[ + f"{source}.{upstream}.{suffix}" + ] + for layer in range(12): + prefix = f"encoder.block.{layer}.layer" + for index, stack in ((0, "time_blocks"), (1, "group_blocks")): + source = f"{prefix}.{index}" + target = f"icl_block.{stack}.{layer}" + norm = ckpt[f"{source}.layer_norm.weight"] + out[f"{target}.query_norm.weight"] = norm + out[f"{target}.key_value_norm.weight"] = norm + out[f"{target}.attn.qkv_lin.weight"] = torch.cat( + [ + ckpt[f"{source}.self_attention.{key}.weight"] + for key in "qkv" + ] + ) + out[f"{target}.attn.out_lin.weight"] = ckpt[ + f"{source}.self_attention.o.weight" + ] + target = f"icl_block.mlps.{layer}" + out[f"{target}.0.weight"] = ckpt[f"{prefix}.2.layer_norm.weight"] + out[f"{target}.1.weight"] = ckpt[f"{prefix}.2.mlp.wi.weight"] + out[f"{target}.3.weight"] = ckpt[f"{prefix}.2.mlp.wo.weight"] + out["icl_block.norm.weight"] = ckpt["encoder.final_layer_norm.weight"] + + freq = torch.arange(0, 64, 2, device=out["reg_token"].device).float() + inv_freq = 1.0 / (10_000 ** (freq / 64)) + for layer in range(12): + for transform in ("query_transform", "key_transform"): + key = f"icl_block.time_blocks.{layer}.attn.{transform}.inv_freq" + out[key] = inv_freq + return out diff --git a/sdm/models/chronos2/icl.py b/sdm/models/chronos2/icl.py new file mode 100644 index 000000000..15e3c8da5 --- /dev/null +++ b/sdm/models/chronos2/icl.py @@ -0,0 +1,90 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: D101, D102 + +from typing import Any + +import torch +from torch import Tensor +from torch.nn import Linear, ModuleList, ReLU, Sequential + +from sdm.models.chronos2.block import Chronos2TransformerBlock, LayerNorm +from sdm.nn import RotaryEmbedding + + +class ICLBlock(torch.nn.Module): + def __init__( + self, + channels: int, + num_layers: int, + num_heads: int, + hidden_channels: int, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + rope = RotaryEmbedding( + channels=channels // num_heads, + layout="split_half", + theta=10_000, + requires_grad=False, + device=device, + ) + self.time_blocks = ModuleList( + Chronos2TransformerBlock( + channels=channels, + num_heads=num_heads, + rope=rope, + **factory_kwargs, + ) + for _ in range(num_layers) + ) + self.group_blocks = ModuleList( + Chronos2TransformerBlock( + channels=channels, + num_heads=num_heads, + **factory_kwargs, + ) + for _ in range(num_layers) + ) + self.mlps = ModuleList( + Sequential( + LayerNorm(channels, eps=1e-6, **factory_kwargs), + Linear( + channels, hidden_channels, bias=False, **factory_kwargs + ), + ReLU(), + Linear( + hidden_channels, channels, bias=False, **factory_kwargs + ), + ) + for _ in range(num_layers) + ) + self.norm = LayerNorm(channels, eps=1e-6, **factory_kwargs) + + def forward( + self, + x: Tensor, # [S, N, C] + observed: Tensor, # [S, N] + group_ids: Tensor, # [S] + ) -> Tensor: + # Finite additive masks preserve upstream's fully masked-row behavior. + minimum = torch.finfo(x.dtype).min + time_mask = (~observed).to(x.dtype)[:, None, None, :] * minimum + group_mask = group_ids[:, None] == group_ids[None, :] + group_mask = group_mask[..., None] & observed[None, :, :] + group_mask = (~group_mask).permute(2, 0, 1).unsqueeze(1) + group_mask = group_mask.to(x.dtype) * minimum + for time_block, group_block, mlp in zip( + self.time_blocks, self.group_blocks, self.mlps + ): + x = time_block(query=x, attn_mask=time_mask) + x = group_block( + query=x.transpose(0, 1), + attn_mask=group_mask, + ).transpose(0, 1) + x = x + mlp(x) + return self.norm(x) diff --git a/sdm/models/chronos2/model.py b/sdm/models/chronos2/model.py new file mode 100644 index 000000000..a0d95ceff --- /dev/null +++ b/sdm/models/chronos2/model.py @@ -0,0 +1,329 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: D205 + +from __future__ import annotations + +from typing import Any, ClassVar, cast + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +from torch import Tensor + +from sdm import Recipe, RelatedTables, Stype, TableTensor, Task +from sdm.cache import Cache +from sdm.models._huggingface import download_checkpoint +from sdm.models.base import ICLModel +from sdm.models.chronos2.block import ResidualBlock +from sdm.models.chronos2.ckpt import remap_ckpt +from sdm.models.chronos2.icl import ICLBlock +from sdm.models.chronos2.recipe import default_recipe +from sdm.models.chronos2.util import rollout_quantiles +from sdm.nn import RevIN + +QUANTILES = (0.01, *[i / 100 for i in range(5, 100, 5)], 0.99) + + +class Chronos2(ICLModel): + """The time series forecasting model from the + `"Chronos-2: From Univariate to Universal Forecasting" + `_ paper. + + Args: + pretrained: Whether to load the pretrained checkpoint. + device: The device. + """ + + supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( + {Stype.numerical} + ) + supported_target_stypes: ClassVar[frozenset[Stype]] = frozenset( + {Stype.numerical} + ) + supports_multi_target: ClassVar[bool] = True + supports_context_only_columns: ClassVar[bool] = True + supports_related_tables: ClassVar[bool] = False + + def __init__( + self, + pretrained: bool = True, + device: torch.device | str | None = None, + ) -> None: + super().__init__(task=Task.regression) + + self.model = _Chronos2(device="meta" if pretrained else device) + + if pretrained: + self._load_from_pretrained(device) + + self.eval() + + def forward( + self, + x_context: Tensor | TableTensor, + y_context: Tensor | TableTensor, + *args: Any, + **kwargs: Any, + ) -> TableTensor: + r""":meta private:""" # noqa: D415 + return super().forward( + x_context[..., -self.model.context_length :, :], + y_context[..., -self.model.context_length :, :], + *args, + **kwargs, + ) + + def fit( + self, + x: Tensor | TableTensor, + y: Tensor | TableTensor, + *args: Any, + **kwargs: Any, + ) -> None: + r""":meta private:""" # noqa: D415 + super().fit( + x[..., -self.model.context_length :, :], + y[..., -self.model.context_length :, :], + *args, + **kwargs, + ) + + @classmethod + def default_recipe(cls) -> Recipe: + r""":meta private:""" # noqa: D415 + return default_recipe() + + def _load_from_pretrained( + self, + device: torch.device | str | None, + ) -> Chronos2: + device = torch.get_default_device() if device is None else device + + path = download_checkpoint( + repo_id="amazon/chronos-2", + filename="model.safetensors", + revision="29ec3766d36d6f73f0696f85560a422f50e8498c", + ) + state = remap_ckpt(load_file(path, device=str(device))) + self.model.load_state_dict(state, strict=True, assign=True) + + return self + + def _forward( + self, + x_context: TableTensor | None, + y_context: TableTensor | None, + x_query: TableTensor | None, + x_context_only: TableTensor | None, + related_context_tables: RelatedTables[TableTensor] | None, + related_query_tables: RelatedTables[TableTensor] | None, + cache: Cache | None, + generator: torch.Generator | None, + **kwargs: Any, + ) -> TableTensor: + + if cache is not None and not cache.is_recording: + x_context = cast(TableTensor, cache["x_context"]) + y_context = cast(TableTensor, cache["y_context"]) + x_context_only = cast(TableTensor | None, cache["x_context_only"]) + assert x_context is not None + assert y_context is not None + if cache is not None and cache.is_recording: + cache["x_context"] = x_context + cache["y_context"] = y_context + cache["x_context_only"] = x_context_only + y = y_context.numerical + columns = [ + f"{name}__q{round(q * 100)}" + for name in y_context.columns[Stype.numerical] + for q in QUANTILES + ] + if x_query is None or x_query.size(-2) == 0: + out = y.new_empty((*y.size()[:-2], 0, len(columns))) + else: + past = ( + y.new_empty((*y.size()[:-1], 0)) + if x_context_only is None + else x_context_only.numerical + ) + context = torch.cat([y, past, x_context.numerical], dim=-1) + *batch, length, variates = context.shape + context = context.transpose(-2, -1).reshape(-1, length).float() + horizon = x_query.size(-2) + future = context.new_full((*batch, horizon, variates), torch.nan) + known = x_query.numerical + future[..., y.size(-1) + past.size(-1) :] = known + future = future.transpose(-2, -1).reshape(-1, horizon) + groups = torch.arange( + context.size(0) // variates, device=context.device + ).repeat_interleave(variates) + out = self.model(context, future, groups, horizon) + out = out.reshape(*batch, variates, len(QUANTILES), horizon) + out = out[..., : y.size(-1), :, :].movedim(-1, -3).flatten(-2) + return TableTensor.from_tensor(out, columns=columns) + + +class _Chronos2(torch.nn.Module): + def __init__( + self, + channels: int = 768, + hidden_channels: int = 3072, + num_layers: int = 12, + num_heads: int = 12, + patch_size: int = 16, + context_length: int = 8192, + max_output_patches: int = 64, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + self.patch_size = patch_size + self.context_length = context_length + self.max_output_patches = max_output_patches + self.norm = RevIN(eps=0) + self.patch_embedding = ResidualBlock( + in_channels=3 * patch_size, + hidden_channels=hidden_channels, + out_channels=channels, + **factory_kwargs, + ) + self.reg_token = torch.nn.Parameter( + torch.zeros(1, channels, **factory_kwargs) + ) + self.icl_block = ICLBlock( + channels=channels, + num_layers=num_layers, + num_heads=num_heads, + hidden_channels=hidden_channels, + **factory_kwargs, + ) + self.head = ResidualBlock( + in_channels=channels, + hidden_channels=hidden_channels, + out_channels=len(QUANTILES) * patch_size, + **factory_kwargs, + ) + + def forward( + self, + context: Tensor, # [S, T] + future: Tensor, # [S, H], NaNs mark unknown future values. + group_ids: Tensor, # [S] + prediction_length: int, + ) -> Tensor: # [S, Q, H] + outputs = [] + remaining = prediction_length + patches = min( + (remaining + self.patch_size - 1) // self.patch_size, + self.max_output_patches, + ) + prediction = self._predict(context, future, group_ids, patches) + outputs.append(prediction) + remaining -= prediction.size(-1) + if remaining > 0: + paths = 9 + # Rollout feedback and refitted statistics use the linear scale. + prediction = prediction.sinh() + context, future = context.sinh(), future.sinh() + context = context[:, None, :].expand(-1, paths, -1) + future = future[:, None, :].expand(-1, paths, -1) + group_ids = group_ids[:, None] * paths + torch.arange( + paths, device=group_ids.device + ) + while remaining > 0: + # Default rollout levels 0.1, ..., 0.9 are every second knot. + feedback = prediction[:, 2:19:2] + context = torch.cat([context, feedback], dim=-1) + context = context[..., -self.context_length :] + known = future[..., : prediction.size(-1)] + context[..., -known.size(-1) :] = torch.where( + known.isnan(), context[..., -known.size(-1) :], known + ) + future = future[..., prediction.size(-1) :] + patches = min( + (remaining + self.patch_size - 1) // self.patch_size, + self.max_output_patches, + ) + # Refit only after extending the recipe-normalized context. + x = context.flatten(0, 1) + mean = x.nanmean(-1, keepdim=True).nan_to_num(nan=0.0) + scale = (x - mean).square().nanmean(-1, keepdim=True).sqrt() + scale = scale.nan_to_num(nan=1.0) + scale = scale.where(scale != 0, 1e-5) + samples = self._predict( + self.norm(x, mean, scale).asinh(), + self.norm(future.flatten(0, 1), mean, scale).asinh(), + group_ids.flatten(), + patches, + ) + samples = self.norm.inverse( + samples.sinh(), mean.unsqueeze(-1), scale.unsqueeze(-1) + ) + samples = samples.unflatten(0, (-1, 9)).flatten(1, 2) + prediction = rollout_quantiles(samples, QUANTILES) + outputs.append(prediction.asinh()) + remaining -= prediction.size(-1) + return torch.cat(outputs, dim=-1)[..., :prediction_length] + + def _predict( + self, + context: Tensor, + future: Tensor, + group_ids: Tensor, + num_output_patches: int, + ) -> Tensor: + context = context[..., -self.context_length :].float() + x = context + observed = ~context.isnan() + padding = -context.size(-1) % self.patch_size + x = F.pad(x, (padding, 0)).unflatten(-1, (-1, self.patch_size)) + observed = F.pad(observed, (padding, 0)) + observed = observed.unflatten(-1, (-1, self.patch_size)) + x = x.masked_fill(~observed, 0) + length = num_output_patches * self.patch_size + future = future[..., :length].float() + future = F.pad(future, (0, length - future.size(-1)), value=torch.nan) + future_mask = ~future.isnan() + future = future.masked_fill(~future_mask, 0) + future = future.unflatten(-1, (num_output_patches, self.patch_size)) + future_mask = future_mask.unflatten( + -1, (num_output_patches, self.patch_size) + ) + context_length = x.size(-2) * self.patch_size + time = ( + torch.arange( + -context_length, length, device=x.device, dtype=torch.float32 + ) + / 8192 + ) + time = time.unflatten(-1, (-1, self.patch_size)) + values = torch.cat([x, future], dim=-2).to(self.reg_token.dtype) + mask = torch.cat([observed, future_mask], dim=-2) + time = time.to(values.dtype).expand_as(values) + x = self.patch_embedding( + torch.cat([time, values, mask.to(values.dtype)], dim=-1) + ) + first = observed.size(-2) + x = torch.cat( + [ + x[:, :first], + self.reg_token.expand(x.size(0), 1, -1), + x[:, first:], + ], + dim=-2, + ) + attention_mask = torch.cat( + [ + observed.any(-1), + observed.new_ones((x.size(0), 1 + num_output_patches)), + ], + dim=-1, + ) + x = self.icl_block(x, attention_mask, group_ids) + out = self.head(x[:, -num_output_patches:]) + out = out.unflatten(-1, (len(QUANTILES), self.patch_size)) + return out.transpose(-3, -2).flatten(-2).float() diff --git a/sdm/models/chronos2/recipe.py b/sdm/models/chronos2/recipe.py new file mode 100644 index 000000000..305a850b5 --- /dev/null +++ b/sdm/models/chronos2/recipe.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sdm.processing as sp + + +def default_recipe() -> sp.Recipe: # noqa: D103 + return sp.Recipe( + features=[ + sp.Standardize(), + sp.Asinh(), + ], + target=[ + sp.Standardize(), + sp.Asinh(), + ], + output=sp.AverageEstimators(), + ) diff --git a/sdm/models/chronos2/util.py b/sdm/models/chronos2/util.py new file mode 100644 index 000000000..80124570a --- /dev/null +++ b/sdm/models/chronos2/util.py @@ -0,0 +1,45 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from collections.abc import Sequence + +import torch +from torch import Tensor + + +def rollout_quantiles(samples: Tensor, levels: Sequence[float]) -> Tensor: + """Combine the nine forecast paths using upstream's weighted empirical CDF. + + Args: + samples: Predictions with shape ``[S, 9 * Q, H]``, where ``S`` is + series, ``Q`` is quantiles and ``H`` is forecast steps. + levels: The model's quantile levels. + """ + quantiles = torch.tensor(levels, device=samples.device) + paths = torch.arange(1, 10, device=samples.device).float() / 10 + + def mass(q: Tensor) -> Tensor: + boundaries = torch.cat([q.new_zeros(1), q, q.new_ones(1)]) + weights = (boundaries[2:] - boundaries[:-2]) / 2 + return weights / weights.sum() + + weights = (mass(paths)[:, None] * mass(quantiles)[None, :]).flatten() + weights = weights / weights.sum() + values, index = samples.transpose(-1, -2).sort(-1) + cdf = weights.expand_as(values).gather(-1, index).cumsum(-1).clamp(0, 1) + # Upstream sorts the CDF too; ties can change sample order in float32. + cdf, index = cdf.sort(-1) + values = values.gather(-1, index) + # The first CDF entry is positive. All model levels are below the final + # CDF entry, so only the lower endpoint is needed for interpolation. + cdf = torch.cat([cdf.new_zeros((*cdf.shape[:-1], 1)), cdf], dim=-1) + values = torch.cat([values[..., :1], values], dim=-1) + query = quantiles.expand(*cdf.shape[:-1], -1).contiguous() + upper = torch.searchsorted(cdf.contiguous(), query, right=True) + upper = upper.clamp(max=cdf.size(-1) - 1) + lower = upper - 1 + lo, hi = cdf.gather(-1, lower), cdf.gather(-1, upper) + weight = ((query - lo) / (hi - lo)).nan_to_num(nan=0.0) + lo, hi = values.gather(-1, lower), values.gather(-1, upper) + return (lo + weight * (hi - lo)).transpose(-1, -2) diff --git a/sdm/processing/__init__.py b/sdm/processing/__init__.py index 91230c732..3d86788b5 100644 --- a/sdm/processing/__init__.py +++ b/sdm/processing/__init__.py @@ -35,6 +35,7 @@ QuantileTransform, RankGaussian, Standardize, + Asinh, RobustScale, MinMaxScale, FlipSign, @@ -90,6 +91,7 @@ "QuantileTransform", "RankGaussian", "Standardize", + "Asinh", "RobustScale", "MinMaxScale", "FlipSign", diff --git a/sdm/processing/numerical/__init__.py b/sdm/processing/numerical/__init__.py index 673cfff32..2970a0da5 100644 --- a/sdm/processing/numerical/__init__.py +++ b/sdm/processing/numerical/__init__.py @@ -12,6 +12,7 @@ from sdm.processing.numerical.quantile import QuantileTransform from sdm.processing.numerical.rank_gaussian import RankGaussian from sdm.processing.numerical.standardize import Standardize +from sdm.processing.numerical.asinh import Asinh from sdm.processing.numerical.robust_scale import RobustScale from sdm.processing.numerical.min_max_scale import MinMaxScale from sdm.processing.numerical.flip_sign import FlipSign @@ -32,6 +33,7 @@ "QuantileTransform", "RankGaussian", "Standardize", + "Asinh", "RobustScale", "MinMaxScale", "FlipSign", diff --git a/sdm/processing/numerical/asinh.py b/sdm/processing/numerical/asinh.py new file mode 100644 index 000000000..3a6545c0b --- /dev/null +++ b/sdm/processing/numerical/asinh.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from sdm import Stype, TableTensor +from sdm.processing import InvertibleMixin, Processor + + +class Asinh(Processor, InvertibleMixin): + """Apply the inverse hyperbolic sine to numerical values. + + Inverse transformation applies the hyperbolic sine. + """ + + handles_stypes = (Stype.numerical,) + requires_fit = False + + def _transform(self, table: TableTensor) -> TableTensor: + return table.replace_blocks(numerical=table.numerical.asinh()) + + def _inverse_transform(self, table: TableTensor) -> TableTensor: + return table.replace_blocks(numerical=table.numerical.sinh()) diff --git a/sdm/processing/numerical/standardize.py b/sdm/processing/numerical/standardize.py index a03c7b380..8be1e7873 100644 --- a/sdm/processing/numerical/standardize.py +++ b/sdm/processing/numerical/standardize.py @@ -21,6 +21,9 @@ class Standardize(Processor, InvertibleMixin): invertible. NaN and infinite values are ignored when fitting statistics and preserved during the transform. + Inverse transformation supports multiple consecutive outputs per fitted + column, such as quantile predictions. + Args: eps: Value added to each fitted standard deviation. """ @@ -84,7 +87,16 @@ def _transform(self, table: TableTensor) -> TableTensor: def _inverse_transform(self, table: TableTensor) -> TableTensor: dtype = table.numerical.dtype - numerical = (table.numerical * self.scale).add_(self.mean) + numerical = table.numerical + if self.scale.ndim and numerical.size(-1) != self.scale.size(-1): + numerical = numerical.unflatten(-1, (self.scale.size(-1), -1)) + numerical = ( + (numerical * self.scale.unsqueeze(-1)) + .add_(self.mean.unsqueeze(-1)) + .flatten(-2) + ) + else: + numerical = (numerical * self.scale).add_(self.mean) return table.replace_blocks(numerical=numerical.to(dtype=dtype)) def __repr__(self, *, indent: int = 0) -> str: