#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright 2019 Chris Sewell
#
# This file is part of aiida-gulp.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms and conditions
# of version 3 of the GNU Lesser General Public License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
import copy
INDEX_SEP = "-"
[docs]def filter_by_species(data, species):
"""filter a potential dict by a subset of species
Parameters
----------
data : dict
a potential or fitting dict
species : list[str]
the species to filter by
Returns
-------
dict
data filtered by species and with all species index keys re-indexed
Raises
------
KeyError
if the data does not adhere to the potential or fitting jsonschema
AssertionError
if the filter set is not a subset of the available species
"""
species = sorted(list(set(species)))
if not set(species).issubset(data["species"]):
raise AssertionError(
"the filter set ({}) is not a subset of the available species ({})".format(
set(species), set(data["species"])
)
)
data = copy.deepcopy(data)
indices = set([str(i) for i, s in enumerate(data["species"]) if s in species])
def convert_indices(key):
return INDEX_SEP.join(
[str(species.index(data["species"][int(k)])) for k in key.split(INDEX_SEP)]
)
for key in ["1body", "2body", "3body", "4body"]:
if key not in data:
continue
data[key] = {
convert_indices(k): v
for k, v in data[key].items()
if indices.issuperset(k.split(INDEX_SEP))
}
data["species"] = species
return data