# torch-friction

**Differentiable Rate-and-State Friction layers for PyTorch.**

This library implements the Rate-and-State friction laws (Dieterich, 1979; Ruina, 1983) using a **Runge-Kutta 4 (RK4)** solver for numerical stability. It is designed for geophysics research where you need to embed physical laws into learnable neural networks.

## Features

*   **Differentiable:** Gradients flow through the ODE solver to physical parameters (`a`, `b`, `Dc`).
*   **Stable:** Uses RK4 integration to handle stiffness in the state evolution equation.
*   **Robust:** Implements **Softplus Parameterization** for physical constants ($a, b, D_c$) to guarantee positivity and prevent divergence during training.
*   **Optimized:** Uses zero-copy pre-allocation for high-performance sequence simulation.

## Installation

```bash
pip install torch-friction

# for development (editable install):
# pip install -e .
```

## Usage

### Single Step
```python
import torch
from torch_friction import FrictionLayer

# Initialize the layer (law='aging' or 'slip')
layer = FrictionLayer(
    mu0=0.6, a=0.015, b=0.02, Dc=0.01, 
    law='slip', 
    seed=42
)

# Simulating a single step
velocity = torch.tensor([1e-6]) # m/s
state = torch.tensor([0.1])     # State variable 'theta' (seconds)
dt = 0.1                        # time step (seconds)

friction, next_state = layer(velocity, state, dt)
```

### High-Performance Trajectory
For simulating long time-series, use `forward_trajectory`:

```python
# Batch=1, Time=1000
velocity_sequence = torch.ones((1, 1000)) * 1e-6 
dt = 0.01

# frictions: [1, 1000], states: [1, 1000]
frictions, states = layer.forward_trajectory(velocity_sequence, dt)
```

## Simulating Earthquakes

To see the classic "Velocity Step" response (spike and decay), run the core verification test:

```bash
python tests/test_core.py
```

This will verify the instantaneous direct effect and long-term steady state decay against analytical solutions.

## Learning Friction (The Inverse Problem)

Can an AI learn friction parameters from noisy data? Yes. See the example in `examples/inverse_problem_demo.py`:

```bash
python examples/inverse_problem_demo.py
```

This script generates synthetic data with hidden parameters and trains a `FrictionLayer` to discover them from scratch.

### Advanced Usage: Freezing Parameters

When solving inverse problems, you may want to freeze some parameters while learning others. Use `freeze_parameters()`:

```python
layer = FrictionLayer(a=0.015, b=0.02, Dc=0.01)

# Freeze 'a' and 'Dc', only learn 'b' and 'mu0'
layer.freeze_parameters(['a', 'Dc'])

# IMPORTANT: Call freeze_parameters BEFORE creating optimizer
optimizer = torch.optim.Adam(layer.parameters(), lr=0.001)

# Now only 'b' and 'mu0' will be updated during training
```

Valid parameter names: `'a'`, `'b'`, `'Dc'`, `'mu0'`.

## Development

Run all tests:
```bash
python -m unittest discover tests
```

Or run specific test files:
```bash
python -m unittest tests/test_edge_cases.py
```

## Citation

If you use this library in your research, please cite:

```bibtex
@software{torch_friction_2026,
  author       = {Aman},
  title        = {AmanSinghNp/torch-friction: v0.1.0 - Initial Release of torch-friction},
  year         = {2026},
  publisher    = {Zenodo},
  version      = {v0.1.0},
  doi          = {10.5281/zenodo.18366445},
  url          = {https://doi.org/10.5281/zenodo.18366445}
}
```

Aman. (2026). AmanSinghNp/torch-friction: v0.1.0 - Initial Release of torch-friction (v0.1.0). Zenodo. https://doi.org/10.5281/zenodo.18366445
