def test_save_wrong_dim(self):
# Write it to the temp directory
test_file = os.path.join(tempfile.gettempdir(),
"save_initial_distribution_wrong_dim.gdf")
with open(test_file, "wb") as f:
with self.assertRaises(ValueError):
easygdf.save_initial_distribution(f,
x=np.linspace(0, 1, 11),
y=np.linspace(0, 1, 7))
def test_uniform_interface(self):
# Try to load the file
with pkg_resources.resource_stream(
"easygdf.tests", "data/initial_distribution.gdf") as f:
all_data = easygdf.load_initial_distribution(f)
# Directly save it
test_file = os.path.join(
tempfile.gettempdir(),
"save_initial_distributionuniform_interface.gdf")
with open(test_file, "wb") as f:
easygdf.save_initial_distribution(f, **all_data)
def test_save_both_speed_and_momentum(self):
"""
Confirms error is thrown if user tries to write both speed and momentum to file
:return:
"""
# Write it to the temp directory
test_file = os.path.join(
tempfile.gettempdir(),
"save_initial_distribution_both_speed_and_momentum.gdf")
with open(test_file, "wb") as f:
with self.assertRaises(ValueError):
easygdf.save_initial_distribution(f,
Bx=np.linspace(0, 1, 11),
GBx=np.linspace(0, 1, 11))
def test_save_length_normalization_B(self):
# Write it to the temp directory
test_file = os.path.join(
tempfile.gettempdir(),
"save_initial_distribution_length_normalization_B.gdf")
with open(test_file, "wb") as f:
easygdf.save_initial_distribution(f, Bx=np.linspace(0, 1, 11))
# Read it back
with open(test_file, "rb") as f:
all_data = easygdf.load_initial_distribution(f)
# Check array lengths
arr_names = ['x', 'y', 'z', 'Bx', 'By', 'Bz']
for a in arr_names:
self.assertEqual(all_data[a].size, 11)
def test_save(self):
# Write it to the temp directory
test_file = os.path.join(tempfile.gettempdir(),
"save_initial_distribution.gdf")
with open(test_file, "wb") as f:
easygdf.save_initial_distribution(f, **self.ref)
# Read it back
with open(test_file, "rb") as f:
all_data = easygdf.load_initial_distribution(f)
# Confirm that the keys are the same
self.assertEqual(self.ref.keys(), all_data.keys())
for k in self.ref:
if isinstance(self.ref[k], np.ndarray):
np.testing.assert_almost_equal(self.ref[k], all_data[k])
else:
self.assertEqual(self.ref[k], all_data[k])
def test_save2(self):
"""
Tests saving another valid input. Difference between this and test_save is that we use velocity as well as all
optional elements.
:return:
"""
# Write it to the temp directory
test_file = os.path.join(tempfile.gettempdir(),
"save_initial_distribution2.gdf")
with open(test_file, "wb") as f:
easygdf.save_initial_distribution(f, **self.ref2)
# Read it back
with open(test_file, "rb") as f:
all_data = easygdf.load_initial_distribution(f)
# Confirm that the keys are the same
self.assertEqual(self.ref2.keys(), all_data.keys())
for k in self.ref2:
if isinstance(self.ref2[k], np.ndarray):
np.testing.assert_almost_equal(self.ref2[k], all_data[k])
else:
self.assertEqual(self.ref2[k], all_data[k])
import easygdf
import numpy as np
# Save some data to an initial distribution file. Unspecified required values are autofilled for us
easygdf.save_initial_distribution(
"initial.gdf",
x=np.random.normal(size=(3, )),
GBx=np.random.normal(size=(3, )),
t=np.random.random((3, )),
)