Source code for megnet.data.qm9

"""
Simple qm9 utils, kept here for historical reasons
"""
from monty.json import MSONable

ATOMNUM2TYPE = {"1": 1, "6": 2, "7": 4, "8": 6, "9": 8}


[docs]class AtomNumberToTypeConverter(MSONable): """ Convert atomic number Z into the atomic type in the QM9 dataset """ def __init__(self, mapping=ATOMNUM2TYPE): self.mapping = mapping
[docs] def convert(self, l): return [self.mapping[str(i)] for i in l]
[docs]def ring_to_vector(l): """ Convert the ring sizes vector to a fixed length vector For example, l can be [3, 5, 5], meaning that the atom is involved in 1 3-sized ring and 2 5-sized ring. This function will convert it into [ 0, 0, 1, 0, 2, 0, 0, 0, 0, 0]. Args: l: (list of integer) ring_sizes attributes Returns: (list of integer) fixed size list with the i-1 th element indicates number of i-sized ring this atom is involved in. """ return_l = [0] * 9 if l: for i in l: return_l[i - 1] += 1 return return_l