diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fd9e860ad..c9e36ccd67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -334,6 +334,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Adds `CrossUnet` photovoltaic-power forecaster + (`physicsnemo.experimental.models.pv_power.CrossUnet`), a hierarchical + patch-Transformer with channel-correlation conditioning. Includes a Hydra + training example with synthetic data + (`examples/weather/pv_power_cross_unet`). +- Adds real-data Cross-Unet paper-reproduction utilities, paper target + comparisons, and English/Chinese usage documentation for the PV-power + example. - Adds GLOBE model (`physicsnemo.experimental.models.globe.model.GLOBE`), including new variant that uses a dual tree traversal algorithm to reduce the complexity of the kernel evaluations from O(N^2) to O(N). diff --git a/examples/weather/pv_power_cross_unet/CROSS_UNET_USAGE.md b/examples/weather/pv_power_cross_unet/CROSS_UNET_USAGE.md new file mode 100644 index 0000000000..de08049981 --- /dev/null +++ b/examples/weather/pv_power_cross_unet/CROSS_UNET_USAGE.md @@ -0,0 +1,134 @@ + + + +# Cross-Unet Real PV Data Usage + +This guide describes how to use Cross-Unet for photovoltaic power prediction. +The workflow reads one CSV file at a 15-minute cadence and uses user-specified +column names for time, power, and weather inputs. + +## 1. Environment + +Install the example requirements in a PhysicsNeMo environment with PyTorch +available: + +```bash +cd /path/to/physicsnemo +pip install -r examples/weather/pv_power_cross_unet/requirements.txt +``` + +## 2. Data Preparation + +The input data should be in a CSV file: + +- one timestamp column +- one historical power column +- one or more weather columns +- regular 15-minute timestamps with no duplicate or missing times + +Set the column names in `conf/real_data.yaml` or through Hydra overrides: + +```yaml +data_file: ./dataset/your_data.csv +time_col: Time +target_col: pv_power +weather_cols: + - irradiation +freq_minutes: 15 +``` + +The name of the data columns should be in accordance with that in the CSV file. + +## 3. Training + +From the example directory: + +```bash +cd examples/weather/pv_power_cross_unet +python train_real_cross_unet.py mode=train +``` + +A quick smoke run: + +```bash +python train_real_cross_unet.py \ + mode=train \ + data_file=./dataset/your_data.csv \ + time_col=Time \ + target_col=pv_power \ + weather_cols='[irradiation]' \ + horizon_label=4h \ + max_epochs=1 \ + max_train_samples=4 \ + max_valid_samples=2 \ + batch_size=2 \ + batch_size_valid=2 \ + d_model=32 \ + d_ff=64 \ + output_dir=./outputs/smoke +``` + +The default model settings follow the Cross-Unet paper: + +```text +d_model=256, d_ff=512, n_heads=4, e_layers=3, +start_lr=1e-4, max_epochs=100, seed=2021, +nonlinear_correlation_proj=True, use_bottleneck_in_decoder=True +``` + +`early_stop_patience` defaults to `5` and can be overridden. + +## 4. Horizons + +Use `horizon_label` for common PV forecasting windows: + +```text +4h -> pred_len=16, seq_len=96, seg_len=12 +1d -> pred_len=96, seq_len=96, seg_len=24 +7d -> pred_len=672, seq_len=672, seg_len=48 +``` + +You can also set `seq_len`, `pred_len`, and `seg_len` directly. + +## 5. Prediction + +After training, run: + +```bash +python train_real_cross_unet.py \ + mode=predict \ + data_file=./dataset/your_data.csv \ + horizon_label=4h \ + output_dir=./outputs/smoke +``` + +The prediction CSV contains: + +```text +data_file,timestamp,horizon_step,prediction,target +``` + +Prediction and target values are saved on the original target scale. + +## 6. Evaluation + +`mode=evaluate` runs prediction and appends one row to the metrics CSV: + +```bash +python train_real_cross_unet.py \ + mode=evaluate \ + data_file=./dataset/your_data.csv \ + horizon_label=4h \ + output_dir=./outputs/smoke \ + metrics_csv_path=./outputs/smoke/metrics.csv +``` + +The metrics CSV contains: + +```text +data_file,horizon_label,seq_len,pred_len,weather_cols,weather_channels, +mae,mse,rmse,r2,best_valid_loss,epoch,checkpoint_path,prediction_path +``` + +MAE, MSE, RMSE, and R2 are computed over all forecast steps after flattening +the prediction and target arrays. diff --git a/examples/weather/pv_power_cross_unet/CROSS_UNET_USAGE.zh.md b/examples/weather/pv_power_cross_unet/CROSS_UNET_USAGE.zh.md new file mode 100644 index 0000000000..97a96fc5a2 --- /dev/null +++ b/examples/weather/pv_power_cross_unet/CROSS_UNET_USAGE.zh.md @@ -0,0 +1,131 @@ + + + +# Cross-Unet 真实光伏数据使用说明 + +本文档说明使用Cross-Unet完成光伏功率预测的工作流。该工作流读取一个 15 分钟频率的 CSV 文件，用户通过 +配置指定时间列、功率目标列和天气输入列。 + +## 1. 环境准备 + +在已安装 PhysicsNeMo 和 PyTorch 的环境中安装示例依赖： + +```bash +cd /path/to/physicsnemo +pip install -r examples/weather/pv_power_cross_unet/requirements.txt +``` + +## 2. 数据准备 + +输入数据应为CSV格式，要求如下： + +- 一列时间戳 +- 一列光伏功率历史数据 +- 至少一列天气数据 +- 时间戳严格按 15 分钟间隔排列，不能有重复时间或缺失 + +可以在 `conf/real_data.yaml` 中修改列名，也可以用 Hydra override： + +```yaml +data_file: ./dataset/your_data.csv +time_col: Time +target_col: pv_power +weather_cols: + - irradiation +freq_minutes: 15 +``` + +数据列名称应与CSV文件中的列名称一致。 + +## 3. 模型训练 + +从示例目录运行： + +```bash +cd examples/weather/pv_power_cross_unet +python train_real_cross_unet.py mode=train +``` + +快速 smoke run： + +```bash +python train_real_cross_unet.py \ + mode=train \ + data_file=./dataset/your_data.csv \ + time_col=Time \ + target_col=pv_power \ + weather_cols='[irradiation]' \ + horizon_label=4h \ + max_epochs=1 \ + max_train_samples=4 \ + max_valid_samples=2 \ + batch_size=2 \ + batch_size_valid=2 \ + d_model=32 \ + d_ff=64 \ + output_dir=./outputs/smoke +``` + +默认模型配置为Cross_Unet论文中的原始配置： + +```text +d_model=256, d_ff=512, n_heads=4, e_layers=3, +start_lr=1e-4, max_epochs=100, seed=2021, +nonlinear_correlation_proj=True, use_bottleneck_in_decoder=True +``` + +`early_stop_patience` 默认是 `5`，可以在命令行或配置文件中修改。 + +## 4. 预测窗口 + +常用预测窗口可通过 `horizon_label` 设置： + +```text +4h -> pred_len=16, seq_len=96, seg_len=12 +1d -> pred_len=96, seq_len=96, seg_len=24 +7d -> pred_len=672, seq_len=672, seg_len=48 +``` + +也可以直接设置 `seq_len`、`pred_len` 和 `seg_len`。 + +## 5. 模型推理 + +训练完成后执行： + +```bash +python train_real_cross_unet.py \ + mode=predict \ + data_file=./dataset/your_data.csv \ + horizon_label=4h \ + output_dir=./outputs/smoke +``` + +预测 CSV 字段为： + +```text +data_file,timestamp,horizon_step,prediction,target +``` + +`prediction` 和 `target` 都是反归一化后的原始目标尺度，便于检查。 + +## 6. 结果评估 + +`mode=evaluate` 会先执行预测，然后向 metrics CSV 追加一行： + +```bash +python train_real_cross_unet.py \ + mode=evaluate \ + data_file=./dataset/your_data.csv \ + horizon_label=4h \ + output_dir=./outputs/smoke \ + metrics_csv_path=./outputs/smoke/metrics.csv +``` + +metrics CSV 字段为： + +```text +data_file,horizon_label,seq_len,pred_len,weather_cols,weather_channels, +mae,mse,rmse,r2,best_valid_loss,epoch,checkpoint_path,prediction_path +``` + +MAE、MSE、RMSE、R2 会在所有预测步展平后统一计算。 diff --git a/examples/weather/pv_power_cross_unet/README.md b/examples/weather/pv_power_cross_unet/README.md new file mode 100644 index 0000000000..4f68809c77 --- /dev/null +++ b/examples/weather/pv_power_cross_unet/README.md @@ -0,0 +1,114 @@ +# Cross_Unet PV-Power Forecasting + +This example contains two CrossUnet workflows: + +- `train_cross_unet.py`: a self-contained synthetic-data training example. +- `train_real_cross_unet.py`: a generic real CSV workflow for 15-minute PV + power forecasting. + +The model is +[`CrossUnet`](../../../physicsnemo/experimental/models/pv_power/cross_unet.py), +which consumes five tensors per sample: + +```text +x_enc, w_enc, seq_w_nwp_hist, seq_x_hist, target +``` + +`x_enc` and `seq_x_hist` carry historical target power. `w_enc` and +`seq_w_nwp_hist` carry weather inputs. Training loss uses the primary power +target channel. + +## Synthetic Example + +```bash +pip install -r requirements.txt +python train_cross_unet.py +``` + +Hydra overrides can change any field: + +```bash +python train_cross_unet.py max_epochs=5 batch_size=16 d_model=128 +``` + +## Real CSV Workflow + +The real-data script reads one CSV file. Configure the file and columns in +`conf/real_data.yaml` or on the command line: + +```yaml +data_file: ./dataset/your_data.csv +time_col: Time +target_col: pv_power +weather_cols: + - irradiation +freq_minutes: 15 +``` + +CSV requirements: + +- regular 15-minute timestamps +- no duplicate or missing timestamps +- one historical power column +- at least one weather column + +Train with the default config: + +```bash +python train_real_cross_unet.py mode=train +``` + +Run prediction: + +```bash +python train_real_cross_unet.py mode=predict +``` + +Run evaluation and append metrics: + +```bash +python train_real_cross_unet.py mode=evaluate +``` + +Prediction CSV columns: + +```text +data_file,timestamp,horizon_step,prediction,target +``` + +Metrics CSV columns: + +```text +data_file,horizon_label,seq_len,pred_len,weather_cols,weather_channels, +mae,mse,rmse,r2,best_valid_loss,epoch,checkpoint_path,prediction_path +``` + +## Common Horizons + +`horizon_label` can set the standard PV forecasting windows: + +```text +4h -> pred_len=16, seq_len=96, seg_len=12 +1d -> pred_len=96, seq_len=96, seg_len=24 +7d -> pred_len=672, seq_len=672, seg_len=48 +``` + +The same values can be set manually with `seq_len`, `pred_len`, and `seg_len`. + +## Default Real-Data Model Settings + +The real-data config follows the Cross-Unet paper defaults used for +actual PV runs: + +```text +d_model=256, d_ff=512, n_heads=4, e_layers=3, +start_lr=1e-4, max_epochs=100, seed=2021, +nonlinear_correlation_proj=True, use_bottleneck_in_decoder=True +``` + +`early_stop_patience` defaults to `5`. + +## More Documentation + +- English: [`CROSS_UNET_USAGE.md`](CROSS_UNET_USAGE.md) +- Chinese: [`CROSS_UNET_USAGE.zh.md`](CROSS_UNET_USAGE.zh.md) diff --git a/examples/weather/pv_power_cross_unet/conf/config.yaml b/examples/weather/pv_power_cross_unet/conf/config.yaml new file mode 100644 index 0000000000..f2aecc37bb --- /dev/null +++ b/examples/weather/pv_power_cross_unet/conf/config.yaml @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +hydra: + job: + chdir: True + run: + dir: ./outputs/ + +# Data +seed: 0 +num_train_samples: 256 +num_valid_samples: 32 +seq_len: 96 # 24 h at 15-min cadence +pred_len: 16 # 4 h forecast horizon +target_channels: 4 # 3 history channels + 1 current power channel +weather_channels: 3 # irradiance, temperature proxy, humidity proxy + +# Model +seg_len: 12 +e_layers: 3 +d_model: 64 +n_heads: 4 +d_ff: 128 +dropout: 0.0 +nonlinear_correlation_proj: True +attention_kind: two_stage +merge_kind: seg_merge +use_bottleneck_in_decoder: True + +# Optimization +batch_size: 8 +batch_size_valid: 8 +max_epochs: 2 +start_lr: 1.0e-3 +lr_scheduler_gamma: 0.99 +checkpoint_save_freq: 1 diff --git a/examples/weather/pv_power_cross_unet/conf/real_data.yaml b/examples/weather/pv_power_cross_unet/conf/real_data.yaml new file mode 100644 index 0000000000..f5fbbc709b --- /dev/null +++ b/examples/weather/pv_power_cross_unet/conf/real_data.yaml @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +hydra: + job: + chdir: False + run: + dir: . + +# train, predict, or evaluate +mode: train + +# Data +data_file: ./dataset/your_data.csv +time_col: Time +target_col: power +weather_cols: + - SWR +freq_minutes: 15 +seq_len: 96 +pred_len: 16 +horizon_label: null # null, 4h, 1d, or 7d; labels override seq_len/pred_len/seg_len +normalization: standard # standard or minmax +max_train_samples: null +max_valid_samples: null +max_predict_samples: null + +# Model +seg_len: 12 +e_layers: 3 +d_model: 256 +n_heads: 4 +d_ff: 512 +dropout: 0.0 +nonlinear_correlation_proj: True +attention_kind: two_stage +merge_kind: seg_merge +use_bottleneck_in_decoder: True + +# Optimization +seed: 2021 +batch_size: 16 +batch_size_valid: 16 +max_epochs: 100 +early_stop_patience: 5 +start_lr: 1.0e-4 +lr_scheduler_gamma: 0.99 + +# Outputs +output_dir: ./outputs/real_cross_unet +checkpoint_path: null +prediction_path: null +metrics_csv_path: ./outputs/real_cross_unet/metrics.csv diff --git a/examples/weather/pv_power_cross_unet/requirements.txt b/examples/weather/pv_power_cross_unet/requirements.txt new file mode 100644 index 0000000000..b4e6b161da --- /dev/null +++ b/examples/weather/pv_power_cross_unet/requirements.txt @@ -0,0 +1,3 @@ +hydra-core>=1.2.0 +pandas>=2.0.0 +termcolor>=2.1.1 diff --git a/examples/weather/pv_power_cross_unet/train_cross_unet.py b/examples/weather/pv_power_cross_unet/train_cross_unet.py new file mode 100644 index 0000000000..452ff528fb --- /dev/null +++ b/examples/weather/pv_power_cross_unet/train_cross_unet.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Train the experimental :class:`CrossUnet` PV-power forecaster on synthetic data. + +This example fabricates a deterministic multivariate time series that mimics +photovoltaic power forecasting (a diurnal irradiance signal plus a noisy +power channel and two lagged history channels). It demonstrates the +PhysicsNeMo training-loop primitives -- ``DistributedManager``, +``LaunchLogger``, ``save_checkpoint``/``load_checkpoint`` -- against the +multi-input ``CrossUnet`` forward signature without requiring any external +dataset download. +""" + +from __future__ import annotations + +import math + +import hydra +import torch +import torch.nn.functional as F +from omegaconf import DictConfig +from torch.utils.data import DataLoader, Dataset + +from physicsnemo.distributed import DistributedManager +from physicsnemo.experimental.models.pv_power import CrossUnet +from physicsnemo.utils import load_checkpoint, save_checkpoint +from physicsnemo.utils.logging import LaunchLogger, PythonLogger + + +def _build_long_series( + *, + n_steps: int, + target_channels: int, + weather_channels: int, + generator: torch.Generator, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate one long deterministic series the dataset slices into windows.""" + t = torch.arange(n_steps, dtype=torch.float32) + + # Diurnal irradiance (96 = 24 h at 15-min cadence). + irradiance = torch.clamp(torch.sin(2 * math.pi * t / 96.0), min=0.0) + # Slow modulation of conversion efficiency. + efficiency = 0.7 + 0.3 * torch.sin(t / 97.0) + # Noisy primary power signal in [0, 1]. + noise = 0.05 * torch.randn(n_steps, generator=generator) + power = torch.clamp(irradiance * efficiency + noise, min=0.0, max=1.0) + + # Weather feature pool (truncated/padded to ``weather_channels``). + temp_proxy = 0.5 + 0.5 * torch.cos(2 * math.pi * t / 96.0) + humidity_proxy = 0.5 + 0.3 * torch.sin(2 * math.pi * t / 97.0) + weather_pool = torch.stack([irradiance, temp_proxy, humidity_proxy], dim=-1) + if weather_channels <= weather_pool.shape[-1]: + weather = weather_pool[:, :weather_channels] + else: + pad = torch.zeros(n_steps, weather_channels - weather_pool.shape[-1]) + weather = torch.cat([weather_pool, pad], dim=-1) + + # Target features = (target_channels - 1) lagged copies of power + power itself. + lagged = [ + torch.roll(power, shifts=lag).unsqueeze(-1) for lag in range(1, target_channels) + ] + targets = torch.cat([*lagged, power.unsqueeze(-1)], dim=-1) + return ( + targets, + weather, + ) # shapes: (n_steps, target_channels), (n_steps, weather_channels) + + +class SyntheticPVDataset(Dataset): + """Sliding-window dataset over a deterministic synthetic series.""" + + def __init__( + self, + *, + num_samples: int, + seq_len: int, + pred_len: int, + target_channels: int, + weather_channels: int, + seed: int, + ) -> None: + # Generate enough data for ``num_samples`` non-overlapping windows. + n_steps = num_samples * (seq_len + pred_len) + generator = torch.Generator(device="cpu").manual_seed(seed) + self.targets, self.weather = _build_long_series( + n_steps=n_steps, + target_channels=target_channels, + weather_channels=weather_channels, + generator=generator, + ) + self.seq_len = seq_len + self.pred_len = pred_len + self.num_samples = num_samples + self.target_channels = target_channels + self.weather_channels = weather_channels + + def __len__(self) -> int: + return self.num_samples + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + start = idx * (self.seq_len + self.pred_len) + window_end = start + self.seq_len + target_end = window_end + self.pred_len + + x_enc = self.targets[start:window_end] # (L, target_channels) + w_enc = self.weather[start:window_end] # (L, weather_channels) + # "Historical weather" reuses the encoder weather window for this + # synthetic example; in real deployments this would be a separate + # measurement stream. + seq_w_nwp_hist = w_enc.clone() + seq_x_hist = x_enc.clone() # (L, target_channels) + target = self.targets[window_end:target_end] # (H, target_channels) + return { + "x_enc": x_enc, + "w_enc": w_enc, + "seq_w_nwp_hist": seq_w_nwp_hist, + "seq_x_hist": seq_x_hist, + "target": target, + } + + +def _move(batch: dict[str, torch.Tensor], device) -> dict[str, torch.Tensor]: + return {k: v.to(device) for k, v in batch.items()} + + +@hydra.main(version_base="1.2", config_path="conf", config_name="config") +def main(cfg: DictConfig) -> None: + """Hydra entrypoint: train Cross-Unet on synthetic PV data.""" + DistributedManager.initialize() + dist = DistributedManager() + LaunchLogger.initialize() + logger = PythonLogger("cross_unet") + + torch.manual_seed(cfg.seed) + + train_ds = SyntheticPVDataset( + num_samples=cfg.num_train_samples, + seq_len=cfg.seq_len, + pred_len=cfg.pred_len, + target_channels=cfg.target_channels, + weather_channels=cfg.weather_channels, + seed=cfg.seed, + ) + valid_ds = SyntheticPVDataset( + num_samples=cfg.num_valid_samples, + seq_len=cfg.seq_len, + pred_len=cfg.pred_len, + target_channels=cfg.target_channels, + weather_channels=cfg.weather_channels, + seed=cfg.seed + 1, + ) + train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True) + valid_loader = DataLoader(valid_ds, batch_size=cfg.batch_size_valid, shuffle=False) + + model = CrossUnet( + target_channels=cfg.target_channels, + weather_channels=cfg.weather_channels, + seq_len=cfg.seq_len, + pred_len=cfg.pred_len, + seg_len=cfg.seg_len, + e_layers=cfg.e_layers, + d_model=cfg.d_model, + n_heads=cfg.n_heads, + d_ff=cfg.d_ff, + dropout=cfg.dropout, + nonlinear_correlation_proj=cfg.nonlinear_correlation_proj, + attention_kind=cfg.attention_kind, + merge_kind=cfg.merge_kind, + use_bottleneck_in_decoder=cfg.use_bottleneck_in_decoder, + ).to(dist.device) + + optimizer = torch.optim.Adam(model.parameters(), lr=cfg.start_lr) + scheduler = torch.optim.lr_scheduler.ExponentialLR( + optimizer, gamma=cfg.lr_scheduler_gamma + ) + + loaded_epoch = load_checkpoint( + "./checkpoints", + models=model, + optimizer=optimizer, + scheduler=scheduler, + device=dist.device, + ) + + target_slice = slice(-cfg.target_channels, None) + + for epoch in range(max(1, loaded_epoch + 1), cfg.max_epochs + 1): + model.train() + with LaunchLogger( + "train", + epoch=epoch, + num_mini_batch=len(train_loader), + epoch_alert_freq=1, + ) as log: + for batch in train_loader: + batch = _move(batch, dist.device) + optimizer.zero_grad() + pred = model( + batch["x_enc"], + batch["w_enc"], + batch["seq_w_nwp_hist"], + batch["seq_x_hist"], + ) + loss = F.mse_loss(pred[..., target_slice], batch["target"]) + loss.backward() + optimizer.step() + log.log_minibatch({"loss": loss.detach()}) + scheduler.step() + log.log_epoch({"learning_rate": optimizer.param_groups[0]["lr"]}) + + # Validation + model.eval() + valid_losses = [] + with torch.no_grad(): + for batch in valid_loader: + batch = _move(batch, dist.device) + pred = model( + batch["x_enc"], + batch["w_enc"], + batch["seq_w_nwp_hist"], + batch["seq_x_hist"], + ) + valid_losses.append( + F.mse_loss(pred[..., target_slice], batch["target"]).item() + ) + with LaunchLogger("valid", epoch=epoch) as log: + log.log_epoch({"loss": sum(valid_losses) / max(1, len(valid_losses))}) + + if epoch % cfg.checkpoint_save_freq == 0: + save_checkpoint( + "./checkpoints", + models=model, + optimizer=optimizer, + scheduler=scheduler, + epoch=epoch, + ) + + logger.info("Training complete.") + + +if __name__ == "__main__": + main() diff --git a/examples/weather/pv_power_cross_unet/train_real_cross_unet.py b/examples/weather/pv_power_cross_unet/train_real_cross_unet.py new file mode 100644 index 0000000000..2cf7df7016 --- /dev/null +++ b/examples/weather/pv_power_cross_unet/train_real_cross_unet.py @@ -0,0 +1,705 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Train, predict, and evaluate CrossUnet on user-provided PV-power CSV data.""" + +from __future__ import annotations + +import csv +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Literal + +import hydra +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F +from hydra.utils import to_absolute_path +from omegaconf import DictConfig, OmegaConf +from torch import Tensor +from torch.utils.data import DataLoader, Dataset + +from physicsnemo.experimental.models.pv_power import CrossUnet +from physicsnemo.utils.logging import PythonLogger + +Split = Literal["train", "valid", "test"] +Normalization = Literal["standard", "minmax"] + +HORIZON_CONFIGS: dict[str, dict[str, int]] = { + "4h": {"pred_len": 16, "seq_len": 96, "seg_len": 12}, + "1d": {"pred_len": 96, "seq_len": 96, "seg_len": 24}, + "7d": {"pred_len": 672, "seq_len": 672, "seg_len": 48}, +} + +METRIC_COLUMNS = [ + "data_file", + "horizon_label", + "seq_len", + "pred_len", + "weather_cols", + "weather_channels", + "mae", + "mse", + "rmse", + "r2", + "best_valid_loss", + "epoch", + "checkpoint_path", + "prediction_path", +] + + +@dataclass +class RealPVDataConfig: + """Configuration for one generic 15-minute PV-power CSV file.""" + + data_file: str | Path + time_col: str + target_col: str + weather_cols: list[str] + freq_minutes: int = 15 + seq_len: int = 96 + pred_len: int = 16 + normalization: Normalization = "standard" + max_samples: int | None = None + + +@dataclass +class ArrayScaler: + """Small serializable scaler for numpy arrays.""" + + mean: np.ndarray + std: np.ndarray + kind: Normalization = "standard" + + @classmethod + def fit(cls, values: np.ndarray, kind: Normalization = "standard") -> "ArrayScaler": + """Fit scaler statistics from a 2-D array of shape (N, C).""" + if kind == "minmax": + mean = values.min(axis=0) + std = values.max(axis=0) - mean + else: + mean = values.mean(axis=0) + std = values.std(axis=0) + std = np.where(std < 1.0e-6, 1.0, std) + return cls( + mean=mean.astype(np.float32), + std=std.astype(np.float32), + kind=kind, + ) + + @classmethod + def from_state(cls, state: dict[str, Any]) -> "ArrayScaler": + """Restore an ArrayScaler from a state dict produced by state_dict().""" + return cls( + mean=np.asarray(state["mean"], dtype=np.float32), + std=np.asarray(state["std"], dtype=np.float32), + kind=state.get("kind", "standard"), + ) + + def transform(self, values: np.ndarray) -> np.ndarray: + """Normalize values using fitted mean and std.""" + return ((values - self.mean) / self.std).astype(np.float32) + + def inverse_transform(self, values: np.ndarray) -> np.ndarray: + """Invert normalization for all channels.""" + return (values * self.std + self.mean).astype(np.float32) + + def inverse_last_channel(self, values: np.ndarray) -> np.ndarray: + """Invert normalization for the last channel only (the power target).""" + return values * self.std[-1] + self.mean[-1] + + def state_dict(self) -> dict[str, Any]: + """Return serializable scaler state for checkpoint saving.""" + return { + "mean": self.mean.tolist(), + "std": self.std.tolist(), + "kind": self.kind, + } + + +@dataclass +class EarlyStopping: + """Track validation loss and stop after stale epochs exceed patience.""" + + patience: int = 5 + min_delta: float = 0.0 + best_loss: float = float("inf") + stale_epochs: int = 0 + + def __post_init__(self) -> None: + if self.patience < 0: + raise ValueError(f"patience must be non-negative, got {self.patience}.") + + def step(self, loss: float) -> tuple[bool, bool]: + """Update state and return ``(improved, should_stop)``.""" + if loss < self.best_loss - self.min_delta: + self.best_loss = loss + self.stale_epochs = 0 + return True, False + + self.stale_epochs += 1 + return False, self.stale_epochs >= self.patience + + +def horizon_config(label: str) -> dict[str, int]: + """Return the default CrossUnet window settings for a named horizon.""" + if label not in HORIZON_CONFIGS: + raise ValueError( + f"Unsupported horizon_label={label!r}; expected one of " + f"{sorted(HORIZON_CONFIGS)}." + ) + return dict(HORIZON_CONFIGS[label]) + + +def _split_ranges(n_rows: int) -> dict[Split, tuple[int, int]]: + num_train = int(n_rows * 0.8) + num_test = int(n_rows * 0.1) + num_valid = n_rows - num_train - num_test + return { + "train": (0, num_train), + "valid": (num_train, num_train + num_valid), + "test": (num_train + num_valid, num_train + num_valid + num_test), + } + + +def _load_generic_csv(cfg: RealPVDataConfig) -> pd.DataFrame: + data_file = Path(cfg.data_file) + if not data_file.is_file(): + raise FileNotFoundError(f"data_file not found: {data_file}") + if not cfg.weather_cols: + raise ValueError("weather_cols must contain at least one column.") + + required_cols = [cfg.time_col, cfg.target_col, *cfg.weather_cols] + frame = pd.read_csv(data_file) + missing = [col for col in required_cols if col not in frame.columns] + if missing: + raise ValueError(f"{data_file} is missing required columns: {missing}") + + frame = frame[required_cols].copy() + frame[cfg.time_col] = pd.to_datetime(frame[cfg.time_col]) + numeric_cols = [cfg.target_col, *cfg.weather_cols] + frame[numeric_cols] = frame[numeric_cols].apply(pd.to_numeric, errors="coerce") + frame = frame.dropna(subset=[cfg.time_col, *numeric_cols]).sort_values(cfg.time_col) + frame = frame.reset_index(drop=True) + + duplicated = frame[cfg.time_col].duplicated() + if bool(duplicated.any()): + raise ValueError(f"{data_file} has duplicate timestamps in {cfg.time_col!r}.") + + expected_delta = pd.Timedelta(minutes=cfg.freq_minutes) + deltas = frame[cfg.time_col].diff().dropna() + if not deltas.empty and bool((deltas != expected_delta).any()): + raise ValueError( + f"{data_file} expected {cfg.freq_minutes}-minute cadence in " + f"{cfg.time_col!r}." + ) + return frame + + +def compute_prediction_metrics( + predictions: np.ndarray, + targets: np.ndarray, +) -> dict[str, float]: + """Compute flattened MAE, MSE, RMSE, and R2.""" + pred = np.asarray(predictions, dtype=np.float64).reshape(-1) + target = np.asarray(targets, dtype=np.float64).reshape(-1) + error = pred - target + mae = float(np.mean(np.abs(error))) + mse = float(np.mean(error**2)) + target_mean = float(np.mean(target)) + ss_res = float(np.sum(error**2)) + ss_tot = float(np.sum((target - target_mean) ** 2)) + r2 = 0.0 if ss_tot <= 1.0e-12 else 1.0 - ss_res / ss_tot + return {"mae": mae, "mse": mse, "rmse": float(np.sqrt(mse)), "r2": float(r2)} + + +class RealPVDataset(Dataset): + """Sliding-window dataset for generic 15-minute PV-power CSV files.""" + + def __init__( + self, + cfg: RealPVDataConfig, + split: Split, + target_scaler: ArrayScaler | None = None, + weather_scaler: ArrayScaler | None = None, + ) -> None: + self.cfg = cfg + self.split = split + self.target_cols = [cfg.target_col] + self.weather_cols = list(cfg.weather_cols) + self.target_channels = 1 + self.weather_channels = len(cfg.weather_cols) + + frame = _load_generic_csv(cfg) + if len(frame) < 2 * cfg.seq_len + max(cfg.seq_len, cfg.pred_len): + raise ValueError( + f"{cfg.data_file} has {len(frame)} usable rows; at least " + f"{2 * cfg.seq_len + max(cfg.seq_len, cfg.pred_len)} are required." + ) + + ranges = _split_ranges(len(frame)) + for split_name in ("valid", "test"): + s, e = ranges[split_name] + if e - s < cfg.pred_len: + raise ValueError( + f"{cfg.data_file}: the {split_name} split has only {e - s} rows " + f"but pred_len={cfg.pred_len} requires at least {cfg.pred_len}. " + f"Use a larger dataset or a shorter pred_len." + ) + train_start, train_end = ranges["train"] + split_start, split_end = ranges[split] + self.split_target_start_time = ( + frame[cfg.time_col].iloc[split_start].to_datetime64() + ) + + target_values = frame[[cfg.target_col]].to_numpy(dtype=np.float32) + weather_values = frame[cfg.weather_cols].to_numpy(dtype=np.float32) + self.target_scaler = target_scaler or ArrayScaler.fit( + target_values[train_start:train_end], kind=cfg.normalization + ) + self.weather_scaler = weather_scaler or ArrayScaler.fit( + weather_values[train_start:train_end], kind=cfg.normalization + ) + + self.times = frame[cfg.time_col].to_numpy() + self.target_data = self.target_scaler.transform(target_values) + self.weather_data = self.weather_scaler.transform(weather_values) + self.raw_target = target_values + first_start = max(0, split_start - 2 * cfg.seq_len) + last_by_split = split_end - cfg.pred_len - 2 * cfg.seq_len + last_by_frame = len(frame) - max(cfg.seq_len, cfg.pred_len) - 2 * cfg.seq_len + last_start = min(last_by_split, last_by_frame) + self.window_starts = list(range(first_start, last_start + 1)) + if cfg.max_samples is not None: + self.window_starts = self.window_starts[: cfg.max_samples] + if not self.window_starts: + raise ValueError( + f"Split {split!r} has no usable windows for seq_len={cfg.seq_len}, " + f"pred_len={cfg.pred_len}." + ) + + def __len__(self) -> int: + return len(self.window_starts) + + def __getitem__(self, index: int) -> dict[str, Tensor]: + start = self.window_starts[index] + seq_len = self.cfg.seq_len + pred_len = self.cfg.pred_len + hist_begin = start + enc_begin = start + seq_len + target_begin = enc_begin + seq_len + weather_end = target_begin + seq_len + target_end = target_begin + pred_len + + return { + "x_enc": torch.from_numpy(self.target_data[enc_begin:target_begin]), + "w_enc": torch.from_numpy(self.weather_data[target_begin:weather_end]), + "seq_w_nwp_hist": torch.from_numpy( + self.weather_data[enc_begin:target_begin] + ), + "seq_x_hist": torch.from_numpy(self.target_data[hist_begin:enc_begin]), + "target": torch.from_numpy(self.target_data[target_begin:target_end]), + } + + def target_times(self, index: int) -> list[np.datetime64]: + """Return the forecast timestamps for sample at index.""" + start = self.window_starts[index] + 2 * self.cfg.seq_len + end = start + self.cfg.pred_len + return list(self.times[start:end]) + + def raw_primary_target(self, index: int) -> np.ndarray: + """Return un-normalized power values for the forecast window at index.""" + start = self.window_starts[index] + 2 * self.cfg.seq_len + end = start + self.cfg.pred_len + return self.raw_target[start:end] + + +def _move(batch: dict[str, Tensor], device: torch.device) -> dict[str, Tensor]: + return {key: value.to(device) for key, value in batch.items()} + + +def _make_model(cfg: DictConfig, dataset: RealPVDataset) -> CrossUnet: + return CrossUnet( + target_channels=dataset.target_channels, + weather_channels=dataset.weather_channels, + seq_len=cfg.seq_len, + pred_len=cfg.pred_len, + seg_len=cfg.seg_len, + e_layers=cfg.e_layers, + d_model=cfg.d_model, + n_heads=cfg.n_heads, + d_ff=cfg.d_ff, + dropout=cfg.dropout, + nonlinear_correlation_proj=cfg.nonlinear_correlation_proj, + attention_kind=cfg.attention_kind, + merge_kind=cfg.merge_kind, + use_bottleneck_in_decoder=cfg.use_bottleneck_in_decoder, + ) + + +def _data_cfg_from_hydra(cfg: DictConfig, max_samples: int | None) -> RealPVDataConfig: + return RealPVDataConfig( + data_file=Path(to_absolute_path(str(cfg.data_file))), + time_col=cfg.time_col, + target_col=cfg.target_col, + weather_cols=list(cfg.weather_cols), + freq_minutes=cfg.freq_minutes, + seq_len=cfg.seq_len, + pred_len=cfg.pred_len, + normalization=cfg.normalization, + max_samples=max_samples, + ) + + +def _data_config_state(cfg: RealPVDataConfig) -> dict[str, Any]: + state = asdict(cfg) + state["data_file"] = str(state["data_file"]) + return state + + +def _checkpoint_name(data_file: str | Path, horizon_label: str | None) -> str: + stem = Path(data_file).stem + if horizon_label: + return f"real_cross_unet_{stem}_{horizon_label}.pt" + return f"real_cross_unet_{stem}.pt" + + +def _output_stem(data_file: str | Path, horizon_label: str | None) -> str: + stem = Path(data_file).stem + if horizon_label: + return f"{stem}_{horizon_label}" + return stem + + +def train(cfg: DictConfig) -> Path: + """Train CrossUnet on one real CSV and save the best validation checkpoint.""" + logger = PythonLogger("real_cross_unet") + torch.manual_seed(cfg.seed) + np.random.seed(cfg.seed) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + train_ds = RealPVDataset( + _data_cfg_from_hydra(cfg, cfg.max_train_samples), split="train" + ) + valid_ds = RealPVDataset( + _data_cfg_from_hydra(cfg, cfg.max_valid_samples), + split="valid", + target_scaler=train_ds.target_scaler, + weather_scaler=train_ds.weather_scaler, + ) + train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True) + valid_loader = DataLoader(valid_ds, batch_size=cfg.batch_size_valid, shuffle=False) + + model = _make_model(cfg, train_ds).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=cfg.start_lr) + scheduler = torch.optim.lr_scheduler.ExponentialLR( + optimizer, gamma=cfg.lr_scheduler_gamma + ) + + early_stopping = EarlyStopping(patience=cfg.early_stop_patience) + output_dir = Path(to_absolute_path(str(cfg.output_dir))) + checkpoint_dir = ( + output_dir / _output_stem(cfg.data_file, cfg.horizon_label) / "checkpoints" + ) + checkpoint_dir.mkdir(parents=True, exist_ok=True) + checkpoint_path = checkpoint_dir / _checkpoint_name( + cfg.data_file, cfg.horizon_label + ) + + last_epoch = 0 + for epoch in range(1, cfg.max_epochs + 1): + last_epoch = epoch + model.train() + train_losses = [] + for batch in train_loader: + batch = _move(batch, device) + optimizer.zero_grad() + pred = model( + batch["x_enc"], + batch["w_enc"], + batch["seq_w_nwp_hist"], + batch["seq_x_hist"], + ) + loss = F.mse_loss(pred[..., -1:], batch["target"][..., -1:]) + loss.backward() + optimizer.step() + train_losses.append(loss.item()) + scheduler.step() + + valid_loss = validation_loss(model, valid_loader, device) + logger.info( + f"epoch={epoch} train_loss={float(np.mean(train_losses)):.6f} " + f"valid_loss={valid_loss:.6f}" + ) + improved, should_stop = early_stopping.step(valid_loss) + if improved: + torch.save( + { + "model_state": model.state_dict(), + "data_config": _data_config_state(train_ds.cfg), + "model_config": { + key: OmegaConf.to_container(cfg)[key] + for key in [ + "seg_len", + "e_layers", + "d_model", + "n_heads", + "d_ff", + "dropout", + "nonlinear_correlation_proj", + "attention_kind", + "merge_kind", + "use_bottleneck_in_decoder", + ] + }, + "target_scaler": train_ds.target_scaler.state_dict(), + "weather_scaler": train_ds.weather_scaler.state_dict(), + "best_valid_loss": early_stopping.best_loss, + "epoch": epoch, + "last_epoch": last_epoch, + "early_stop_patience": cfg.early_stop_patience, + "horizon_label": cfg.horizon_label, + }, + checkpoint_path, + ) + if should_stop: + logger.info( + f"early stopping at epoch={epoch}; " + f"best_valid_loss={early_stopping.best_loss:.6f}" + ) + break + logger.info(f"saved checkpoint to {checkpoint_path}") + return checkpoint_path + + +@torch.no_grad() +def validation_loss( + model: CrossUnet, loader: DataLoader, device: torch.device +) -> float: + """Evaluate scaled primary-target MSE.""" + model.eval() + losses = [] + for batch in loader: + batch = _move(batch, device) + pred = model( + batch["x_enc"], + batch["w_enc"], + batch["seq_w_nwp_hist"], + batch["seq_x_hist"], + ) + losses.append(F.mse_loss(pred[..., -1:], batch["target"][..., -1:]).item()) + return float(np.mean(losses)) + + +def save_prediction_csv( + path: Path, + data_file: str, + target_times: list[list[np.datetime64]], + predictions: np.ndarray, + targets: np.ndarray, +) -> None: + """Save primary-target predictions in a simple long CSV format.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="") as fp: + writer = csv.writer(fp) + writer.writerow( + ["data_file", "timestamp", "horizon_step", "prediction", "target"] + ) + for sample_times, sample_pred, sample_target in zip( + target_times, predictions, targets + ): + for step, (timestamp, pred, target) in enumerate( + zip(sample_times, sample_pred[:, 0], sample_target[:, 0]) + ): + writer.writerow( + [ + data_file, + np.datetime_as_string(timestamp, unit="m"), + step, + float(pred), + float(target), + ] + ) + + +def _checkpoint_path_from_cfg(cfg: DictConfig) -> Path: + if cfg.checkpoint_path is not None: + return Path(to_absolute_path(str(cfg.checkpoint_path))) + return ( + Path(to_absolute_path(str(cfg.output_dir))) + / _output_stem(cfg.data_file, cfg.horizon_label) + / "checkpoints" + / _checkpoint_name(cfg.data_file, cfg.horizon_label) + ) + + +def _prediction_path_from_cfg(cfg: DictConfig, data_file: str | Path) -> Path: + if cfg.prediction_path is not None: + return Path(to_absolute_path(str(cfg.prediction_path))) + return ( + Path(to_absolute_path(str(cfg.output_dir))) + / _output_stem(data_file, cfg.horizon_label) + / "predictions" + / f"{_output_stem(data_file, cfg.horizon_label)}_predictions.csv" + ) + + +@torch.no_grad() +def _run_prediction(cfg: DictConfig) -> tuple[Path, dict[str, float], dict[str, Any]]: + checkpoint_path = _checkpoint_path_from_cfg(cfg) + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + + saved_data_cfg = checkpoint["data_config"] + # Prefer the data_file from the current config so users can redirect to a + # new path; fall back to the checkpoint's saved path if the override does + # not resolve to an existing file (e.g. the placeholder default). + data_file_override = Path(to_absolute_path(str(cfg.data_file))) + data_file = ( + data_file_override + if data_file_override.is_file() + else Path(to_absolute_path(str(saved_data_cfg["data_file"]))) + ) + data_cfg = RealPVDataConfig( + data_file=data_file, + time_col=saved_data_cfg["time_col"], + target_col=saved_data_cfg["target_col"], + weather_cols=list(saved_data_cfg["weather_cols"]), + freq_minutes=saved_data_cfg.get("freq_minutes", 15), + seq_len=saved_data_cfg["seq_len"], + pred_len=saved_data_cfg["pred_len"], + normalization=saved_data_cfg.get("normalization", "standard"), + max_samples=cfg.max_predict_samples, + ) + target_scaler = ArrayScaler.from_state(checkpoint["target_scaler"]) + weather_scaler = ArrayScaler.from_state(checkpoint["weather_scaler"]) + test_ds = RealPVDataset( + data_cfg, + split="test", + target_scaler=target_scaler, + weather_scaler=weather_scaler, + ) + + merged_cfg = OmegaConf.create({**checkpoint["model_config"], **asdict(data_cfg)}) + model = _make_model(merged_cfg, test_ds) + model.load_state_dict(checkpoint["model_state"]) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device).eval() + + loader = DataLoader(test_ds, batch_size=cfg.batch_size_valid, shuffle=False) + predictions = [] + targets = [] + target_times = [] + sample_offset = 0 + for batch in loader: + batch_size = batch["x_enc"].shape[0] + batch = _move(batch, device) + pred = model( + batch["x_enc"], + batch["w_enc"], + batch["seq_w_nwp_hist"], + batch["seq_x_hist"], + ) + pred_primary = pred[..., -1:].detach().cpu().numpy() + target_primary = batch["target"][..., -1:].detach().cpu().numpy() + predictions.append(target_scaler.inverse_last_channel(pred_primary)) + targets.append(target_scaler.inverse_last_channel(target_primary)) + target_times.extend( + test_ds.target_times(idx) + for idx in range(sample_offset, sample_offset + batch_size) + ) + sample_offset += batch_size + + pred_array = np.concatenate(predictions, axis=0) + target_array = np.concatenate(targets, axis=0) + prediction_path = _prediction_path_from_cfg(cfg, data_cfg.data_file) + save_prediction_csv( + prediction_path, + data_file=Path(data_cfg.data_file).name, + target_times=target_times, + predictions=pred_array, + targets=target_array, + ) + metrics = compute_prediction_metrics(pred_array, target_array) + return prediction_path, metrics, checkpoint + + +def predict(cfg: DictConfig) -> Path: + """Load a checkpoint, run test-set inference, and save a prediction CSV.""" + logger = PythonLogger("real_cross_unet") + prediction_path, metrics, _ = _run_prediction(cfg) + logger.info(f"saved predictions to {prediction_path}") + logger.info( + f"test MAE={metrics['mae']:.6f} MSE={metrics['mse']:.6f} " + f"RMSE={metrics['rmse']:.6f} R2={metrics['r2']:.6f}" + ) + return prediction_path + + +def evaluate(cfg: DictConfig) -> Path: + """Run prediction and append one row to the metrics CSV.""" + logger = PythonLogger("real_cross_unet") + prediction_path, metrics, checkpoint = _run_prediction(cfg) + metrics_path = Path(to_absolute_path(str(cfg.metrics_csv_path))) + metrics_path.parent.mkdir(parents=True, exist_ok=True) + row = { + "data_file": Path(checkpoint["data_config"]["data_file"]).name, + "horizon_label": checkpoint.get("horizon_label") or cfg.horizon_label or "", + "seq_len": checkpoint["data_config"]["seq_len"], + "pred_len": checkpoint["data_config"]["pred_len"], + "weather_cols": "|".join(checkpoint["data_config"]["weather_cols"]), + "weather_channels": len(checkpoint["data_config"]["weather_cols"]), + "mae": metrics["mae"], + "mse": metrics["mse"], + "rmse": metrics["rmse"], + "r2": metrics["r2"], + "best_valid_loss": checkpoint.get("best_valid_loss", ""), + "epoch": checkpoint.get("epoch", ""), + "checkpoint_path": str(_checkpoint_path_from_cfg(cfg)), + "prediction_path": str(prediction_path), + } + write_header = not metrics_path.exists() + with metrics_path.open("a", newline="") as fp: + writer = csv.DictWriter(fp, fieldnames=METRIC_COLUMNS) + if write_header: + writer.writeheader() + writer.writerow(row) + logger.info(f"saved predictions to {prediction_path}") + logger.info(f"appended metrics to {metrics_path}") + return metrics_path + + +@hydra.main(version_base="1.2", config_path="conf", config_name="real_data") +def main(cfg: DictConfig) -> None: + """Hydra entrypoint.""" + if cfg.horizon_label: + for key, value in horizon_config(str(cfg.horizon_label)).items(): + cfg[key] = value + if cfg.mode == "train": + train(cfg) + elif cfg.mode == "predict": + predict(cfg) + elif cfg.mode == "evaluate": + evaluate(cfg) + else: + raise ValueError( + f"Unsupported mode={cfg.mode!r}; expected train, predict, or evaluate." + ) + + +if __name__ == "__main__": + main() diff --git a/physicsnemo/experimental/models/pv_power/__init__.py b/physicsnemo/experimental/models/pv_power/__init__.py new file mode 100644 index 0000000000..ae83bfe674 --- /dev/null +++ b/physicsnemo/experimental/models/pv_power/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Photovoltaic-power forecasting models (experimental).""" + +from .cross_unet import CrossUnet, CrossUnetMetaData + +__all__ = ["CrossUnet", "CrossUnetMetaData"] diff --git a/physicsnemo/experimental/models/pv_power/attention.py b/physicsnemo/experimental/models/pv_power/attention.py new file mode 100644 index 0000000000..d02caa374d --- /dev/null +++ b/physicsnemo/experimental/models/pv_power/attention.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Attention building blocks for the Cross_Unet PV-power model. + +The mask branch from upstream ``FullAttention`` is removed because every +``Cross_Unet`` instantiation uses ``mask_flag=False``; modules such as the +upstream ``space_attn`` / ``dim_sender`` / ``dim_receiver`` / ``router`` that are +allocated but never read in the forward pass are also dropped here for +clarity. +""" + +from __future__ import annotations + +from math import sqrt + +import torch +import torch.nn as nn +from einops import rearrange + + +class FullAttention(nn.Module): + r"""Scaled dot-product attention without masking. + + Parameters + ---------- + attention_dropout : float, optional, default=0.0 + Dropout probability applied to the attention weights. + """ + + def __init__(self, attention_dropout: float = 0.0) -> None: + super().__init__() + self.dropout = nn.Dropout(attention_dropout) + + def forward( + self, + queries: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + ) -> torch.Tensor: + _, _, _, e = queries.shape + scale = 1.0 / sqrt(e) + scores = torch.einsum("blhe,bshe->bhls", queries, keys) + attn = self.dropout(torch.softmax(scale * scores, dim=-1)) + out = torch.einsum("bhls,bshd->blhd", attn, values) + return out.contiguous() + + +class AttentionLayer(nn.Module): + r"""Multi-head attention layer wrapping :class:`FullAttention`. + + Parameters + ---------- + d_model : int + Hidden dimension of the input and output tokens. + n_heads : int + Number of attention heads. + attention_dropout : float, optional, default=0.0 + Dropout applied within the attention scores. + d_keys : int or None, optional, default=None + Per-head key dimension. Defaults to ``d_model // n_heads``. + d_values : int or None, optional, default=None + Per-head value dimension. Defaults to ``d_model // n_heads``. + """ + + def __init__( + self, + d_model: int, + n_heads: int, + attention_dropout: float = 0.0, + d_keys: int | None = None, + d_values: int | None = None, + ) -> None: + super().__init__() + d_keys = d_keys or (d_model // n_heads) + d_values = d_values or (d_model // n_heads) + + self.inner_attention = FullAttention(attention_dropout=attention_dropout) + self.query_projection = nn.Linear(d_model, d_keys * n_heads) + self.key_projection = nn.Linear(d_model, d_keys * n_heads) + self.value_projection = nn.Linear(d_model, d_values * n_heads) + self.out_projection = nn.Linear(d_values * n_heads, d_model) + self.n_heads = n_heads + + def forward( + self, + queries: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + ) -> torch.Tensor: + b, lq, _ = queries.shape + _, lk, _ = keys.shape + h = self.n_heads + q = self.query_projection(queries).view(b, lq, h, -1) + k = self.key_projection(keys).view(b, lk, h, -1) + v = self.value_projection(values).view(b, lk, h, -1) + out = self.inner_attention(q, k, v) + out = out.view(b, lq, -1) + return self.out_projection(out) + + +class TwoStageAttentionLayer(nn.Module): + r"""Two-stage attention block: time attention + correlation channel mixing. + + The block first applies multi-head self-attention along the patch axis of + each channel independently (the *time* stage). It then mixes channels with + a precomputed correlation matrix (the *channel* stage). The result is + fused with a residual connection, layer-normalised, and refined by a + feed-forward MLP. + + Parameters + ---------- + n_vars : int + Total number of channels :math:`C` flowing through the block. + d_model : int + Per-token hidden dimension. + n_heads : int + Number of attention heads. + d_ff : int or None, optional, default=None + Feed-forward inner width. Defaults to ``4 * d_model``. + dropout : float, optional, default=0.1 + Dropout probability for attention and MLP residuals. + swap_corr_axis : bool, optional, default=False + If True, multiply the channel-mixing correlation matrix on the right + of the channel-flattened activations (``[B, L, C] @ [B, C, C]``) + rather than the left. + """ + + def __init__( + self, + n_vars: int, + d_model: int, + n_heads: int, + d_ff: int | None = None, + dropout: float = 0.1, + swap_corr_axis: bool = False, + ) -> None: + super().__init__() + d_ff = d_ff or 4 * d_model + self.time_attention = AttentionLayer( + d_model=d_model, n_heads=n_heads, attention_dropout=dropout + ) + self.dropout = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.mlp = nn.Sequential( + nn.Linear(d_model, d_ff), + nn.GELU(), + nn.Linear(d_ff, d_model), + ) + self.swap_corr_axis = swap_corr_axis + self.n_vars = n_vars + + def _apply_correlation( + self, x_flat: torch.Tensor, corr: torch.Tensor + ) -> torch.Tensor: + # ``x_flat``: (B, n_vars, seg_num * d_model); ``corr``: (B, n_vars, n_vars). + if self.swap_corr_axis: + return torch.bmm(x_flat.permute(0, 2, 1), corr).permute(0, 2, 1) + return torch.bmm(corr, x_flat) + + def forward( + self, x: torch.Tensor, corr: torch.Tensor + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + # ``x``: (B, n_vars, seg_num, d_model). + b, n_vars, seg_num, d_model = x.shape + + # Time-axis attention applied independently per channel. + time_in = rearrange(x, "b ts_d seg_num d_model -> (b ts_d) seg_num d_model") + time_enc = self.time_attention(time_in, time_in, time_in) + time_reco = time_enc.view(b, n_vars, seg_num, d_model) + + # Channel-axis mixing using the supplied correlation matrix. + x_flat = rearrange(x, "b ts_d seg_num d_model -> b ts_d (seg_num d_model)") + corr_matrix = self._apply_correlation(x_flat, corr).view( + b, n_vars, seg_num, d_model + ) + attn_trace = [x_flat.detach(), corr.detach(), corr_matrix.detach()] + + # Residual fusion + post-norm + feed-forward block. + dim_in = self.norm1(x + self.dropout(time_reco + corr_matrix)) + dim_in = self.norm2(dim_in + self.dropout(self.mlp(dim_in))) + return dim_in, attn_trace + + +class ParallelTwoStageAttentionLayer(nn.Module): + r"""Parallel two-stage attention: simultaneous time and channel attention. + + Differs from :class:`TwoStageAttentionLayer` in that it replaces the + correlation-based channel mixer with a second self-attention pass over the + channel axis. The two stages are then fused by addition. + + Parameters + ---------- + n_vars : int + Total number of channels :math:`C`. + d_model : int + Per-token hidden dimension. + n_heads : int + Number of attention heads. + d_ff : int or None, optional, default=None + Feed-forward inner width. Defaults to ``4 * d_model``. + dropout : float, optional, default=0.1 + Dropout probability. + """ + + def __init__( + self, + n_vars: int, + d_model: int, + n_heads: int, + d_ff: int | None = None, + dropout: float = 0.1, + ) -> None: + super().__init__() + d_ff = d_ff or 4 * d_model + self.time_attention = AttentionLayer( + d_model=d_model, n_heads=n_heads, attention_dropout=dropout + ) + self.dropout = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.mlp = nn.Sequential( + nn.Linear(d_model, d_ff), + nn.GELU(), + nn.Linear(d_ff, d_model), + ) + self.n_vars = n_vars + + def forward( + self, x: torch.Tensor, corr: torch.Tensor + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + # ``x``: (B, n_vars, seg_num, d_model); ``corr`` accepted but unused. + del corr + b, n_vars, seg_num, d_model = x.shape + + # Time-axis attention. + time_in = rearrange(x, "b ts_d seg_num d_model -> (b ts_d) seg_num d_model") + time_enc = self.time_attention(time_in, time_in, time_in) + time_reco = time_enc.view(b, n_vars, seg_num, d_model) + + # Channel-axis attention reuses the same head weights. + x_perm = x.permute(0, 2, 1, 3).contiguous().view(b * seg_num, n_vars, d_model) + space_enc = self.time_attention(x_perm, x_perm, x_perm) + space_reco = space_enc.view(b, seg_num, n_vars, d_model).permute(0, 2, 1, 3) + attn_trace = [time_reco.detach(), space_reco.detach()] + + # Residual fusion + post-norm + feed-forward block. + dim_in = self.norm1(x + self.dropout(time_reco + space_reco)) + dim_in = self.norm2(dim_in + self.dropout(self.mlp(dim_in))) + return dim_in, attn_trace diff --git a/physicsnemo/experimental/models/pv_power/cross_unet.py b/physicsnemo/experimental/models/pv_power/cross_unet.py new file mode 100644 index 0000000000..9dbd008b93 --- /dev/null +++ b/physicsnemo/experimental/models/pv_power/cross_unet.py @@ -0,0 +1,495 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Cross_Unet photovoltaic-power forecasting model. + +Adapted from the upstream open-source PV-power benchmark suite (Apache-2.0, +2023 Yunhao Zhang & Junchi Yan). The PhysicsNeMo port refactors the upstream +``Model(configs)`` API to explicit typed kwargs and drops dead modules while +preserving the upstream channel-correlation construction: correlations are +computed from historical weather channels, the full historical target window, +and the current primary target channel, then reduced to the model's +:math:`C \times C` channel-mixing grid. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from math import ceil +from typing import Literal, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat +from jaxtyping import Float +from torch import Tensor +from torch.func import vmap + +from physicsnemo.core.meta import ModelMetaData +from physicsnemo.core.module import Module +from physicsnemo.experimental.models.pv_power.attention import ( + AttentionLayer, + ParallelTwoStageAttentionLayer, + TwoStageAttentionLayer, +) +from physicsnemo.experimental.models.pv_power.embedding import PatchEmbedding +from physicsnemo.experimental.models.pv_power.encoder_decoder import ( + Decoder, + DecoderLayer, + Encoder, + ScaleBlock, +) + + +@dataclass +class CrossUnetMetaData(ModelMetaData): + """Metadata for :class:`CrossUnet`.""" + + # Optimization + jit: bool = False # vmap + dynamic ceil-padding break TorchScript trace + cuda_graphs: bool = False # dynamic per-call seg padding + amp_cpu: bool = False + amp_gpu: bool = True + bf16: bool = True + torch_fx: bool = False + # Inference + onnx: bool = False + onnx_runtime: bool = False + # Physics informed + func_torch: bool = True # uses torch.func.vmap internally + auto_grad: bool = True + + +def _pearson_correlation(sample: Tensor) -> Tensor: + r"""Per-sample Pearson correlation of the channel axis of ``sample``. + + Parameters + ---------- + sample : torch.Tensor + Single-sample slice of shape :math:`(L, C)`. + + Returns + ------- + torch.Tensor + Channel-correlation matrix of shape :math:`(C, C)`. + """ + return torch.corrcoef(sample.T) + + +class CrossUnet(Module): + r"""Cross_Unet: hierarchical patch-Transformer for PV-power forecasting. + + The model embeds a multi-channel time series as patches, applies a + U-Net-style encoder/decoder of two-stage (time + channel) attention + blocks, and conditions the attention on a Pearson-correlation matrix + computed from the historical weather and target channels. It is designed + for short-to-medium-horizon photovoltaic power forecasting at 15-minute + cadence (the default ``seg_len=12`` corresponds to 3-hour patches), but + works for any tabular multivariate forecasting problem with a single + designated *target* channel. + + Channels are split into two semantic groups: + + * ``target_channels``: the target time series plus any history channels + that should be reconstructed. The **last** channel of ``x_enc`` is the + primary signal whose correlation with the other channels conditions + the attention. + * ``weather_channels``: exogenous weather forecast features that are + concatenated to ``x_enc`` before the patch embedding (set to ``0`` to + disable the weather branch entirely). + + Internally the network operates on ``total_channels = target_channels + + weather_channels`` channels. The channel-correlation matrix is computed + from one extra correlation input channel: the full historical target + window plus the current primary target channel. + + Adapted from the upstream `PV-power Cross_Unet + `_ (Apache-2.0). + + Parameters + ---------- + target_channels : int + Number of target/history channels :math:`C_{tgt}` (must be + :math:`\geq 1`). The last channel of ``x_enc`` is the primary signal + whose feature correlations condition the attention. + weather_channels : int, optional, default=0 + Number of exogenous weather channels :math:`C_{wx}` concatenated to + ``x_enc``. Set to ``0`` to disable the weather branch. + seq_len : int, optional, default=96 + Input window length :math:`L`. + pred_len : int, optional, default=16 + Forecast horizon :math:`H`. + seg_len : int, optional, default=12 + Patch length along the time axis. Defaults to ``12`` (= 3 h at + 15-min cadence). + e_layers : int, optional, default=3 + Number of encoder/decoder levels (the decoder has ``e_layers + 1`` + levels including the bottleneck-fed level). + d_model : int, optional, default=128 + Per-token hidden dimension :math:`D`. + n_heads : int, optional, default=4 + Number of attention heads. + d_ff : int, optional, default=256 + Feed-forward inner width inside attention layers. + dropout : float, optional, default=0.0 + Dropout probability shared across attention and MLP layers. + nonlinear_correlation_proj : bool, optional, default=False + If True, run the channel-correlation vector through a small learned + nonlinear projection before broadcasting it into a mixing matrix. + swap_corr_axis : bool, optional, default=False + If True, multiply the correlation matrix on the right of the + channel-flattened activations. See + :class:`~physicsnemo.experimental.models.pv_power.attention.TwoStageAttentionLayer`. + merge_kind : Literal["seg_merge", "cnn_merge"], optional, default="seg_merge" + Selects the segment-downsampling layer used between encoder levels. + attention_kind : Literal["two_stage", "parallel"], optional, default="two_stage" + Selects between correlation-conditioned attention (``"two_stage"``) + and pure self-attention along both axes (``"parallel"``). + use_bottleneck_in_decoder : bool, optional, default=True + If True, the encoder bottleneck output is fused into the deepest + decoder level (U-Net skip-connection variant). + + Forward + ------- + x_enc : torch.Tensor + Target window of shape :math:`(B, L, C_{tgt})`. The last channel + is treated as the primary forecast target. + w_enc : torch.Tensor or None + Weather forecast window of shape :math:`(B, L, C_{wx})`. May be + ``None`` only when ``weather_channels == 0``. + seq_w_nwp_hist : torch.Tensor or None + Historical weather signal of shape :math:`(B, L, C_{wx})` used to + compute the channel-correlation matrix. May be ``None`` only when + ``weather_channels == 0``. + seq_x_hist : torch.Tensor + Historical target/history channels of shape + :math:`(B, L, C_{tgt})`. Concatenated with ``seq_w_nwp_hist`` and + the last channel of ``x_enc`` to produce a + ``(B, L, C_{tgt} + C_{wx} + 1)`` matrix on which Pearson + correlations are evaluated. + + Outputs + ------- + torch.Tensor + Forecast tensor of shape :math:`(B, H, C_{tgt} + C_{wx})`. The last + ``target_channels`` columns correspond to the target channels. + + Examples + -------- + >>> import torch + >>> from physicsnemo.experimental.models.pv_power import CrossUnet + >>> model = CrossUnet( + ... target_channels=4, weather_channels=3, + ... seq_len=96, pred_len=16, seg_len=12, + ... e_layers=2, d_model=32, n_heads=4, d_ff=64, + ... ) + >>> x_enc = torch.randn(2, 96, 4) + >>> w_enc = torch.randn(2, 96, 3) + >>> hist_w = torch.randn(2, 96, 3) + >>> hist_x = torch.randn(2, 96, 4) + >>> out = model(x_enc, w_enc, hist_w, hist_x) + >>> out.shape + torch.Size([2, 16, 7]) + """ + + def __init__( + self, + target_channels: int, + weather_channels: int = 0, + seq_len: int = 96, + pred_len: int = 16, + seg_len: int = 12, + e_layers: int = 3, + d_model: int = 128, + n_heads: int = 4, + d_ff: int = 256, + dropout: float = 0.0, + nonlinear_correlation_proj: bool = False, + swap_corr_axis: bool = False, + merge_kind: Literal["seg_merge", "cnn_merge"] = "seg_merge", + attention_kind: Literal["two_stage", "parallel"] = "two_stage", + use_bottleneck_in_decoder: bool = True, + ) -> None: + if target_channels < 1: + raise ValueError(f"target_channels must be >= 1, got {target_channels}.") + if weather_channels < 0: + raise ValueError(f"weather_channels must be >= 0, got {weather_channels}.") + if seq_len <= 0 or pred_len <= 0 or seg_len <= 0: + raise ValueError( + f"seq_len/pred_len/seg_len must be positive, got " + f"({seq_len}, {pred_len}, {seg_len})." + ) + if e_layers < 1: + raise ValueError(f"e_layers must be >= 1, got {e_layers}.") + + super().__init__(meta=CrossUnetMetaData()) + + self.target_channels = target_channels + self.weather_channels = weather_channels + self.total_channels = target_channels + weather_channels + self.seq_len = seq_len + self.pred_len = pred_len + self.seg_len = seg_len + self.e_layers = e_layers + self.d_model = d_model + self.dropout_p = dropout + self.nonlinear_correlation_proj = nonlinear_correlation_proj + self.use_weather = weather_channels > 0 + + # Patch / segment grid sizes (mirrors upstream: pad to a multiple of seg_len). + win_size = 2 + self.win_size = win_size + pad_in_len = ceil(seq_len / seg_len) * seg_len + pad_out_len = ceil(pred_len / seg_len) * seg_len + in_seg_num = pad_in_len // seg_len + out_seg_num = ceil(in_seg_num / (win_size ** (e_layers - 1))) + dec_seg_num = pad_out_len // seg_len + # Decoder must be at least as long as the encoder bottleneck only when + # the bottleneck skip-add is enabled in ``Decoder.forward``. + if use_bottleneck_in_decoder and dec_seg_num < out_seg_num: + raise ValueError( + f"pred_len={pred_len} (decoder segments={dec_seg_num}) is shorter " + f"than the encoder bottleneck (segments={out_seg_num}) implied by " + f"seq_len={seq_len}, seg_len={seg_len}, e_layers={e_layers}. " + f"Either increase pred_len, decrease e_layers, or decrease seg_len." + ) + self.in_seg_num = in_seg_num + self.out_seg_num = out_seg_num + + # Embedding stack. + self.enc_value_embedding = PatchEmbedding( + d_model=d_model, + patch_len=seg_len, + stride=seg_len, + padding=pad_in_len - seq_len, + dropout=0.0, + ) + self.enc_pos_embedding = nn.Parameter( + torch.randn(1, self.total_channels, in_seg_num, d_model) + ) + self.dec_pos_embedding = nn.Parameter( + torch.randn(1, self.total_channels, pad_out_len // seg_len, d_model) + ) + self.pre_norm = nn.LayerNorm(d_model) + + # Encoder. + scale_blocks = [ + ScaleBlock( + win_size=1 if level == 0 else win_size, + n_vars=self.total_channels, + seg_num=in_seg_num + if level == 0 + else ceil(in_seg_num / win_size**level), + d_model=d_model, + n_heads=n_heads, + d_ff=d_ff, + dropout=dropout, + merge_kind=merge_kind, + attention_kind=attention_kind, + swap_corr_axis=swap_corr_axis, + ) + for level in range(e_layers) + ] + self.encoder = Encoder( + scale_blocks=scale_blocks, + d_model=d_model, + n_heads=n_heads, + d_ff=d_ff, + dropout=dropout, + ) + + # Decoder (one extra level: e_layers + 1 layers). + decoder_layers: list[DecoderLayer] = [] + for _ in range(e_layers + 1): + if attention_kind == "two_stage": + self_attn: nn.Module = TwoStageAttentionLayer( + n_vars=self.total_channels, + d_model=d_model, + n_heads=n_heads, + d_ff=d_ff, + dropout=dropout, + swap_corr_axis=swap_corr_axis, + ) + else: + self_attn = ParallelTwoStageAttentionLayer( + n_vars=self.total_channels, + d_model=d_model, + n_heads=n_heads, + d_ff=d_ff, + dropout=dropout, + ) + cross_attn = AttentionLayer( + d_model=d_model, + n_heads=n_heads, + attention_dropout=dropout, + ) + decoder_layers.append( + DecoderLayer( + self_attention=self_attn, + cross_attention=cross_attn, + seg_len=seg_len, + d_model=d_model, + dropout=dropout, + ) + ) + self.decoder = Decoder( + layers=decoder_layers, + use_bottleneck=use_bottleneck_in_decoder, + ) + + # Optional non-linear correlation projection. + if nonlinear_correlation_proj: + corr_dim = self.total_channels + self.channel_proj1 = nn.Sequential( + nn.Linear(corr_dim, corr_dim * 4, bias=True), + nn.Sigmoid(), + nn.Linear(corr_dim * 4, corr_dim, bias=True), + ) + self.channel_proj2 = nn.Sequential( + nn.Linear(corr_dim * corr_dim, corr_dim * corr_dim * 4, bias=True), + nn.Sigmoid(), + nn.Linear(corr_dim * corr_dim * 4, corr_dim * corr_dim, bias=True), + ) + + def _compute_channel_correlation(self, samples: Tensor) -> Tensor: + r"""Pearson correlations reduced to a source-style mixing matrix. + + Parameters + ---------- + samples : torch.Tensor + Tensor of shape :math:`(B, L, C + 1)` whose final channel is the + primary target signal used as the correlation reference. + + Returns + ------- + torch.Tensor + Mixing matrix of shape :math:`(B, C, C)` for the model channels. + """ + b, _, c = samples.shape + if c <= 1: + return torch.ones(b, 1, 1, dtype=samples.dtype, device=samples.device) + + # Pearson correlation matrix per sample, shape (B, C, C). + corr = vmap(_pearson_correlation)(samples) + # Pull "feature -> target" column (excluding the target's diagonal). + corr_to_target = corr[:, :-1, -1] # (B, C - 1) + corr_to_target = torch.nan_to_num( + corr_to_target, nan=0.0, posinf=0.0, neginf=0.0 + ) + corr_dim = c - 1 + + if self.nonlinear_correlation_proj: + # Project raw correlations, then mask negative projected values. + v = corr_to_target.unsqueeze(-1) # (B, C - 1, 1) + v = self.channel_proj1(v.permute(0, 2, 1)).permute(0, 2, 1) + v = v.repeat(1, 1, corr_dim) # (B, C - 1, C - 1) + flat = v.view(b, corr_dim * corr_dim) + flat = self.channel_proj2(flat) + mixing = flat.view(b, corr_dim, corr_dim).clamp(min=0.0) + else: + # Source-style rank-1 mixing: each feature-target softmax weight + # is broadcast across destination rows for ``corr @ x``. + corr_to_target = corr_to_target.clamp(min=0.0) + row = F.softmax(corr_to_target, dim=-1) # (B, C - 1) + mixing = row.unsqueeze(1).repeat(1, corr_dim, 1) + + row_sum = mixing.sum(dim=-1, keepdim=True) + uniform = torch.full_like(mixing, 1.0 / corr_dim) + return torch.where( + row_sum > 1.0e-12, mixing / row_sum.clamp_min(1.0e-12), uniform + ) + + def _forecast(self, x: Tensor, corr: Tensor) -> Tensor: + # ``x``: (B, L, C); ``corr``: (B, C, C). + x_emb, n_vars = self.enc_value_embedding(x.permute(0, 2, 1)) + x_emb = rearrange( + x_emb, "(b d) seg_num d_model -> b d seg_num d_model", d=n_vars + ) + x_emb = x_emb + self.enc_pos_embedding + x_emb = self.pre_norm(x_emb) + enc_out, _ = self.encoder(x_emb, corr) + dec_in = repeat( + self.dec_pos_embedding, + "b ts_d l d -> (repeat b) ts_d l d", + repeat=x_emb.shape[0], + ) + return self.decoder(dec_in, enc_out, corr) + + def forward( + self, + x_enc: Float[Tensor, "batch seq_len target_channels"], + w_enc: Optional[Float[Tensor, "batch seq_len weather_channels"]], + seq_w_nwp_hist: Optional[Float[Tensor, "batch seq_len weather_channels"]], + seq_x_hist: Float[Tensor, "batch seq_len target_channels"], + ) -> Float[Tensor, "batch pred_len total_channels"]: + # Validate input shapes (skip under torch.compile per MOD-005). + if not torch.compiler.is_compiling(): + self._validate_forward_inputs(x_enc, w_enc, seq_w_nwp_hist, seq_x_hist) + + # Build the correlation-conditioning input: weather history + target + # history + full target history + the primary target channel from the + # current window. + if self.use_weather and w_enc is not None and seq_w_nwp_hist is not None: + corr_input = torch.cat( + [seq_w_nwp_hist, seq_x_hist, x_enc[:, :, -1:]], dim=-1 + ) + x = torch.cat([w_enc, x_enc], dim=-1) + else: + corr_input = torch.cat([seq_x_hist, x_enc[:, :, -1:]], dim=-1) + x = x_enc + + corr = self._compute_channel_correlation(corr_input) + dec_out = self._forecast(x, corr) + return dec_out[:, : self.pred_len, :] + + def _validate_forward_inputs( + self, + x_enc: Tensor, + w_enc: Optional[Tensor], + seq_w_nwp_hist: Optional[Tensor], + seq_x_hist: Tensor, + ) -> None: + batch = x_enc.shape[0] + if x_enc.shape != (batch, self.seq_len, self.target_channels): + raise ValueError( + f"Expected x_enc of shape (B, {self.seq_len}, {self.target_channels}) " + f"but got tensor of shape {tuple(x_enc.shape)}." + ) + if seq_x_hist.shape != (batch, self.seq_len, self.target_channels): + raise ValueError( + f"Expected seq_x_hist of shape " + f"(B, {self.seq_len}, {self.target_channels}) " + f"but got tensor of shape {tuple(seq_x_hist.shape)}." + ) + if self.use_weather: + if w_enc is None or seq_w_nwp_hist is None: + raise ValueError( + "w_enc and seq_w_nwp_hist are required when " + f"weather_channels={self.weather_channels} > 0." + ) + if w_enc.shape != (batch, self.seq_len, self.weather_channels): + raise ValueError( + f"Expected w_enc of shape (B, {self.seq_len}, " + f"{self.weather_channels}) but got tensor of shape " + f"{tuple(w_enc.shape)}." + ) + if seq_w_nwp_hist.shape != (batch, self.seq_len, self.weather_channels): + raise ValueError( + f"Expected seq_w_nwp_hist of shape (B, {self.seq_len}, " + f"{self.weather_channels}) but got tensor of shape " + f"{tuple(seq_w_nwp_hist.shape)}." + ) diff --git a/physicsnemo/experimental/models/pv_power/embedding.py b/physicsnemo/experimental/models/pv_power/embedding.py new file mode 100644 index 0000000000..44b93b76d4 --- /dev/null +++ b/physicsnemo/experimental/models/pv_power/embedding.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Patch and positional embedding modules for the Cross_Unet PV-power model.""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + + +class PositionalEmbedding(nn.Module): + r"""Sinusoidal positional embedding with a fixed maximum length. + + Computes the standard transformer sin/cos positional encoding once at + construction and exposes the prefix matching the input length on each + forward call. + + Parameters + ---------- + d_model : int + Embedding dimension :math:`D`. + max_len : int, optional, default=5000 + Maximum sequence length supported by the cached table. + """ + + def __init__(self, d_model: int, max_len: int = 5000) -> None: + super().__init__() + pe = torch.zeros(max_len, d_model).float() + position = torch.arange(0, max_len).float().unsqueeze(1) + div_term = ( + torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model) + ).exp() + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]]) + self.register_buffer("pe", pe.unsqueeze(0)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.pe[:, : x.size(1)] + + +class PatchEmbedding(nn.Module): + r"""Per-channel patch embedding for time-series. + + Splits the input along the time axis into overlapping patches with the + given ``patch_len`` and ``stride`` (after replication padding), then + projects each patch to ``d_model`` and adds a sinusoidal positional + embedding. + + Parameters + ---------- + d_model : int + Output embedding dimension :math:`D`. + patch_len : int + Patch length along the time axis. + stride : int + Stride between consecutive patches. + padding : int + Trailing replication padding applied before patching. + dropout : float + Dropout probability applied to the embedded patches. + + Forward + ------- + x : torch.Tensor + Input tensor of shape :math:`(B, C, L)` with channel dimension first. + + Outputs + ------- + torch.Tensor + Patch embeddings of shape :math:`(B \cdot C, P, D)` where + :math:`P` is the number of patches. + int + The original number of channels :math:`C`, returned so that + downstream code can rearrange the merged batch-channel dimension. + """ + + def __init__( + self, + d_model: int, + patch_len: int, + stride: int, + padding: int, + dropout: float, + ) -> None: + super().__init__() + self.patch_len = patch_len + self.stride = stride + self.padding_patch_layer = nn.ReplicationPad1d((0, padding)) + self.value_embedding = nn.Linear(patch_len, d_model, bias=False) + self.position_embedding = PositionalEmbedding(d_model) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int]: + n_vars = x.shape[1] + x = self.padding_patch_layer(x) + x = x.unfold(dimension=-1, size=self.patch_len, step=self.stride) + x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3])) + x = self.value_embedding(x) + self.position_embedding(x) + return self.dropout(x), n_vars diff --git a/physicsnemo/experimental/models/pv_power/encoder_decoder.py b/physicsnemo/experimental/models/pv_power/encoder_decoder.py new file mode 100644 index 0000000000..1509abaaa6 --- /dev/null +++ b/physicsnemo/experimental/models/pv_power/encoder_decoder.py @@ -0,0 +1,397 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Encoder, decoder and merge utilities for the Cross_Unet PV-power model.""" + +from __future__ import annotations + +from typing import Literal + +import torch +import torch.nn as nn +from einops import rearrange + +from physicsnemo.experimental.models.pv_power.attention import ( + AttentionLayer, + ParallelTwoStageAttentionLayer, + TwoStageAttentionLayer, +) + + +class SegMerging(nn.Module): + r"""Token-wise segment merging that halves the patch axis (linear). + + Parameters + ---------- + d_model : int + Per-token hidden dimension :math:`D`. + win_size : int + Number of consecutive segments combined into one. + """ + + def __init__(self, d_model: int, win_size: int) -> None: + super().__init__() + self.d_model = d_model + self.win_size = win_size + self.linear_trans = nn.Linear(win_size * d_model, d_model) + self.norm = nn.LayerNorm(win_size * d_model) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # ``x``: (B, ts_d, seg_num, d_model). + _, _, seg_num, _ = x.shape + pad_num = seg_num % self.win_size + if pad_num != 0: + pad_num = self.win_size - pad_num + x = torch.cat((x, x[:, :, -pad_num:, :]), dim=-2) + seg_to_merge = [x[:, :, i :: self.win_size, :] for i in range(self.win_size)] + x = torch.cat(seg_to_merge, -1) + x = self.norm(x) + return self.linear_trans(x) + + +class CNNMerging(nn.Module): + r"""Convolutional segment merging using a 1D conv with stride ``win_size``. + + Parameters + ---------- + d_model : int + Per-token hidden dimension. + win_size : int + Conv stride (segments merged per output token). + """ + + def __init__(self, d_model: int, win_size: int) -> None: + super().__init__() + self.win_size = win_size + self.conv = nn.Conv1d( + in_channels=d_model, + out_channels=d_model, + kernel_size=3, + stride=win_size, + padding=1, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # ``x``: (B, ts_d, seg_num, d_model) -> (B, ts_d, seg_num // win, d_model). + batch_size, ts_d, seg_num, d_model = x.shape + x = x.permute(0, 1, 3, 2).reshape(batch_size * ts_d, d_model, seg_num) + x = self.conv(x) + new_seg_num = x.shape[-1] + return ( + x.reshape(batch_size, ts_d, d_model, new_seg_num) + .permute(0, 1, 3, 2) + .contiguous() + ) + + +class ScaleBlock(nn.Module): + r"""One U-Net level: optional segment merging followed by attention layers. + + Parameters + ---------- + win_size : int + Patch-axis downsampling factor for this level. Use ``1`` to skip the + merge layer entirely. + n_vars : int + Total channels :math:`C` flowing through the model. + seg_num : int + Number of patches at this scale (input to the attention layers). + d_model : int + Per-token hidden dimension. + n_heads : int + Number of attention heads. + d_ff : int + Feed-forward inner width inside attention layers. + dropout : float + Dropout probability. + depth : int, optional, default=1 + Number of attention layers stacked at this level. + merge_kind : Literal["seg_merge", "cnn_merge"], optional, default="seg_merge" + Selects the merge layer used to compress consecutive segments. + attention_kind : Literal["two_stage", "parallel"], optional, default="two_stage" + Selects between :class:`TwoStageAttentionLayer` (channel mixing via + correlation) and :class:`ParallelTwoStageAttentionLayer` (channel + mixing via self-attention). + swap_corr_axis : bool, optional, default=False + Forwarded to :class:`TwoStageAttentionLayer`. + """ + + def __init__( + self, + win_size: int, + n_vars: int, + seg_num: int, + d_model: int, + n_heads: int, + d_ff: int, + dropout: float, + depth: int = 1, + merge_kind: Literal["seg_merge", "cnn_merge"] = "seg_merge", + attention_kind: Literal["two_stage", "parallel"] = "two_stage", + swap_corr_axis: bool = False, + ) -> None: + super().__init__() + if win_size > 1: + if merge_kind == "seg_merge": + self.merge_layer: nn.Module | None = SegMerging(d_model, win_size) + elif merge_kind == "cnn_merge": + self.merge_layer = CNNMerging(d_model, win_size) + else: + raise ValueError( + f"Unknown merge_kind {merge_kind!r}; expected 'seg_merge' or 'cnn_merge'." + ) + else: + self.merge_layer = None + + self.encode_layers = nn.ModuleList() + for _ in range(depth): + if attention_kind == "two_stage": + self.encode_layers.append( + TwoStageAttentionLayer( + n_vars=n_vars, + d_model=d_model, + n_heads=n_heads, + d_ff=d_ff, + dropout=dropout, + swap_corr_axis=swap_corr_axis, + ) + ) + elif attention_kind == "parallel": + self.encode_layers.append( + ParallelTwoStageAttentionLayer( + n_vars=n_vars, + d_model=d_model, + n_heads=n_heads, + d_ff=d_ff, + dropout=dropout, + ) + ) + else: + raise ValueError( + f"Unknown attention_kind {attention_kind!r}; " + "expected 'two_stage' or 'parallel'." + ) + + def forward( + self, x: torch.Tensor, corr: torch.Tensor + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + if self.merge_layer is not None: + x = self.merge_layer(x) + atten_new: list[torch.Tensor] = [] + for layer in self.encode_layers: + x, atten_new = layer(x, corr) + return x, atten_new + + +class BottleneckLayer(nn.Module): + r"""Self-attention + feed-forward block applied at the encoder bottleneck. + + Parameters + ---------- + d_model : int + Per-token hidden dimension. + n_heads : int + Number of attention heads. + d_ff : int + Feed-forward inner width. + dropout : float + Dropout probability. + """ + + def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float) -> None: + super().__init__() + self.self_attn = nn.MultiheadAttention( + d_model, n_heads, dropout=dropout, batch_first=True + ) + self.feed_forward = nn.Sequential( + nn.Linear(d_model, d_ff), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(d_ff, d_model), + ) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # ``x``: (B, ts_d, seg_num, d_model). + batch_size, ts_d, seg_num, d_model = x.shape + x_flat = x.reshape(batch_size * ts_d, seg_num, d_model) + attn_output, _ = self.self_attn(x_flat, x_flat, x_flat) + x_flat = self.norm1(x_flat + self.dropout(attn_output)) + x_flat = self.norm2(x_flat + self.dropout(self.feed_forward(x_flat))) + return x_flat.reshape(batch_size, ts_d, seg_num, d_model) + + +class Encoder(nn.Module): + r"""Stack of :class:`ScaleBlock` levels followed by a bottleneck. + + Parameters + ---------- + scale_blocks : list of ScaleBlock + Encoder levels in coarse-to-fine order. + d_model : int + Per-token hidden dimension. + n_heads : int + Number of attention heads in the bottleneck layer. + d_ff : int + Feed-forward inner width in the bottleneck layer. + dropout : float + Dropout probability. + """ + + def __init__( + self, + scale_blocks: list[ScaleBlock], + d_model: int, + n_heads: int, + d_ff: int, + dropout: float, + ) -> None: + super().__init__() + self.encode_blocks = nn.ModuleList(scale_blocks) + self.bottleneck = BottleneckLayer( + d_model=d_model, n_heads=n_heads, d_ff=d_ff, dropout=dropout + ) + + def forward( + self, x: torch.Tensor, corr: torch.Tensor + ) -> tuple[list[torch.Tensor], list[list[torch.Tensor]]]: + # ``x``: (B, n_vars, seg_num, d_model). + encode_x: list[torch.Tensor] = [x] + atten_weights: list[list[torch.Tensor]] = [] + for block in self.encode_blocks: + x, attns = block(x, corr) + encode_x.append(x) + atten_weights.append(attns) + encode_x.append(self.bottleneck(x)) + return encode_x, atten_weights + + +class DecoderLayer(nn.Module): + r"""One decoder level: self-attention, encoder cross-attention, and prediction head. + + Parameters + ---------- + self_attention : nn.Module + Self-attention layer that consumes the decoder state and channel + correlation matrix (typically :class:`TwoStageAttentionLayer` or + :class:`ParallelTwoStageAttentionLayer`). + cross_attention : AttentionLayer + Cross-attention layer used to attend to one encoder level. + seg_len : int + Length of one output segment in the time domain (used to size the + per-segment linear prediction head). + d_model : int + Per-token hidden dimension. + dropout : float, optional, default=0.1 + Dropout probability. + """ + + def __init__( + self, + self_attention: nn.Module, + cross_attention: AttentionLayer, + seg_len: int, + d_model: int, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.self_attention = self_attention + self.cross_attention = cross_attention + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + self.mlp = nn.Sequential( + nn.Linear(d_model, d_model), + nn.GELU(), + nn.Linear(d_model, d_model), + ) + self.linear_pred = nn.Linear(d_model, seg_len) + + def forward( + self, x: torch.Tensor, cross: torch.Tensor, corr: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + # ``x``: (B, ts_d, out_seg_num, d_model); ``cross``: same as ``x`` from + # the encoder side (or the bottleneck output). + batch = x.shape[0] + x, _ = self.self_attention(x, corr) + x = rearrange(x, "b ts_d out_seg_num d_model -> (b ts_d) out_seg_num d_model") + cross = rearrange( + cross, "b ts_d in_seg_num d_model -> (b ts_d) in_seg_num d_model" + ) + cross_attn_out = self.cross_attention(x, cross, cross) + x = self.norm1(x + self.dropout(cross_attn_out)) + x = self.norm2(x + self.mlp(x)) + dec_output = rearrange( + x, + "(b ts_d) seg_dec_num d_model -> b ts_d seg_dec_num d_model", + b=batch, + ) + layer_predict = self.linear_pred(dec_output) + layer_predict = rearrange( + layer_predict, + "b out_d seg_num seg_len -> b (out_d seg_num) seg_len", + ) + return dec_output, layer_predict + + +class Decoder(nn.Module): + r"""Stack of :class:`DecoderLayer` blocks producing a multi-segment forecast. + + Parameters + ---------- + layers : list of DecoderLayer + Decoder levels in coarse-to-fine order. The first layer attends to + the encoder bottleneck when ``use_bottleneck`` is True. + use_bottleneck : bool, optional, default=True + If True, the encoder's bottleneck output is consumed by the deepest + decoder layer (and skip connections start from one level shallower). + """ + + def __init__(self, layers: list[DecoderLayer], use_bottleneck: bool = True) -> None: + super().__init__() + self.decode_layers = nn.ModuleList(layers) + self.use_bottleneck = use_bottleneck + + def forward( + self, x: torch.Tensor, cross_all: list[torch.Tensor], corr: torch.Tensor + ) -> torch.Tensor: + if self.use_bottleneck: + bottleneck_output = cross_all[-1] + cross = cross_all[:-1] + else: + bottleneck_output = None + cross = cross_all + + ts_d = x.shape[1] + layer_num = len(self.decode_layers) + # Initialise ``final_predict`` from the first layer; ``Decoder`` is built + # with at least one decoder layer (``e_layers + 1 >= 2``). + final_predict = torch.zeros(0, dtype=x.dtype, device=x.device) + for i, layer in enumerate(self.decode_layers): + if self.use_bottleneck and i == 0: + cross_in = bottleneck_output + x[:, :, : bottleneck_output.size(2), :] + else: + cross_in = cross[layer_num - i - 1] if self.use_bottleneck else cross[i] + x, layer_predict = layer(x, cross_in, corr) + final_predict = layer_predict if i == 0 else final_predict + layer_predict + + return rearrange( + final_predict, + "b (out_d seg_num) seg_len -> b (seg_num seg_len) out_d", + out_d=ts_d, + ) diff --git a/test/examples/weather/test_pv_power_cross_unet_real_data.py b/test/examples/weather/test_pv_power_cross_unet_real_data.py new file mode 100644 index 0000000000..51f318577b --- /dev/null +++ b/test/examples/weather/test_pv_power_cross_unet_real_data.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from omegaconf import OmegaConf + +from examples.weather.pv_power_cross_unet.train_real_cross_unet import ( + ArrayScaler, + EarlyStopping, + RealPVDataConfig, + RealPVDataset, + compute_prediction_metrics, + horizon_config, + save_prediction_csv, +) + +CONFIG_PATH = Path("examples/weather/pv_power_cross_unet/conf/real_data.yaml") + + +def _write_csv( + path: Path, + *, + rows: int = 80, + freq_minutes: int = 15, + time_col: str = "timestamp", + target_col: str = "power", + weather_cols: tuple[str, ...] = ("irradiance", "temperature"), +) -> pd.DataFrame: + times = pd.date_range("2020-01-01", periods=rows, freq=f"{freq_minutes}min") + frame = pd.DataFrame({time_col: times, target_col: np.arange(rows, dtype=float)}) + for offset, col in enumerate(weather_cols, start=1): + frame[col] = np.arange(rows, dtype=float) + offset * 100.0 + frame["ignored"] = "unused" + frame.to_csv(path, index=False) + return frame + + +def _dataset_cfg( + path: Path, + *, + seq_len: int = 8, + pred_len: int = 4, + weather_cols: tuple[str, ...] = ("irradiance", "temperature"), + max_samples: int | None = 3, +) -> RealPVDataConfig: + return RealPVDataConfig( + data_file=path, + time_col="timestamp", + target_col="power", + weather_cols=list(weather_cols), + freq_minutes=15, + seq_len=seq_len, + pred_len=pred_len, + max_samples=max_samples, + ) + + +def test_real_pv_dataset_accepts_arbitrary_column_names_and_ignores_extras(tmp_path): + data_file = tmp_path / "custom_station.csv" + _write_csv(data_file, rows=80) + + dataset = RealPVDataset(_dataset_cfg(data_file), split="train") + sample = dataset[0] + + assert dataset.target_channels == 1 + assert dataset.weather_channels == 2 + assert dataset.weather_cols == ["irradiance", "temperature"] + assert sample["x_enc"].shape == (8, 1) + assert sample["seq_x_hist"].shape == (8, 1) + assert sample["w_enc"].shape == (8, 2) + assert sample["seq_w_nwp_hist"].shape == (8, 2) + assert sample["target"].shape == (4, 1) + + +def test_real_pv_dataset_rejects_missing_required_columns(tmp_path): + data_file = tmp_path / "missing.csv" + _write_csv(data_file, rows=80, weather_cols=("irradiance",)) + + with pytest.raises(ValueError, match="missing required columns"): + RealPVDataset( + _dataset_cfg( + data_file, + weather_cols=("irradiance", "temperature"), + ), + split="train", + ) + + +def test_real_pv_dataset_rejects_duplicate_times(tmp_path): + data_file = tmp_path / "duplicate.csv" + frame = _write_csv(data_file, rows=80) + frame.loc[10, "timestamp"] = frame.loc[9, "timestamp"] + frame.to_csv(data_file, index=False) + + with pytest.raises(ValueError, match="duplicate timestamps"): + RealPVDataset(_dataset_cfg(data_file), split="train") + + +def test_real_pv_dataset_rejects_non_15_minute_cadence(tmp_path): + data_file = tmp_path / "hourly.csv" + _write_csv(data_file, rows=80, freq_minutes=60) + + with pytest.raises(ValueError, match="expected 15-minute cadence"): + RealPVDataset(_dataset_cfg(data_file), split="train") + + +def test_real_pv_dataset_uses_real_prior_history_without_target_leakage(tmp_path): + data_file = tmp_path / "ordered.csv" + _write_csv(data_file, rows=120) + + dataset = RealPVDataset(_dataset_cfg(data_file, max_samples=1), split="valid") + sample = dataset[0] + + assert not np.array_equal(sample["seq_x_hist"].numpy(), sample["x_enc"].numpy()) + target_start = dataset.target_times(0)[0] + assert target_start >= dataset.split_target_start_time + + +def test_compute_prediction_metrics_flattens_predictions_and_includes_rmse(): + pred = np.array([[[1.0], [3.0]], [[2.0], [4.0]]], dtype=np.float32) + target = np.array([[[1.0], [1.0]], [[3.0], [5.0]]], dtype=np.float32) + + metrics = compute_prediction_metrics(pred, target) + + assert metrics["mae"] == 1.0 + assert metrics["mse"] == 1.5 + assert np.isclose(metrics["rmse"], np.sqrt(1.5)) + assert np.isclose(metrics["r2"], 0.4545454545454546) + + +def test_horizon_config_matches_requested_windows(): + assert horizon_config("4h") == {"pred_len": 16, "seq_len": 96, "seg_len": 12} + assert horizon_config("1d") == {"pred_len": 96, "seq_len": 96, "seg_len": 24} + assert horizon_config("7d") == {"pred_len": 672, "seq_len": 672, "seg_len": 48} + + +def test_array_scaler_supports_minmax_normalization(): + values = np.array([[2.0, 10.0], [4.0, 14.0], [6.0, 18.0]], dtype=np.float32) + + scaler = ArrayScaler.fit(values, kind="minmax") + scaled = scaler.transform(values) + restored = scaler.inverse_transform(scaled) + + assert scaler.kind == "minmax" + assert np.allclose(scaled[0], [0.0, 0.0]) + assert np.allclose(scaled[-1], [1.0, 1.0]) + assert np.allclose(restored, values) + + +def test_early_stopping_respects_patience_and_resets_on_improvement(): + early_stop = EarlyStopping(patience=2) + + assert early_stop.step(1.0) == (True, False) + assert early_stop.step(1.1) == (False, False) + assert early_stop.step(0.9) == (True, False) + assert early_stop.step(0.95) == (False, False) + assert early_stop.step(0.96) == (False, True) + + +def test_real_data_config_uses_cross_unet_defaults_and_patience_five(): + cfg = OmegaConf.load(CONFIG_PATH) + + assert cfg.early_stop_patience == 5 + assert cfg.d_model == 256 + assert cfg.d_ff == 512 + assert cfg.n_heads == 4 + assert cfg.e_layers == 3 + assert cfg.start_lr == 1.0e-4 + assert cfg.max_epochs == 100 + assert cfg.seed == 2021 + assert cfg.nonlinear_correlation_proj is True + assert cfg.use_bottleneck_in_decoder is True + + +def test_save_prediction_csv_writes_generic_columns(tmp_path): + target_times = [ + [np.datetime64("2020-01-01T00:00"), np.datetime64("2020-01-01T00:15")] + ] + predictions = np.array([[[1.0], [2.0]]], dtype=np.float32) + targets = np.array([[[1.5], [2.5]]], dtype=np.float32) + out_path = tmp_path / "predictions.csv" + + save_prediction_csv( + out_path, + data_file="custom_station.csv", + target_times=target_times, + predictions=predictions, + targets=targets, + ) + + text = out_path.read_text() + assert "data_file,timestamp,horizon_step,prediction,target" in text + assert "custom_station.csv,2020-01-01T00:00,0,1.0,1.5" in text diff --git a/test/experimental/models/pv_power/cross_unet/data/cross_unet_default_output.pth b/test/experimental/models/pv_power/cross_unet/data/cross_unet_default_output.pth new file mode 100644 index 0000000000..1ad8e2d26c Binary files /dev/null and b/test/experimental/models/pv_power/cross_unet/data/cross_unet_default_output.pth differ diff --git a/test/experimental/models/pv_power/cross_unet/test_cross_unet.py b/test/experimental/models/pv_power/cross_unet/test_cross_unet.py new file mode 100644 index 0000000000..fe65255e27 --- /dev/null +++ b/test/experimental/models/pv_power/cross_unet/test_cross_unet.py @@ -0,0 +1,540 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# NOTE on regression-data generation: +# ``validate_forward_accuracy`` creates the reference ``.pth`` if missing and +# then errors. To regenerate references after an intentional model change, +# delete the corresponding ``data/cross_unet_*_output.pth`` and re-run the +# test twice (the first pass writes; the second passes). + +import random +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F +from omegaconf import OmegaConf + +from physicsnemo.core.module import Module +from physicsnemo.experimental.models.pv_power import CrossUnet +from physicsnemo.experimental.models.pv_power.embedding import PositionalEmbedding +from test.common import validate_checkpoint, validate_forward_accuracy + + +def _make_inputs( + device, + *, + batch_size: int = 4, + seq_len: int = 96, + target_channels: int = 4, + weather_channels: int = 3, +): + """Deterministic sample inputs for the standard CrossUnet config.""" + g = torch.Generator(device="cpu").manual_seed(0) + x_enc = torch.randn(batch_size, seq_len, target_channels, generator=g).to(device) + w_enc = torch.randn(batch_size, seq_len, weather_channels, generator=g).to(device) + hist_w = torch.randn(batch_size, seq_len, weather_channels, generator=g).to(device) + hist_x = torch.randn(batch_size, seq_len, target_channels, generator=g).to(device) + return x_enc, w_enc, hist_w, hist_x + + +@pytest.mark.parametrize( + "config", + ["default", "custom"], + ids=["with_defaults", "with_custom_args"], +) +def test_cross_unet_constructor(config): + """Exercise the CrossUnet constructor over default and custom configurations.""" + if config == "default": + model = CrossUnet( + target_channels=4, + weather_channels=3, + seq_len=96, + pred_len=16, + seg_len=12, + e_layers=3, + d_model=32, + n_heads=4, + d_ff=64, + ) + # Default attribute values + assert model.target_channels == 4 + assert model.weather_channels == 3 + assert model.total_channels == 7 + assert model.seq_len == 96 + assert model.pred_len == 16 + assert model.use_weather is True + assert model.nonlinear_correlation_proj is False + assert len(model.encoder.encode_blocks) == 3 + assert len(model.decoder.decode_layers) == 4 # e_layers + 1 + else: + model = CrossUnet( + target_channels=3, + weather_channels=0, + seq_len=48, + pred_len=24, + seg_len=6, + e_layers=2, + d_model=24, + n_heads=3, + d_ff=48, + dropout=0.1, + nonlinear_correlation_proj=True, + swap_corr_axis=True, + merge_kind="cnn_merge", + attention_kind="parallel", + use_bottleneck_in_decoder=False, + ) + assert model.target_channels == 3 + assert model.weather_channels == 0 + assert model.total_channels == 3 + assert model.use_weather is False + assert model.nonlinear_correlation_proj is True + assert hasattr(model, "channel_proj1") + assert hasattr(model, "channel_proj2") + assert len(model.encoder.encode_blocks) == 2 + assert len(model.decoder.decode_layers) == 3 + assert model.decoder.use_bottleneck is False + + # Common invariants + assert isinstance(model, Module) + assert hasattr(model, "meta") + assert model.meta.jit is False + assert model.meta.cuda_graphs is False + + +def test_cross_unet_forward_regression(device): + """Compare forward output against committed reference data.""" + torch.manual_seed(0) + model = CrossUnet( + target_channels=4, + weather_channels=3, + seq_len=96, + pred_len=16, + seg_len=12, + e_layers=3, + d_model=32, + n_heads=4, + d_ff=64, + ).to(device) + + x_enc, w_enc, hist_w, hist_x = _make_inputs(device) + assert validate_forward_accuracy( + model, + (x_enc, w_enc, hist_w, hist_x), + file_name="experimental/models/pv_power/cross_unet/data/cross_unet_default_output.pth", + atol=2e-3, + ) + + +def test_cross_unet_forward_no_weather(device): + """Forward path with ``weather_channels=0`` (correlation built from target only).""" + torch.manual_seed(0) + model = CrossUnet( + target_channels=4, + weather_channels=0, + seq_len=48, + pred_len=24, + seg_len=6, + e_layers=2, + d_model=24, + n_heads=3, + d_ff=48, + ).to(device) + + g = torch.Generator(device="cpu").manual_seed(1) + x_enc = torch.randn(4, 48, 4, generator=g).to(device) + hist_x = torch.randn(4, 48, 4, generator=g).to(device) + + out = model(x_enc, None, None, hist_x) + assert out.shape == (4, 24, 4) + + +def test_cross_unet_source_compatible_correlation_shape(device): + """Correlation input has one extra target column and returns model-channel mixing.""" + model = CrossUnet( + target_channels=4, + weather_channels=3, + seq_len=96, + pred_len=16, + seg_len=12, + e_layers=3, + d_model=32, + n_heads=4, + d_ff=64, + ).to(device) + + g = torch.Generator(device="cpu").manual_seed(2) + samples = torch.randn(2, 96, model.total_channels + 1, generator=g).to(device) + + corr = model._compute_channel_correlation(samples) + raw_corr = torch.vmap(lambda sample: torch.corrcoef(sample.T))(samples) + corr_to_target = torch.nan_to_num(raw_corr[:, :-1, -1], nan=0.0).clamp(min=0.0) + expected = ( + F.softmax(corr_to_target, dim=-1) + .unsqueeze(1) + .repeat(1, model.total_channels, 1) + ) + + assert corr.shape == (2, model.total_channels, model.total_channels) + assert torch.allclose(corr, expected) + + +def test_cross_unet_correlation_single_channel_fallback(device): + """Direct helper calls with a single input channel should not crash.""" + model = CrossUnet( + target_channels=1, + weather_channels=0, + seq_len=4, + pred_len=4, + seg_len=2, + e_layers=1, + d_model=8, + n_heads=2, + d_ff=16, + ).to(device) + + samples = torch.randn(3, 4, 1, device=device) + + corr = model._compute_channel_correlation(samples) + assert corr.shape == (3, 1, 1) + assert torch.allclose(corr, torch.ones_like(corr)) + + +def test_cross_unet_forward_single_target_channel_no_weather(device): + """Forward should work for the minimal target-only channel configuration.""" + model = CrossUnet( + target_channels=1, + weather_channels=0, + seq_len=12, + pred_len=6, + seg_len=6, + e_layers=2, + d_model=16, + n_heads=4, + d_ff=32, + ).to(device) + + g = torch.Generator(device="cpu").manual_seed(3) + x_enc = torch.randn(2, 12, 1, generator=g).to(device) + hist_x = torch.randn(2, 12, 1, generator=g).to(device) + + out = model(x_enc, None, None, hist_x) + assert out.shape == (2, 6, 1) + + +def test_cross_unet_short_horizon_without_bottleneck_decoder(device): + """Short decoder grids are allowed when the decoder does not add the bottleneck.""" + model = CrossUnet( + target_channels=2, + weather_channels=0, + seq_len=96, + pred_len=6, + seg_len=12, + e_layers=2, + d_model=16, + n_heads=4, + d_ff=32, + use_bottleneck_in_decoder=False, + ).to(device) + + g = torch.Generator(device="cpu").manual_seed(4) + x_enc = torch.randn(2, 96, 2, generator=g).to(device) + hist_x = torch.randn(2, 96, 2, generator=g).to(device) + + out = model(x_enc, None, None, hist_x) + assert out.shape == (2, 6, 2) + + +def test_cross_unet_nonlinear_correlation_projection_forward(device): + """Nonlinear source-style correlation projection should match model channels.""" + model = CrossUnet( + target_channels=3, + weather_channels=2, + seq_len=24, + pred_len=12, + seg_len=6, + e_layers=2, + d_model=16, + n_heads=4, + d_ff=32, + nonlinear_correlation_proj=True, + ).to(device) + + x_enc, w_enc, hist_w, hist_x = _make_inputs( + device, + batch_size=2, + seq_len=24, + target_channels=3, + weather_channels=2, + ) + + out = model(x_enc, w_enc, hist_w, hist_x) + assert out.shape == (2, 12, 5) + + +def test_cross_unet_nonlinear_correlation_projection_masks_after_projection(device): + """P-corr should project raw correlations, mask negatives, then row-normalize.""" + model = CrossUnet( + target_channels=2, + weather_channels=0, + seq_len=4, + pred_len=4, + seg_len=2, + e_layers=1, + d_model=8, + n_heads=2, + d_ff=16, + nonlinear_correlation_proj=True, + ).to(device) + + with torch.no_grad(): + proj1_a = model.channel_proj1[0] + proj1_b = model.channel_proj1[2] + proj1_a.weight.zero_() + proj1_a.bias.copy_( + torch.tensor([8.0, -8.0, 8.0, -8.0, 0.0, 0.0, 0.0, 0.0], device=device) + ) + proj1_b.weight.copy_( + torch.tensor( + [ + [1.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, -1.0, 0.0, 0.0, 0.0, 0.0], + ], + device=device, + ) + ) + proj1_b.bias.zero_() + + proj2_a = model.channel_proj2[0] + proj2_b = model.channel_proj2[2] + proj2_a.weight.zero_() + proj2_a.bias.copy_( + torch.tensor( + [ + 8.0, + -8.0, + -8.0, + 8.0, + -8.0, + 8.0, + 8.0, + -8.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + device=device, + ) + ) + proj2_b.weight.copy_( + torch.tensor( + [ + [ + 1.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + [ + 0.0, + 0.0, + 1.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + [ + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 1.0, + -1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + ], + device=device, + ) + ) + proj2_b.bias.zero_() + + t = torch.arange(4, dtype=torch.float32, device=device) + samples = torch.stack([t, -t, t], dim=-1).unsqueeze(0) + + corr = model._compute_channel_correlation(samples) + + assert corr.shape == (1, 2, 2) + assert torch.all(corr >= 0) + assert torch.allclose(corr.sum(dim=-1), torch.ones(1, 2, device=device)) + assert corr[0, 0, 1] == 0 + assert corr[0, 1, 0] == 0 + + +def test_cross_unet_positional_embedding_odd_d_model(): + """Odd-width sinusoidal embeddings should construct successfully.""" + embedding = PositionalEmbedding(d_model=3, max_len=8) + out = embedding(torch.zeros(2, 5, 3)) + assert out.shape == (1, 5, 3) + + +def test_cross_unet_example_default_config_constructs_model(): + """The example's default Hydra config should describe a runnable model.""" + config_path = ( + Path(__file__).parents[5] + / "examples/weather/pv_power_cross_unet/conf/config.yaml" + ) + cfg = OmegaConf.load(config_path) + + model = CrossUnet( + target_channels=cfg.target_channels, + weather_channels=cfg.weather_channels, + seq_len=cfg.seq_len, + pred_len=cfg.pred_len, + seg_len=cfg.seg_len, + e_layers=cfg.e_layers, + d_model=cfg.d_model, + n_heads=cfg.n_heads, + d_ff=cfg.d_ff, + dropout=cfg.dropout, + nonlinear_correlation_proj=cfg.nonlinear_correlation_proj, + attention_kind=cfg.attention_kind, + merge_kind=cfg.merge_kind, + use_bottleneck_in_decoder=cfg.use_bottleneck_in_decoder, + ) + + assert model.nonlinear_correlation_proj is True + + +def test_cross_unet_checkpoint(device): + """Save model_1 to .mdlus and verify model_2 reproduces its forward output.""" + kwargs = dict( + target_channels=4, + weather_channels=3, + seq_len=96, + pred_len=16, + seg_len=12, + e_layers=3, + d_model=32, + n_heads=4, + d_ff=64, + ) + torch.manual_seed(0) + model_1 = CrossUnet(**kwargs).to(device) + torch.manual_seed(1) + model_2 = CrossUnet(**kwargs).to(device) + + bsize = random.randint(1, 4) + x_enc, w_enc, hist_w, hist_x = _make_inputs(device, batch_size=bsize) + + assert validate_checkpoint(model_1, model_2, (x_enc, w_enc, hist_w, hist_x)) + + +def test_cross_unet_forward_validation(device): + """Forward should raise on shape mismatches.""" + model = CrossUnet( + target_channels=4, + weather_channels=3, + seq_len=96, + pred_len=16, + seg_len=12, + e_layers=3, + d_model=32, + n_heads=4, + d_ff=64, + ).to(device) + + x_enc, w_enc, hist_w, hist_x = _make_inputs(device) + + # Wrong target_channels + with pytest.raises(ValueError, match="x_enc"): + model(x_enc[:, :, :3], w_enc, hist_w, hist_x) + + # Missing weather inputs when weather_channels > 0 + with pytest.raises(ValueError, match="weather_channels"): + model(x_enc, None, None, hist_x) + + # Wrong seq_x_hist channel count + with pytest.raises(ValueError, match="seq_x_hist"): + model(x_enc, w_enc, hist_w, hist_x[:, :, :3]) + + +def test_cross_unet_constructor_validation(): + """Constructor should reject obviously invalid inputs.""" + with pytest.raises(ValueError, match="target_channels"): + CrossUnet(target_channels=0, weather_channels=3, seq_len=96, pred_len=16) + with pytest.raises(ValueError, match="weather_channels"): + CrossUnet(target_channels=4, weather_channels=-1, seq_len=96, pred_len=16) + with pytest.raises(ValueError, match="seq_len|pred_len|seg_len"): + CrossUnet(target_channels=4, weather_channels=3, seq_len=0, pred_len=16) + with pytest.raises(ValueError, match="e_layers"): + CrossUnet( + target_channels=4, weather_channels=3, seq_len=96, pred_len=16, e_layers=0 + )