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

pip install torch-friction

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

Usage

Single Step

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:

# 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:

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:

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():

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:

python -m unittest discover tests

Or run specific test files:

python -m unittest tests/test_edge_cases.py

Citation

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

@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