discontinuum 1.0.4__py3-none-any.whl → 1.0.6__py3-none-any.whl

discontinuum/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.0.4'
-__version_tuple__ = version_tuple = (1, 0, 4)
+__version__ = version = '1.0.6'
+__version_tuple__ = version_tuple = (1, 0, 6)

discontinuum/engines/gpytorch.py

@@ -49,10 +49,11 @@ class MarginalGPyTorch(BaseModel):
             target: Dataset,
             target_unc: Dataset = None,
             iterations: int = 100,
-            optimizer: str = "adam",
+            optimizer: str = "adamw",
             learning_rate: float = None,
             early_stopping: bool = False,
-            early_stopping_patience: int = 100,
+            patience: int = 30,
+            gradient_noise: bool = False,
             ):
         """Fit the model to data.
 
@@ -67,13 +68,15 @@ class MarginalGPyTorch(BaseModel):
         iterations : int, optional
             Number of iterations for optimization. The default is 100.
         optimizer : str, optional
-            Optimization method. The default is "adam".
+            Optimization method. Supported: "adam", "adamw". The default is "adamw".
         learning_rate : float, optional
             Learning rate for optimization. If None, uses adaptive defaults.
         early_stopping : bool, optional
             Whether to use early stopping. The default is False.
-        early_stopping_patience : int, optional
-            Number of iterations to wait without improvement before stopping. The default is 100.
+        patience : int, optional
+            Number of iterations to wait without improvement before stopping. The default is 60.
+        gradient_noise : bool, optional
+            Whether to inject Gaussian noise into gradients each step (std = 0.1 × current learning rate). The default is False.
         """
         self.is_fitted = True
         # setup data manager (self.dm)
@@ -95,31 +98,38 @@ class MarginalGPyTorch(BaseModel):
         self.model.train()
         self.likelihood.train()
 
-        # Adaptive learning rate selection for faster convergence
         if learning_rate is None:
-            if optimizer == "adam":
-                learning_rate = 0.1  # More aggressive default for faster convergence
-            elif optimizer == "lbfgs":
-                learning_rate = 1.0   # L-BFGS doesn't use learning rate the same way
+            # More conservative starting LR
+            learning_rate = 0.05
         
-        # Use the specified optimizer with stabilization
-        if optimizer != "adam":
-            raise NotImplementedError(f"Only 'adam' optimizer is supported. Got '{optimizer}'.")
-        optimizer = torch.optim.Adam(
-            self.model.parameters(), 
-            lr=learning_rate,
-            betas=(0.9, 0.999),    # Slightly more conservative momentum
-            eps=1e-8,              # Numerical stability
-            weight_decay=1e-4      # Small L2 regularization
-        )
-        # More responsive learning rate scheduler for faster adaptation
+        if optimizer == "adamw":
+            optimizer_obj = torch.optim.AdamW(
+                self.model.parameters(),
+                lr=learning_rate,
+                betas=(0.9, 0.999),
+                eps=1e-8,
+                weight_decay=1e-2      # Stronger regularization for AdamW
+            )
+        elif optimizer == "adam":
+            optimizer_obj = torch.optim.Adam(
+                self.model.parameters(),
+                lr=learning_rate,
+                betas=(0.9, 0.999),
+                eps=1e-8,
+                weight_decay=1e-4      # Lighter regularization for Adam
+            )
+        else:
+            raise NotImplementedError(f"Only 'adam' and 'adamw' optimizers are supported. Got '{optimizer}'.")
+
+        # Use ReduceLROnPlateau for more stable learning rate adaptation
         scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optimizer, 
-            mode='min', 
-            factor=0.6,    # Reduce LR by 40% when loss plateaus (more aggressive)
-            patience=40,   # Reduce sooner for faster adaptation
-            min_lr=1e-5,   # Higher minimum learning rate
-            threshold=1e-4 # Less sensitive to plateaus
+            optimizer_obj,
+            mode='min',
+            factor=0.5,                      # Reduce LR by half
+            patience=max(2, patience),
+            threshold=5e-1,                  # Aggressive plateau detection
+            #threshold_mode='rel',            # Use relative threshold
+            min_lr=1e-5
         )
 
         # "Loss" for GPs - the marginal log likelihood
@@ -133,98 +143,97 @@ class MarginalGPyTorch(BaseModel):
         min_lr_for_early_stop = 2e-5  # Stop if patience is exceeded and LR is below this
         
         for i in pbar:
