import numpy as np
import h5py
import include.phase_imprinting

"""Turbulence simulation for a scalar BEC. Vortices are generated by via phase imprinting."""

# --------------------------------------------------------------------------------------------------------------------
# Controlled variables:
# --------------------------------------------------------------------------------------------------------------------
Nx, Ny = 1024, 1024
Mx, My = Nx // 2, Ny // 2  # Number of grid pts
dx = dy = 1  # Grid spacing
dkx = np.pi / (Mx * dx)
dky = np.pi / (My * dy)  # K-space spacing
len_x = Nx * dx  # Box length
len_y = Ny * dy
x = np.arange(-Mx, Mx) * dx
y = np.arange(-My, My) * dy
X, Y = np.meshgrid(x, y)  # Spatial meshgrid

# k-space arrays and meshgrid:
kx = np.arange(-Mx, Mx) * dkx
ky = np.arange(-My, My) * dky
Kx, Ky = np.meshgrid(kx, ky)  # K-space meshgrid
Kx, Ky = np.fft.fftshift(Kx), np.fft.fftshift(Ky)

# Controlled variables
V = 0.  # Doubly periodic box
c0 = 3e-5

# Time steps, number and wavefunction save variables
Nt = 10000000
Nframe = 20000   # Save data every Nframe number of timesteps
dt = 1e-2  # Imaginary time timestep
t = 0.
k = 0   # Array index

filename = 'scalar_imp_noise_cpu'    # Name of file to save data to
data_path = '../scratch/data/scalar/{}.hdf5'.format(filename)
backup_data_path = '../scratch/data/scalar/backups/{}_backup.hdf5'.format(filename)

fresh_simulation = True  # Boolean that corresponds to a fresh simulation if True or a continued simulation if False
loading_vortex_pos = False

# --------------------------------------------------------------------------------------------------------------------
# Generating initial state:
# --------------------------------------------------------------------------------------------------------------------
# If it is a continued simulation, load the previous data and continue to evolution:
if not fresh_simulation:
    previous_data = h5py.File(backup_data_path, 'r')
    psi_k = np.array(previous_data['wavefunction/psi_k'])
    t = np.round(previous_data['time'][...])
    k = previous_data['array_index'][...]
    previous_data.close()

# If it is a fresh simulation, generate the initial state:
else:
    n_0 = 3.2e9 / (Nx * dx * Ny * dy)
    xi = 1 / np.sqrt(n_0 * c0)  # Healing length
    cvals = np.linspace(0, n_0 * 1.2, 100)
    # Store parameters in dictionary for saving
    parameters = {
        "c0": c0,
        "n_0": n_0,
        "xi": xi,
        "m": 1 / 2,
    }

    # Generate vortex positions:
    N_vort = 48 ** 2

    if loading_vortex_pos:
        with h5py.File('vortex_positions/vortex_pos_uniform.hdf5', 'r') as data:
            vort_pos = iter(data['positions'][...])
    else:
        vort_pos = include.phase_imprinting.get_positions_cpu(N_vort, 2 * xi, len_x, len_y)

    with h5py.File('vortex_positions/theta_uniform.hdf5', 'r') as file:
        theta = file['theta'][...]

    # Construct wavefunction
    psi = np.sqrt(n_0) * np.exp(1j * theta)

    atom_num = dx * dy * np.sum(np.abs(psi) ** 2)

    theta_fix = np.angle(psi)
    psi_k = np.fft.fft2(psi)

    # ------------------------------------------------------------------------------------------------------------------
    # Imaginary time evolution
    # ------------------------------------------------------------------------------------------------------------------
    for i in range(2000):

        # Kinetic energy:
        psi_k *= np.exp(-0.25 * dt * (Kx ** 2 + Ky ** 2))

        # Backward FFT:
        psi = np.fft.ifft2(psi_k)

        # Interaction term:
        psi *= np.exp(-dt * (c0 * np.abs(psi) ** 2))

        # Forward FFT:
        psi_k = np.fft.fft2(psi)

        # Kinetic energy:
        psi_k *= np.exp(-0.25 * dt * (Kx ** 2 + Ky ** 2))

        # Re-normalising:
        atom_num_new = dx * dy * np.sum(np.abs(np.fft.ifft2(psi_k)) ** 2)
        psi_k = np.fft.fft2(np.sqrt(atom_num) * np.fft.ifft2(psi_k) / np.sqrt(atom_num_new))

        # Fixing phase:
        psi = np.fft.ifft2(psi_k)
        psi *= np.exp(1j * theta_fix) / np.exp(1j * np.angle(psi))
        psi_k = np.fft.fft2(psi)
        print(f"On iteration {i}")

    # Add noise to k!=0 modes:
    for i in range(Nx):
        for j in range(Ny):
            if i == 0 and j == 0:
                continue
            else:
                psi_k[i, j] += np.sqrt(Nx * Ny) * np.sqrt(1 / 2) * np.exp(1j * np.random.uniform(0, 2 * np.pi))

    # Creating file to save to:
    with h5py.File(data_path, 'w') as data:
        # Saving spatial data:
        data.create_dataset('grid/x', x.shape, data=x)
        data.create_dataset('grid/y', y.shape, data=y)

        # Saving time variables:
        data.create_dataset('time/Nt', data=Nt)
        data.create_dataset('time/dt', data=dt)
        data.create_dataset('time/Nframe', data=Nframe)

        # Save parameters:
        data.create_dataset('parameters', data=str(parameters))

        # Creating empty wavefunction datasets to store data:
        data.create_dataset('wavefunction/psi', (Nx, Ny, 1), maxshape=(Nx, Ny, None), dtype='complex64')

        # Stores initial state:
        data.create_dataset('initial_state/psi', data=np.fft.ifft2(psi_k))

