cgcnn2 0.2.1__tar.gz

cgcnn2-0.2.1/PKG-INFO

@@ -0,0 +1,83 @@
+Metadata-Version: 2.3
+Name: cgcnn2
+Version: 0.2.1
+Summary: Crystal Graph Convolutional Neural Networks
+License: MIT
+Keywords: python,gnn,vasp,crystal
+Author: Jiacheng Wang
+Author-email: jiachengwang@umass.edu
+Maintainer: Jiacheng Wang
+Requires-Python: >=3.10,<3.13
+Classifier: Development Status :: 1 - Planning
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: ase (>=3.23.0,<4.0.0)
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: pymatgen (>=2025.1.0,<2026.0.0)
+Requires-Dist: pymatviz (>=0.15.0,<0.16.0)
+Requires-Dist: scikit-learn
+Requires-Dist: torch (>=2.5.1,<3.0.0)
+Project-URL: Homepage, https://github.com/jcwang587/cgcnn2/
+Project-URL: Repository, https://github.com/jcwang587/cgcnn2/
+Description-Content-Type: text/markdown
+
+# CGCNN2
+
+As the original Crystal Graph Convolutional Neural Networks (CGCNN) repository is no longer actively maintained, this repository is a reproduction of [CGCNN](https://github.com/txie-93/cgcnn) by Xie et al. It includes necessary updates for deprecated components and a few additional functions to ensure smooth operation. Despite its age, CGCNN remains a straightforward and fast deep learning framework that is easy to learn and use.
+
+The package provides following major functions:
+
+- **Training** a CGCNN model with a customized dataset.
+- **Predicting** material properties with a pre-trained CGCNN model.
+- **Fine-tuning** a pre-trained CGCNN model on a new dataset.
+- **Extracting** atomic features as descriptors for the downstream task.
+
+## Installation
+
+Make sure you have a Python interpreter, preferably version 3.10 or higher. Then, you can simply install xdatbus from
+PyPI using `pip`:
+
+```bash
+pip install cgcnn2
+```
+
+If you'd like to use the latest unreleased version on the main branch, you can install it directly from GitHub:
+
+```bash
+pip install git+https://github.com/jcwang587/cgcnn2
+```
+
+## Get Started
+
+```bash
+cgcnn-ft --help
+```
+
+
+## References
+
+The original paper describes the details of the CGCNN framework:
+
+```bibtex
+@article{PhysRevLett.120.145301,
+  title = {Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties},
+  author = {Xie, Tian and Grossman, Jeffrey C.},
+  journal = {Phys. Rev. Lett.},
+  volume = {120},
+  issue = {14},
+  pages = {145301},
+  numpages = {6},
+  year = {2018},
+  month = {Apr},
+  publisher = {American Physical Society},
+  doi = {10.1103/PhysRevLett.120.145301},
+  url = {https://link.aps.org/doi/10.1103/PhysRevLett.120.145301}
+}
+```
+

cgcnn2-0.2.1/README.md

@@ -0,0 +1,53 @@
+# CGCNN2
+
+As the original Crystal Graph Convolutional Neural Networks (CGCNN) repository is no longer actively maintained, this repository is a reproduction of [CGCNN](https://github.com/txie-93/cgcnn) by Xie et al. It includes necessary updates for deprecated components and a few additional functions to ensure smooth operation. Despite its age, CGCNN remains a straightforward and fast deep learning framework that is easy to learn and use.
+
+The package provides following major functions:
+
+- **Training** a CGCNN model with a customized dataset.
+- **Predicting** material properties with a pre-trained CGCNN model.
+- **Fine-tuning** a pre-trained CGCNN model on a new dataset.
+- **Extracting** atomic features as descriptors for the downstream task.
+
+## Installation
+
+Make sure you have a Python interpreter, preferably version 3.10 or higher. Then, you can simply install xdatbus from
+PyPI using `pip`:
+
+```bash
+pip install cgcnn2
+```
+
+If you'd like to use the latest unreleased version on the main branch, you can install it directly from GitHub:
+
+```bash
+pip install git+https://github.com/jcwang587/cgcnn2
+```
+
+## Get Started
+
+```bash
+cgcnn-ft --help
+```
+
+
+## References
+
+The original paper describes the details of the CGCNN framework:
+
+```bibtex
+@article{PhysRevLett.120.145301,
+  title = {Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties},
+  author = {Xie, Tian and Grossman, Jeffrey C.},
+  journal = {Phys. Rev. Lett.},
+  volume = {120},
+  issue = {14},
+  pages = {145301},
+  numpages = {6},
+  year = {2018},
+  month = {Apr},
+  publisher = {American Physical Society},
+  doi = {10.1103/PhysRevLett.120.145301},
+  url = {https://link.aps.org/doi/10.1103/PhysRevLett.120.145301}
+}
+```

cgcnn2-0.2.1/cgcnn2/__init__.py

@@ -0,0 +1,4 @@
+from .cgcnn_data import *
+from .cgcnn_model import *
+from .cgcnn_utils import *
+from .cgcnn_ft import *

cgcnn2-0.2.1/cgcnn2/cgcnn_data.py

@@ -0,0 +1,392 @@
+import csv
+import functools
+import json
+import os
+import random
+import warnings
+
+import numpy as np
+import torch
+from pymatgen.core.structure import Structure
+from torch.utils.data import Dataset
+
+
+def collate_pool(dataset_list):
+    """
+    Collate a list of data and return a batch for predicting crystal
+    properties.
+
+    Parameters
+    ----------
+
+    dataset_list: list of tuples for each data point.
+      (atom_fea, nbr_fea, nbr_fea_idx, target)
+
+      atom_fea: torch.Tensor shape (n_i, atom_fea_len)
+      nbr_fea: torch.Tensor shape (n_i, M, nbr_fea_len)
+      nbr_fea_idx: torch.LongTensor shape (n_i, M)
+      target: torch.Tensor shape (1, )
+      cif_id: str or int
+
+    Returns
+    -------
+    N = sum(n_i); N0 = sum(i)
+
+    batch_atom_fea: torch.Tensor shape (N, orig_atom_fea_len)
+      Atom features from atom type
+    batch_nbr_fea: torch.Tensor shape (N, M, nbr_fea_len)
+      Bond features of each atom's M neighbors
+    batch_nbr_fea_idx: torch.LongTensor shape (N, M)
+      Indices of M neighbors of each atom
+    crystal_atom_idx: list of torch.LongTensor of length N0
+      Mapping from the crystal idx to atom idx
+    target: torch.Tensor shape (N, 1)
+      Target value for prediction
+    batch_cif_ids: list
+    """
+    batch_atom_fea, batch_nbr_fea, batch_nbr_fea_idx = [], [], []
+    crystal_atom_idx, batch_target = [], []
+    batch_cif_ids = []
+    base_idx = 0
+    for i, ((atom_fea, nbr_fea, nbr_fea_idx), target, cif_id) in enumerate(
+        dataset_list
+    ):
+        n_i = atom_fea.shape[0]  # number of atoms for this crystal
+        batch_atom_fea.append(atom_fea)
+        batch_nbr_fea.append(nbr_fea)
+        batch_nbr_fea_idx.append(nbr_fea_idx + base_idx)
+        new_idx = torch.LongTensor(np.arange(n_i) + base_idx)
+        crystal_atom_idx.append(new_idx)
+        batch_target.append(target)
+        batch_cif_ids.append(cif_id)
+        base_idx += n_i
+    return (
+        (
+            torch.cat(batch_atom_fea, dim=0),
+            torch.cat(batch_nbr_fea, dim=0),
+            torch.cat(batch_nbr_fea_idx, dim=0),
+            crystal_atom_idx,
+        ),
+        torch.stack(batch_target, dim=0),
+        batch_cif_ids,
+    )
+
+
+class GaussianDistance(object):
+    """
+    Expands the distance by Gaussian basis.
+
+    Unit: angstrom
+    """
+
+    def __init__(self, dmin, dmax, step, var=None):
+        """
+        Parameters
+        ----------
+
+        dmin: float
+          Minimum interatomic distance
+        dmax: float
+          Maximum interatomic distance
+        step: float
+          Step size for the Gaussian filter
+        """
+        assert dmin < dmax
+        assert dmax - dmin > step
+        self.filter = np.arange(dmin, dmax + step, step)
+        if var is None:
+            var = step
+        self.var = var
+
+    def expand(self, distances):
+        """
+        Apply Gaussian distance filter to a numpy distance array
+
+        Parameters
+        ----------
+
+        distances: np.ndarray
+          A distance matrix of any shape
+
+        Returns
+        -------
+        expanded_distance: shape (n+1)-d array
+          Expanded distance matrix with the last dimension of length
+          len(self.filter)
+        """
+        return np.exp(-((distances[..., np.newaxis] - self.filter) ** 2) / self.var**2)
+
+
+class AtomInitializer(object):
+    """
+    Base class for initializing the vector representation for atoms.
+
+    !!! Use one AtomInitializer per dataset !!!
+    """
+
+    def __init__(self, atom_types):
+        self.atom_types = set(atom_types)
+        self._embedding = {}
+
+    def get_atom_fea(self, atom_type):
+        assert atom_type in self.atom_types
+        return self._embedding[atom_type]
+
+    def load_state_dict(self, state_dict):
+        self._embedding = state_dict
+        self.atom_types = set(self._embedding.keys())
+        self._decodedict = {
+            idx: atom_type for atom_type, idx in self._embedding.items()
+        }
+
+    def state_dict(self):
+        return self._embedding
+
+    def decode(self, idx):
+        if not hasattr(self, "_decodedict"):
+            self._decodedict = {
+                idx: atom_type for atom_type, idx in self._embedding.items()
+            }
+        return self._decodedict[idx]
+
+
+class AtomCustomJSONInitializer(AtomInitializer):
+    """
+    Initialize atom feature vectors using a JSON file, which is a python
+    dictionary mapping from element number to a list representing the
+    feature vector of the element.
+
+    Parameters
+    ----------
+
+    elem_embedding_file: str
+        The path to the .json file
+    """
+
+    def __init__(self, elem_embedding_file):
+        with open(elem_embedding_file) as f:
+            elem_embedding = json.load(f)
+        elem_embedding = {int(key): value for key, value in elem_embedding.items()}
+        atom_types = set(elem_embedding.keys())
+        super(AtomCustomJSONInitializer, self).__init__(atom_types)
+        for key, value in elem_embedding.items():
+            self._embedding[key] = np.array(value, dtype=float)
+
+
+class CIFData(Dataset):
+    """
+    The CIFData dataset is a wrapper for a dataset where the crystal structures
+    are stored in the form of CIF files. The dataset should have the following
+    directory structure:
+
+    root_dir
+    ├── id_prop.csv
+    ├── atom_init.json
+    ├── id0.cif
+    ├── id1.cif
+    ├── ...
+
+    id_prop.csv: a CSV file with two columns. The first column recodes a
+    unique ID for each crystal, and the second column recodes the value of
+    target property.
+
+    atom_init.json: a JSON file that stores the initialization vector for each
+    element.
+
+    ID.cif: a CIF file that recodes the crystal structure, where ID is the
+    unique ID for the crystal.
+
+    Parameters
+    ----------
+
+    root_dir: str
+        The path to the root directory of the dataset
+    max_num_nbr: int
+        The maximum number of neighbors while constructing the crystal graph
+    radius: float
+        The cutoff radius for searching neighbors
+    dmin: float
+        The minimum distance for constructing GaussianDistance
+    step: float
+        The step size for constructing GaussianDistance
+    random_seed: int
+        Random seed for shuffling the dataset
+
+    Returns
+    -------
+
+    atom_fea: torch.Tensor shape (n_i, atom_fea_len)
+    nbr_fea: torch.Tensor shape (n_i, M, nbr_fea_len)
+    nbr_fea_idx: torch.LongTensor shape (n_i, M)
+    target: torch.Tensor shape (1, )
+    cif_id: str or int
+    """
+
+    def __init__(
+        self, root_dir, max_num_nbr=12, radius=8, dmin=0, step=0.2, random_seed=123
+    ):
+        self.root_dir = root_dir
+        self.max_num_nbr, self.radius = max_num_nbr, radius
+        assert os.path.exists(root_dir), "root_dir does not exist!"
+        id_prop_file = os.path.join(self.root_dir, "id_prop.csv")
+        assert os.path.exists(id_prop_file), "id_prop.csv does not exist!"
+        with open(id_prop_file) as f:
+            reader = csv.reader(f)
+            self.id_prop_data = [row for row in reader]
+        random.seed(random_seed)
+        random.shuffle(self.id_prop_data)
+        atom_init_file = os.path.join(self.root_dir, "atom_init.json")
+        assert os.path.exists(atom_init_file), "atom_init.json does not exist!"
+        self.ari = AtomCustomJSONInitializer(atom_init_file)
+        self.gdf = GaussianDistance(dmin=dmin, dmax=self.radius, step=step)
+
+    def __len__(self):
+        return len(self.id_prop_data)
+
+    @functools.lru_cache(maxsize=None)  # Cache loaded structures
+    def __getitem__(self, idx):
+        cif_id, target = self.id_prop_data[idx]
+        crystal = Structure.from_file(os.path.join(self.root_dir, cif_id + ".cif"))
+        atom_fea = np.vstack(
+            [
+                self.ari.get_atom_fea(crystal[i].specie.number)
+                for i in range(len(crystal))
+            ]
+        )
+        atom_fea = torch.Tensor(atom_fea)
+        all_nbrs = crystal.get_all_neighbors(self.radius, include_index=True)
+        all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs]
+        nbr_fea_idx, nbr_fea = [], []
+        for nbr in all_nbrs:
+            if len(nbr) < self.max_num_nbr:
+                warnings.warn(
+                    "{} not find enough neighbors to build graph. "
+                    "If it happens frequently, consider increase "
+                    "radius.".format(cif_id)
+                )
+                nbr_fea_idx.append(
+                    list(map(lambda x: x[2], nbr)) + [0] * (self.max_num_nbr - len(nbr))
+                )
+                nbr_fea.append(
+                    list(map(lambda x: x[1], nbr))
+                    + [self.radius + 1.0] * (self.max_num_nbr - len(nbr))
+                )
+            else:
+                nbr_fea_idx.append(list(map(lambda x: x[2], nbr[: self.max_num_nbr])))
+                nbr_fea.append(list(map(lambda x: x[1], nbr[: self.max_num_nbr])))
+        nbr_fea_idx, nbr_fea = np.array(nbr_fea_idx), np.array(nbr_fea)
+        nbr_fea = self.gdf.expand(nbr_fea)
+        atom_fea = torch.Tensor(atom_fea)
+        nbr_fea = torch.Tensor(nbr_fea)
+        nbr_fea_idx = torch.LongTensor(nbr_fea_idx)
+        target = torch.Tensor([float(target)])
+        return (atom_fea, nbr_fea, nbr_fea_idx), target, cif_id
+
+
+class CIFData_pred(Dataset):
+    """
+    The CIFData dataset is a wrapper for a dataset where the crystal structures
+    are stored in the form of CIF files. The dataset should have the following
+    directory structure:
+
+    root_dir
+    ├── id_prop.csv
+    ├── atom_init.json
+    ├── id0.cif
+    ├── id1.cif
+    ├── ...
+
+    id_prop.csv: a CSV file with two columns. The first column recodes a
+    unique ID for each crystal, and the second column recodes the value of
+    target property.
+
+    atom_init.json: a JSON file that stores the initialization vector for each
+    element.
+
+    ID.cif: a CIF file that recodes the crystal structure, where ID is the
+    unique ID for the crystal.
+
+    Parameters
+    ----------
+
+    root_dir: str
+        The path to the root directory of the dataset
+    max_num_nbr: int
+        The maximum number of neighbors while constructing the crystal graph
+    radius: float
+        The cutoff radius for searching neighbors
+    dmin: float
+        The minimum distance for constructing GaussianDistance
+    step: float
+        The step size for constructing GaussianDistance
+    random_seed: int
+        Random seed for shuffling the dataset
+
+    Returns
+    -------
+
+    atom_fea: torch.Tensor shape (n_i, atom_fea_len)
+    nbr_fea: torch.Tensor shape (n_i, M, nbr_fea_len)
+    nbr_fea_idx: torch.LongTensor shape (n_i, M)
+    target: torch.Tensor shape (1, )
+    cif_id: str or int
+    """
+
+    def __init__(
+        self, root_dir, max_num_nbr=12, radius=8, dmin=0, step=0.2, random_seed=123
+    ):
+        self.root_dir = root_dir
+        self.max_num_nbr, self.radius = max_num_nbr, radius
+        assert os.path.exists(root_dir), "root_dir does not exist!"
+        id_prop_file = os.path.join(self.root_dir, "id_prop.csv")
+        assert os.path.exists(id_prop_file), "id_prop.csv does not exist!"
+        with open(id_prop_file) as f:
+            reader = csv.reader(f)
+            self.id_prop_data = [row for row in reader]
+        random.seed(random_seed)
+        atom_init_file = os.path.join(self.root_dir, "atom_init.json")
+        assert os.path.exists(atom_init_file), "atom_init.json does not exist!"
+        self.ari = AtomCustomJSONInitializer(atom_init_file)
+        self.gdf = GaussianDistance(dmin=dmin, dmax=self.radius, step=step)
+
+    def __len__(self):
+        return len(self.id_prop_data)
+
+    @functools.lru_cache(maxsize=None)  # Cache loaded structures
+    def __getitem__(self, idx):
+        cif_id, target = self.id_prop_data[idx]
+        crystal = Structure.from_file(os.path.join(self.root_dir, cif_id + ".cif"))
+        atom_fea = np.vstack(
+            [
+                self.ari.get_atom_fea(crystal[i].specie.number)
+                for i in range(len(crystal))
+            ]
+        )
+        atom_fea = torch.Tensor(atom_fea)
+        all_nbrs = crystal.get_all_neighbors(self.radius, include_index=True)
+        all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs]
+        nbr_fea_idx, nbr_fea = [], []
+        for nbr in all_nbrs:
+            if len(nbr) < self.max_num_nbr:
+                warnings.warn(
+                    "{} not find enough neighbors to build graph. "
+                    "If it happens frequently, consider increase "
+                    "radius.".format(cif_id)
+                )
+                nbr_fea_idx.append(
+                    list(map(lambda x: x[2], nbr)) + [0] * (self.max_num_nbr - len(nbr))
+                )
+                nbr_fea.append(
+                    list(map(lambda x: x[1], nbr))
+                    + [self.radius + 1.0] * (self.max_num_nbr - len(nbr))
+                )
+            else:
+                nbr_fea_idx.append(list(map(lambda x: x[2], nbr[: self.max_num_nbr])))
+                nbr_fea.append(list(map(lambda x: x[1], nbr[: self.max_num_nbr])))
+        nbr_fea_idx, nbr_fea = np.array(nbr_fea_idx), np.array(nbr_fea)
+        nbr_fea = self.gdf.expand(nbr_fea)
+        atom_fea = torch.Tensor(atom_fea)
+        nbr_fea = torch.Tensor(nbr_fea)
+        nbr_fea_idx = torch.LongTensor(nbr_fea_idx)
+        target = torch.Tensor([float(target)])
+        return (atom_fea, nbr_fea, nbr_fea_idx), target, cif_id