HSGP with Neumann Boundary Conditions

[juanitorduz]

Add Neumann Boundary Conditions to HSGP Approximation

Based on discussions with @maresb :)

Implementation: I guided agents with Opus 4.5 (and manual checks)

Summary

This PR extends the Hilbert Space Gaussian Process (HSGP) approximation to support Neumann boundary conditions in addition to the existing Dirichlet boundary conditions. This provides users with more flexibility when the default Dirichlet conditions (which force the GP approximation toward zero at domain boundaries) are not appropriate for their application.

The notebook is for testing purposes, I will remove it from the PR whenever we are ready to merge.

Motivation

The current HSGP implementation uses Dirichlet boundary conditions exclusively, which impose the constraint f(±L) = 0 at the boundaries of the approximation domain. This forces the approximation toward zero at the edges of [-L, L], which can introduce undesirable artifacts in regression problems where boundary effects should not influence the solution.

Neumann boundary conditions enforce f'(±L) = 0 (flat derivatives at boundaries), which is often more natural for regression applications and avoids forcing values toward zero at domain edges.

Changes

Core Implementation

1. Added boundary condition type definition:

• New BoundaryConditionType = Literal["dirichlet", "neumann"] type alias for type safety and extensibility

2. Refactored eigenvalue calculations:

• calc_eigenvalues_dirichlet(L, m): Extracted existing implementation for Dirichlet conditions
• calc_eigenvalues_neumann(L, m): New implementation for Neumann conditions
• calc_eigenvalues(L, m, boundary): Dispatcher function using Python match statement

3. Refactored eigenvector calculations:

• calc_eigenvectors_dirichlet(Xs, L, eigvals, m): Sine-based eigenfunctions for Dirichlet
• calc_eigenvectors_neumann(Xs, L, eigvals, m): Cosine-based eigenfunctions for Neumann
• calc_eigenvectors(Xs, L, eigvals, m, boundary): Dispatcher function

4. Updated HSGP class:

• Added boundary parameter to __init__ (default: "dirichlet" for backward compatibility)
• Parameter validation for boundary condition values
• Updated prior_linearized and _build_conditional to propagate boundary condition

Key Design Decisions

Identifiability: Both Dirichlet and Neumann implementations start their eigenfunction indices at j=1 (excluding j=0). For Neumann conditions, this avoids the constant eigenfunction which would create non-identifiable models when an intercept or mean function is present. This ensures both boundary conditions use the same number of basis vectors prod(m).

Normalization: Both boundary conditions use the same normalization factor 1/√L for consistency and simplicity.

Eigenvalue formula: Both use λ_j = (πj/(2L))² for j = 1, 2, ..., m.

Documentation

All new and modified functions follow numpy-style docstrings (numpydoc) with:

• Detailed parameter descriptions
• Return value specifications
• "See Also" sections linking related functions
• "Notes" sections explaining boundary condition behavior
• References to Solin & Särkkä (2019)

The HSGP class docstring includes comprehensive documentation of the new boundary parameter explaining when to use each option.

Tests

Unit Tests (TestEigenvaluesFunctions and TestEigenvectorsFunctions)

• Verify eigenvalues match theoretical formulas
• Test that Neumann and Dirichlet eigenvalues are identical (both start at j=1)
• Verify Dirichlet eigenvectors are zero at boundaries
• Verify Neumann eigenvector derivatives are zero at boundaries (theoretical check)
• Test approximate orthonormality of eigenvectors via numerical integration
• Verify output shapes for multi-dimensional cases
• Test dispatcher functions route to correct implementations

Integration Tests (TestHSGPBoundaryConditions)

• Parameter validation for invalid boundary conditions
• Default boundary condition is Dirichlet (backward compatibility)
• Prior matching between HSGP and unapproximated GP (MMD two-sample test) for both boundaries
• Conditional matching prior when no data observed (both boundaries)
• Verify Neumann does NOT suppress variance at domain edges
• Verify Dirichlet DOES suppress variance at domain edges (as expected)
• Test prior_linearized works correctly with boundary parameter