# ---------------------------------------------------------------------------------------------------------------------
# Real time evolution
# ---------------------------------------------------------------------------------------------------------------------
for i in range(Nt):

    # Kinetic energy:
    psi_k *= np.exp(-0.25 * 1j * dt * (Kx ** 2 + Ky ** 2))

    # Backward FFT:
    psi = np.fft.ifft2(psi_k)

    # Interaction term:
    psi *= np.exp(-1j * dt * (c0 * np.abs(psi) ** 2))

    # Forward FFT:
    psi_k = np.fft.fft2(psi)

    # Kinetic energy:
    psi_k *= np.exp(-0.25 * 1j * dt * (Kx ** 2 + Ky ** 2))

    # Saves data
    if np.mod(i + 1, Nframe) == 0:
        with h5py.File(data_path, 'r+') as data:
            new_psi = data['wavefunction/psi']
            new_psi.resize((Nx, Ny, k + 1))
            new_psi[:, :, k] = np.fft.ifft2(psi_k)
        k += 1

    # # Saves 'backup' wavefunction we can use to continue simulations if ended:
    # if np.mod(i + 1, 50000) == 0:
    #     with h5py.File(backup_data_path, 'w') as backup:
    #         backup.create_dataset('time', data=t)
    #         backup.create_dataset('wavefunction/psi_k', shape=psi_k.shape, dtype='complex64', data=(psi_k))
    #         backup.create_dataset('array_index', data=k)

    # Prints current time
    if np.mod(i, 50) == 0:
        print('t = %1.4f' % t)

    t += dt

import numpy as np
import h5py
import include.phase_imprinting

"""Turbulence simulation for a scalar BEC. Vortices are generated by via phase imprinting."""

# --------------------------------------------------------------------------------------------------------------------
# Controlled variables:
# --------------------------------------------------------------------------------------------------------------------
Nx, Ny = 1024, 1024
Mx, My = Nx // 2, Ny // 2  # Number of grid pts
dx = dy = 1  # Grid spacing
dkx = np.pi / (Mx * dx)
dky = np.pi / (My * dy)  # K-space spacing
len_x = Nx * dx  # Box length
len_y = Ny * dy
x = np.arange(-Mx, Mx) * dx
y = np.arange(-My, My) * dy
X, Y = np.meshgrid(x, y)  # Spatial meshgrid

# k-space arrays and meshgrid:
kx = np.arange(-Mx, Mx) * dkx
ky = np.arange(-My, My) * dky
Kx, Ky = np.meshgrid(kx, ky)  # K-space meshgrid
Kx, Ky = np.fft.fftshift(Kx), np.fft.fftshift(Ky)

