Metadata-Version: 2.4
Name: nato-opt
Version: 0.1.3
Summary: NATO Optimizer with Fourier Spectral Penalty (FSP), Kakeya Directional Penalty, and N-D FFT gradient filtering
Author: Malhar Pangarkar, Atharva Khambete
Author-email: malharpangarkar19@gmail.com, atharvakhambete1@gmail.com
License: MIT
Project-URL: Homepage, https://github.com/Malhar1912/NATO
Project-URL: Repository, https://github.com/Malhar1912/NATO
Project-URL: Issues, https://github.com/Malhar1912/NATO/issues
Keywords: pytorch,optimizer,deep-learning,fourier,spectral-penalty,gradient-filtering
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
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
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: numpy
Dynamic: license-file

# NATO: Fourier Spectral Regularization

[![PyPI version](https://img.shields.io/pypi/v/nato-opt.svg)](https://pypi.org/project/nato-opt/)
[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

**NATO** — Fourier Spectral Regularization (FSR) with selective high-frequency penalty, N-D FFT gradient filtering, and experimental directional regularization for PyTorch.

## Features

- 📊 **Fourier Spectral Penalty (FSP)**: Selective high-frequency penalty (Definition 3.2) with configurable τ cutoff, hypercube & radial masks
- 🔧 **Low-Pass Gradient Filtering**: N-D FFT-based gradient smoothing with correct frequency-domain masking
- 🚀 **NATOOptimizer**: Custom Adam-variant optimizer with tethered updates *(experimental)*
- 🎯 **Kakeya Directional Penalty**: Gradient direction consistency regularization *(experimental)*
- ⚡ **GPU Accelerated**: Full CUDA support for all operations

## Installation

### From PyPI
```bash
pip install nato-opt
```

### From Source (Editable)
```bash
git clone https://github.com/Malhar1912/NATO.git
cd NATO
pip install -e .
```

## Quick Start

```python
import torch
import torch.nn as nn
from nato_opt import fourier_spectral_penalty, low_pass_filter_gradients

model = nn.Sequential(nn.Conv2d(3, 16, 3), nn.ReLU(), nn.Linear(16, 10))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for inputs, targets in dataloader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # Selective high-frequency penalty (Definition 3.2)
    # tau=2 penalizes only high-frequency weight components
    fsp = fourier_spectral_penalty(model, lambda_fsp=1e-6, tau=2, mask_mode="hypercube")
    total_loss = loss + fsp
    total_loss.backward()

    # Low-pass filter gradients before optimizer step (Theorem 9.1)
    low_pass_filter_gradients(model, tau_ratio=0.5)

    optimizer.step()
```

### Full-spectrum mode (ℓ₂ weight decay, Proposition 3.2)

```python
# tau=None (default) gives P_FSP^full = ||W||_F² — equivalent to weight decay
fsp = fourier_spectral_penalty(model, lambda_fsp=1e-6, tau=None)
```

## API Reference

### fourier_spectral_penalty

```python
fourier_spectral_penalty(
    model,
    lambda_fsp=1e-6,
    tau=None,             # None = full-spectrum (weight decay), float = selective HF
    mask_mode="hypercube", # "hypercube" (ℓ∞, Def 2.2) or "radial" (ℓ₂, §3.4)
    include_conv=True,
    include_linear=True,
    module_whitelist=None,
    module_blacklist=None,
    device=None
) -> torch.Tensor
```

Compute Fourier Spectral Penalty on model weights. The returned tensor retains `grad_fn` so autograd can differentiate through it.

**Parameters:**
- `model`: PyTorch model
- `lambda_fsp`: Penalty coefficient (default: 1e-6)
- `tau`: Frequency cutoff. `None` = full-spectrum (≡ weight decay). `float >= 0` = selective HF penalty.
- `mask_mode`: `"hypercube"` (ℓ∞ cutoff) or `"radial"` (ℓ₂ cutoff)
- `include_conv`: Include Conv layers (default: True)
- `include_linear`: Include Linear layers (default: True)
- `module_whitelist`: Only include these module names
- `module_blacklist`: Exclude these module names

**Returns:** Scalar penalty tensor (with grad_fn)

---

### low_pass_filter_gradients

```python
low_pass_filter_gradients(
    model,
    tau_ratio=0.5,    # ρ(τ): fraction of frequencies to keep per dimension
    skip_bias=True,
    in_place=True
)
```

Apply low-pass FFT filtering to gradients, smoothing high-frequency noise. Uses `fftshift`/`ifftshift` internally so the centered mask correctly selects near-DC frequencies.

**Parameters:**
- `model`: PyTorch model with computed gradients
- `tau_ratio`: Frequency retention ratio ρ(τ), 0 < tau_ratio ≤ 1. Lower = more filtering.

---

### adjust_learning_rate

```python
adjust_learning_rate(optimizer, epoch, ...)
```

Utility function for learning rate scheduling.

---

## Experimental Components

> **Note:** The following components are part of the broader DSR (Directional–Spectral Regularization) research hypothesis described in `DSR_Concept_Note.md`. They are **not** covered by the FSR paper's theoretical guarantees.

### NATOOptimizer

```python
NATOOptimizer(params, lr=1e-3, beta1=0.9, beta2=0.999,
              epsilon=1e-8, gamma=0.01, tether_interval=100, ...)
```

Custom Adam-variant optimizer with a tether term that penalizes drift from a periodic parameter checkpoint.

### kakeya_directional_penalty

```python
kakeya_directional_penalty(
    model,
    state,          # persistent dict for gradient history
    lambda_k=1e-4
) -> torch.Tensor
```

Penalizes gradients that maintain high cosine similarity with previous gradients.

**Usage (must be added to loss *before* `.backward()`):**
```python
kakeya_state = {}  # persistent across steps

for inputs, targets in dataloader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # Kakeya uses stored gradients from the PREVIOUS step
    k_penalty = kakeya_directional_penalty(model, kakeya_state)

    total_loss = loss + k_penalty
    total_loss.backward()
    optimizer.step()
```

## Requirements

- Python >= 3.8
- PyTorch
- NumPy

## License

MIT License - see [LICENSE](LICENSE) for details.

## Authors

- **Malhar Pangarkar** - [malharpangarkar19@gmail.com](mailto:malharpangarkar19@gmail.com)
- **Atharva Khambete** - [atharvakhambete1@gmail.com](mailto:atharvakhambete1@gmail.com)

## Citation

If you use this in your research, please cite:

```bibtex
@software{nato_opt,
  title = {NATO: Fourier Spectral Regularization},
  author = {Pangarkar, Malhar and Khambete, Atharva},
  year = {2026},
  url = {https://github.com/Malhar1912/NATO}
}
```

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

1. Fork the repository
2. Create your feature branch (`git checkout -b feature/AmazingFeature`)
3. Commit your changes (`git commit -m 'Add some AmazingFeature'`)
4. Push to the branch (`git push origin feature/AmazingFeature`)
5. Open a Pull Request
