def test_batching_5D():
"""Test the batching with high dimensional data (5D)"""
x0 = np.random.uniform(0, 1, (50, 4, 4, 4, 2))
y0 = np.random.uniform(0, 1, (50, 4, 1, 1, 1))
p0 = x0.copy()
x_batches, y_batches, p_batches = PhysicsGuidedNeuralNetwork.make_batches(
x0, y0, p0, n_batch=6, shuffle=False)
assert len(x_batches) == 6
assert len(y_batches) == 6
assert len(p_batches) == 6
assert len(x0.shape) == 5
assert len(y0.shape) == 5
assert len(p0.shape) == 5
assert len(x_batches[0].shape) == 5
assert len(y_batches[0].shape) == 5
assert len(p_batches[0].shape) == 5
assert (x_batches[0] == x0[:len(x_batches[0])]).all()
assert (y_batches[0] == y0[:len(y_batches[0])]).all()
assert (p_batches[0] == p0[:len(p_batches[0])]).all()
assert (x_batches[-1] == x0[-(len(x_batches[0]) - 1):]).all()
assert (y_batches[-1] == y0[-(len(y_batches[0]) - 1):]).all()
assert (p_batches[-1] == p0[-(len(p_batches[0]) - 1):]).all()
def test_batching_shuffle():
"""Test the batching operation with shuffling"""
x_batches, y_batches, p_batches = PhysicsGuidedNeuralNetwork.make_batches(
X, Y, P, n_batch=4, shuffle=True)
assert len(x_batches) == 4
assert len(y_batches) == 4
assert len(p_batches) == 4
assert ~(x_batches[0] == X[0:len(x_batches[0]), :]).all()
assert ~(y_batches[0] == Y[0:len(y_batches[0]), :]).all()
assert ~(p_batches[0] == P[0:len(p_batches[0]), :]).all()
for i, x_b in enumerate(x_batches):
assert (x_b == p_batches[i]).all()
truth = np.sqrt(x_b[:, 0]**2 + x_b[:, 1]**2).reshape((len(x_b), 1))
y_check = y_batches[i]
assert np.allclose(truth, y_check)