# Controlled variables
V = 0.  # Doubly periodic box
c0 = 3e-5

# Time steps, number and wavefunction save variables
Nt = 10000000
Nframe = 20000   # Save data every Nframe number of timesteps
dt = 1e-2  # Imaginary time timestep
t = 0.
k = 0   # Array index

filename = 'scalar_imp_noise_cpu'    # Name of file to save data to
data_path = '../scratch/data/scalar/{}.hdf5'.format(filename)
backup_data_path = '../scratch/data/scalar/backups/{}_backup.hdf5'.format(filename)

fresh_simulation = True  # Boolean that corresponds to a fresh simulation if True or a continued simulation if False
loading_vortex_pos = False

# --------------------------------------------------------------------------------------------------------------------
# Generating initial state:
# --------------------------------------------------------------------------------------------------------------------
# If it is a continued simulation, load the previous data and continue to evolution:
if not fresh_simulation:
    previous_data = h5py.File(backup_data_path, 'r')
    psi_k = np.array(previous_data['wavefunction/psi_k'])
    t = np.round(previous_data['time'][...])
    k = previous_data['array_index'][...]
    previous_data.close()

# If it is a fresh simulation, generate the initial state:
else:
    n_0 = 3.2e9 / (Nx * dx * Ny * dy)
    xi = 1 / np.sqrt(n_0 * c0)  # Healing length
    cvals = np.linspace(0, n_0 * 1.2, 100)
    # Store parameters in dictionary for saving
    parameters = {
        "c0": c0,
        "n_0": n_0,
        "xi": xi,
        "m": 1 / 2,
    }

    # Generate vortex positions:
    N_vort = 48 ** 2

    if loading_vortex_pos:
        with h5py.File('vortex_positions/vortex_pos_uniform.hdf5', 'r') as data:
            vort_pos = iter(data['positions'][...])
    else:
        vort_pos = include.phase_imprinting.get_positions_cpu(N_vort, 2 * xi, len_x, len_y)

    with h5py.File('vortex_positions/theta_uniform.hdf5', 'r') as file:
        theta = file['theta'][...]

    # Construct wavefunction
    psi = np.sqrt(n_0) * np.exp(1j * theta)

    atom_num = dx * dy * np.sum(np.abs(psi) ** 2)

    theta_fix = np.angle(psi)
    psi_k = np.fft.fft2(psi)

    # ------------------------------------------------------------------------------------------------------------------
    # Imaginary time evolution
    # ------------------------------------------------------------------------------------------------------------------
    for i in range(2000):

        # Kinetic energy:
        psi_k *= np.exp(-0.25 * dt * (Kx ** 2 + Ky ** 2))

        # Backward FFT:
        psi = np.fft.ifft2(psi_k)

        # Interaction term:
        psi *= np.exp(-dt * (c0 * np.abs(psi) ** 2))

        # Forward FFT:
        psi_k = np.fft.fft2(psi)

        # Kinetic energy:
        psi_k *= np.exp(-0.25 * dt * (Kx ** 2 + Ky ** 2))

        # Re-normalising:
        atom_num_new = dx * dy * np.sum(np.abs(np.fft.ifft2(psi_k)) ** 2)
        psi_k = np.fft.fft2(np.sqrt(atom_num) * np.fft.ifft2(psi_k) / np.sqrt(atom_num_new))

        # Fixing phase:
        psi = np.fft.ifft2(psi_k)
        psi *= np.exp(1j * theta_fix) / np.exp(1j * np.angle(psi))
        psi_k = np.fft.fft2(psi)
        print(f"On iteration {i}")

    # Add noise to k!=0 modes:
    for i in range(Nx):
        for j in range(Ny):
            if i == 0 and j == 0:
                continue
            else:
                psi_k[i, j] += np.sqrt(Nx * Ny) * np.sqrt(1 / 2) * np.exp(1j * np.random.uniform(0, 2 * np.pi))

    # Creating file to save to:
    with h5py.File(data_path, 'w') as data:
        # Saving spatial data:
        data.create_dataset('grid/x', x.shape, data=x)
        data.create_dataset('grid/y', y.shape, data=y)

        # Saving time variables:
        data.create_dataset('time/Nt', data=Nt)
        data.create_dataset('time/dt', data=dt)
        data.create_dataset('time/Nframe', data=Nframe)

        # Save parameters:
        data.create_dataset('parameters', data=str(parameters))

        # Creating empty wavefunction datasets to store data:
        data.create_dataset('wavefunction/psi', (Nx, Ny, 1), maxshape=(Nx, Ny, None), dtype='complex64')

        # Stores initial state:
        data.create_dataset('initial_state/psi', data=np.fft.ifft2(psi_k))

