patchfm 2.0.1__tar.gz → 2.1.1__tar.gz

{patchfm-2.0.1/src/patchfm.egg-info → patchfm-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchfm
-Version: 2.0.1
+Version: 2.1.1
 Summary: a Foundation Model for Univariate Time Series Forecasting
 Author-email: Samy-Melwan Vilhes <samy-melwan.vilhes@insa-rouen.fr>
 License: MIT License

{patchfm-2.0.1 → patchfm-2.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "patchfm"
-version = "2.0.1"
+version = "2.1.1"
 authors = [
   { name="Samy-Melwan Vilhes", email="samy-melwan.vilhes@insa-rouen.fr" },
 ]

patchfm-2.1.1/src/patchfm/__init__.py

@@ -0,0 +1,2 @@
+from patchfm.configs.model_config import PatchFMConfig
+from patchfm.inference.forecaster import Forecaster

{patchfm-2.0.1 → patchfm-2.1.1}/src/patchfm/configs/model_config.py

@@ -1,11 +1,12 @@
-from dataclasses import dataclass, field, asdict
+from dataclasses import asdict, dataclass, field
+
 
 @dataclass
 class PatchFMConfig:
     max_seq_len: int = 1024
     patch_len: int = 32
-    d_model: int = 2048
-    n_heads: int = 64
+    d_model: int = 1024
+    n_heads: int = 32
     n_layers_encoder: int = 6
     quantiles: list[float] = field(default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
     

{patchfm-2.0.1 → patchfm-2.1.1}/src/patchfm/inference/forecaster.py

@@ -1,7 +1,10 @@
+import numpy as np
 import torch
 import torch.nn as nn
 from einops import rearrange
-from patchfm.inference.modules import CausalRevIN, ResidualBlock, TransformerEncoder, PatchFM, SeqTypeConverter
+
+from patchfm.inference.modules import (CausalRevIN, PatchFM, ResidualBlock,
+                                       SeqTypeConverter, TransformerEncoder)
 
 
 # --- Forecaster Model ---
@@ -20,7 +23,7 @@ class Forecaster(nn.Module):
         self.max_patches = self.max_seq_len // self.patch_len
 
         print("Loading base model from HuggingFace Hub...")
-        base_model = PatchFM.from_pretrained("vilhess/PatchFM")
+        base_model = PatchFM.from_pretrained("vilhess/PatchFM-CausalRevIN-asinh")
         self._init_from_base(base_model)
 
         self.eval()
@@ -59,91 +62,13 @@ class Forecaster(nn.Module):
         self.proj_output = base_model.proj_output
     
     @torch.inference_mode()
-    def forecast(self, x: torch.Tensor, forecast_horizon: int | None = None, quantiles: list[float] | None = None) -> torch.Tensor: 
-        x = self.converter.convert(x)
-        assert x.ndim in (1, 2), f"Input dimension must be 1D (time) or 2D (batch, time), got {x.ndim}D."
-
-        batch_dim=True
-        if x.ndim != 2:
-            x = x.unsqueeze(0)
-            batch_dim=False
-        bs, ws = x.size()
-
-        x = x.to(self.device)
-
-        if ws > self.max_seq_len:
-            print(f"Warning: Input length {ws} exceeds max_seq_len {self.max_seq_len}. Truncating input.")
-            x = x[:, -self.max_seq_len:]
-            ws = self.max_seq_len
+    def auto_regressive_quantile_decoding(self, x: torch.Tensor, forecast_horizon: int | None = None, quantiles: list[float] | None = None) -> torch.Tensor: 
 
-        # Pad so length is divisible by patch_len
-        pad = (self.patch_len - ws % self.patch_len) % self.patch_len
-        if pad > 0:
-            x = torch.cat([x[:, :1].repeat(1, pad), x], dim=1)
+        q = torch.tensor(self.quantiles, device=self.device)
 
         # Default horizon = patch_len
         forecast_horizon = forecast_horizon or self.patch_len
 
-        # Reshape into patches
-        x = rearrange(x, "b (pn pl) -> b pn pl", pl=self.patch_len)  
-
-        rollouts = -(-forecast_horizon // self.patch_len)  # ceil division
-        predictions = []
-
-        for _ in range(rollouts):
-
-            if x.size(1) > self.max_patches:
-                x = x[:, -self.max_patches:, :]
-
-            # Forward pass
-            x = self.revin(x, mode="norm")
-            x = self.proj_embedding(x)
-            x = self.transformer_encoder(x)
-            x = x[:, -1:, :]  # Keep only the last patch for autoregressive forecasting
-            forecasting = self.proj_output(x)
-            forecasting = self.revin(forecasting, mode="denorm")
-
-            # Reshape to (bs, patch_num, patch_len, n_quantiles)
-            forecasting = rearrange(
-                forecasting, "b 1 (pl q) -> b 1 pl q", 
-                pl=self.patch_len, q=self.n_quantiles
-            )
-            
-            # Take median quantile (index 4)
-            patch_median = forecasting[:, -1:, :, 4].detach()
-            predictions.append(forecasting[:, -1, :, :])
-
-            # Append median patch for next rollout
-            x = patch_median.clone()
-        
-        pred_quantiles = torch.cat(predictions, dim=1)
-        pred_quantiles = pred_quantiles[:, :forecast_horizon, :]
-        pred_median = pred_quantiles[:, :, 4]
-
-        pred_quantiles = pred_quantiles[..., [self.quantiles.index(q) for q in quantiles]] if quantiles is not None else pred_quantiles
-
-        self.clear_cache()
-
-        if torch.any(torch.isnan(pred_median)) or torch.any(torch.isinf(pred_median)):
-            print("Warning: NaN or Inf values detected in predictions. Returning zeros.")
-            pred_median = torch.zeros_like(pred_median)
-            pred_quantiles = torch.zeros_like(pred_quantiles)
-        
-        if not batch_dim:
-            pred_median = pred_median.squeeze(0)
-            pred_quantiles = pred_quantiles.squeeze(0)   
-
-        pred_median, pred_quantiles = self.converter.deconvert(pred_median, pred_quantiles)
-        return pred_median, pred_quantiles
-
-    @torch.inference_mode()
-    def auto_regressive_quantile_decoding(self, x: torch.Tensor, forecast_horizon: int | None = None, quantiles: list[float] | None = None) -> torch.Tensor: 
-
-        if forecast_horizon<=self.patch_len:
-            print("Warning: forecast_horizon is less than or equal to patch_len. Using standard forecast method instead.")
-            return self.forecast(x, forecast_horizon, quantiles)
-
-        print("Using Auto-Regressive Quantile Decoding... This may be slower because no KV-Caching.")
         x = self.converter.convert(x)
         assert x.ndim in (1, 2), f"Input dimension must be 1D (time) or 2D (batch, time), got {x.ndim}D."
 
@@ -165,12 +90,8 @@ class Forecaster(nn.Module):
         if pad > 0:
             x = torch.cat([x[:, :1].repeat(1, pad), x], dim=1)
 
-        # Default horizon = patch_len
-        forecast_horizon = forecast_horizon or self.patch_len
-
         # Reshape into patches
-        x = rearrange(x, "b (pn pl) -> b pn pl", pl=self.patch_len)
-        context = x.clone()  
+        x = rearrange(x, "b (pn pl) -> b pn pl", pl=self.patch_len) 
 
         rollouts = -(-forecast_horizon // self.patch_len)  # ceil division
         predictions = []
@@ -188,18 +109,10 @@ class Forecaster(nn.Module):
         forecasting = rearrange(
             forecasting, "b 1 (pl q) -> b 1 pl q", 
             pl=self.patch_len, q=self.n_quantiles
-        )
+        )     
+        x = forecasting.permute(0, 3, 1, 2).reshape(forecasting.size(0)*self.n_quantiles, 1, self.patch_len)
 
-        context_expanded = torch.repeat_interleave(context.unsqueeze(-1), repeats=self.n_quantiles, dim=-1) # batch x patch_num x patch_len x n_quantiles
-        base_context_expanded = torch.cat((context_expanded, forecasting), dim=1)
-        context_expanded = base_context_expanded.permute(0, 3, 1, 2).reshape(bs*self.n_quantiles, base_context_expanded.size(1), self.patch_len)       
-
-        x = context_expanded
-        q = torch.tensor(self.quantiles, device=self.device)
-
-        predictions.append(forecasting[:, -1, :, :])
-
-        self.clear_cache()
+        predictions.append(forecasting[:, -1, :, :].detach())
 
         for _ in range(rollouts-1):
                 
@@ -223,15 +136,11 @@ class Forecaster(nn.Module):
             )
             forecasting = forecasting.permute(0, 2, 3, 1, 4).flatten(start_dim=-2)  # batch x 1 x patch_len x n_quantiles**2
             forecasting = torch.quantile(forecasting, q, dim=-1) # n_quantiles x batch x 1 x patch_len
-            forecasting = forecasting.permute(1, 2, 3, 0) # batch x 1 x patch_len x n_quantiles
-
-            base_context_expanded = torch.cat((base_context_expanded, forecasting), dim=1)
-            context_expanded = base_context_expanded.permute(0, 3, 1, 2).reshape(bs*self.n_quantiles, base_context_expanded.size(1), self.patch_len)       
-
-            x = context_expanded
-            predictions.append(forecasting[:, -1, :, :])
+            
+            x = forecasting.permute(1, 0, 2, 3).reshape(-1, 1, self.patch_len)
+            predictions.append(forecasting.permute(1, 2, 3, 0)[:, 0].detach())
 
-            self.clear_cache()
+        self.clear_cache()
 
         pred_quantiles = torch.cat(predictions, dim=1)
         pred_quantiles = pred_quantiles[:, :forecast_horizon, :]
@@ -251,13 +160,30 @@ class Forecaster(nn.Module):
         pred_median, pred_quantiles = self.converter.deconvert(pred_median, pred_quantiles)
         return pred_median, pred_quantiles
 
-    def __call__(self, context: torch.Tensor, forecast_horizon: int | None = None, quantiles: list[float] | None = None, quantile_decoding: bool = False) -> torch.Tensor:
-        if quantile_decoding:
-            return self.auto_regressive_quantile_decoding(context, forecast_horizon, quantiles)
+    def __call__(self, context: torch.Tensor, forecast_horizon: int | None = None, quantiles: list[float] | None = None, flip_equivariance: bool = False) -> torch.Tensor:
+        if flip_equivariance:
+            print("Flip equivariance enabled: forecast = (f(x) - f(-x)) / 2. This requires multiplying by 2 the batch size (Reverso: Efficient Time Series Foundation Models for Zero-shot Forecasting).")
+            bs = context.size(0)
+            context_flipped = -context
+            concat_context = torch.cat([context, context_flipped], dim=0)
+            pred_median_full, pred_quantiles_full = self.auto_regressive_quantile_decoding(concat_context, forecast_horizon, quantiles)
+            pred_median, pred_quantiles = pred_median_full[:bs], pred_quantiles_full[:bs]
+            pred_median2, pred_quantiles2 = pred_median_full[bs:], pred_quantiles_full[bs:]
+            pred_median = (pred_median - pred_median2) / 2
+            pred_quantiles = (pred_quantiles - flip_last_dim(pred_quantiles2)) / 2
         else:
-            return self.forecast(context, forecast_horizon, quantiles)
+            pred_median, pred_quantiles = self.auto_regressive_quantile_decoding(context, forecast_horizon, quantiles)
+        return pred_median, pred_quantiles
     
     def clear_cache(self):
         self.revin.clear_cache()    
         for layer in self.transformer_encoder.layers:
-            layer.attn.clear_cache()  
+            layer.attn.clear_cache()  
+
+def flip_last_dim(x):
+    if isinstance(x, torch.Tensor):
+        return torch.flip(x, dims=[-1])
+    elif isinstance(x, np.ndarray):
+        return np.flip(x, axis=-1)
+    else:
+        raise TypeError(f"Unsupported type: {type(x)}")

{patchfm-2.0.1 → patchfm-2.1.1}/src/patchfm/inference/modules.py

@@ -1,11 +1,12 @@
 # Modules efficient for inference with caching
 
+import numpy as np
 import torch
 import torch.nn as nn
 from einops import rearrange
-from rotary_embedding_torch import RotaryEmbedding
 from huggingface_hub import PyTorchModelHubMixin
-import numpy as np
+from rotary_embedding_torch import RotaryEmbedding
+
 
 class SeqTypeConverter:
     def __init__(self):
@@ -136,11 +137,14 @@ class CausalRevIN(nn.Module):
         counts = counts - nan_counts
     
         if self.cached_counts is not None:
+            factor = B//self.cached_counts.size(0)
+            self.cached_counts = self.cached_counts.repeat_interleave(factor, dim=0)
             counts += self.cached_counts
         self.cached_counts = counts[:, -1:, :]
 
         cumsum_x = torch.cumsum(x.nansum(dim=-1, keepdim=True), dim=1)
         if self.cached_cumsum_x is not None:
+            self.cached_cumsum_x = self.cached_cumsum_x.repeat_interleave(factor, dim=0)
             cumsum_x += self.cached_cumsum_x
         self.cached_cumsum_x = cumsum_x[:, -1:, :]
 
@@ -149,6 +153,7 @@ class CausalRevIN(nn.Module):
 
         cumsum_x2 = torch.cumsum((x**2).nansum(dim=-1, keepdim=True), dim=1)
         if self.cached_cumsum_x2 is not None:
+            self.cached_cumsum_x2 = self.cached_cumsum_x2.repeat_interleave(factor, dim=0)
             cumsum_x2 += self.cached_cumsum_x2
         self.cached_cumsum_x2 = cumsum_x2[:, -1:, :]
 
@@ -161,6 +166,8 @@ class CausalRevIN(nn.Module):
         self.cached_cumsum_x = None
         self.cached_cumsum_x2 = None
         self.cached_counts = None
+        self.cached_mean = None
+        self.cached_std = None
 
 class ResidualBlock(nn.Module):
     def __init__(self, in_dim, hid_dim, out_dim):
@@ -218,6 +225,9 @@ class MultiHeadAttention(nn.Module):
         if self.k_cache is not None and self.v_cache is not None:
             offset += self.k_cache.size(2)
             is_causal = False
+            factor = q.size(0) // self.k_cache.size(0)
+            self.k_cache = self.k_cache.repeat_interleave(factor, dim=0)
+            self.v_cache = self.v_cache.repeat_interleave(factor, dim=0)
             k = torch.cat([self.k_cache, k], dim=2)
             v = torch.cat([self.v_cache, v], dim=2)
 

{patchfm-2.0.1 → patchfm-2.1.1/src/patchfm.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchfm
-Version: 2.0.1
+Version: 2.1.1
 Summary: a Foundation Model for Univariate Time Series Forecasting
 Author-email: Samy-Melwan Vilhes <samy-melwan.vilhes@insa-rouen.fr>
 License: MIT License

patchfm-2.0.1/src/patchfm/__init__.py

@@ -1,2 +0,0 @@
-from patchfm.inference.forecaster import Forecaster
-from patchfm.configs.model_config import PatchFMConfig