-            if optimizer.__class__.__name__ == "LBFGS":
-                # L-BFGS requires a closure function
-                def closure():
-                    optimizer.zero_grad()
-                    output = self.model(train_x)
-                    with gpytorch.settings.cholesky_jitter(jitter):
-                        loss = -mll(output, train_y).sum()
-                    loss.backward()
-                    return loss
-                
-                loss = optimizer.step(closure)
-                pbar.set_postfix(loss=loss.item())
-            else:
-                # Adam optimizer with stability features
-                optimizer.zero_grad()
-                output = self.model(train_x)
-                
-                # Attempt loss calculation with dynamic jitter
-                try:
-                    with gpytorch.settings.cholesky_jitter(jitter):
-                        loss = -mll(output, train_y)
-                except Exception as e:
-                    # Increase jitter if numerical issues occur
-                    jitter = min(jitter * 10, 1e-2)
-                    current_lr = optimizer.param_groups[0]['lr']
-                    pbar.set_postfix_str(
-                        f'lr={current_lr:.1e} jitter={jitter:.1e} | Numerical issue - increasing jitter'
-                    )
-                    continue
-                
-                # Check for NaN loss
-                if torch.isnan(loss) or torch.isinf(loss):
-                    current_lr = optimizer.param_groups[0]['lr']
-                    pbar.set_postfix_str(
-                        f'lr={current_lr:.1e} jitter={jitter:.1e} | NaN/Inf loss detected - skipping step'
-                    )
-                    continue
-                
-                loss.backward()
-                
-                # Gradient clipping for stability
-                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
-                
-                # Check for NaN gradients
-                has_nan_grad = False
+            # Adam/AdamW optimizer with stability features
+            optimizer_obj.zero_grad()
+            output = self.model(train_x)
+
+            # Attempt loss calculation with dynamic jitter
+            try:
+                with gpytorch.settings.cholesky_jitter(jitter):
+                    loss = -mll(output, train_y)
+            except Exception as e:
+                # Increase jitter if numerical issues occur
+                jitter = min(jitter * 10, 1e-2)
+                current_lr = optimizer_obj.param_groups[0]['lr']
+                pbar.set_postfix_str(
+                    f'lr={current_lr:.1e} jitter={jitter:.1e} | Numerical issue - increasing jitter'
+                )
+                continue
+
+            # Check for NaN loss
+            if torch.isnan(loss) or torch.isinf(loss):
+                current_lr = optimizer_obj.param_groups[0]['lr']
+                pbar.set_postfix_str(
+                    f'lr={current_lr:.1e} jitter={jitter:.1e} | NaN/Inf loss detected - skipping step'
+                )
+                continue
+
+            loss.backward()
+
+            # Get current learning rate before gradient noise injection
+            current_lr = optimizer_obj.param_groups[0]['lr']
+
+            # Gradient noise injection (if enabled)
+            if gradient_noise:
+                gradient_noise_scale = 0.1
+                adaptive_noise = gradient_noise_scale * current_lr
                 for param in self.model.parameters():
-                    if param.grad is not None and torch.isnan(param.grad).any():
-                        has_nan_grad = True
-                        break
-                
-                if has_nan_grad:
-                    # Don't update scheduler on NaN gradients - this prevents rapid LR decay
-                    # The scheduler should only respond to actual optimization progress
-                    current_lr = optimizer.param_groups[0]['lr']
-                    
-                    # Update best loss tracking (loss is still valid, just gradients are NaN)
-                    if loss.item() < best_loss:
-                        best_loss = loss.item()
-                        patience_counter = 0
-                    else:
-                        patience_counter += 1
-                    
-                    # Display comprehensive info even with NaN gradients
-                    pbar.set_postfix_str(
-                        f'loss={loss.item():.4f} lr={current_lr:.1e} jitter={jitter:.1e} best={best_loss:.4f} | NaN gradients - skipping step'
-                    )
-                    continue
-                
-                optimizer.step()
-                
-                # Update learning rate scheduler for Adam
-                scheduler.step(loss)
-                current_lr = optimizer.param_groups[0]['lr']
-                
-                # Early stopping check (more aggressive)
+                    if param.grad is not None:
+                        noise = torch.normal(mean=0.0, std=adaptive_noise, size=param.grad.shape, device=param.grad.device)
+                        param.grad.add_(noise)
+
+            # Gradient clipping for stability
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+
+            # Check for NaN gradients
+            has_nan_grad = False
+            for param in self.model.parameters():
+                if param.grad is not None and torch.isnan(param.grad).any():
+                    has_nan_grad = True
+                    break
+
+            if has_nan_grad:
+                # Don't update scheduler on NaN gradients - this prevents rapid LR decay
+                # The scheduler should only respond to actual optimization progress
+                current_lr = optimizer_obj.param_groups[0]['lr']
+
+                # Update best loss tracking (loss is still valid, just gradients are NaN)
                 if loss.item() < best_loss:
                     best_loss = loss.item()
                     patience_counter = 0
                 else:
                     patience_counter += 1
 