# ---------------------------------------------------------------------------------------------------------------------
# Real time evolution
# ---------------------------------------------------------------------------------------------------------------------
for i in range(Nt):

    # Kinetic energy:
    psi_k *= np.exp(-0.25 * 1j * dt * (Kx ** 2 + Ky ** 2))

    # Backward FFT:
    psi = np.fft.ifft2(psi_k)

    # Interaction term:
    psi *= np.exp(-1j * dt * (c0 * np.abs(psi) ** 2))

    # Forward FFT:
    psi_k = np.fft.fft2(psi)

    # Kinetic energy:
    psi_k *= np.exp(-0.25 * 1j * dt * (Kx ** 2 + Ky ** 2))

    # Saves data
    if np.mod(i + 1, Nframe) == 0:
        with h5py.File(data_path, 'r+') as data:
            new_psi = data['wavefunction/psi']
            new_psi.resize((Nx, Ny, k + 1))
            new_psi[:, :, k] = np.fft.ifft2(psi_k)
        k += 1

    # # Saves 'backup' wavefunction we can use to continue simulations if ended:
    # if np.mod(i + 1, 50000) == 0:
    #     with h5py.File(backup_data_path, 'w') as backup:
    #         backup.create_dataset('time', data=t)
    #         backup.create_dataset('wavefunction/psi_k', shape=psi_k.shape, dtype='complex64', data=(psi_k))
    #         backup.create_dataset('array_index', data=k)

    # Prints current time
    if np.mod(i, 50) == 0:
        print('t = %1.4f' % t)

    t += dt

Turbulence simulation for a scalar BEC. Vortices are generated by via phase imprinting. # -------------------------------------------------------------------------------------------------------------------- # Controlled variables: # -------------------------------------------------------------------------------------------------------------------- # Number of grid pts # Grid spacing # K-space spacing # Box length # Spatial meshgrid # k-space arrays and meshgrid: # K-space meshgrid # Controlled variables # Doubly periodic box # Time steps, number and wavefunction save variables # Save data every Nframe number of timesteps # Imaginary time timestep # Array index # Name of file to save data to # Boolean that corresponds to a fresh simulation if True or a continued simulation if False # -------------------------------------------------------------------------------------------------------------------- # Generating initial state: # -------------------------------------------------------------------------------------------------------------------- # If it is a continued simulation, load the previous data and continue to evolution: # If it is a fresh simulation, generate the initial state: # Healing length # Store parameters in dictionary for saving # Generate vortex positions: # Construct wavefunction # ------------------------------------------------------------------------------------------------------------------ # Imaginary time evolution # ------------------------------------------------------------------------------------------------------------------ # Kinetic energy: # Backward FFT: # Interaction term: # Forward FFT: # Kinetic energy: # Re-normalising: # Fixing phase: # Add noise to k!=0 modes: # Creating file to save to: # Saving spatial data: # Saving time variables: # Save parameters: # Creating empty wavefunction datasets to store data: # Stores initial state: # --------------------------------------------------------------------------------------------------------------------- # Real time evolution # --------------------------------------------------------------------------------------------------------------------- # Kinetic energy: # Backward FFT: # Interaction term: # Forward FFT: # Kinetic energy: # Saves data # # Saves 'backup' wavefunction we can use to continue simulations if ended: # if np.mod(i + 1, 50000) == 0: #     with h5py.File(backup_data_path, 'w') as backup: #         backup.create_dataset('time', data=t) #         backup.create_dataset('wavefunction/psi_k', shape=psi_k.shape, dtype='complex64', data=(psi_k)) #         backup.create_dataset('array_index', data=k) # Prints current time