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.
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.
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()If you encounter any issues, please open an issue or submit a pull request.