import numpy as np
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
from os import listdir
from tensorflow.keras.callbacks import ModelCheckpoint

dataDir = "./data/trainSmallFA/"
files = listdir(dataDir)
files.sort()
totalLength = len(files)
inputs = np.empty((len(files), 3, 64, 64))
targets = np.empty((len(files), 3, 64, 64))

for i, file in enumerate(files):
    npfile = np.load(dataDir + file)
    d = npfile['a']
    inputs[i] = d[0:3]  # inx, iny, mask
    targets[i] = d[3:6]  # p, velx, vely

# print("inputs shape = ", inputs.shape)
print(np.shape(targets[:, 1, :, :].flatten()))
maxvel = np.amax(np.sqrt(targets[:, 1, :, :]* targets[:, 1, :, :]
                         + targets[:, 2, :, :]* targets[:, 2, :, :]))
print(maxvel)
targets[:, 1:3, :, :] /= maxvel
targets[:, 0, :, :] /= np.amax(targets[:, 0, :, :])
for input in inputs:
    plt.figure(num=None, figsize=(20, 10), dpi=80, facecolor='w', edgecolor='k')

    # predicted data
    plt.subplot(331)
    plt.title('x vel')
    plt.imshow(input[0, :, :], cmap='jet')  # vmin=-100,vmax=100, cmap='jet')
    plt.colorbar()
    plt.subplot(332)
    plt.title('y vel')
    plt.imshow(input[1, :, :], cmap='jet')
    plt.colorbar()

    plt.show()
[((9, 8, 9, 24), 'os.listdir', 'listdir', ({(9, 16, 9, 23): 'dataDir'}, {}), '(dataDir)', False, 'from os import listdir\n'), ((27, 23, 27, 51), 'numpy.amax', 'np.amax', ({(27, 31, 27, 50): 'targets[:, (0), :, :]'}, {}), '(targets[:, (0), :, :])', True, 'import numpy as np\n'), ((16, 13, 16, 36), 'numpy.load', 'np.load', ({(16, 21, 16, 35): 'dataDir + file'}, {}), '(dataDir + file)', True, 'import numpy as np\n'), ((23, 17, 24, 68), 'numpy.sqrt', 'np.sqrt', ({(23, 25, 24, 67): 'targets[:, (1), :, :] * targets[:, (1), :, :] + targets[:, (2), :, :\n    ] * targets[:, (2), :, :]'}, {}), '(targets[:, (1), :, :] * targets[:, (1), :, :] + targets[:, (2), :,\n    :] * targets[:, (2), :, :])', True, 'import numpy as np\n'), ((29, 4, 29, 80), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((32, 4, 32, 20), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(32, 16, 32, 19): '(331)'}, {}), '(331)', True, 'import matplotlib.pyplot as plt\n'), ((33, 4, 33, 22), 'matplotlib.pyplot.title', 'plt.title', ({(33, 14, 33, 21): '"""x vel"""'}, {}), "('x vel')", True, 'import matplotlib.pyplot as plt\n'), ((34, 4, 34, 42), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((35, 4, 35, 18), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((36, 4, 36, 20), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(36, 16, 36, 19): '(332)'}, {}), '(332)', True, 'import matplotlib.pyplot as plt\n'), ((37, 4, 37, 22), 'matplotlib.pyplot.title', 'plt.title', ({(37, 14, 37, 21): '"""y vel"""'}, {}), "('y vel')", True, 'import matplotlib.pyplot as plt\n'), ((38, 4, 38, 42), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((39, 4, 39, 18), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((41, 4, 41, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n')]