-                # Display progress with comprehensive metadata
+                # Display comprehensive info even with NaN gradients, skip normal progress update
+                pbar.set_postfix_str(
+                    f'loss={loss.item():.4f} lr={current_lr:.1e} jitter={jitter:.1e} best={best_loss:.4f} | NaN gradients - skipping step'
+                )
+                continue
+
+            optimizer_obj.step()
+
+            # Update learning rate scheduler for Adam/AdamW
+            scheduler.step(loss.item())
+            current_lr = optimizer_obj.param_groups[0]['lr']
+
+            # Early stopping check (more aggressive)
+            if loss.item() < best_loss:
+                best_loss = loss.item()
+                patience_counter = 0
+            else:
+                patience_counter += 1
+
+            # Only update progress bar if not skipped above
+            if not has_nan_grad:
                 progress_info = f'loss={loss.item():.4f} lr={current_lr:.1e} jitter={jitter:.1e} best={best_loss:.4f}'
                 if early_stopping:
-                    progress_info += f' patience={patience_counter}/25'
+                    progress_info += f' patience={patience_counter}/{patience}'
                 pbar.set_postfix_str(progress_info)
 
-                # More aggressive early stopping: patience=25 and require LR to be low
-                if early_stopping and patience_counter >= 25 and current_lr <= min_lr_for_early_stop:
-                    print(f"\nEarly stopping triggered after {i+1} iterations (patience exceeded and LR low)")
-                    print(f"Best loss: {best_loss:.6f}")
-                    break
+            if early_stopping and patience_counter >= patience and current_lr <= min_lr_for_early_stop:
+                print(f"\nEarly stopping triggered after {i+1} iterations")
+                print(f"Best loss: {best_loss:.6f}")
+                break
 
     @is_fitted
     def predict(self,

{discontinuum-1.0.4.dist-info → discontinuum-1.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: discontinuum
-Version: 1.0.4
+Version: 1.0.6
 Summary: Estimate discontinuous timeseries from continuous covariates.
 Maintainer-email: Timothy Hodson <thodson@usgs.gov>
 License: License

{discontinuum-1.0.4.dist-info → discontinuum-1.0.6.dist-info}/RECORD

@@ -1,17 +1,17 @@
 discontinuum/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-discontinuum/_version.py,sha256=rXTOeD0YpRo_kJ2LqUiMnTKEFf43sO_PBvJHDh0SQUA,511
+discontinuum/_version.py,sha256=B7xX94ww8E2YsJCo2PT7LI1Lp5224NjscDIhXgKzj3U,511
 discontinuum/data_manager.py,sha256=LiZoPR0nnu7YAUfh5L1ZDRfaS3dgfVIELXIHkzUKyBg,4416
 discontinuum/pipeline.py,sha256=1avuZnFai-b3HmihcpZ8M3WFNQ8lXAFSNTrnfl2NrY0,10074
 discontinuum/plot.py,sha256=eZQS6-Ydq8FFcEukPtNuDVB-weV6lHyWMyJ1hqTkVrU,2969
 discontinuum/utils.py,sha256=07hIHQk_oDlkjz7tasgBjqqPOC6D0iNcy0eu-88aNbM,1540
 discontinuum/engines/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 discontinuum/engines/base.py,sha256=OlHd4ssIQoWvYHKoVqk5fKAVBcKsIIkR4ul9iNBvaYg,2396
-discontinuum/engines/gpytorch.py,sha256=36TxE_qfRUjuOB16eXmyrxPlicKzXkdQ7xnfqL2ucy0,14539
+discontinuum/engines/gpytorch.py,sha256=05x7Ha0g2vywM_moL18fMFDGeh0CF3vJpF-mDImrIx8,14387
 discontinuum/engines/pymc.py,sha256=phbtE-3UCSVcP1MhbXwAHIWDZWDr56wK9U7aRt-w-2o,5961
 discontinuum/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 discontinuum/providers/base.py,sha256=Yn2EHS1b4fYl09-m2MYuf2P9VRUXAP-WDpSoZrCbRvY,720
 discontinuum/tests/test_pipeline.py,sha256=_FhkGxbFIxNb35lGaIdZk7Zjgs6CkxEF3gFUX3PE8EU,918
-discontinuum-1.0.4.dist-info/licenses/LICENSE.md,sha256=XElVHHnS2uQ15M_Z2giPH1vmeWMzdpGQ48ItkuZurVA,1650
+discontinuum-1.0.6.dist-info/licenses/LICENSE.md,sha256=XElVHHnS2uQ15M_Z2giPH1vmeWMzdpGQ48ItkuZurVA,1650
 loadest_gp/__init__.py,sha256=YISfvbc7Zy2y0BOxS1A2KzqxyoNJTz0EnLMnRW6iVT8,740
 loadest_gp/plot.py,sha256=x2PK7vBCc44dX9lu5YV-rvw1u4pvXSLdcrTSvYLiHMA,2595
 loadest_gp/utils.py,sha256=m5QaqR_0JiuRXPfryH8nI5lODp8PqvQla5C05WDN3LY,2772
@@ -22,14 +22,14 @@ loadest_gp/models/pymc.py,sha256=ShP-XNoEwNAf62yCvTuS6Q8iAiB9NQk0dVS69WgkPsE,344
 loadest_gp/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 loadest_gp/providers/usgs.py,sha256=LJ5uh0g3nrZ-8I4poGwCdVqcXcpyZkroixwwt99vBcI,10885
 rating_gp/pipeline.py,sha256=1HgxN6DD3ZL5lhUb3DK2in2IXiml7W4Ja272GBMTc08,1884
-rating_gp/plot.py,sha256=CJphwqWWAfIY22j5Oz5DRwj7TcQCRyIQvM79_3KEdlc,9635
+rating_gp/plot.py,sha256=_XaeNLYXEcJxg7B4UCxyYccSNzNow0e4dV1z93_THaQ,10899
 rating_gp/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rating_gp/models/base.py,sha256=e2Kq644I88YLHWPNA0qyRgitF5wimdLW4618vKX-o_s,1474
-rating_gp/models/gpytorch.py,sha256=4SqOdWIvI93kDq9S4cDPHXX25EHNjT_hKwZijhAR4C0,7121
+rating_gp/models/gpytorch.py,sha256=bNFJFT13DyVqlKhMRD3W0r6-Y72E3S2fw9E-0houyoM,7068
 rating_gp/models/kernels.py,sha256=3xg2mhY3aEgjI3r5vyAll9MA4c3M5UKqRi3FApNhJJQ,11579
 rating_gp/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rating_gp/providers/usgs.py,sha256=KmKYN3c8Mi-ly2l6X80WT3taEhqCPXeEcRNi9HvbJmY,8134
-discontinuum-1.0.4.dist-info/METADATA,sha256=A6T6BQocZmIox600f7nU5Tb9r7x5YthC5ba1WRET2XM,6302
-discontinuum-1.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-discontinuum-1.0.4.dist-info/top_level.txt,sha256=mwU_PSFrZYSJrBgqIuTJTo7Pp9ODDv6XdDed7kAagXM,34
-discontinuum-1.0.4.dist-info/RECORD,,
+discontinuum-1.0.6.dist-info/METADATA,sha256=3h3AhrQZ3eNviDGAsbcKEd3yd_cYv8XaaIxMlCQdK0s,6302
+discontinuum-1.0.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+discontinuum-1.0.6.dist-info/top_level.txt,sha256=mwU_PSFrZYSJrBgqIuTJTo7Pp9ODDv6XdDed7kAagXM,34
+discontinuum-1.0.6.dist-info/RECORD,,

rating_gp/models/gpytorch.py

@@ -5,9 +5,9 @@ from discontinuum.engines.gpytorch import MarginalGPyTorch, NoOpMean
 
 from gpytorch.kernels import (
     MaternKernel,
-    RBFKernel,
     RQKernel,
     ScaleKernel,
+    PeriodicKernel,
 )
 from gpytorch.priors import (
     GammaPrior,
@@ -15,10 +15,9 @@ from gpytorch.priors import (
     NormalPrior,
 )
 
-from linear_operator.operators import MatmulLinearOperator
 from rating_gp.models.base import RatingDataMixin, ModelConfig
 from rating_gp.plot import RatingPlotMixin
-from rating_gp.models.kernels import StageTimeKernel, SigmoidKernel, LogWarp, TanhWarp
+from rating_gp.models.kernels import SigmoidKernel
 
 
 class PowerLawTransform(torch.nn.Module):
@@ -94,53 +93,55 @@ class ExactGPModel(gpytorch.models.ExactGP):
 
         self.powerlaw = PowerLawTransform()
 
-        # self.mean_module = gpytorch.means.ConstantMean()
-        # self.mean_module = gpytorch.means.LinearMean(input_size=1)
         self.mean_module = NoOpMean()
 
-        #self.warp_stage_dim = TanhWarp()
-        #self.warp_stage_dim = LogWarp()
-
-        # self.covar_module = (
-        #     (self.cov_stage() * self.cov_stagetime())
-        #     + self.cov_residual()
-        # )
-         
-        # Stage * time kernel with large time length
-        # + stage * time kernel only at low stage with smaller time length.
-        # Note that stage gets transformed to q, so the kernel is actually
-        # q * time
-        b_min = np.quantile(train_y, 0.30)
-        b_max = np.quantile(train_y, 0.90)
+        # Use stage (not y) for sigmoid kernel constraint
+        stage = train_x[:, self.stage_dim[0]]#.cpu().numpy()
+        b_min = np.quantile(stage, 0.10)
+        b_max = np.quantile(stage, 0.90)
         self.covar_module = (
-            (self.cov_stage(ls_prior=GammaPrior(concentration=2,  rate=1))
-             * self.cov_time(ls_prior=GammaPrior(concentration=1,  rate=1)))
-             + (self.cov_stage(ls_prior=GammaPrior(concentration=5, rate=1))
-               * self.cov_time(ls_prior=GammaPrior(concentration=1, rate=5))
-               * SigmoidKernel(
-                   active_dims=self.stage_dim,
-                   # a_prior=NormalPrior(loc=20, scale=1),
-                   # b_prior=NormalPrior(loc=0.5, scale=0.2),
-                   b_constraint=gpytorch.constraints.Interval(
-                       b_min,
-                       b_max,
-                   ),
-               )
-              )
+            # core time kernel
+            (
+                 self.cov_time(
+                     #ls_prior=GammaPrior(concentration=2, rate=1),
+                     ls_prior=GammaPrior(concentration=3, rate=1),
+                     eta_prior=HalfNormalPrior(scale=0.3),
+                 ) 
+                 *
+                 self.cov_stage(ls_prior=GammaPrior(concentration=3, rate=2))
+                 #self.cov_stage(ls_prior=GammaPrior(concentration=2, rate=1))
+             )
+             # gated shift component
+             + (
+                self.cov_time(
+                    ls_prior=GammaPrior(concentration=2, rate=5),
+                    eta_prior=HalfNormalPrior(scale=1),
+                    )
+                * SigmoidKernel(
+                    active_dims=self.stage_dim,
+                    # b_prior=NormalPrior(loc=0.7, scale=0.001),
+                    b_constraint=gpytorch.constraints.Interval(
+                        b_min,
+                        b_max,
+                    ),
+                )
+            )
+            # additive periodic component for seasonal effects
+            + self.cov_periodic()
         )
 
 
     def forward(self, x):
-        self.powerlaw.b.data.clamp_(1.5, 2.5)
+        self.powerlaw.b.data.clamp_(1.2, 2.5)
         #x = x.clone()
         #q = self.powerlaw(x[:, self.stage_dim])
         #x_t[:, self.stage_dim] = self.warp_stage_dim(x_t[:, self.stage_dim])
         x_t = x.clone()
         x_t[:, self.stage_dim] = self.powerlaw(x_t[:, self.stage_dim])
         q = x_t[:, self.stage_dim]
-
         mean_x = self.mean_module(q)
-        covar_x = self.covar_module(x_t)
+        #covar_x = self.covar_module(x_t)
+        covar_x = self.covar_module(x)
         return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
 
     def cov_stage(self, ls_prior=None):
@@ -155,61 +156,57 @@ class ExactGPModel(gpytorch.models.ExactGP):
             outputscale_prior=eta,
         )
 
-    def cov_time(self, ls_prior=None):
-        eta = HalfNormalPrior(scale=1)
+    def cov_time(self, ls_prior=None, eta_prior=None):
+        if eta_prior is None:
+            eta_prior = HalfNormalPrior(scale=1)
 
         # Base Matern kernel for long-term trends
-        base_kernel = ScaleKernel(
+        return ScaleKernel(
             MaternKernel(
                 active_dims=self.time_dim,
                 lengthscale_prior=ls_prior,
-                nu=1.5, # was 2.5
-            ),
-            outputscale_prior=eta,
-        )
-        
-        # Periodic kernel for annual seasonality
-        # Locally periodic kernel: Periodic * Matern
-        periodic_kernel = ScaleKernel(
-            gpytorch.kernels.PeriodicKernel(
-                active_dims=self.time_dim,
-                period_length_prior=NormalPrior(loc=1.0, scale=0.05),  # ~1 year
-                lengthscale_prior=GammaPrior(concentration=6, rate=1),
-            ) * MaternKernel(
-                active_dims=self.time_dim,
-                nu=2.5,
-                lengthscale_prior=GammaPrior(concentration=4, rate=3),
+                nu=1.5, # was 1.5 XXX
             ),
-            outputscale_prior=HalfNormalPrior(scale=0.2),
+            outputscale_prior=eta_prior,
         )
-        
-        return base_kernel + periodic_kernel
     
+    def cov_periodic(self, ls_prior=None, eta_prior=None):
+        """
+        Smooth, time-dependent periodic kernel for seasonal effects.
+        """
+        if eta_prior is None:
+            eta_prior = HalfNormalPrior(scale=0.5) 
 
-
-
-    def cov_stagetime(self):
-        eta = HalfNormalPrior(scale=1)
-        ls = GammaPrior(concentration=2, rate=1)
+        if ls_prior is None:
+            ls_prior = GammaPrior(concentration=3, rate=1)
 
         return ScaleKernel(
-            StageTimeKernel(
-                active_dims=self.dims,
-                # lengthscale_prior=ls,
+            PeriodicKernel(
+                active_dims=self.time_dim,
+                period_length_prior=NormalPrior(loc=1.0, scale=0.1),  # ~1 year
+                # lengthscale_prior=GammaPrior(concentration=2, rate=4),
             ),
-            # outputscale_prior=eta,
+            # *
+            # MaternKernel(
+            #     active_dims=self.stage_dim,
+            #     lengthscale_prior=ls_prior,
+            #     nu=2.5,  # Smoother kernel (was nu=1.5)
+            # ),
+            outputscale_prior=HalfNormalPrior(scale=0.5),
         )
-
-    def cov_residual(self):
-        eta = HalfNormalPrior(scale=0.2)
-        ls = GammaPrior(concentration=2, rate=10)
-
+    
+    def cov_base(self):
+        """
+        Smooth, time-independent base rating curve using a Matern kernel on stage.
+        """
+        # Base should capture most variation
+        eta = HalfNormalPrior(scale=1)
+        ls = GammaPrior(concentration=3, rate=1)
         return ScaleKernel(
             MaternKernel(
-                ard_num_dims=2,
-                nu=1.5,
-                active_dims=self.dims,
+                active_dims=self.stage_dim,
                 lengthscale_prior=ls,
             ),
             outputscale_prior=eta,
         )
+ 

rating_gp/plot.py

@@ -12,6 +12,7 @@ import pandas as pd
 from discontinuum.engines.base import is_fitted
 from discontinuum.plot import BasePlotMixin
 from scipy.stats import norm
+from rating_gp.models.kernels import SigmoidKernel
 import xarray as xr
 from xarray import DataArray
 from xarray.plot.utils import label_from_attrs
@@ -117,6 +118,28 @@ class RatingPlotMixin(BasePlotMixin):
             zorder=1,
             **kwargs
         )
+        # Plot switch point if sigmoid kernel is in model
+        try:
+            # find first SigmoidKernel in covar_module
+            sig_kernels = [m for m in self.model.covar_module.modules() if isinstance(m, SigmoidKernel)]
+            if sig_kernels:
+                sig = sig_kernels[0]
+                # b_sig is in normalized stage space: inverse-transform to original stage units
+                b_sig = sig.b.item()
+                # use scaler step directly to inverse-transform normalized stage
+                pipeline = self.dm.covariate_pipelines['stage']
+                scaler = pipeline.named_steps['scaler']
+                stage_switch = float(scaler.inverse_transform(b_sig))
+                # Draw switch point on every call; label only once
+                if not getattr(ax, 'switch_point_plotted', False):
+                    ax.axvline(stage_switch, linestyle='--', color='gray', label='switch point')
+                    ax.switch_point_plotted = True
+                    ax.legend()
+                else:
+                    # subsequent calls, draw without label
+                    ax.axvline(stage_switch, linestyle='--', color='gray')
+        except Exception:
+            pass
 
         # self.plot_observed_rating(ax, zorder=3)