Skip to content

common

Common module.

harm_osc_wavefunction(n, x, l_osc)

For given quantum number n=0,1,2,... return the value of the harmonic oscillator wave function :math:\psi_n(x) = N H_n(x/l_{osc}) \exp(-x^2/2l_\text{ osc}), N being the proper normalization factor.

Parameters

n: index of wave function, n=0 is ground state x: coordinate(s) where wave function is evaluated l_osc: oscillator length, defined via <0|x^2|0> = l_osc^2/2

Returns

value of harmonic oscillator wave function
Source code in jaxquantum/devices/common/utils.py
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def harm_osc_wavefunction(n, x, l_osc):
    r"""
    For given quantum number n=0,1,2,... return the value of the harmonic
    oscillator wave function :math:`\psi_n(x) = N H_n(x/l_{osc}) \exp(-x^2/2l_\text{
    osc})`, N being the proper normalization factor.

    Parameters
    ----------
    n:
        index of wave function, n=0 is ground state
    x:
        coordinate(s) where wave function is evaluated
    l_osc:
        oscillator length, defined via <0|x^2|0> = l_osc^2/2

    Returns
    -------
        value of harmonic oscillator wave function
    """
    return harm_osc_wavefunctions(n + 1, x, l_osc)[n]

harm_osc_wavefunctions(num_levels, x, l_osc)

Evaluate the first num_levels normalized oscillator wavefunctions.

Source code in jaxquantum/devices/common/utils.py
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def harm_osc_wavefunctions(num_levels, x, l_osc):
    """Evaluate the first ``num_levels`` normalized oscillator wavefunctions."""
    if num_levels < 1:
        raise ValueError("num_levels must be positive")

    coordinate = 2 * jnp.pi * jnp.asarray(x) / l_osc
    psi0 = jnp.exp(-(coordinate**2) / 2) / jnp.sqrt(l_osc * jnp.sqrt(jnp.pi))
    if num_levels == 1:
        return psi0[None]

    psi1 = jnp.sqrt(2.0) * coordinate * psi0

    def next_level(carry, level):
        previous, current = carry
        following = (
            jnp.sqrt(2.0 / (level + 1)) * coordinate * current
            - jnp.sqrt(level / (level + 1)) * previous
        )
        return (current, following), following

    _, remaining = lax.scan(
        next_level,
        (psi0, psi1),
        jnp.arange(1, num_levels - 1),
    )
    return jnp.concatenate((psi0[None], psi1[None], remaining), axis=0)