Folders and files

Name		Name	Last commit message	Last commit date
Latest commit History 16 Commits
.github/workflows		.github/workflows
examples		examples
tests		tests
thermox		thermox
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
LICENSE		LICENSE
README.md		README.md
pyproject.toml		pyproject.toml

Repository files navigation

thermox

This package provides a very simple interface toexactly simulateOrnstein-Uhlenbeck (OU) processes of the form

$$ dx = - A(x - b) dt + \mathcal{N}(0, D dt) $$

To collect samples from this process, define sampling timests, initial statex0, drift matrixA, displacement vectorb, diffusion matrixD and a JAX random key. Then runthermox.sample:

thermox.sample(key,ts,x0,A,b,D)

Samples are then collected by exact diagonalization (therefore there is no discretization error) and JAX scans.

You can access log-probabilities of the OU process by runningthermox.log_prob:

thermox.log_prob(ts,xs,A,b,D)

which can be useful for e.g. maximum likelihood estimation of the parametersA,b andD by composing withjax.grad.

Additionallythermox provides ascipy style suit ofthermodynamic linear algebra primitives:thermox.linalg.solve,thermox.linalg.inv,thermox.linalg.expm andthermox.linalg.negexpm which all simulate an OU process under the hood. More details can be found in thethermo_linear_algebra.ipynb notebook.

Contributing

Before submitting any pull request, make sure to runpre-commit run --all-files.

Example usage

Here is a simple code example for a 5-dimensional OU process:

importthermoximportjaximportjax.numpyasjnpimportmatplotlib.pyplotasplt# Set random seedkey=jax.random.PRNGKey(0)# Timeframedt=0.01ts=jnp.arange(0,1,dt)# System parameters for a 5-dimensional OU processA=jnp.array([[2.0,0.5,0.0,0.0,0.0],               [0.5,2.0,0.5,0.0,0.0],               [0.0,0.5,2.0,0.5,0.0],               [0.0,0.0,0.5,2.0,0.5],               [0.0,0.0,0.0,0.5,2.0]])b,x0=jnp.zeros(5),jnp.zeros(5)# Zero drift displacement vector and initial state# Diffusion matrix with correlations between x_1 and x_2D=jnp.array([[2,1,0,0,0],               [1,2,0,0,0],               [0,0,2,0,0],               [0,0,0,2,0],               [0,0,0,0,2]])# Collect samplessamples=thermox.sample(key,ts,x0,A,b,D)plt.figure(figsize=(12,5))plt.plot(ts,samples,label=[f'Dimension{i+1}'foriinrange(5)])plt.xlabel('Time')plt.ylabel('Value')plt.title('Trajectories of 5-Dimensional OU Process')plt.legend()plt.show()

Citation

If you usethermox in your research, please cite the library using the following BibTeX entry:

@misc{duffield2024thermox,title={thermox: Exact OU processes with JAX},author={Duffield, Samuel and Donatella, Kaelan and Melanson, Denis},howpublished={\url{https://github.com/normal-computing/thermox}},year={2024}}