Backward Compatibility

✅ Fully backward compatible. The boundary parameter defaults to "dirichlet", preserving existing behavior. All existing code will continue to work without modification.

Usage Example

import pymc as pm
import numpy as np

X = np.linspace(-5, 5, 100)[:, None]

with pm.Model() as model:
    cov_func = pm.gp.cov.ExpQuad(1, ls=1.0)
    
    # Dirichlet (default): forces approximation toward zero at edges
    gp_dirichlet = pm.gp.HSGP(m=[50], c=2.0, boundary="dirichlet", cov_func=cov_func)
    
    # Neumann: flat derivatives at edges, no forced zero values
    gp_neumann = pm.gp.HSGP(m=[50], c=2.0, boundary="neumann", cov_func=cov_func)
    
    f = gp_neumann.prior("f", X=X)

References

• Solin, A., & Särkkä, S. (2019). Hilbert Space Methods for Reduced-Rank Gaussian Process Regression. Statistics and Computing, 30, 419-446. https://link.springer.com/article/10.1007/s11222-019-09886-w

Future Extensions

The implementation uses a match statement pattern that makes it straightforward to add additional boundary conditions in the future (e.g., periodic, Robin, or mixed boundary conditions) by:

1. Extending the BoundaryConditionType literal
2. Adding corresponding calc_eigenvalues_* and calc_eigenvectors_* functions
3. Adding new cases to the dispatcher match statements

[juanitorduz]

Added in 1187747 , before this test was failing. I need to investigate more when the Neumann boundary conditions might fail bad (and if this tests is too restrictive)

[maresb]

Just took a very quick look and I haven't quite parsed this test, but I think it may be trying to be too clever. I'm pretty sure that Dirichlet and Neumann both have simple closed-form covariances, so maybe we should just be testing this instead? I'm not sure.

[juanitorduz]

Initial TODOs:

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[maresb]

This is very exciting, thanks so much @juanitorduz!!!

I'm surprised to see such a small difference in the test notebook. Maybe if we decrease m and c the distinction will become more obvious?

[maresb]

Very nice! I'm about to take off, but wanted to share a few quick thoughts from a very initial skimming.

[maresb]

There's no way this is correct. It should be the opposite.

[maresb]

I find this highly unexpected.

Maybe we should be a separate mean-zero-neumann BC?

[jessegrabowski]

My understanding was that the API is moving away from this "skip-first basis" API, because it is not relevant for periodic kernels and vector-valued GPs. There's a depreciation notice on the drop_first argument in the HSGP class with a link to some discussion about it between @bwengals and @theorashid

[maresb]

Oh ya, I have a vague memory that skip-first was being applied in some strange way like to the Dirichlet basis. I think it should make sense for periodic and Neumann cases, but I had indeed forgotten about that part of the API.

[jessegrabowski]

Well my point was that we're moving away from it. I think the user should drop the first basis function himself if he doesn't want it -- internally, we shouldn't assume there's going to be an intercept in the model.

[maresb]

Just took a very quick look and I haven't quite parsed this test, but I think it may be trying to be too clever. I'm pretty sure that Dirichlet and Neumann both have simple closed-form covariances, so maybe we should just be testing this instead? I'm not sure.

[maresb]

Ping also @AlexAndorra.

I think it'd also make sense (probably best for a follow-up PR) to implement mixed Dirichlet/Neumann boundary conditions, so Dirichlet at the left endpoint and Neumann at the right endpoint. Perhaps we could call it "dirichlet-neumann" and "neumann-dirichlet"?

Does this work properly for dim>1?

[maresb]

Side point, probably best for a follow-up PR, is sorting out periodicity. All the current boundary conditions are actually periodic when extended, and I think the division between HSGP and Periodic is artificial. It'd be great to move towards unifying them sometime soon.