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 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__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 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 test_minimal_xyz_uneven(self, data_dir, ext, dim):
Nxyz = np.arange(5, 7)[:dim]
xyz = np.meshgrid(*[np.exp(np.linspace(0, 1, _N)) for _N in Nxyz],
indexing="ij",
sparse=True)
t = np.exp(np.linspace(0, 2, 4))
data = io.WData(prefix="test1", data_dir=data_dir, xyz=xyz, t=t)
self.check_roundtrip(data)
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_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_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_extra_abscissa(self, data_dir):
Nxyz = np.arange(5, 7)
xyz = np.meshgrid(*[np.exp(np.linspace(0, 1, _N)) for _N in Nxyz],
indexing="ij",
sparse=True)
t = np.exp(np.linspace(0, 1, 4))
data = io.WData(prefix="test1",
data_dir=data_dir,
xyz=xyz,
t=t,
ext="npy")
data.save()
Path(os.path.join(data_dir, f"{data.prefix}__t.wdat")).touch()
with pytest.warns(
UserWarning,
match=r"Multiple files found for varying abscissa t: .*"):
data1 = io.WData.load(infofile=data.infofile)
def test_missing_abscissa(self, data_dir, ext):
Nxyz = np.arange(5, 7)
xyz = np.meshgrid(*[np.exp(np.linspace(0, 1, _N)) for _N in Nxyz],
indexing="ij",
sparse=True)
t = np.exp(np.linspace(0, 1, 4))
data = io.WData(prefix="test1",
data_dir=data_dir,
xyz=xyz,
t=t,
ext=ext)
data.save()
os.remove(os.path.join(data_dir, f"{data.prefix}__t.{ext}"))
with pytest.raises(
ValueError,
match=r"Abscissa t has varying dt but no files .* found"):
data1 = io.WData.load(infofile=data.infofile)
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_issue5(self, data_dir, ext):
x = np.array([1, 2, 4, 5])
y = np.array([1, 2, 3, 4, 5])
xyz = [x[:, np.newaxis], y[np.newaxis, :]]
t = [0]
prefix = "tmp"
full_prefix = os.path.join(data_dir, prefix)
data = io.WData(prefix=prefix,
data_dir=data_dir,
xyz=xyz,
t=t,
ext=ext)
assert np.isnan(data.dt)
assert np.isnan(data.dxyz[0])
assert np.allclose(1.0, data.dxyz[1])
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 varying # Time interval between frames\n"
)
assert len(found) == 3
wdata = io.WData.load(infofile=infofile)
assert np.isnan(wdata.dt)
assert np.isnan(wdata.dxyz[0])
assert np.allclose(1.0, wdata.dxyz[1])
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_dts(self):
data = io.WData(Nxyz=(4, 5, 6), t=[1, 2, 3, 4, 5])
assert data.dt == 1
data = io.WData(Nxyz=(4, 5, 6), t=[1, 2, 3, 5])
assert np.isnan(data.dt)
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
def test_minimal(self, data_dir, ext, dim):
Nxyz = np.arange(5, 7)[:dim]
data = io.WData(prefix="test1", data_dir=data_dir, Nxyz=Nxyz, Nt=1)
self.check_roundtrip(data)
def test_empty(self):
data = io.WData(Nxyz=(4, 5), t=[1, 2])
assert data.variables == []
assert data.constants == {}
assert data.aliases == {}
assert data.keys() == []
def test_getattr(self):
data = io.WData(Nxyz=(4, 5), t=[1, 2])
with pytest.raises(AttributeError) as excinfo:
data.unknown
assert str(
excinfo.value) == "'WData' object has no attribute 'unknown'"