discontinuum 1.0.2__py3-none-any.whl → 1.0.4__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.2'
-__version_tuple__ = version_tuple = (1, 0, 2)
+__version__ = version = '1.0.4'
+__version_tuple__ = version_tuple = (1, 0, 4)

discontinuum/engines/gpytorch.py

@@ -50,6 +50,9 @@ class MarginalGPyTorch(BaseModel):
             target_unc: Dataset = None,
             iterations: int = 100,
             optimizer: str = "adam",
+            learning_rate: float = None,
+            early_stopping: bool = False,
+            early_stopping_patience: int = 100,
             ):
         """Fit the model to data.
 
@@ -65,6 +68,12 @@ class MarginalGPyTorch(BaseModel):
             Number of iterations for optimization. The default is 100.
         optimizer : str, optional
             Optimization method. The default is "adam".
+        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.
         """
         self.is_fitted = True
         # setup data manager (self.dm)
@@ -86,26 +95,136 @@ class MarginalGPyTorch(BaseModel):
         self.model.train()
         self.likelihood.train()
 
-        # Use the adam optimizer
-        if optimizer == "adam":
-            optimizer = torch.optim.Adam(self.model.parameters(), lr=0.05) # default previously lr=0.1
-        else:
-            raise NotImplementedError("Only Adam optimizer is implemented")
+        # 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
+        
+        # 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
+        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
+        )
 
         # "Loss" for GPs - the marginal log likelihood
         mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.likelihood, self.model)
 
-        pbar = tqdm.tqdm(range(iterations), ncols=70)
+        # Training loop with stability features
+        pbar = tqdm.tqdm(range(iterations), ncols=100)  # Wider progress bar
+        jitter = 1e-6  # Dynamic jitter for numerical stability
+        best_loss = float('inf')
+        patience_counter = 0
+        min_lr_for_early_stop = 2e-5  # Stop if patience is exceeded and LR is below this
+        
         for i in pbar:
