discontinuum 1.0.3__tar.gz → 1.0.4__tar.gz

{discontinuum-1.0.3 → discontinuum-1.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: discontinuum
-Version: 1.0.3
+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.3 → discontinuum-1.0.4}/README.md

@@ -36,11 +36,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.3 → discontinuum-1.0.4}/docs/source/notebooks/loadest-gp-demo.ipynb

@@ -14,7 +14,9 @@
     "# loadest-gp (prototype)\n",
     "LOAD ESTimator (LOADEST) is a software program for estimating some constituent using surrogate variables (covariates).\n",
     "However, LOADEST has several serious limitations, and it has been all but replaced by another model known as Weighted Regressions on Time, Discharge, and Season (WRTDS).\n",
-    "`loadest-gp` essentially reimplements WRTDS as a Gaussian process."
+    "`loadest-gp` essentially reimplements WRTDS as a Gaussian process.\n",
+    "\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/thodson-usgs/discontinuum/blob/main/docs/source/notebooks/loadest-gp-demo.ipynb)"
    ]
   },
   {
@@ -25,7 +27,7 @@
    "outputs": [],
    "source": [
     "# install the latest version of discontinuum\n",
-    "# !pip install discontinuum\n",
+    "# !pip install discontinuum[loadest_gp]\n",
     "%load_ext autoreload\n",
     "%autoreload 2"
    ]

{discontinuum-1.0.3 → discontinuum-1.0.4}/docs/source/notebooks/rating-gp-demo.ipynb

@@ -13,7 +13,9 @@
    "source": [
     "# rating-gp (prototype)\n",
     "`rating-gp` is a prototype model that can fit rating curves (stage-discharge relationship) using a Gaussian process.\n",
-    "This model seeks to expand the typical rating curve fitting process to include shifts in the rating curve with time such that the time evolution in the rating curve can be included in the model."
+    "This model seeks to expand the typical rating curve fitting process to include shifts in the rating curve with time such that the time evolution in the rating curve can be included in the model.\n",
+    "\n",
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/thodson-usgs/discontinuum/blob/main/docs/source/notebooks/rating-gp-demo.ipynb)"
    ]
   },
   {
@@ -24,7 +26,9 @@
    "outputs": [],
    "source": [
     "%load_ext autoreload\n",
-    "%autoreload 2"
+    "%autoreload 2\n",
+    "\n",
+    "#!pip install discontinuum[rating_gp]"
    ]
   },
   {
@@ -116,7 +120,10 @@
     "model = RatingGP()\n",
     "model.fit(target=training_data['discharge'],\n",
     "          covariates=training_data[['stage']],\n",
-    "          target_unc=training_data['discharge_unc'])"
+    "          target_unc=training_data['discharge_unc'],\n",
+    "          iterations=2000,\n",
+    "          early_stopping=True,\n",
+    "          )"
    ]
   },
   {
@@ -242,7 +249,9 @@
     "    model[site].fit(target=training_data['discharge'],\n",
     "                    covariates=training_data[['stage']],\n",
     "                    target_unc=training_data['discharge_unc'],\n",
-    "                    iterations=200)"
+    "                    iterations=2000,\n",
+    "                    early_stopping=True,\n",
+    "                    )"
    ]
   },
   {

{discontinuum-1.0.3 → discontinuum-1.0.4}/src/discontinuum/_version.py

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

{discontinuum-1.0.3 → discontinuum-1.0.4}/src/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.3 → discontinuum-1.0.4}/src/discontinuum.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: discontinuum
-Version: 1.0.3
+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.3 → discontinuum-1.0.4}/src/rating_gp/models/gpytorch.py

@@ -71,8 +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=True,
-            noise_prior=gpytorch.priors.HalfNormalPrior(scale=0.01),
+            noise_prior=gpytorch.priors.HalfNormalPrior(scale=0.005),
         )
 
         model = ExactGPModel(X, y, self.likelihood)
@@ -109,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_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=1, rate=2))
-               * self.cov_time(ls_prior=GammaPrior(concentration=2, rate=5))
+             + (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,
                    ),
                )
               )
@@ -141,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,
         )
@@ -153,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)