Chemiscope NaN to sparse environments

chemiscope-nan-to-sparse.py

import gzip
import json
import numpy as np

input_path = "TODO"
output_path = "TODO"


if input_path.endswith(".gz"):
    with gzip.open(input_path) as fd:
        data = json.load(fd)
else:
    with open(input_path) as fd:
        data = json.load(fd)


atom_properties = [
    name
    for name, property in data["properties"].items()
    if property["target"] == "atom"
]

valid = np.ones(len(data["properties"][atom_properties[0]]["values"]), dtype=np.bool_)
for name in atom_properties:
    # only take the non-NaN entries
    mask = np.isfinite(data["properties"][name]["values"])
    valid = np.logical_and(valid, mask)

environments = []
for i, environment in enumerate(data["environments"]):
    if valid[i]:
        environments.append(environment)

old_environments = data["environments"]
data["environments"] = environments

if "settings" in data and "pinned" in data["settings"]:
    cleaned_pinned = []
    for pinned in data["settings"]["pinned"]:
        structure = old_environments[pinned]["structure"]
        center = old_environments[pinned]["center"]

        for new_index, env in enumerate(environments):
            if env["structure"] == structure and env["center"] == center:
                cleaned_pinned.append(new_index)

    data["settings"]["pinned"] = cleaned_pinned


for name, property in data["properties"].items():
    if name in atom_properties:
        values = np.array(property["values"])[valid]
        property["values"] = values.tolist()


if output_path.endswith(".gz"):
    with gzip.open(output_path, "w") as fd:
        fd.write(json.dumps(data).encode("utf8"))
else:
    with open(output_path, "w") as fd:
        fd.write(json.dumps(data))