A PyTorch library for learning multivariate Hawkes processes, with a focus on extensibility, speed, and scalability.
Overview · Quick Start · Organization · Documentation · Future Work · Citation
This library is designed around multivariate linear exponential Hawkes processes, which have intensity functions of the form:
| Symbol | Description |
|---|---|
| Base intensity of node |
|
| Influence of node |
|
| Decay rate of timescale |
|
| Number of events in the sequence | |
| Number of dimensions (nodes) in the process | |
| Number of kernels (timescales of influence) | |
| Time of the |
|
| Dimension (node type) of event |
We implement a highly parallel algorithm for fitting Hawkes processes via exact maximum likelihood estimation. This allows scaling to massive datasets; e.g., ~1,000 nodes and ~10,000,000 events. Therefore, a GPU or many CPU cores are recommended to realize these speedups.
Per-epoch time (left) and memory usage (right) on an Nvidia A100 GPU with
See our paper for experiment results and more details on the optimizations done for this library.
For a detailed worked example of simulation, fitting, and evaluation, see examples/full_example.ipynb.
The main prerequisite is PyTorch, ideally with GPU support, and matplotlib to run the plotting code. The numpy package is also required to run some of the examples.
Since this library is under active development, it will be updated frequently to add features and fix bugs. We recommend installing via pip directly from GitHub so you can easily pull in the latest updates:
pip install git+https://github.com/ahmrr/HawkesTorchAlternatively, you can add this repo as a submodule to your project if you want to modify the source.
Here is an example of simulating a Hawkes process and fitting a model to it. See examples/full_example.ipynb for a more detailed walkthrough.
import torch
from hawkes import models
from hawkes.utils import config
# Simulate data
M = 10 # number of nodes
N = 100000 # number of events
sim_gamma = torch.tensor([1.5]) # decay rate
sim_mu = torch.tensor([0.1] * M) # constant base rate
sim_alpha = torch.rand(1, M, M) * 0.5 # influence (K x M x M)
sim_base_process = models.Poisson(M, mu_init=sim_mu)
sim_model = models.Hawkes(
gamma=sim_gamma,
gamma_param=False,
base_process=sim_base_process,
alpha_init=sim_alpha,
)
seq = sim_model.simulate(max_events=N)
# Fit model
est_gamma = torch.tensor([1.0]) # initial decay rate
est_base_process = models.Poisson(M)
est_model = models.Hawkes(
gamma=est_gamma,
gamma_param=True, # allow learning the decay rate
base_process=est_base_process,
)
fit_config = config.FitConfig(
num_steps=4000, # epochs to train for
monitor_interval=400, # how often to print training progress
learning_rate=0.1, # learning rate for Adam optimizer
)
stats = est_model.fit(seq, fit_config) # returns dict of training detailsArrows indicate inheritance, diamonds indicate composition.
Most functionality is implemented in three main abstract base classes. Subclasses of these implement specific Hawkes and Poisson parameterizations and penalizations, which you can also extend to your own use cases.
-
HawkesBase: Provides intensity and likelihood computation, as well asfitandsimulatemethods. Subclasses implement custom parameterizations of$\textcolor{green}{\alpha_{p, q, k}}$ and$\textcolor{blue}{\gamma_k}$ , e.g., low-rank, sparse, etc. -
PoissonBase: Represents the base rate$\textcolor{red}{\mu_p(t)}$ and is used inHawkesBasevia thebase_processattribute. Also provides standalonefitandsimulatemethods, if raw Poisson fitting is desired. Subclasses implement custom parameterizations of the time-varying base rate. -
Penalty: An abstractdataclassthat represents a regularization term added to the log-likelihood during fitting. Extend this to create custom penalties beyond those already implemented. Penalties are attached to a model via thepenalizationattribute, which accepts adataclassthat storesPenaltyinstances for each parameter.
To implement your own Hawkes or Poisson process, or to add custom penalization, please reference:
- Files in
examples/templatesfor templates you can copy and modify. - The
HawkesBase,PoissonBase, andPenaltybase classes for core methods that are already implemented. - The
Hawkes,HawkesLowRank,Poisson,PoissonFourier, andNormPenaltysubclasses for concrete implementation examples.
There will be an API Reference available soon with detailed documentation and examples. In the meantime, please refer to the code, docstrings, examples, and paper.
-
Even faster likelihood computation via a custom CUDA implementation (using CUB
DeviceScan) -
More
PoissonBasemodels (e.g., piecewise constant, splines, etc.) -
API documentation
If you encounter a bug or have a feature request, please feel free to open an issue. Suggestions are also much appreciated!
If you find this library to be useful in your research, please cite the following paper:
@article{raza2026hawkestorch,
title = {Massively Parallel Exact Inference for Hawkes Processes},
author = {Raza, Ahmer and Smith, Hudson},
journal = {arXiv preprint arXiv:2604.01342},
year = 2026,
}
