Using precompute tables

Note: this example expects precomputed lookup-table files named single_bragg_sig0.260.h5 and multi_bragg_sig0.260.h5 to be present in the notebook working directory. Generate or copy those tables before executing the notebook.

First lets define our imports, the lookup table file paths, and the parameters of the precompute Bragg pulses. Then we will load in the actual precompute tables.

[1]:
import numpy as np
from numba import jit, float64
from mwave.integrate import make_kvec, make_phi, propagate, phase_fnc_constant
from mwave.precompute import load_precomputed_gbragg

# Define parameters used to generate precompute table, these can optionally be passed into the load_precomputed_gbragg function
nbragg = 4
sigma = 0.259658916
tau_factor = 5

# Make a kvector
kvec, _, _ = make_kvec(0,nbragg)

# Try loading just the single frequency precompute table, then the combined single and multifrequency one
lookup_single_fname = 'single_bragg_sig0.260.h5'
lookup_multi_fname = 'multi_bragg_sig0.260.h5'

pgbragg_single = load_precomputed_gbragg(lookup_single_fname, table_sigma=sigma)
pgbragg_both = load_precomputed_gbragg(lookup_single_fname, lookup_multi_fname, table_sigma=sigma, table_modulation_frequency=8*nbragg)
Loading single frequency Bragg precompute table, this could take a while...
Precompute table loaded! Performing checks...
Checks passed!
Loading single frequency Bragg precompute table, this could take a while...
Precompute table loaded! Performing checks...
Checks passed!
Loading multifrequency Bragg precompute table, this could take a while...
Precompute table loaded! Performing checks...
Checks passed!

Now that we have loaded in our precompute tables, lets construct a function that takes the same arguments as pgbragg_single and pgbragg_both but directly computes the output wavefunction instead of using the lookup table.

[2]:
def gbragg(kvec, k0, sigma, omega, delta, delta_phase, mod_freq=None, mod_phase=0.0):

    # To stop edge effects effecting us when we compare to the precompute table we should actually use a much larger kvec
    inner_kvec = make_kvec(0,0,npad=50)[0]

    tfinal = 2*tau_factor*sigma

    if mod_freq:
        # mod_freq not None, use multi-Bragg

        @jit(float64(float64, float64[:]))
        def multi_omega_fnc(t, args):
            omega, sigma, t0, mod_freq, mod_phase = args
            return 2*np.cos(mod_freq*t + mod_phase)*omega*np.exp(-np.square(t-t0)/(2*(sigma**2)))

        sol = propagate(inner_kvec, make_phi(inner_kvec, k0/2), tfinal, delta, multi_omega_fnc, np.array([omega, sigma, tfinal/2, mod_freq, mod_phase]), phase_fnc_constant, np.array([delta_phase]))

        return sol.phi_final[np.isin(inner_kvec, kvec)]

    else:
        # mod_freq is None, use single Bragg

        @jit(float64(float64, float64[:]))
        def omega_fnc_gaussian(t, args):
            omega, sigma, t0 = args
            return omega*np.exp(-np.square(t-t0)/(2*(sigma**2)))

        sol = propagate(inner_kvec, make_phi(inner_kvec, k0/2), tfinal, delta, omega_fnc_gaussian, np.array([omega, sigma, tfinal/2]), phase_fnc_constant, np.array([delta_phase]))

        return sol.phi_final[np.isin(inner_kvec, kvec)]

Now that we have the direct and precomputed functions we and print out the differences between them. Lets start with a single frequency pulse.

[11]:
# Randomly sample parameters to cehck precompute table against
omega = np.random.rand()*45
k0 = 2*int(np.random.rand()*10-5)
delta = 4*nbragg + np.random.rand()*8 - 4
delta_phase = np.random.rand()*2*np.pi

# Compare direct computation to precomputed values
gb = gbragg(kvec, k0, sigma, omega, delta, delta_phase, 0.0, 0.0)
gbp = pgbragg_single(kvec, k0, sigma, np.array([omega]), np.array([delta]), delta_phase)
gbbp = pgbragg_both(kvec, k0, sigma, np.array([omega]), np.array([delta]), delta_phase)

# Print absolute differences
print('directly computed populations')
print(np.abs(gb[:])**2)
print('pgbragg_single')
print(f'phase deviations err {np.angle(gb[:])-np.angle(gbp[0,:])}')
print(f'population err {np.abs(gb[:])**2-np.abs(gbp[0,:])**2}')
print('\npgbragg_both')
print(f'phase deviations err {np.angle(gb[:])-np.angle(gbbp[0,:])}')
print(f'population err {np.abs(gb[:])**2-np.abs(gbbp[0,:])**2}')
directly computed populations
[7.73693526e-28 1.64408211e-26 6.69012041e-26 5.04494003e-25
 1.32733681e-24 6.56453256e-25 4.77292244e-24 3.74553216e-15
 5.11763100e-06 1.71478211e-03 4.99429309e-01 1.22978914e-01
 6.17643911e-02 5.66803520e-02 2.42007799e-01 1.54192069e-02
 1.26301855e-07 4.66194478e-17 4.93461050e-26 3.71251858e-26
 3.61496358e-26 2.61121489e-26 2.13290728e-26 3.63420981e-27
 7.39412202e-28]
