import numpy as np
'''
sin_train = np.load('./data/3d_sin_5_5_train.npy')
sin_test = np.load('./data/3d_sin_5_5_test.npy')
helix_train = np.load('./data/helix_3D_train.npy')
helix_test = np.load('./data/helix_3D_test.npy')
sin_bounds = np.load('./data/3d_sin_5_5_bounds.npy')
helix_bounds = np.load('./data/helix_3D_bounds.npy')

merged_train = np.concatenate((sin_train, helix_train), axis=0)
merged_test = np.concatenate((sin_test, helix_test), axis=0)
merged_bounds = np.concatenate((sin_bounds, helix_bounds), axis=0)

np.save('./data/3d_sin_helix_train.npy', merged_train)
np.save('./data/3d_sin_helix_test.npy', merged_test)
np.save('./data/3d_sin_helix_bounds.npy', merged_bounds)
'''
from generate import *

train_data = get_cylinder_helix('train')
print(train_data.shape)
test_data = get_cylinder_helix('test')
print(test_data.shape)

np.save('./data/3d_cylinder_helix_train.npy',train_data)
np.save('./data/3d_cylinder_helix_test.npy',test_data)

helix_3d_train = np.load('./data/helix_3D_train.npy')
helix_3d_test = np.load('./data/helix_3D_test.npy')

helix_3d_train[:,2] += 2
helix_3d_test[:,2] += 2

np.save('./data/shifted_helix_3D_train.npy',helix_3d_train)
np.save('./data/shifted_helix_3D_test.npy',helix_3d_test)