def test_sum_all(self, collection_data, collection_info, expected):
dc = DOSCollection([
RawDOSData(*item, info=info)
for item, info in zip(collection_data, collection_info)
])
summed_dc = dc.sum_all()
energies, weights, ref_info = expected
assert np.allclose(summed_dc.get_energies(), energies)
assert np.allclose(summed_dc.get_weights(), weights)
assert summed_dc.info == ref_info
def test_from_data(self, x, weights, bad_info):
dc = DOSCollection.from_data(x, weights)
for i, dos_data in enumerate(dc):
assert dos_data.info == {}
assert np.allclose(dos_data.get_energies(), x)
assert np.allclose(dos_data.get_weights(), weights[i])
with pytest.raises(ValueError):
dc = DOSCollection.from_data(x, weights, info=bad_info)
def test_total(self, collection_data, collection_info):
dc = DOSCollection([
RawDOSData(*item, info=info)
for item, info in zip(collection_data, collection_info)
])
summed = dc.sum_all()
total = dc.total()
assert np.allclose(summed.get_energies(), total.get_energies())
assert np.allclose(summed.get_weights(), total.get_weights())
assert (set(total.info.items()) - set(summed.info.items()) == set([
('label', 'Total')
]))
def test_select(self, select_info, select_query, select_result,
select_not_result, sum_by_result):
dc = DOSCollection(
[RawDOSData([0.], [0.], info=info) for info in select_info])
if select_result is None:
assert dc.select(**select_query)._almost_equals(DOSCollection([]))
else:
assert select_result == [
data.info for data in dc.select(**select_query)
]
if select_not_result is None:
assert (dc.select_not(**select_query)._almost_equals(
DOSCollection([])))
else:
assert select_not_result == [
data.info for data in dc.select_not(**select_query)
]
assert sum_by_result == [
data.info for data in dc.sum_by(*sorted(select_query.keys()))
]
def get_pdos(self) -> DOSCollection:
"""Phonon DOS, including atomic contributions"""
energies = self.get_energies()
masses = self._atoms[self.get_mask()].get_masses()
# Get weights as N_moving_atoms x N_modes array
vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis]**-0.5
all_weights = (np.linalg.norm(vectors, axis=-1)**2).T
mask = self.get_mask()
all_info = [{
'index': i,
'symbol': a.symbol
} for i, a in enumerate(self._atoms) if mask[i]]
return DOSCollection([
RawDOSData(energies, weights, info=info)
for weights, info in zip(all_weights, all_info)
])
def test_sum_empty(self):
dc = DOSCollection([])
with pytest.raises(IndexError):
dc.sum_all()