ladam 0.2.0__tar.gz

ladam-0.2.0/.gitignore

@@ -0,0 +1,16 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+*.egg
+dist/
+build/
+.eggs/
+*.whl
+.pytest_cache/
+.tox/
+.coverage
+htmlcov/
+*.so
+.env
+.venv

ladam-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Greg Partin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ladam-0.2.0/PKG-INFO

@@ -0,0 +1,246 @@
+Metadata-Version: 2.4
+Name: ladam
+Version: 0.2.0
+Summary: LAdam: Laplacian Adam — Adam with spatially-coupled variance estimates via discrete Laplacian
+Project-URL: Homepage, https://github.com/gpartin/ladam
+Project-URL: Documentation, https://github.com/gpartin/ladam#usage
+Project-URL: Issues, https://github.com/gpartin/ladam/issues
+Author: Greg Partin
+License: MIT
+License-File: LICENSE
+Keywords: adam,deep-learning,laplacian,optimizer,pinn,pytorch,scientific-ml,spatial-regularization
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.8
+Requires-Dist: torch>=1.10.0
+Provides-Extra: dev
+Requires-Dist: pytest-cov>=4.0; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# LAdam
+
+**Laplacian Adam — spatially-aware adaptive optimizer for PyTorch**
+
+[![PyPI](https://img.shields.io/pypi/v/ladam.svg)](https://pypi.org/project/ladam/)
+[![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://python.org)
+
+LAdam is a drop-in Adam replacement that applies **discrete Laplacian regularization** to Adam's second-moment estimate (v_t). This couples neighboring weight learning rates, producing spatially-smoothed adaptive optimization.
+
+## Why LAdam?
+
+Adam computes independent per-parameter learning rates. But adjacent weights in trained networks are often functionally correlated — the per-parameter variance estimates should reflect this structure.
+
+LAdam adds **one operation** to Adam: a Laplacian diffusion step on v_t, controlled by a single scalar `c2`. The Laplacian allows each weight's learning rate to be informed by its neighbors, smoothing the optimization landscape.
+
+## Results
+
+| Task | Architecture | Metric | Adam | LAdam | Improvement |
+|------|-------------|--------|------|-------|-------------|
+| **Wave Equation PINN** | 5×128 MLP | L2 Error | 0.0310 | **0.0172** | **-44.6%** |
+| **FashionMNIST** | Transformer | Accuracy | 89.46% | **89.66%** | **+0.20%** (p=0.0005) |
+| FashionMNIST | MLP | Accuracy | 89.10% | 89.12% | +0.02% (n.s.) |
+| FashionMNIST | CNN | Accuracy | 91.15% | 91.14% | -0.01% (tie) |
+
+> **LAdam excels on architectures with spatially-correlated weight structure** — particularly PINNs and transformers. For CNNs (whose conv filters are already spatial detectors), the Laplacian is redundant.
+
+## Installation
+
+```bash
+pip install ladam
+```
+
+## Optimizers
+
+LAdam ships three Laplacian-enhanced optimizers:
+
+| Optimizer | Base | Laplacian target | Best for |
+|-----------|------|------------------|----------|
+| **LAdam** | Adam | Second moment v_t | PINNs, transformers, CNNs |
+| **LAdaGrad** | AdaGrad | Cumulative sum G_t | Sparse features, NLP |
+| **LRMSProp** | RMSProp | Running average v_t | RNNs, non-stationary losses |
+
+All three share the same Laplacian kernel infrastructure and `c2` parameter.
+
+## Usage
+
+### Basic — Drop-in Adam replacement
+
+```python
+from ladam import LAdam
+
+optimizer = LAdam(model.parameters(), lr=1e-3, c2=1e-4)
+
+# Training loop is identical to Adam
+for batch in dataloader:
+    loss = criterion(model(batch))
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+```
+
+### LAdaGrad and LRMSProp
+
+```python
+from ladam import LAdaGrad, LRMSProp
+
+# AdaGrad with Laplacian smoothing on cumulative squared gradients
+optimizer = LAdaGrad(model.parameters(), lr=1e-2, c2=1e-4)
+
+# RMSProp with Laplacian smoothing on running variance
+optimizer = LRMSProp(model.parameters(), lr=1e-2, alpha=0.99, c2=1e-4)
+```
+
+### Per-layer c2 with parameter groups
+
+```python
+optimizer = LAdam([
+    {'params': model.attention.parameters(), 'c2': 1e-4},   # Transformer attention
+    {'params': model.ffn.parameters(), 'c2': 1e-5},         # Feed-forward
+    {'params': model.norm.parameters(), 'c2': 0.0},         # Skip for norms
+], lr=3e-4)
+```
+
+### Architecture-aware defaults
+
+```python
+from ladam import LAdam, suggest_c2
+
+c2 = suggest_c2('pinn')         # Returns 1e-5
+c2 = suggest_c2('transformer')  # Returns 1e-4
+
+optimizer = LAdam(model.parameters(), lr=1e-3, c2=c2)
+```
+
+## Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `lr` | 1e-3 | Learning rate |
+| `betas` | (0.9, 0.999) | EMA coefficients (same as Adam) |
+| `eps` | 1e-8 | Numerical stability (same as Adam) |
+| `weight_decay` | 0 | L2 regularization (same as AdamW behavior) |
+| `c2` | 1e-4 | **Laplacian coupling strength.** Controls how much neighboring variance estimates influence each other. |
+| `mode` | 'variance_lap' | Which quantity to smooth. `'variance_lap'` is best. |
+| `stencil` | '9point' | **Discrete Laplacian stencil.** `'9point'` (isotropic, 0.46% anisotropy) or `'5point'` (legacy, 12.3% anisotropy). |
+| `min_spatial_size` | 16 | Skip Laplacian for params with fewer elements (biases, LayerNorm). |
+
+### Stencil Selection
+
+The `stencil` parameter controls the discrete Laplacian kernel used for spatial coupling:
+
+- **`'9point'` (default)**: Isotropic stencil with face + edge neighbors. Treats diagonal neighbors with 1/6 weight vs 4/6 for face neighbors.
+- **`'5point'`**: Standard cross-pattern stencil (faces only). Slightly faster but 25× more anisotropic.
+
+At typical `c2` values (1e-5 to 1e-3), the effective learning rate difference between stencils is <0.3%. The 9-point default is recommended for correctness.
+
+### Choosing c2
+
+`c2` is the only new hyperparameter. It's robust across 3 orders of magnitude:
+
+| c2 | Best For | Notes |
+|----|----------|-------|
+| `1e-5` | PINNs, scientific ML | Gentle coupling, biggest error reduction |
+| `1e-4` | Transformers, general | **Safe default** |
+| `1e-3` | Aggressive smoothing | Works but slightly less stable |
+| `0` | Disable | Reduces to standard Adam |
+
+All 7 values tested in [1e-6, 1e-3] outperformed Adam on transformers (B12 sweep).
+
+## How It Works
+
+Standard Adam computes per-parameter adaptive learning rates from the second moment:
+
+```
+v_t = β₂·v_{t-1} + (1-β₂)·g_t²     # Variance estimate
+lr_effective = lr / (√v_t + ε)        # Per-parameter learning rate
+```
+
+LAdam adds a Laplacian coupling step:
+
+```
+v_smooth = v_t + c2 · ∇²v_t           # Spatial smoothing
+lr_effective = lr / (√v_smooth + ε)    # Coupled learning rate
+```
+
+Where `\nabla^2` is the discrete Laplacian computed via a single `F.conv2d` kernel (9-point isotropic by default) -- efficient and GPU-friendly. The Laplacian treats weight matrices as 2D fields, coupling each weight's learning rate with its spatial neighbors.
+
+**Overhead**: ~2-5% wall-clock time increase per step. The Laplacian is a single fused convolution kernel, not point-wise iteration.
+
+## Benchmarks
+
+### PINN: Wave Equation (u_tt = c^2 u_xx)
+
+5-layer, 128-unit tanh MLP trained for 5000 steps on the 1D wave equation.
+
+| Optimizer | L2 Error | vs Adam |
+|-----------|----------|---------|
+| Adam (lr=1e-3) | 0.0310 | — |
+| LAdam c²=1e-4 | 0.0240 | -22.8% |
+| **LAdam c²=1e-5** | **0.0172** | **-44.6%** |
+| LAdam c²=1e-3 | 0.0185 | -40.3% |
+
+### Transformer: FashionMNIST Classification
+
+4-head, 128-dim, 2-layer transformer, 30 epochs, 5 independent seeds.
+
+| Optimizer | Accuracy (mean ± std) | p-value (vs Adam) |
+|-----------|----------------------|-------------------|
+| Adam | 89.46 ± 0.10% | — |
+| **LAdam c²=1e-4** | **89.66 ± 0.06%** | **0.0005** |
+
+### c² Robustness Sweep
+
+7 c² values on the same transformer task. **All 7 beat Adam:**
+
+| c² | Accuracy | Δ vs Adam |
+|----|----------|-----------|
+| 1e-6 | 89.62% | +0.16% |
+| 5e-6 | 89.73% | +0.27% |
+| 1e-5 | 89.79% | +0.33% |
+| 5e-5 | 89.75% | +0.29% |
+| 1e-4 | 89.67% | +0.21% |
+| 5e-4 | 89.64% | +0.18% |
+| 1e-3 | 89.66% | +0.20% |
+
+## FAQ
+
+**Q: Does this work for LLMs / GPT-scale models?**
+A: No. LAdam **hurts** LLM training (tested on GPT-2/WikiText-2). Attention weight matrices encode semantic structure, not spatial structure — the Laplacian destroys per-feature specialization. Use standard Adam/AdamW for LLMs.
+
+**Q: Why not smooth the gradient instead of the variance?**
+A: [Osher et al. (2018)](https://arxiv.org/abs/1806.06317) explored Laplacian smoothing of gradients. We found that smoothing the *variance estimate* is more effective because it smooths the *learning rate landscape* rather than the *descent direction*. These are mathematically distinct: ∇²(EMA(g²)) ≠ (∇²g)².
+
+**Q: Why does this help PINNs so much?**
+A: PDE-based loss landscapes have inherent spatial structure from the differential operators in the loss function. The Laplacian on v_t aligns the optimizer's internal representation with this structure.
+
+**Q: Can I use this with learning rate schedulers?**
+A: Yes. LAdam is fully compatible with any `torch.optim.lr_scheduler`.
+
+## Citation
+
+If you use LAdam in your research, please cite:
+
+```bibtex
+@software{partin2026ladam,
+  author = {Partin, Greg},
+  title = {LAdam: Spatially-Aware Adaptive Optimization via Laplacian-Regularized Variance Estimates},
+  year = {2026},
+  url = {https://github.com/gpartin/ladam}
+}
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE) for details.

ladam-0.2.0/README.md

@@ -0,0 +1,216 @@
+# LAdam
+
+**Laplacian Adam — spatially-aware adaptive optimizer for PyTorch**
+
+[![PyPI](https://img.shields.io/pypi/v/ladam.svg)](https://pypi.org/project/ladam/)
+[![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)
+[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://python.org)
+
+LAdam is a drop-in Adam replacement that applies **discrete Laplacian regularization** to Adam's second-moment estimate (v_t). This couples neighboring weight learning rates, producing spatially-smoothed adaptive optimization.
+
+## Why LAdam?
+
+Adam computes independent per-parameter learning rates. But adjacent weights in trained networks are often functionally correlated — the per-parameter variance estimates should reflect this structure.
+
+LAdam adds **one operation** to Adam: a Laplacian diffusion step on v_t, controlled by a single scalar `c2`. The Laplacian allows each weight's learning rate to be informed by its neighbors, smoothing the optimization landscape.
+
+## Results
+
+| Task | Architecture | Metric | Adam | LAdam | Improvement |
+|------|-------------|--------|------|-------|-------------|
+| **Wave Equation PINN** | 5×128 MLP | L2 Error | 0.0310 | **0.0172** | **-44.6%** |
+| **FashionMNIST** | Transformer | Accuracy | 89.46% | **89.66%** | **+0.20%** (p=0.0005) |
+| FashionMNIST | MLP | Accuracy | 89.10% | 89.12% | +0.02% (n.s.) |
+| FashionMNIST | CNN | Accuracy | 91.15% | 91.14% | -0.01% (tie) |
+
+> **LAdam excels on architectures with spatially-correlated weight structure** — particularly PINNs and transformers. For CNNs (whose conv filters are already spatial detectors), the Laplacian is redundant.
+
+## Installation
+
+```bash
+pip install ladam
+```
+
+## Optimizers
+
+LAdam ships three Laplacian-enhanced optimizers:
+
+| Optimizer | Base | Laplacian target | Best for |
+|-----------|------|------------------|----------|
+| **LAdam** | Adam | Second moment v_t | PINNs, transformers, CNNs |
+| **LAdaGrad** | AdaGrad | Cumulative sum G_t | Sparse features, NLP |
+| **LRMSProp** | RMSProp | Running average v_t | RNNs, non-stationary losses |
+
+All three share the same Laplacian kernel infrastructure and `c2` parameter.
+
+## Usage
+
+### Basic — Drop-in Adam replacement
+
+```python
+from ladam import LAdam
+
+optimizer = LAdam(model.parameters(), lr=1e-3, c2=1e-4)
+
+# Training loop is identical to Adam
+for batch in dataloader:
+    loss = criterion(model(batch))
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+```
+
+### LAdaGrad and LRMSProp
+
+```python
+from ladam import LAdaGrad, LRMSProp
+
+# AdaGrad with Laplacian smoothing on cumulative squared gradients
+optimizer = LAdaGrad(model.parameters(), lr=1e-2, c2=1e-4)
+
+# RMSProp with Laplacian smoothing on running variance
+optimizer = LRMSProp(model.parameters(), lr=1e-2, alpha=0.99, c2=1e-4)
+```
+
+### Per-layer c2 with parameter groups
+
+```python
+optimizer = LAdam([
+    {'params': model.attention.parameters(), 'c2': 1e-4},   # Transformer attention
+    {'params': model.ffn.parameters(), 'c2': 1e-5},         # Feed-forward
+    {'params': model.norm.parameters(), 'c2': 0.0},         # Skip for norms
+], lr=3e-4)
+```
+
+### Architecture-aware defaults
+
+```python
+from ladam import LAdam, suggest_c2
+
+c2 = suggest_c2('pinn')         # Returns 1e-5
+c2 = suggest_c2('transformer')  # Returns 1e-4
+
+optimizer = LAdam(model.parameters(), lr=1e-3, c2=c2)
+```
+
+## Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `lr` | 1e-3 | Learning rate |
+| `betas` | (0.9, 0.999) | EMA coefficients (same as Adam) |
+| `eps` | 1e-8 | Numerical stability (same as Adam) |
+| `weight_decay` | 0 | L2 regularization (same as AdamW behavior) |
+| `c2` | 1e-4 | **Laplacian coupling strength.** Controls how much neighboring variance estimates influence each other. |
+| `mode` | 'variance_lap' | Which quantity to smooth. `'variance_lap'` is best. |
+| `stencil` | '9point' | **Discrete Laplacian stencil.** `'9point'` (isotropic, 0.46% anisotropy) or `'5point'` (legacy, 12.3% anisotropy). |
+| `min_spatial_size` | 16 | Skip Laplacian for params with fewer elements (biases, LayerNorm). |
+
+### Stencil Selection
+
+The `stencil` parameter controls the discrete Laplacian kernel used for spatial coupling:
+
+- **`'9point'` (default)**: Isotropic stencil with face + edge neighbors. Treats diagonal neighbors with 1/6 weight vs 4/6 for face neighbors.
+- **`'5point'`**: Standard cross-pattern stencil (faces only). Slightly faster but 25× more anisotropic.
+
+At typical `c2` values (1e-5 to 1e-3), the effective learning rate difference between stencils is <0.3%. The 9-point default is recommended for correctness.
+
+### Choosing c2
+
+`c2` is the only new hyperparameter. It's robust across 3 orders of magnitude:
+
+| c2 | Best For | Notes |
+|----|----------|-------|
+| `1e-5` | PINNs, scientific ML | Gentle coupling, biggest error reduction |
+| `1e-4` | Transformers, general | **Safe default** |
+| `1e-3` | Aggressive smoothing | Works but slightly less stable |
+| `0` | Disable | Reduces to standard Adam |
+
+All 7 values tested in [1e-6, 1e-3] outperformed Adam on transformers (B12 sweep).
+
+## How It Works
+
+Standard Adam computes per-parameter adaptive learning rates from the second moment:
+
+```
+v_t = β₂·v_{t-1} + (1-β₂)·g_t²     # Variance estimate
+lr_effective = lr / (√v_t + ε)        # Per-parameter learning rate
+```
+
+LAdam adds a Laplacian coupling step:
+
+```
+v_smooth = v_t + c2 · ∇²v_t           # Spatial smoothing
+lr_effective = lr / (√v_smooth + ε)    # Coupled learning rate
+```
+
+Where `\nabla^2` is the discrete Laplacian computed via a single `F.conv2d` kernel (9-point isotropic by default) -- efficient and GPU-friendly. The Laplacian treats weight matrices as 2D fields, coupling each weight's learning rate with its spatial neighbors.
+
+**Overhead**: ~2-5% wall-clock time increase per step. The Laplacian is a single fused convolution kernel, not point-wise iteration.
+
+## Benchmarks
+
+### PINN: Wave Equation (u_tt = c^2 u_xx)
+
+5-layer, 128-unit tanh MLP trained for 5000 steps on the 1D wave equation.
+
+| Optimizer | L2 Error | vs Adam |
+|-----------|----------|---------|
+| Adam (lr=1e-3) | 0.0310 | — |
+| LAdam c²=1e-4 | 0.0240 | -22.8% |
+| **LAdam c²=1e-5** | **0.0172** | **-44.6%** |
+| LAdam c²=1e-3 | 0.0185 | -40.3% |
+
+### Transformer: FashionMNIST Classification
+
+4-head, 128-dim, 2-layer transformer, 30 epochs, 5 independent seeds.
+
+| Optimizer | Accuracy (mean ± std) | p-value (vs Adam) |
+|-----------|----------------------|-------------------|
+| Adam | 89.46 ± 0.10% | — |
+| **LAdam c²=1e-4** | **89.66 ± 0.06%** | **0.0005** |
+
+### c² Robustness Sweep
+
+7 c² values on the same transformer task. **All 7 beat Adam:**
+
+| c² | Accuracy | Δ vs Adam |
+|----|----------|-----------|
+| 1e-6 | 89.62% | +0.16% |
+| 5e-6 | 89.73% | +0.27% |
+| 1e-5 | 89.79% | +0.33% |
+| 5e-5 | 89.75% | +0.29% |
+| 1e-4 | 89.67% | +0.21% |
+| 5e-4 | 89.64% | +0.18% |
+| 1e-3 | 89.66% | +0.20% |
+
+## FAQ
+
+**Q: Does this work for LLMs / GPT-scale models?**
+A: No. LAdam **hurts** LLM training (tested on GPT-2/WikiText-2). Attention weight matrices encode semantic structure, not spatial structure — the Laplacian destroys per-feature specialization. Use standard Adam/AdamW for LLMs.
+
+**Q: Why not smooth the gradient instead of the variance?**
+A: [Osher et al. (2018)](https://arxiv.org/abs/1806.06317) explored Laplacian smoothing of gradients. We found that smoothing the *variance estimate* is more effective because it smooths the *learning rate landscape* rather than the *descent direction*. These are mathematically distinct: ∇²(EMA(g²)) ≠ (∇²g)².
+
+**Q: Why does this help PINNs so much?**
+A: PDE-based loss landscapes have inherent spatial structure from the differential operators in the loss function. The Laplacian on v_t aligns the optimizer's internal representation with this structure.
+
+**Q: Can I use this with learning rate schedulers?**
+A: Yes. LAdam is fully compatible with any `torch.optim.lr_scheduler`.
+
+## Citation
+
+If you use LAdam in your research, please cite:
+
+```bibtex
+@software{partin2026ladam,
+  author = {Partin, Greg},
+  title = {LAdam: Spatially-Aware Adaptive Optimization via Laplacian-Regularized Variance Estimates},
+  year = {2026},
+  url = {https://github.com/gpartin/ladam}
+}
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE) for details.

ladam-0.2.0/pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "ladam"
+version = "0.2.0"
+description = "LAdam: Laplacian Adam — Adam with spatially-coupled variance estimates via discrete Laplacian"
+readme = "README.md"
+license = {text = "MIT"}
+requires-python = ">=3.8"
+authors = [
+    {name = "Greg Partin"},
+]
+keywords = [
+    "optimizer", "adam", "deep-learning", "pytorch",
+    "laplacian", "pinn", "scientific-ml", "spatial-regularization",
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Mathematics",
+]
+dependencies = [
+    "torch>=1.10.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "pytest-cov>=4.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/gpartin/ladam"
+Documentation = "https://github.com/gpartin/ladam#usage"
+Issues = "https://github.com/gpartin/ladam/issues"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/ladam"]

ladam-0.2.0/src/ladam/__init__.py

@@ -0,0 +1,12 @@
+"""
+LAdam — Laplacian Adam Optimizer
+"""
+
+__version__ = "0.2.0"
+
+from .optimizer import LAdam, LAdaGrad, LRMSProp, suggest_c2
+
+# Backward compatibility alias
+WaveAdam = LAdam
+
+__all__ = ["LAdam", "LAdaGrad", "LRMSProp", "WaveAdam", "suggest_c2"]

ladam-0.2.0/src/ladam/optimizer.py

@@ -0,0 +1,442 @@
+"""
+LAdam — Laplacian Adam Optimizer
+================================
+
+A drop-in Adam replacement that applies discrete Laplacian regularization
+to the second moment estimate (v_t), producing spatially-smoothed adaptive
+learning rates.
+
+Key insight: adjacent weights in neural networks are often functionally
+correlated, but Adam computes independent per-parameter learning rates.
+LAdam restores spatial coherence by applying a Laplacian diffusion
+operator to Adam's variance estimate, allowing neighboring weights to
+share curvature information.
+
+This is particularly effective for:
+  - **PINNs** (physics-informed neural networks): -44.6% L2 error vs Adam
+  - **Transformers**: +0.20% accuracy, p=0.0005 across 5 seeds
+  - Any architecture with spatially-correlated weight structure
+
+Usage:
+    from ladam import LAdam
+
+    optimizer = LAdam(model.parameters(), lr=1e-3, c2=1e-4)
+
+    for batch in dataloader:
+        loss = model(batch)
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+"""
+
+import torch
+import torch.nn.functional as F
+from torch.optim import Optimizer
+
+
+class LAdam(Optimizer):
+    """
+    LAdam: Laplacian Adam — Adam with Laplacian-regularized variance estimates.
+
+    Applies a discrete Laplacian to Adam's second-moment estimate v_t,
+    coupling adjacent weight learning rates. This smooths the adaptive
+    learning rate landscape, improving convergence for architectures with
+    spatially-structured parameters.
+
+    Modes:
+        'variance_lap':  Laplacian on v_t (default, best overall)
+        'update_lap':    Laplacian on the update direction
+        'weight_lap':    Standard Adam step + weight-space diffusion
+
+    Args:
+        params: Model parameters or param groups.
+        lr (float): Learning rate. Default: 1e-3.
+        betas (Tuple[float, float]): EMA coefficients for (m_t, v_t).
+            Default: (0.9, 0.999).
+        eps (float): Numerical stability. Default: 1e-8.
+        weight_decay (float): L2 weight decay. Default: 0.
+        c2 (float): Laplacian coupling strength. Controls how much
+            neighboring variance estimates influence each other.
+            Recommended: 1e-5 for PINNs, 1e-4 for transformers.
+            Default: 1e-4.
+        mode (str): Which internal quantity to apply the Laplacian to.
+            Default: 'variance_lap'.
+        stencil (str): Laplacian stencil type for 2D parameters.
+            '9point': Isotropic stencil [1,4,1;4,-20,4;1,4,1]/6 — 0.46%
+                anisotropy. Default. Based on LFM lattice geometry
+                (face weight 2/3, diagonal weight 1/6).
+            '5point': Standard stencil [0,1,0;1,-4,1;0,1,0] — 12.3%
+                anisotropy. Legacy option for reproducing prior results.
+        min_spatial_size (int): Skip Laplacian for parameters with fewer
+            elements than this threshold. Small params (biases, norms)
+            lack meaningful spatial structure. Default: 16.
+    """
+
+    # Pre-built Laplacian kernels (cached on first use per device/stencil)
+    _lap_kernels_2d = {}
+    _lap_kernel_1d = None
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, c2=1e-4, mode='variance_lap',
+                 stencil='9point', min_spatial_size=16):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if eps < 0.0:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if c2 < 0.0:
+            raise ValueError(f"Invalid coupling strength: {c2}")
+        if mode not in ('variance_lap', 'update_lap', 'weight_lap'):
+            raise ValueError(f"Invalid mode: {mode}. Choose from: "
+                             "'variance_lap', 'update_lap', 'weight_lap'")
+        if stencil not in ('5point', '9point'):
+            raise ValueError(f"Invalid stencil: {stencil}. Choose from: "
+                             "'5point', '9point'")
+
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
+                        c2=c2, mode=mode, stencil=stencil,
+                        min_spatial_size=min_spatial_size)
+        super().__init__(params, defaults)
+
+    @classmethod
+    def _get_lap_kernel_2d(cls, device, dtype, stencil='9point'):
+        """Lazily create and cache the 2D discrete Laplacian kernel.
+
+        stencil='9point' (default): Isotropic [1,4,1;4,-20,4;1,4,1]/6
+            Isotropic stencil — faces weighted 2/3, diagonals 1/6.
+            Only 0.46% anisotropy vs 12.3% for 5-point.
+        stencil='5point': Standard [0,1,0;1,-4,1;0,1,0].
+            Standard cross-pattern stencil.
+        """
+        key = (stencil, device, dtype)
+        if key not in cls._lap_kernels_2d:
+            if stencil == '9point':
+                # Isotropic 2D Laplacian: face weight 4/6, diagonal weight 1/6
+                # Minimizes angular anisotropy (0.46% vs 12.3% for 5-point)
+                k = torch.tensor([[1/6., 4/6., 1/6.],
+                                  [4/6., -20/6., 4/6.],
+                                  [1/6., 4/6., 1/6.]], device=device, dtype=dtype)
+            else:  # 5point
+                k = torch.tensor([[0., 1., 0.],
+                                  [1., -4., 1.],
+                                  [0., 1., 0.]], device=device, dtype=dtype)
+            cls._lap_kernels_2d[key] = k.reshape(1, 1, 3, 3)
+        return cls._lap_kernels_2d[key]
+
+    @classmethod
+    def _get_lap_kernel_1d(cls, device, dtype):
+        """Lazily create and cache the 1D discrete Laplacian kernel."""
+        if cls._lap_kernel_1d is None or cls._lap_kernel_1d.device != device:
+            k = torch.tensor([1., -2., 1.], device=device, dtype=dtype)
+            cls._lap_kernel_1d = k.reshape(1, 1, 3)
+        return cls._lap_kernel_1d.to(dtype=dtype)
+
+    @staticmethod
+    def _laplacian(W, kernel_2d, kernel_1d):
+        """
+        Apply discrete Laplacian via fused conv2d — single GPU kernel launch.
+
+        Uses circular padding to treat weight matrices as periodic fields.
+        High-dimensional tensors are reshaped to 2D while preserving spatial
+        structure.
+        """
+        if W.dim() == 1:
+            n = W.numel()
+            if n < 4:
+                return torch.zeros_like(W)
+            x = W.reshape(1, 1, n)
+            x = F.pad(x, (1, 1), mode='circular')
+            return F.conv1d(x, kernel_1d).reshape(W.shape)
+        else:
+            shape = W.shape
+            W2 = W.reshape(shape[0], -1) if W.dim() > 2 else W
+            h, w = W2.shape
+            if h < 3 or w < 3:
+                return torch.zeros_like(W)
+            x = W2.reshape(1, 1, h, w)
+            x = F.pad(x, (1, 1, 1, 1), mode='circular')
+            result = F.conv2d(x, kernel_2d).reshape(W2.shape)
+            if W.dim() > 2:
+                result = result.reshape(shape)
+            return result
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Args:
+            closure (Callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group['lr']
+            beta1, beta2 = group['betas']
+            eps = group['eps']
+            c2 = group['c2']
+            mode = group['mode']
+            wd = group['weight_decay']
+            min_sp = group['min_spatial_size']
+
+            stencil = group['stencil']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if wd != 0:
+                    grad = grad.add(p.data, alpha=wd)
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['m'] = torch.zeros_like(p.data)
+                    state['v'] = torch.zeros_like(p.data)
+                    state['use_lap'] = (c2 > 0 and p.data.numel() >= min_sp)
+
+                state['step'] += 1
+                m, v = state['m'], state['v']
+                step = state['step']
+                use_lap = state['use_lap']
+
+                # Standard Adam momentum updates
+                m.mul_(beta1).add_(grad, alpha=1 - beta1)
+                v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+                # Bias correction
+                m_hat = m / (1 - beta1 ** step)
+                v_hat = v / (1 - beta2 ** step)
+
+                if mode == 'variance_lap':
+                    # Core innovation: couple neighboring variance estimates
+                    if use_lap:
+                        k2d = self._get_lap_kernel_2d(p.device, v_hat.dtype, stencil)
+                        k1d = self._get_lap_kernel_1d(p.device, v_hat.dtype)
+                        v_smooth = v_hat + c2 * self._laplacian(v_hat, k2d, k1d)
+                        v_smooth = v_smooth.clamp(min=eps)
+                    else:
+                        v_smooth = v_hat
+                    p.data.addcdiv_(m_hat, v_smooth.sqrt().add_(eps), value=-lr)
+
+                elif mode == 'update_lap':
+                    update = m_hat / (v_hat.sqrt() + eps)
+                    if use_lap:
+                        k2d = self._get_lap_kernel_2d(p.device, update.dtype, stencil)
+                        k1d = self._get_lap_kernel_1d(p.device, update.dtype)
+                        update = update + c2 * self._laplacian(update, k2d, k1d)
+                    p.data.add_(update, alpha=-lr)
+
+                elif mode == 'weight_lap':
+                    p.data.addcdiv_(m_hat, v_hat.sqrt().add_(eps), value=-lr)
+                    if use_lap:
+                        k2d = self._get_lap_kernel_2d(p.device, p.data.dtype, stencil)
+                        k1d = self._get_lap_kernel_1d(p.device, p.data.dtype)
+                        p.data.add_(self._laplacian(p.data, k2d, k1d), alpha=c2)
+
+        return loss
+
+
+def suggest_c2(architecture: str = 'auto') -> float:
+    """
+    Suggest a c2 value based on architecture type.
+
+    Based on systematic sweeps across 7 c2 values and 4 architecture types.
+
+    Args:
+        architecture: One of 'pinn', 'transformer', 'mlp', 'cnn', or 'auto'.
+            'auto' returns a safe default.
+
+    Returns:
+        Recommended c2 value.
+    """
+    recommendations = {
+        'pinn': 1e-5,         # -44.6% L2 error vs Adam
+        'transformer': 1e-4,  # +0.20% acc, p=0.0005
+        'mlp': 1e-5,          # Marginal improvement, safe value
+        'cnn': 0.0,           # No benefit; skip Laplacian
+        'auto': 1e-4,         # Safe general-purpose default
+    }
+    arch = architecture.lower()
+    if arch not in recommendations:
+        raise ValueError(f"Unknown architecture: {arch}. "
+                         f"Choose from: {list(recommendations.keys())}")
+    return recommendations[arch]
+
+
+class LAdaGrad(Optimizer):
+    """
+    LAdaGrad: AdaGrad with Laplacian-regularized accumulator.
+
+    Standard AdaGrad accumulates squared gradients independently per parameter.
+    LAdaGrad applies discrete Laplacian diffusion to the accumulator, letting
+    neighboring weights share curvature information.
+
+    Args:
+        params: Model parameters or param groups.
+        lr (float): Learning rate. Default: 1e-2.
+        lr_decay (float): Learning rate decay. Default: 0.
+        eps (float): Numerical stability. Default: 1e-10.
+        weight_decay (float): L2 weight decay. Default: 0.
+        c2 (float): Laplacian coupling strength. Default: 1e-4.
+        stencil (str): '9point' or '5point'. Default: '9point'.
+        min_spatial_size (int): Skip Laplacian below this size. Default: 16.
+    """
+
+    def __init__(self, params, lr=1e-2, lr_decay=0, eps=1e-10,
+                 weight_decay=0, c2=1e-4, stencil='9point',
+                 min_spatial_size=16):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if c2 < 0.0:
+            raise ValueError(f"Invalid coupling strength: {c2}")
+        defaults = dict(lr=lr, lr_decay=lr_decay, eps=eps,
+                        weight_decay=weight_decay, c2=c2,
+                        stencil=stencil, min_spatial_size=min_spatial_size)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group['lr']
+            lr_decay = group['lr_decay']
+            eps = group['eps']
+            c2 = group['c2']
+            wd = group['weight_decay']
+            min_sp = group['min_spatial_size']
+            stencil = group['stencil']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if wd != 0:
+                    grad = grad.add(p.data, alpha=wd)
+
+                state = self.state[p]
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['sum'] = torch.zeros_like(p.data)
+                    state['use_lap'] = (c2 > 0 and p.data.numel() >= min_sp)
+
+                state['step'] += 1
+                state['sum'].addcmul_(grad, grad, value=1)
+
+                clr = lr / (1 + (state['step'] - 1) * lr_decay)
+
+                acc = state['sum']
+                if state['use_lap']:
+                    k2d = LAdam._get_lap_kernel_2d(p.device, acc.dtype, stencil)
+                    k1d = LAdam._get_lap_kernel_1d(p.device, acc.dtype)
+                    acc_smooth = acc + c2 * LAdam._laplacian(acc, k2d, k1d)
+                    acc_smooth = acc_smooth.clamp(min=0)
+                else:
+                    acc_smooth = acc
+
+                p.data.addcdiv_(grad, acc_smooth.sqrt().add_(eps), value=-clr)
+
+        return loss
+
+
+class LRMSProp(Optimizer):
+    """
+    LRMSProp: RMSProp with Laplacian-regularized running average.
+
+    Standard RMSProp maintains a decaying average of squared gradients per
+    parameter. LRMSProp applies discrete Laplacian diffusion to this average,
+    coupling neighbors.
+
+    Args:
+        params: Model parameters or param groups.
+        lr (float): Learning rate. Default: 1e-2.
+        alpha (float): Smoothing constant (decay). Default: 0.99.
+        eps (float): Numerical stability. Default: 1e-8.
+        weight_decay (float): L2 weight decay. Default: 0.
+        momentum (float): Momentum factor. Default: 0.
+        c2 (float): Laplacian coupling strength. Default: 1e-4.
+        stencil (str): '9point' or '5point'. Default: '9point'.
+        min_spatial_size (int): Skip Laplacian below this size. Default: 16.
+    """
+
+    def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8,
+                 weight_decay=0, momentum=0, c2=1e-4, stencil='9point',
+                 min_spatial_size=16):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if alpha < 0.0 or alpha >= 1.0:
+            raise ValueError(f"Invalid alpha: {alpha}")
+        if c2 < 0.0:
+            raise ValueError(f"Invalid coupling strength: {c2}")
+        defaults = dict(lr=lr, alpha=alpha, eps=eps, weight_decay=weight_decay,
+                        momentum=momentum, c2=c2, stencil=stencil,
+                        min_spatial_size=min_spatial_size)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group['lr']
+            alpha = group['alpha']
+            eps = group['eps']
+            c2 = group['c2']
+            wd = group['weight_decay']
+            mom = group['momentum']
+            min_sp = group['min_spatial_size']
+            stencil = group['stencil']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad.data
+                if wd != 0:
+                    grad = grad.add(p.data, alpha=wd)
+
+                state = self.state[p]
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['square_avg'] = torch.zeros_like(p.data)
+                    if mom > 0:
+                        state['momentum_buffer'] = torch.zeros_like(p.data)
+                    state['use_lap'] = (c2 > 0 and p.data.numel() >= min_sp)
+
+                state['step'] += 1
+                sq_avg = state['square_avg']
+                sq_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha)
+
+                if state['use_lap']:
+                    k2d = LAdam._get_lap_kernel_2d(p.device, sq_avg.dtype, stencil)
+                    k1d = LAdam._get_lap_kernel_1d(p.device, sq_avg.dtype)
+                    avg_smooth = sq_avg + c2 * LAdam._laplacian(sq_avg, k2d, k1d)
+                    avg_smooth = avg_smooth.clamp(min=0)
+                else:
+                    avg_smooth = sq_avg
+
+                if mom > 0:
+                    buf = state['momentum_buffer']
+                    buf.mul_(mom).addcdiv_(grad, avg_smooth.sqrt().add_(eps))
+                    p.data.add_(buf, alpha=-lr)
+                else:
+                    p.data.addcdiv_(grad, avg_smooth.sqrt().add_(eps), value=-lr)
+
+        return loss

ladam-0.2.0/tests/test_optimizer.py

@@ -0,0 +1,182 @@
+"""Tests for LAdam optimizer."""
+
+import pytest
+import torch
+import torch.nn as nn
+from ladam import LAdam, suggest_c2
+
+
+class SimpleMLP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(20, 64)
+        self.fc2 = nn.Linear(64, 64)
+        self.fc3 = nn.Linear(64, 1)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        return self.fc3(x)
+
+
+@pytest.fixture
+def model():
+    return SimpleMLP()
+
+
+@pytest.fixture
+def data():
+    torch.manual_seed(42)
+    X = torch.randn(100, 20)
+    y = torch.randn(100, 1)
+    return X, y
+
+
+def test_basic_step(model, data):
+    """LAdam can perform a basic optimization step."""
+    X, y = data
+    opt = LAdam(model.parameters(), lr=1e-3, c2=1e-4)
+    loss_fn = nn.MSELoss()
+
+    loss_before = loss_fn(model(X), y).item()
+    for _ in range(10):
+        opt.zero_grad()
+        loss = loss_fn(model(X), y)
+        loss.backward()
+        opt.step()
+    loss_after = loss_fn(model(X), y).item()
+
+    assert loss_after < loss_before, "Loss should decrease after optimization"
+
+
+def test_c2_zero_equals_adam(model, data):
+    """With c2=0, LAdam should behave identically to Adam."""
+    X, y = data
+    loss_fn = nn.MSELoss()
+
+    # LAdam with c2=0
+    torch.manual_seed(0)
+    m1 = SimpleMLP()
+    opt1 = LAdam(m1.parameters(), lr=1e-3, c2=0.0)
+    for _ in range(5):
+        opt1.zero_grad()
+        loss_fn(m1(X), y).backward()
+        opt1.step()
+
+    # Standard Adam
+    torch.manual_seed(0)
+    m2 = SimpleMLP()
+    opt2 = torch.optim.Adam(m2.parameters(), lr=1e-3)
+    for _ in range(5):
+        opt2.zero_grad()
+        loss_fn(m2(X), y).backward()
+        opt2.step()
+
+    for p1, p2 in zip(m1.parameters(), m2.parameters()):
+        assert torch.allclose(p1, p2, atol=1e-6), \
+            "c2=0 LAdam should match standard Adam"
+
+
+def test_all_modes(model, data):
+    """All three modes should run without error and reduce loss."""
+    X, y = data
+    loss_fn = nn.MSELoss()
+
+    for mode in ['variance_lap', 'update_lap', 'weight_lap']:
+        torch.manual_seed(42)
+        m = SimpleMLP()
+        opt = LAdam(m.parameters(), lr=1e-3, c2=1e-4, mode=mode)
+
+        loss_before = loss_fn(m(X), y).item()
+        for _ in range(20):
+            opt.zero_grad()
+            loss_fn(m(X), y).backward()
+            opt.step()
+        loss_after = loss_fn(m(X), y).item()
+
+        assert loss_after < loss_before, f"Mode '{mode}' should reduce loss"
+
+
+def test_invalid_params():
+    """Invalid parameters should raise ValueError."""
+    m = SimpleMLP()
+    with pytest.raises(ValueError):
+        LAdam(m.parameters(), lr=-1)
+    with pytest.raises(ValueError):
+        LAdam(m.parameters(), c2=-1)
+    with pytest.raises(ValueError):
+        LAdam(m.parameters(), mode='invalid')
+    with pytest.raises(ValueError):
+        LAdam(m.parameters(), betas=(1.5, 0.999))
+
+
+def test_suggest_c2():
+    """Architecture-specific c2 suggestions."""
+    assert suggest_c2('pinn') == 1e-5
+    assert suggest_c2('transformer') == 1e-4
+    assert suggest_c2('cnn') == 0.0
+    assert suggest_c2('auto') == 1e-4
+    with pytest.raises(ValueError):
+        suggest_c2('unknown_arch')
+
+
+def test_param_groups(model, data):
+    """Per-group c2 values should work."""
+    X, y = data
+    loss_fn = nn.MSELoss()
+
+    opt = LAdam([
+        {'params': model.fc1.parameters(), 'c2': 1e-4},
+        {'params': model.fc2.parameters(), 'c2': 1e-5},
+        {'params': model.fc3.parameters(), 'c2': 0.0},
+    ], lr=1e-3)
+
+    for _ in range(10):
+        opt.zero_grad()
+        loss_fn(model(X), y).backward()
+        opt.step()
+
+
+def test_gpu_if_available(data):
+    """LAdam should work on GPU if available."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+    X, y = data
+    X, y = X.cuda(), y.cuda()
+    m = SimpleMLP().cuda()
+    opt = LAdam(m.parameters(), lr=1e-3, c2=1e-4)
+    loss_fn = nn.MSELoss()
+
+    for _ in range(5):
+        opt.zero_grad()
+        loss_fn(m(X), y).backward()
+        opt.step()
+
+
+def test_with_weight_decay(model, data):
+    """Weight decay should work alongside Laplacian."""
+    X, y = data
+    opt = LAdam(model.parameters(), lr=1e-3, c2=1e-4, weight_decay=0.01)
+    loss_fn = nn.MSELoss()
+
+    for _ in range(10):
+        opt.zero_grad()
+        loss_fn(model(X), y).backward()
+        opt.step()
+
+
+def test_closure(model, data):
+    """Closure-based step should work (for LBFGS-style usage)."""
+    X, y = data
+    opt = LAdam(model.parameters(), lr=1e-3, c2=1e-4)
+    loss_fn = nn.MSELoss()
+
+    def closure():
+        opt.zero_grad()
+        loss = loss_fn(model(X), y)
+        loss.backward()
+        return loss
+
+    loss = opt.step(closure)
+    assert loss is not None