# Wigner Kernels

This a collection of programs to benchmark Wigner kernels in JAX. Most notably, they have been used to produce benchmark results on rMD17.


## Prerequisites

The python requirements are listed in requirements.txt

In our programs, we also make use of a Fortran subroutine to calculate scaled modified spherical Bessel functions.
To compile the Fortran file: 

`python -m numpy.f2py -c fortran.f -m fortran`

(A Fortran compiler is required)


## Usage
This directory contains four scripts: `train_wk.py`, `train_wk_forces.py`, `train_wk_forces_hypers.py`, and `measure_latency.py`.

### train_wk.py
This script trains Wigner kernels without forces (only on energies).

### train_wk_forces_hypers.py
This script is used to explore the behavior of Wigner kernels as the two optimizable hyperparameters of the model are changed.

### train_wk_forces.py
This script executes a single fit with forces. It must be invoked with three command line arguments: the first is the dataset file,
the other two are the optimizable hyperparameters of the Wigner kernel model.

### measure_latency.py
This script measures the average execution time of energy + force calculations.


## Datasets

We provide the rMD17 datasets, but we cannot provide files for all 5 splits due to size constraints.
We only provide the first of these for each molecule in the dataset.
The fully rMD17 dataset can be downloaded from its online repository, along with its official
train/test splits.