pgbragg_single
phase deviations err [ 3.46198814e+00  4.40103642e+00  3.09117904e+00  6.20975992e-01
  2.51852691e-01  5.67958471e-02  1.43584660e-01 -2.71233819e-05
  2.57268540e-06 -2.72409143e-06  1.04484705e-06 -1.43832199e-07
  1.42969206e-06 -2.52064104e-06 -1.46709459e-07 -4.04860789e-06
  4.88389526e-05  4.74291736e-05  2.93782326e+00  2.19176450e+00
 -4.17334862e+00  1.85944392e+00 -6.60228378e-01 -4.38346450e+00
 -1.69580328e+00]
population err [-1.30244564e-27 -3.65248032e-24 -2.10536952e-25  1.11149705e-25
  1.83284151e-25  7.33429647e-26 -2.94150736e-25 -1.90206410e-19
  3.87883412e-11  3.40942213e-09 -2.40721744e-06 -2.54747354e-07
  9.49739348e-08  3.87778559e-07  3.19901486e-06  4.14241384e-07
  8.88351312e-12 -4.70473960e-21 -7.01075715e-26 -3.66908404e-28
  3.45642530e-27  1.40626575e-26  2.09893035e-26 -1.13502310e-26
  7.39412202e-28]

pgbragg_both
phase deviations err [ 3.46198814e+00  4.40103642e+00  3.09117904e+00  6.20975992e-01
  2.51852691e-01  5.67958471e-02  1.43584660e-01 -2.71233819e-05
  2.57268540e-06 -2.72409143e-06  1.04484705e-06 -1.43832199e-07
  1.42969206e-06 -2.52064104e-06 -1.46709459e-07 -4.04860789e-06
  4.88389526e-05  4.74291736e-05  2.93782326e+00  2.19176450e+00
 -4.17334862e+00  1.85944392e+00 -6.60228378e-01 -4.38346450e+00
 -1.69580328e+00]
population err [-1.30244564e-27 -3.65248032e-24 -2.10536952e-25  1.11149705e-25
  1.83284151e-25  7.33429647e-26 -2.94150736e-25 -1.90206410e-19
  3.87883412e-11  3.40942213e-09 -2.40721744e-06 -2.54747354e-07
  9.49739348e-08  3.87778559e-07  3.19901486e-06  4.14241384e-07
  8.88351312e-12 -4.70473960e-21 -7.01075715e-26 -3.66908404e-28
  3.45642530e-27  1.40626575e-26  2.09893035e-26 -1.13502310e-26
  7.39412202e-28]

This looks good. We have small population errors across the board, and our phase errors are only large when we have zero (i.e. 1e-12) population.

Now lets test a multi-frequency pulse.

[4]:
# Randomly sample parameters to cehck precompute table against
omega = np.random.rand()*45
k0 = 2*int(np.random.rand()*10-5)
nbloch = 0
delta = 4*nbragg + np.random.rand()*8 - 4
delta_phase = np.random.rand()*2*np.pi
mod_freq = 8*nbragg
mod_phase = 0

# Compare direct computation to precomputed values# Randomly sample parameters to cehck precompute table against

gb = gbragg(kvec, k0, sigma, omega, delta, delta_phase, mod_freq, mod_phase)
gbp = pgbragg_both(kvec, k0, sigma, np.array([omega]), np.array([delta]), delta_phase, mod_freq, mod_phase)

# Print absolute differences
print('directly computed populations')
print(np.abs(gb[:])**2)
print(f'phase deviations err {np.angle(gb[:])-np.angle(gbp[0,:])}')
print(f'population err {np.abs(gb[:])**2-np.abs(gbp[0,:])**2}')
directly computed populations
[6.51229483e-27 1.42337890e-25 1.03835926e-23 8.99007802e-17
 4.97042019e-11 1.74890545e-06 2.98260346e-02 5.18788503e-01
 1.27193768e-01 2.50302236e-01 6.70009950e-02 2.90726271e-03
 2.32270533e-03 1.48546635e-03 4.89464087e-05 2.55653484e-05
 3.62237932e-06 9.45748110e-06 6.93521332e-05 1.42954031e-05
 3.76057016e-08 2.84426433e-12 9.74848077e-18 1.02040308e-22
 3.00026708e-29]
phase deviations err [-4.26092646e+00  4.57348857e-01 -2.04220612e-01  1.27582364e-03
  6.75143828e-04  1.06487428e-05  5.11554111e-05 -3.62033632e-06
 -7.65706056e-07 -2.72753603e-05  1.37881045e-04  3.85593508e-04
 -5.90606717e-04  3.06754976e-03  4.53087003e-03  5.34724102e-03
 -1.09391735e-02  6.21637377e-02  1.34753156e-02  1.47683890e-01
  3.39391490e-01  3.71658534e-01  8.64316677e-01 -5.05147400e+00
  1.13114902e+00]
population err [-8.07988713e-28  4.01502118e-26 -2.50159970e-25  4.02479076e-19
  7.30890560e-14  6.05804583e-10  2.53960916e-06  6.50810047e-06
  2.56293970e-07  2.21832720e-05  6.17576041e-05  7.59193636e-06
  1.02596731e-05 -9.32421292e-07  1.11665033e-06  1.22015972e-06
  3.83786712e-07  2.04471964e-06  2.19419212e-05  5.15005244e-06
  1.19979719e-08  2.12826809e-12  2.16218882e-18  5.71925680e-23
  6.70548038e-30]

Same deal–our phase errors are only large when we have zero (i.e. 1e-12) population.