# 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>
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/cc/11/167b4cff07c20065e2d21f2ca2c47ae7d65b7bbfd9a952a8a9a83bf3bbda/thermox-0.0.1.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",
"bugtrack_url": null,
"license": "Apache-2.0",
"summary": "OU Processes and Linear Algebra with JAX",
"version": "0.0.1",
"project_urls": null,
"split_keywords": [
"jax",
" linear algebra",
" ou process",
" stochastic process"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "f7c390266ab73e1adc51988ac76c95bfbfcc297325e5b61e57b7ab3c8df6bb6c",
"md5": "31044b1eb47ac081816eb69ba255d6f2",
"sha256": "fcfb720469705822aeac0f7dc902d41441c7251fe1612d98c1e8952fe23d547a"
},
"downloads": -1,
"filename": "thermox-0.0.1-py3-none-any.whl",
"has_sig": false,
"md5_digest": "31044b1eb47ac081816eb69ba255d6f2",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.9",
"size": 11885,
"upload_time": "2024-05-29T16:26:16",
"upload_time_iso_8601": "2024-05-29T16:26:16.016521Z",
"url": "https://files.pythonhosted.org/packages/f7/c3/90266ab73e1adc51988ac76c95bfbfcc297325e5b61e57b7ab3c8df6bb6c/thermox-0.0.1-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "cc11167b4cff07c20065e2d21f2ca2c47ae7d65b7bbfd9a952a8a9a83bf3bbda",
"md5": "d3399e1f937c870f47c61175cbb7423c",
"sha256": "a483d514893464bb825fbe0c00facf1634d0fc2c5c87f521e76aff2513d04d66"
},
"downloads": -1,
"filename": "thermox-0.0.1.tar.gz",
"has_sig": false,
"md5_digest": "d3399e1f937c870f47c61175cbb7423c",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.9",
"size": 13255,
"upload_time": "2024-05-29T16:26:17",
"upload_time_iso_8601": "2024-05-29T16:26:17.047732Z",
"url": "https://files.pythonhosted.org/packages/cc/11/167b4cff07c20065e2d21f2ca2c47ae7d65b7bbfd9a952a8a9a83bf3bbda/thermox-0.0.1.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-05-29 16:26:17",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "thermox"
}
JSON