# thermox
This package provides a very simple interface to **exactly** simulate [Ornstein-Uhlenbeck](https://en.wikipedia.org/wiki/Ornstein%E2%80%93Uhlenbeck_process) (OU) processes of the form
$$ dx = - A(x - b) dt + \mathcal{N}(0, D dt) $$
To collect samples from this process, define sampling times `ts`, initial state `x0`, drift matrix `A`, displacement vector `b`, diffusion matrix `D` and a JAX random key. Then run `thermox.sample`:
```python
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 running `thermox.log_prob`:
```python
thermox.log_prob(ts, xs, A, b, D)
```
which can be useful for e.g. maximum likelihood estimation of the parameters `A`, `b` and `D` by composing with `jax.grad`.
Additionally `thermox` provides a [`scipy`](https://docs.scipy.org/doc/scipy/reference/linalg.html) style suit of [**thermodynamic linear algebra**](https://arxiv.org/abs/2308.05660) primitives: `thermox.linalg.solve`, `thermox.linalg.inv`, `thermox.linalg.expm` and `thermox.linalg.negexpm` which all simulate an OU process under the hood. More details can be found in the [`thermo_linear_algebra.ipynb`](/thermo_linear_algebra.ipynb) notebook.
## Contributing
Before submitting any pull request, make sure to run `pre-commit run --all-files`.
## Example usage
Here is a simple code example for a 5-dimensional OU process:
```python
import thermox
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
# Set random seed
key = jax.random.PRNGKey(0)
# Timeframe
dt = 0.01
ts = jnp.arange(0, 1, dt)
# System parameters for a 5-dimensional OU process
A = 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_2
D = 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 samples
samples = thermox.sample(key, ts, x0, A, b, D)
plt.figure(figsize=(12, 5))
plt.plot(ts, samples, label=[f'Dimension {i+1}' for i in range(5)])
plt.xlabel('Time')
plt.ylabel('Value')
plt.title('Trajectories of 5-Dimensional OU Process')
plt.legend()
plt.show()
```
<p align="center">
<img src="https://storage.googleapis.com/normal-blog-artifacts/thermox/ou_trajectories.png" width="800" lineheight = -10%/>
<br>
</p>
# Citation
If you use `thermox` in your research, please cite the library using the following BibTeX entry:
```bibtex
@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}
}
```
Raw data
{
"_id": null,
"home_page": null,
"name": "thermox",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "jax, linear algebra, ou process, stochastic process",
"author": null,
"author_email": "Sam Duffield <sam@normalcomputing.ai>, Kaelan Donatella <kaelan@normalcomputing.ai>",
"download_url": "https://files.pythonhosted.org/packages/d2/22/4cfb4ca9af51e5ecb65aca5af919420284b9a9549aab9f2b2c1f0b0b1108/thermox-0.0.3.tar.gz",
"platform": null,
"description": "# thermox\n\nThis package provides a very simple interface to **exactly** simulate [Ornstein-Uhlenbeck](https://en.wikipedia.org/wiki/Ornstein%E2%80%93Uhlenbeck_process) (OU) processes of the form \n\n$$ dx = - A(x - b) dt + \\mathcal{N}(0, D dt) $$\n\nTo collect samples from this process, define sampling times `ts`, initial state `x0`, drift matrix `A`, displacement vector `b`, diffusion matrix `D` and a JAX random key. Then run `thermox.sample`:\n\n```python\nthermox.sample(key, ts, x0, A, b, D) \n```\nSamples are then collected by exact diagonalization (therefore there is no discretization error) and JAX scans.\n\nYou can access log-probabilities of the OU process by running `thermox.log_prob`:\n\n```python\nthermox.log_prob(ts, xs, A, b, D)\n```\n\nwhich can be useful for e.g. maximum likelihood estimation of the parameters `A`, `b` and `D` by composing with `jax.grad`.\n\nAdditionally `thermox` provides a [`scipy`](https://docs.scipy.org/doc/scipy/reference/linalg.html) style suit of [**thermodynamic linear algebra**](https://arxiv.org/abs/2308.05660) primitives: `thermox.linalg.solve`, `thermox.linalg.inv`, `thermox.linalg.expm` and `thermox.linalg.negexpm` which all simulate an OU process under the hood. More details can be found in the [`thermo_linear_algebra.ipynb`](/thermo_linear_algebra.ipynb) notebook.\n\n## Contributing\n\nBefore submitting any pull request, make sure to run `pre-commit run --all-files`.\n\n\n## Example usage\n\nHere is a simple code example for a 5-dimensional OU process:\n```python\nimport thermox\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\n\n# Set random seed\nkey = jax.random.PRNGKey(0)\n\n# Timeframe\ndt = 0.01\nts = jnp.arange(0, 1, dt)\n\n# System parameters for a 5-dimensional OU process\nA = jnp.array([[2.0, 0.5, 0.0, 0.0, 0.0],\n [0.5, 2.0, 0.5, 0.0, 0.0],\n [0.0, 0.5, 2.0, 0.5, 0.0],\n [0.0, 0.0, 0.5, 2.0, 0.5],\n [0.0, 0.0, 0.0, 0.5, 2.0]])\n\nb, x0 = jnp.zeros(5), jnp.zeros(5) # Zero drift displacement vector and initial state\n\n # Diffusion matrix with correlations between x_1 and x_2\nD = jnp.array([[2, 1, 0, 0, 0],\n [1, 2, 0, 0, 0],\n [0, 0, 2, 0, 0],\n [0, 0, 0, 2, 0],\n [0, 0, 0, 0, 2]])\n\n# Collect samples\nsamples = thermox.sample(key, ts, x0, A, b, D)\n\nplt.figure(figsize=(12, 5))\nplt.plot(ts, samples, label=[f'Dimension {i+1}' for i in range(5)])\nplt.xlabel('Time')\nplt.ylabel('Value')\nplt.title('Trajectories of 5-Dimensional OU Process')\nplt.legend()\nplt.show()\n```\n\n<p align=\"center\">\n <img src=\"https://storage.googleapis.com/normal-blog-artifacts/thermox/ou_trajectories.png\" width=\"800\" lineheight = -10%/>\n <br>\n</p>\n\n\n# Citation\nIf you use `thermox` in your research, please cite the library using the following BibTeX entry:\n\n```bibtex\n@misc{duffield2024thermox,\n title={thermox: Exact OU processes with JAX},\n author={Duffield, Samuel and Donatella, Kaelan and Melanson, Denis},\n howpublished={\\url{https://github.com/normal-computing/thermox}},\n year={2024}\n}\n```\n",
"bugtrack_url": null,
"license": "Apache-2.0",
"summary": "Exact OU processes with JAX",
"version": "0.0.3",
"project_urls": null,
"split_keywords": [
"jax",
" linear algebra",
" ou process",
" stochastic process"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "4050b1d0b87611d46be37acaa3b523246e5991b75d6e9697a62eeaa62d3edbbe",
"md5": "cedae3c6151cc10e844c79c1833dd080",
"sha256": "05983204441506021ce90af4890a8d731fe7b4a58216d682f1df3cf2ab02bbd8"
},
"downloads": -1,
"filename": "thermox-0.0.3-py3-none-any.whl",
"has_sig": false,
"md5_digest": "cedae3c6151cc10e844c79c1833dd080",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 13388,
"upload_time": "2024-09-24T09:57:00",
"upload_time_iso_8601": "2024-09-24T09:57:00.971512Z",
"url": "https://files.pythonhosted.org/packages/40/50/b1d0b87611d46be37acaa3b523246e5991b75d6e9697a62eeaa62d3edbbe/thermox-0.0.3-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "d2224cfb4ca9af51e5ecb65aca5af919420284b9a9549aab9f2b2c1f0b0b1108",
"md5": "ccbc2430fc1af6b9f805540b5956000b",
"sha256": "1e6c66b601dc5f909a6a7ce5f101c25a9d0b7192f766cba3e2548910756966d8"
},
"downloads": -1,
"filename": "thermox-0.0.3.tar.gz",
"has_sig": false,
"md5_digest": "ccbc2430fc1af6b9f805540b5956000b",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 14700,
"upload_time": "2024-09-24T09:57:02",
"upload_time_iso_8601": "2024-09-24T09:57:02.176997Z",
"url": "https://files.pythonhosted.org/packages/d2/22/4cfb4ca9af51e5ecb65aca5af919420284b9a9549aab9f2b2c1f0b0b1108/thermox-0.0.3.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-09-24 09:57:02",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "thermox"
}
JSON