jTWA package#

Subpackages#

Submodules#

jTWA.integrate module#

jTWA.integrate.append_observables(key_to_use, samples, hamiltonian, spin_operators, stored_observables, t, cfg)#

Append observables to an existing array of observables in stored_observables.

Args:
  • key_to_use: A jax.random.PRNGKey that is used to mix each sample with Gaussian noise to simulate simultaneous readout.

  • samples: A list of samples. The array should be of shape (\(N_{samples}\), \(N_{wells}\), \(N_{internal}\)).

  • hamiltonian: The hamiltonian function that describes the system. Used to obtain energy expectation values.

  • spin_operators: A dictionary of the matrices that describe spin observables of interest along with their names.

  • stored_observables: A dictionary that is either empty or contains observable values from previous time steps.

  • t: The current simulation time.

  • cfg: The dictionary that contains the settings of the current run.

Returns:
  • stored_observables: The dictionary to which the new observables have been added.

jTWA.integrate.check_obs_save_condition(t, cfg, eps=0.1)#

Check whether observables should be computed and stored at the current time step.

Args:
  • t: The current simulation time.

  • cfg: The dictionary that contains the settings of the current run.

Returns:
  • A boolean indicating whether observables should be computed and stored or not.

jTWA.integrate.integrate(samples, flow, dt)#

(Compiled) Integrator that propagates the samples along their flow lines by the time step dt using a 4-th order RK-scheme described in integrate_single_sample().

Args:
  • samples: A list of samples. The array should be of shape (\(N_{samples}\), \(N_{wells}\), \(N_{internal}\)).

  • flow: The flow that is induced by the Hamiltonian.

  • dt: A small time step.

Returns:
  • samples: The propagated samples.

jTWA.integrate.integrate_single_sample(sample, flow, dt)#

4-th order RK-integrator that propagates a single sample along its flow by the time step dt.

Args:
  • sample: A single sample. The array should be of shape (\(N_{wells}\), \(N_{internal}\)).

  • flow: The flow that is induced by the Hamiltonian.

  • dt: A small time step.

Returns:
  • sample: The propagated sample.

jTWA.integrate.obtain_evolution(samples, hamiltonian, spin_operators, cfg)#

Stepper functionality to integrate the coupled equations of motion given by the Hamiltonian to the final time specified in cfg. First, the flow that is generated by the Hamiltonian is computed. Then, the samples are integrated along their flow lines using jTWA.integrate.integrate(). At times of interest that are separated by the interval cfg["simulationParameters"]["dtObs"], observables are read out and stored to a dictionary. This dictionary is returned once the final integration time is reached.

Args:
  • samples: A list of samples. The array should be of shape (\(N_{samples}\), \(N_{wells}\), \(N_{internal}\)).

  • hamiltonian: A function that returns the energy of a given single sample, given its complex conjugate, the sample itself as well as cfg.

  • spin_operators: A dictionary of the matrices that describe spin observables of interest along with their names.

  • cfg: The dictionary that contains the settings of the current run.

Returns:
  • obs: Dictionary that holds observables at specified times of interest.

jTWA.util module#

jTWA.util.read_data(cfg)#

Read the observables that are stored in the specified folder within cfg in pickle format.

Args:
  • cfg: The dictionary that contains the settings of the current run, including the working directory.

Returns:
  • obs: The dictionary of stored observables.

jTWA.util.write_data(obs, cfg)#

Write the computed observables to a pickle file in the dictionary that is specified in cfg.

Args:
  • obs: A dictionary holding all observables that are to be stored.

  • cfg: The dictionary that contains the settings of the current run, including the working directory.

jTWA.visualization module#

jTWA.visualization.animate(obs, cfg, fps=15)#

animate generates a video of the evolution of the Wigner distribution given by \(S_x\) and \(Q_{yz}\) over time.

Args:
  • obs: The observables that are to be used for visualizations.

  • cfg: The dictionary that contains the settings of the current run, including the working directory.

jTWA.visualization.create_visuals(obs, cfg)#

create_visuals is a wrapper function to call other plotting functionality.

Args:
  • obs: The observables that are to be used for visualizations.

  • cfg: The dictionary that contains the settings of the current run, including the working directory.

jTWA.visualization.plot_correlation_matrices(obs, cfg)#

plot_correlation_matrices plots correlation matrices of the two observables in obs["spin_obs"] at nine evenly distributed points in time. Currently this function assumes that obs["spin_obs"] only contains the observables \(S_x\) and \(Q_{yz}\).

Args:
  • obs: The observables that are to be used for visualizations.

  • cfg: The dictionary that contains the settings of the current run, including the working directory.

jTWA.visualization.plot_deviations(obs, cfg)#

plot_deviations plots sanity checks, making sure that conservations of particle number and energy are fulfilled.

Args:
  • obs: The observables that are to be used for visualizations.

  • cfg: The dictionary that contains the settings of the current run, including the working directory.

jTWA.visualization.plot_spin_obs(obs, cfg)#

plot_spin_obs plots the total side- and zero-mode populations as well as the different momentum-mode populations.

Args:
  • obs: The observables that are to be used for visualizations.

  • cfg: The dictionary that contains the settings of the current run, including the working directory.

Module contents#