jTWA 0.0.4

Truncated Wigner on GPUs based on JAX.

jTWA

jTWA implements the semiclassical Truncated Wigner Approximation (TWA) in python, relying on Google's jax library. This allows to write easily understandable python code without compromising on speed, as all calculations are compiled and executed on GPUs, if available.

jTWA is designed to work with bosonic systems with potentially different numbers of internal degrees of freedom. However, basic functionalities for spin-1 systems are already available in jTWA.spin1.

Documentation & Installation

All code is documented on readthedocs. Working with jTWA is designed to be straightforward. For the installation procedure, have a look at the corresponding section in the documentation.

Getting Started: A Minimal Example

A minimal working example of the codebase can be found in the main.py located in the root of the directory. To execute, run python main.py config.json as described in the quickstart section of the documentation. To demonstrate the elementary features, here are the contents of the main.py:

import jax
import sys
import json
import matplotlib.pyplot as plt

jax.config.update("jax_enable_x64", True)

import jTWA


if __name__ == "__main__":
    configuration_file = sys.argv[1]

    with open(configuration_file) as f:
        cfg = json.load(f)

    spin_operators = jTWA.spin1.observables.get_spin_operators(cfg)
    samples = jTWA.spin1.initState.getPolarState(cfg)

    cfg = jTWA.spin1.hamiltonian.update_cfg(cfg)
    hamiltonian = jTWA.spin1.hamiltonian.hamiltonian

    obs = jTWA.integrate.obtain_evolution(samples, hamiltonian, spin_operators, cfg)
    jTWA.util.write_data(obs, cfg)

    obs = jTWA.util.read_data(cfg)
    jTWA.visualization.create_visuals(obs, cfg)
    plt.show()

Issues & Requests

If you encounter any issues, please open an issue or submit a pull request.

License

Apache License 2.0