def test_wdata_errors(self):
with pytest.raises(ValueError) as excinfo:
io.WData()
assert str(excinfo.value) == "Must provide one of xyz or Nxyz"
with pytest.raises(ValueError) as excinfo:
io.WData(Nxyz=(3, 4, 5), Nt=1)
assert str(
excinfo.value) == "First dimension of Nxyz==(3, 4, 5) must be > 3."
with pytest.raises(ValueError) as excinfo:
x = io.Var(x=np.zeros((2, 4, 5)))
io.WData(Nxyz=(4, 5), Nt=1, variables=[x])
assert (str(excinfo.value) ==
"Variable 'x' has incompatible Nt=1: data.shape[0] = 2")
with pytest.raises(ValueError) as excinfo:
x = io.Var(x=np.zeros((2, 4)))
io.WData(Nxyz=(5, ), Nt=2, variables=[x])
assert (str(
excinfo.value
) == "Variable 'x' has incompatible Nxyz=(5,): data.shape[-1:] = (4,)")
with pytest.raises(ValueError) as excinfo:
x = io.Var(x=np.zeros((2, 4, 4, 6)))
io.WData(Nxyz=(4, 6), dim=2, Nt=2, variables=[x])
assert (str(
excinfo.value
) == "Variable 'x' has incompatible dim=2: data.shape = (2, 4, 4, 6)")
def infofile(data_dir, ext, dim):
"""Reasonable datasets for testing."""
Nt = 4
Nxyz = (4, 8, 16)
dxyz = (0.1, 0.2, 0.3)
variables = [
io.Var(density=np.random.random((Nt, ) + Nxyz[:dim])),
io.Var(current1=np.random.random((Nt, 1) + Nxyz[:dim])),
io.Var(current2=np.random.random((Nt, 2) + Nxyz[:dim])),
io.Var(current3=np.random.random((Nt, 3) + Nxyz[:dim])),
]
data = io.WData(
prefix="tmp",
data_dir=data_dir,
ext=ext,
dim=dim,
Nxyz=Nxyz,
dxyz=dxyz,
variables=variables,
Nt=Nt,
aliases={"n": "density"},
constants=dict(hbar=1.23),
)
data.save()
infofile = data.infofile
del data
yield infofile
def infofile0(data_dir):
"""Reasonable infofile for testing with no data."""
Nt = 4
Nxyz = (4, 8)
dxyz = (0.1, 0.2)
dim = 2
variables = [
io.Var(name="density", filename="density.npy", descr=float,
shape=Nxyz),
io.Var(name="delta", filename="delta.npy", descr=complex, shape=Nxyz),
]
data = io.WData(
prefix="tmp",
data_dir=data_dir,
dim=dim,
Nxyz=Nxyz,
dxyz=dxyz,
variables=variables,
Nt=Nt,
aliases={"n": "density"},
constants=dict(hbar=1.23),
check_data=False,
)
data.save()
infofile = data.infofile
del data
yield infofile
def test_descr(self):
"""Test setting descr."""
v = io.Var(x=np.ones(10), descr=complex)
assert v.data.dtype == complex
assert io.WData._get_type(v) == "abscissa"
v = io.Var(x=np.ones((3, 10), dtype="<f4"))
assert io.WData._get_type(v) == "<f4"
assert io.WData._get_descr("<f4") == "<f4"
def test_var_errors(self):
with pytest.raises(ValueError) as excinfo:
io.Var()
assert str(
excinfo.value) == "Must provide `name` or data as a kwarg: got {}"
with pytest.raises(ValueError) as excinfo:
io.Var(data=np.ones((4, 4)))
assert str(excinfo.value) == "Got data but no name."
v = io.Var(name="x", filename="tmp.unknown")
with pytest.raises(NotImplementedError) as excinfo:
data = v.data
assert str(
excinfo.value) == "Data format of 'tmp.unknown' not supported."
def test__dir__(self):
x = io.Var(x=np.zeros((2, 4, 5)))
data = io.WData(Nxyz=(4, 5),
t=[1, 2],
variables=[x],
constants=dict(x=3))
assert dir(data) == ["x"]
def test_readonly(self, data_dir, ext):
"""Issue 13: fail on read-only file systems."""
Nxyz = (4, 8)
dxyz = (0.1, 0.2)
prefix = "tmp"
full_prefix = os.path.join(data_dir, prefix)
data = io.WData(prefix=prefix,
data_dir=data_dir,
Nxyz=Nxyz,
dxyz=dxyz,
Nt=1,
ext=ext)
xyz = data.xyz
np.random.seed(2)
psi = np.random.random(Nxyz + (2, )) - 0.5
psi = psi.view(dtype=complex)[..., 0]
assert psi.shape == Nxyz
psis = [psi]
densities = [abs(psi)**2]
data = io.WData(
data_dir=data_dir,
xyz=xyz,
variables=[io.Var(density=densities),
io.Var(delta=psis)],
ext=ext,
)
data.save()
# Make files read-only
for (dirpath, dirnames, filenames) in os.walk(data_dir):
for file in filenames:
os.chmod(os.path.join(dirpath, file), stat.S_IRUSR)
# Try loading
res = io.WData.load(full_prefix=full_prefix)
assert np.allclose(res.delta, psi)
assert np.allclose(res.density, densities)
def test_metadata(self, data_dir, ext):
x = np.array([1, 2, 3, 5])
y = np.array([1, 2, 3, 4, 5])
xyz = [x[:, np.newaxis], y[np.newaxis, :]]
t = [0, 1.2]
prefix = "tmp"
full_prefix = os.path.join(data_dir, prefix)
Nxyz = sum(xyz).shape
x = io.Var(x=np.ones((len(t), ) + Nxyz, dtype=complex))
data = io.WData(
prefix=prefix,
data_dir=data_dir,
xyz=xyz,
t=t,
ext=ext,
variables=[x],
aliases={"y": "x"},
constants=dict(hbar=1.23),
)
data.save()
infofile = f"{full_prefix}.wtxt"
with open(infofile, "r") as f:
found = set()
for line in f.readlines():
if line.startswith("dx"):
found.add("dx")
assert line == "dx varying # Spacing in x direction\n"
if line.startswith("dy"):
found.add("dy")
assert line == "dy 1.0 # ... y ...\n"
if line.startswith("dt"):
found.add("dt")
assert (
line ==
"dt 1.2 # Time interval between frames\n"
)
if line.startswith("var"):
found.add("var")
assert (
line ==
f"var x complex none {ext:>4} # x\n"
)
if line.startswith("link"):
found.add("link")
assert line == "link y x\n"
if line.startswith("const"):
found.add("const")
assert line == "const hbar 1.23"
assert len(found) == 6
def test_eq(self, data_dir):
args = dict(prefix="test", data_dir=data_dir, Nxyz=(4, 5, 6))
data1 = io.WData(Nt=1, **args)
data2 = io.WData(Nt=2, **args)
assert data1 != data2
np.random.seed(3)
Nxyz = args["Nxyz"]
X1 = np.random.random(Nxyz)
X2 = np.random.random(Nxyz)
# Nans should be considered equal.
X1[0, 0] = np.nan
variables = [io.Var(X1=[X1])]
data1 = io.WData(variables=variables, **args)
data2 = io.WData(variables=variables, **args)
assert data1 == data2
data2 = io.WData(variables=[io.Var(X1=[X2])], **args)
assert data1 != data2
args["Nt"] = 1
args["variables"] = variables + [io.Var(X2=[X2])]
data1 = io.WData(constants={"eF": 1}, **args)
data2 = io.WData(constants={"eF": 2}, **args)
assert data1 != data2
# Aliases and constants should be case-sensitive
data2 = io.WData(constants={"ef": 1}, **args)
assert data1 != data2
data1 = io.WData(aliases={"B": "X1"}, **args)
data2 = io.WData(aliases={"C": "X1"}, **args)
assert data1 != data2
data1 = io.WData(aliases={"A": "X1"}, **args)
data2 = io.WData(aliases={"A": "X2"}, **args)
assert data1 != data2
def test_write_data(self, data_dir):
filename = os.path.join(data_dir, "tmp.wdat")
v = io.Var(x=np.arange(10))
with pytest.raises(ValueError) as excinfo:
v.write_data()
assert str(excinfo.value) == "No filename specified in Var."
v.write_data(filename)
assert os.path.exists(f"{filename}")
with pytest.raises(IOError) as excinfo:
v.write_data(filename)
assert str(excinfo.value) == f"File '{filename}' already exists!"
with pytest.raises(NotImplementedError) as excinfo:
v.write_data(filename="x.unknown")
assert str(excinfo.value) == "Unsupported extension for 'x.unknown'"
v = io.Var(name="x", filename=filename)
with pytest.raises(ValueError) as excinfo:
v.write_data()
assert str(excinfo.value) == "Missing data for 'x'!"
def test_wdata1(self, data_dir, ext):
Nxyz = (4, 8, 16)
dxyz = (0.1, 0.2, 0.3)
prefix = "tmp"
full_prefix = os.path.join(data_dir, prefix)
data = io.WData(prefix=prefix,
data_dir=data_dir,
Nxyz=Nxyz,
dxyz=dxyz,
Nt=1,
ext=ext)
xyz = data.xyz
np.random.seed(2)
psi = np.random.random(Nxyz + (2, )) - 0.5
psi = psi.view(dtype=complex)[..., 0]
assert psi.shape == Nxyz
psis = [psi]
densities = [abs(psi)**2]
data = io.WData(
data_dir=data_dir,
xyz=xyz,
variables=[io.Var(density=densities),
io.Var(delta=psis)],
ext=ext,
)
data.save()
res = io.WData.load(full_prefix=full_prefix)
assert np.allclose(res.delta, psi)
assert np.allclose(res.density, densities)
def test_set(self):
"""Test setting attributes."""
v = io.Var(x=np.ones(10))
assert v.data.dtype == float
assert v.shape == (10, )
v.data = 1j * np.arange(12)
assert v.data.dtype == complex
assert v.shape == (12, )
v.shape = (3, 4)
assert v.data.shape == (3, 4)
with pytest.raises(ValueError) as excinfo:
v.shape = (3, 5)
assert (str(excinfo.value) ==
"Property shape=(3, 5) incompatible with data.shape=(3, 4)")
def test_getitem(self):
x = io.Var(x=np.zeros((2, 4, 5)))
data = io.WData(Nxyz=(4, 5),
t=[1, 2],
variables=[x],
constants=dict(x=3))
with pytest.warns(UserWarning,
match="Variable x hides constant of the same name"):
data.x
data = io.WData(Nxyz=(4, 5),
t=[1, 2],
variables=[x],
constants=dict(y=3))
with pytest.raises(AttributeError) as excinfo:
data.q
assert str(excinfo.value) == "'WData' object has no attribute 'q'"
data.keys()
def test_interfaces(self, data_dir):
assert verifyClass(io.IVar, io.Var)
assert verifyClass(io.IWData, io.WData)
Nt = 1
Nxyz = (4, 8, 16)
dxyz = (0.1, 0.2, 0.3)
var = io.Var(density=np.random.random((Nt, ) + Nxyz))
data = io.WData(
prefix="tmp",
data_dir=data_dir,
Nxyz=Nxyz,
dxyz=dxyz,
variables=[var],
Nt=Nt,
)
assert verifyObject(io.IVar, var)
assert verifyObject(io.IWData, data)
def test__get_ext(self, ext):
x = io.Var(x=np.zeros((1, 4, 5)), filename=f"x.{ext}")
data = io.WData(Nxyz=x.shape[1:], variables=[x])
info = data.get_metadata()
assert ext in info