Skip to content

Instantly share code, notes, and snippets.

Created January 24, 2022 14:05
What would you like to do?
Chemiscope NaN to sparse environments
import gzip
import json
import numpy as np
input_path = "TODO"
output_path = "TODO"
if input_path.endswith(".gz"):
with as fd:
data = json.load(fd)
with open(input_path) as fd:
data = json.load(fd)
atom_properties = [
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]:
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:
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, "w") as fd:
with open(output_path, "w") as fd:
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment