Skip to main content

A JAX library for solving block tri-diagonal matrix systems

Project description

blocktrix

A JAX library for efficiently solving block tri-diagonal matrix systems

Author: Philip Mocz (@pmocz)

Implements both a block Thomas (serial) solver and a block cyclic reduction B-cyclic (parallel) solver.

⚠️ Warning: Work in Progress

This library is still under active development and is not guaranteed to work at this point XXX.

Installation

pip install blocktrix

Usage

import jax
from blocktrix import solve_block_tridiagonal_bcyclic, random_block_tridiagonal

# Generate a random test system
key = jax.random.PRNGKey(42)
n_blocks, block_size = 8, 4

lower, diag, upper, rhs = random_block_tridiagonal(key, n_blocks, block_size)

# Solve the system
x = solve_block_tridiagonal_bcyclic(n_blocks, lower, diag, upper, rhs)

Quickstart

It's fast!

Timing comparison

Speedup

Links

Cite this repository

If you use this software, please cite it as below.

XXX

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

blocktrix-0.0.0.tar.gz (8.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

blocktrix-0.0.0-py3-none-any.whl (8.0 kB view details)

Uploaded Python 3

File details

Details for the file blocktrix-0.0.0.tar.gz.

File metadata

  • Download URL: blocktrix-0.0.0.tar.gz
  • Upload date:
  • Size: 8.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.8

File hashes

Hashes for blocktrix-0.0.0.tar.gz
Algorithm Hash digest
SHA256 67c0e01ed287581b5760ad0ad5c89d774eaf982d6af5a03931f7c2f211e02f64
MD5 f75c5926cc26d7471aeb3e092ac1c8d0
BLAKE2b-256 8c5c27b2abc9fd1d5103b1639b47675bcc1fe05ffaa0299c3b4a3c339d952266

See more details on using hashes here.

File details

Details for the file blocktrix-0.0.0-py3-none-any.whl.

File metadata

  • Download URL: blocktrix-0.0.0-py3-none-any.whl
  • Upload date:
  • Size: 8.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.8

File hashes

Hashes for blocktrix-0.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 2cdda472fb49915f1f35e9d1542dc2ad557509a348fda8c6bdf9ab5ddc39c598
MD5 113f7225b4c07165424140be49c3e97b
BLAKE2b-256 b330f7f477429b82741eacae119c1c339c747578efd76f92a4e88f4fc3d82a02

See more details on using hashes here.

Supported by

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