-            # Zero gradients from previous iteration
-            optimizer.zero_grad()
-            # Output from model
-            output = self.model(train_x)
-            # Calc loss and backprop gradients
-            loss = -mll(output, train_y)
-            loss.backward()
-            pbar.set_postfix(loss=loss.item())
-            optimizer.step()
+            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
+                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 loss.item() < best_loss:
+                    best_loss = loss.item()
+                    patience_counter = 0
+                else:
+                    patience_counter += 1
+
+                # Display progress with comprehensive metadata
+                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'
+                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
 
     @is_fitted
     def predict(self,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: discontinuum
-Version: 1.0.2
+Version: 1.0.4
 Summary: Estimate discontinuous timeseries from continuous covariates.
 Maintainer-email: Timothy Hodson <thodson@usgs.gov>
 License: License
@@ -124,11 +124,12 @@ However, LOADEST has several serious limitations
 the more flexible Weighted Regression on Time Discharge and Season (WRTDS),
 which allows the relation between target and covariate to vary through time.
 `loadest-gp` takes the WRTDS idea and reimplements it as a GP.
-Try it out in the [loadest-gp demo](https://code.usgs.gov/wma/uncertainty/discontinuum/-/blob/main/docs/source/notebooks/loadest-gp-demo.ipynb).
+github/thodson-usgs/discontinuum/blob/main/docs/source/notebooks/loadest-gp-demo.ipynb
+Try it out in the [loadest-gp demo](https://github.com/thodson-usgs/discontinuum/blob/main/docs/source/notebooks/loadest-gp-demo.ipynb).
 
 ### rating-gp
 `rating-gp` is a Gaussian-process model for estimating river flow from stage time series.
-Try it out in the [rating-gp demo](https://code.usgs.gov/wma/uncertainty/discontinuum/-/blob/main/docs/source/notebooks/rating-gp-demo.ipynb).
+Try it out in the [rating-gp demo](https://github.com/thodson-usgs/discontinuum/blob/main/docs/source/notebooks/rating-gp-demo.ipynb).
 
 ## Engines
 Currently, the only supported engines are the marginal likelihood implementation in `pymc` and `gpytorch`.

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

@@ -1,17 +1,17 @@
 discontinuum/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-discontinuum/_version.py,sha256=98Mtoja-WpqDUnppClywkGhX6oxX6l-ZHnSYMYYbBUE,511
+discontinuum/_version.py,sha256=rXTOeD0YpRo_kJ2LqUiMnTKEFf43sO_PBvJHDh0SQUA,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=oJMNvNAwKwxQyt3j-QyRE-pjkYDv4i-qqhQfimNQ2HQ,8654
+discontinuum/engines/gpytorch.py,sha256=36TxE_qfRUjuOB16eXmyrxPlicKzXkdQ7xnfqL2ucy0,14539
 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.2.dist-info/licenses/LICENSE.md,sha256=XElVHHnS2uQ15M_Z2giPH1vmeWMzdpGQ48ItkuZurVA,1650
+discontinuum-1.0.4.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
@@ -25,11 +25,11 @@ rating_gp/pipeline.py,sha256=1HgxN6DD3ZL5lhUb3DK2in2IXiml7W4Ja272GBMTc08,1884
 rating_gp/plot.py,sha256=CJphwqWWAfIY22j5Oz5DRwj7TcQCRyIQvM79_3KEdlc,9635
 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=rSxuTMoSeGK2LwqANMDMrmxDsOCXyAWSmENm6KFjRZ0,5930
+rating_gp/models/gpytorch.py,sha256=4SqOdWIvI93kDq9S4cDPHXX25EHNjT_hKwZijhAR4C0,7121
 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.2.dist-info/METADATA,sha256=0oCct8VOrQib5sCzdi2cBNI8KEcpYTYTKj1oCL3DCsc,6231
-discontinuum-1.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-discontinuum-1.0.2.dist-info/top_level.txt,sha256=mwU_PSFrZYSJrBgqIuTJTo7Pp9ODDv6XdDed7kAagXM,34
-discontinuum-1.0.2.dist-info/RECORD,,
+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,,

rating_gp/models/gpytorch.py

@@ -71,7 +71,9 @@ class RatingGPMarginalGPyTorch(
         # noise, *and* you did not specify noise. This is treated as a no-op."
         self.likelihood = gpytorch.likelihoods.FixedNoiseGaussianLikelihood(
             noise=noise,
-            learn_additional_noise=False,
+            #learn_additional_noise=False,
+            learn_additional_noise=True,
+            noise_prior=gpytorch.priors.HalfNormalPrior(scale=0.005),
         )
 
         model = ExactGPModel(X, y, self.likelihood)
@@ -108,17 +110,20 @@ class ExactGPModel(gpytorch.models.ExactGP):
         # + 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)
         self.covar_module = (
-            (self.cov_stage()
-             * self.cov_time(ls_prior=GammaPrior(concentration=10, rate=5)))
-            + (self.cov_stage()
-               * self.cov_time(ls_prior=GammaPrior(concentration=2, rate=5))
+            (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(
-                       train_y.min(),
-                       train_y.max(),
+                       b_min,
+                       b_max,
                    ),
                )
               )
@@ -140,11 +145,12 @@ class ExactGPModel(gpytorch.models.ExactGP):
 
     def cov_stage(self, ls_prior=None):
         eta = HalfNormalPrior(scale=1)
-
+        
         return ScaleKernel(
             MaternKernel(
                 active_dims=self.stage_dim,
                 lengthscale_prior=ls_prior,
+                nu=2.5,  # Smoother kernel (was nu=1.5)
             ),
             outputscale_prior=eta,
         )
@@ -152,13 +158,35 @@ class ExactGPModel(gpytorch.models.ExactGP):
     def cov_time(self, ls_prior=None):
         eta = HalfNormalPrior(scale=1)
 
-        return ScaleKernel(
+        # Base Matern kernel for long-term trends
+        base_kernel = 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),
+            ),
+            outputscale_prior=HalfNormalPrior(scale=0.2),
+        )
+        
+        return base_kernel + periodic_kernel
+    
+
+
 
     def cov_stagetime(self):
         eta = HalfNormalPrior(scale=1)