Skip to main content

Solving tridiagonal systems.

Project description

Tests

tridiax

tridiax implements solvers for tridiagonal systems in jax. All solvers support CPU and GPU, are compatible with jit compilation and can be differentiated with grad.

Implemented solvers

Generally, Thomas algorithm will be faster on CPU whereas the divide and conquer algorithm and Stone's algorithm will be faster on GPU.

Known limitations

Currently, all solvers are only tested for systems whose dimensionality is an exponential of 2.

Usage

from tridiax import thomas_solve, divide_conquer_solve, stone_solve

dim = 1024
diag = jnp.asarray(np.random.randn(dim))
upper = jnp.asarray(np.random.randn(dim - 1))
lower = jnp.asarray(np.random.randn(dim - 1))
solve = jnp.asarray(np.random.randn(dim))
solution = thomas_solve(lower, diag, upper, solve)

If many systems of the same size are solved and the divide and conquer algorithm is used, it helps to precompute the reordering indizes:

from tridiax import divide_conquer_solve, divide_conquer_index

dim = 1024
diag = jnp.asarray(np.random.randn(dim))
upper = jnp.asarray(np.random.randn(dim - 1))
lower = jnp.asarray(np.random.randn(dim - 1))
solve = jnp.asarray(np.random.randn(dim))

indexing = divide_conquer_index(dim)
solution = divide_conquer_solve(lower, diag, upper, solve, indexing=indexing)

Installation

git clone https://github.com/mackelab/tridiax.git
cd tridiax
pip install -e .

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

tridiax-0.1.0.tar.gz (11.0 kB view details)

Uploaded Source

Built Distribution

tridiax-0.1.0-py3-none-any.whl (11.5 kB view details)

Uploaded Python 3

File details

Details for the file tridiax-0.1.0.tar.gz.

File metadata

  • Download URL: tridiax-0.1.0.tar.gz
  • Upload date:
  • Size: 11.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.11.8

File hashes

Hashes for tridiax-0.1.0.tar.gz
Algorithm Hash digest
SHA256 1eb84e1dcff299f2dfd440c229098500f745558209f7e25c3b3c95b6193e2411
MD5 f51f333b0cf48fd88b8948a322c9f07b
BLAKE2b-256 a5d77c18868dc684516b1853961a22ce1d77e4833f86e9403fd5e4e0f107b0b1

See more details on using hashes here.

File details

Details for the file tridiax-0.1.0-py3-none-any.whl.

File metadata

  • Download URL: tridiax-0.1.0-py3-none-any.whl
  • Upload date:
  • Size: 11.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.1 CPython/3.11.8

File hashes

Hashes for tridiax-0.1.0-py3-none-any.whl
Algorithm Hash digest
SHA256 9748032ed963a9c2c465d595559f5670aaf8fcea4040f3f45edcca89f3cbd16d
MD5 18fa69214d38ddde056f3e47bfd0adad
BLAKE2b-256 cbac8712692ce2478472eaef18ac710e7f0e4abb7ac60b3edd25a40f7a48208b

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page