Merge pull request #172 from usnistgov/develop

Develop
usnistgov · Dec 2, 2024 · 7f3f0e2 · 7f3f0e2
2 parents 983ca4f + c17395c
commit 7f3f0e2
Show file tree

Hide file tree

Showing 26 changed files with 1,580 additions and 900 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -9,8 +9,8 @@ jobs:
         matrix:
             os: ["ubuntu-latest"]
     steps:
-      - uses: actions/checkout@v2
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: actions/checkout@v3
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           activate-environment: test
           environment-file: environment.yml
@@ -24,8 +24,8 @@ jobs:
         shell: bash -l {0}
         run: |
             conda install flake8 pycodestyle pydocstyle
-            flake8 --ignore E203,W503 --exclude=examples,tests,scripts --statistics --count --exit-zero alignn
-            pycodestyle --ignore E203,W503 --exclude=examples,tests,scripts alignn
+            flake8 --ignore E203,W503,W605 --exclude=examples,tests,scripts --statistics --count --exit-zero alignn
+            pycodestyle --ignore E203,W503,W605 --exclude=examples,tests,scripts alignn
             pydocstyle --match-dir=core --match-dir=io --match-dir=io --match-dir=ai --match-dir=analysis --match-dir=db --match-dir=tasks --count alignn
       
       - name: Run pytest
@@ -39,6 +39,8 @@ jobs:
             export DGLBACKEND=pytorch
             export CUDA_VISIBLE_DEVICES="-1"
             pip install phonopy flake8 pytest pycodestyle pydocstyle codecov pytest-cov coverage
+            #pip uninstall jarvis-tools -y
+            #pip install -q git+https://github.com/usnistgov/jarvis.git@develop
             python setup.py develop
             echo 'environment.yml'
             conda env export

diff --git a/README.md b/README.md
diff --git a/alignn/__init__.py b/alignn/__init__.py
@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2024.8.30"
+__version__ = "2024.10.30"
diff --git a/alignn/config.py b/alignn/config.py
@@ -8,16 +8,7 @@
 from alignn.models.alignn import ALIGNNConfig
 from alignn.models.alignn_atomwise import ALIGNNAtomWiseConfig
 
-# from alignn.models.modified_cgcnn import CGCNNConfig
-# from alignn.models.icgcnn import ICGCNNConfig
-# from alignn.models.gcn import SimpleGCNConfig
-# from alignn.models.densegcn import DenseGCNConfig
-# from pydantic import model_validator
-# from alignn.models.dense_alignn import DenseALIGNNConfig
-# from alignn.models.alignn_cgcnn import ACGCNNConfig
-# from alignn.models.alignn_layernorm import ALIGNNConfig as ALIGNN_LN_Config
-
-# from typing import List
+# import torch
 
 try:
     VERSION = (
@@ -162,14 +153,13 @@ class TrainingConfig(BaseSettings):
     ] = "dft_3d"
     target: TARGET_ENUM = "exfoliation_energy"
     atom_features: Literal["basic", "atomic_number", "cfid", "cgcnn"] = "cgcnn"
-    neighbor_strategy: Literal["k-nearest", "voronoi", "radius_graph"] = (
-        "k-nearest"
-    )
+    neighbor_strategy: Literal[
+        "k-nearest", "voronoi", "radius_graph", "radius_graph_jarvis"
+    ] = "k-nearest"
     id_tag: Literal["jid", "id", "_oqmd_entry_id"] = "jid"
 
-    # logging configuration
-
     # training configuration
+    dtype: str = "float32"
     random_seed: Optional[int] = 123
     classification_threshold: Optional[float] = None
     # target_range: Optional[List] = None
@@ -220,26 +210,4 @@ class TrainingConfig(BaseSettings):
     model: Union[
         ALIGNNConfig,
         ALIGNNAtomWiseConfig,
-        # CGCNNConfig,
-        # ICGCNNConfig,
-        # SimpleGCNConfig,
-        # DenseGCNConfig,
-        # ALIGNN_LN_Config,
-        # DenseALIGNNConfig,
-        # ACGCNNConfig,
     ] = ALIGNNAtomWiseConfig(name="alignn_atomwise")
-
-    # @root_validator()
-    # @model_validator(mode='before')
-    # def set_input_size(cls, values):
-    #    """Automatically configure node feature dimensionality."""
-    #    values["model"].atom_input_features = FEATURESET_SIZE[
-    #        values["atom_features"]
-    #    ]
-
-    #    return values
-
-    # @property
-    # def atom_input_features(self):
-    #     """Automatically configure node feature dimensionality."""
-    #     return FEATURESET_SIZE[self.atom_features]
diff --git a/alignn/data.py b/alignn/data.py
@@ -122,6 +122,7 @@ def get_train_val_loaders(
     target_atomwise: str = "",
     target_grad: str = "",
     target_stress: str = "",
+    target_additional_output: str = "",
     atom_features: str = "cgcnn",
     neighbor_strategy: str = "k-nearest",
     n_train=None,
@@ -153,13 +154,15 @@ def get_train_val_loaders(
     world_size=0,
     rank=0,
     use_lmdb: bool = True,
+    dtype="float32",
 ):
     """Help function to set up JARVIS train and val dataloaders."""
     if use_lmdb:
         print("Using LMDB dataset.")
         from alignn.lmdb_dataset import get_torch_dataset
     else:
         print("Not using LMDB dataset, memory footprint maybe high.")
+        print("WARNING: not using LMDB might result errors.")
         from alignn.dataset import get_torch_dataset
     train_sample = filename + "_train.data"
     val_sample = filename + "_val.data"
@@ -372,6 +375,7 @@ def get_train_val_loaders(
             target_atomwise=target_atomwise,
             target_grad=target_grad,
             target_stress=target_stress,
+            target_additional_output=target_additional_output,
             neighbor_strategy=neighbor_strategy,
             use_canonize=use_canonize,
             name=dataset,
@@ -383,6 +387,7 @@ def get_train_val_loaders(
             output_dir=output_dir,
             sampler=train_sampler,
             tmp_name=tmp_name,
+            dtype=dtype,
             # tmp_name="train_data",
         )
         tmp_name = filename + "val_data"
@@ -395,6 +400,7 @@ def get_train_val_loaders(
                 target_atomwise=target_atomwise,
                 target_grad=target_grad,
                 target_stress=target_stress,
+                target_additional_output=target_additional_output,
                 neighbor_strategy=neighbor_strategy,
                 use_canonize=use_canonize,
                 name=dataset,
@@ -406,6 +412,7 @@ def get_train_val_loaders(
                 classification=classification_threshold is not None,
                 output_dir=output_dir,
                 tmp_name=tmp_name,
+                dtype=dtype,
                 # tmp_name="val_data",
             )
             if len(dataset_val) > 0
@@ -421,6 +428,7 @@ def get_train_val_loaders(
                 target_atomwise=target_atomwise,
                 target_grad=target_grad,
                 target_stress=target_stress,
+                target_additional_output=target_additional_output,
                 neighbor_strategy=neighbor_strategy,
                 use_canonize=use_canonize,
                 name=dataset,
@@ -431,6 +439,7 @@ def get_train_val_loaders(
                 classification=classification_threshold is not None,
                 output_dir=output_dir,
                 tmp_name=tmp_name,
+                dtype=dtype,
                 # tmp_name="test_data",
             )
             if len(dataset_test) > 0

diff --git a/alignn/dataset.py b/alignn/dataset.py
@@ -14,6 +14,8 @@
 tqdm.pandas()
 
 
+# NOTE: Use lmd_dataset,
+# need to fix adding lattice in dataloader
 def load_graphs(
     dataset=[],
     name: str = "dft_3d",
@@ -26,6 +28,7 @@ def load_graphs(
     id_tag="jid",
     # extra_feats_json=None,
     map_size=1e12,
+    dtype="float32",
 ):
     """Construct crystal graphs.
 
@@ -54,6 +57,7 @@ def atoms_to_graph(atoms):
             compute_line_graph=False,
             use_canonize=use_canonize,
             neighbor_strategy=neighbor_strategy,
+            dtype=dtype,
         )
 
     if cachedir is not None:
@@ -84,6 +88,7 @@ def atoms_to_graph(atoms):
                 use_canonize=use_canonize,
                 neighbor_strategy=neighbor_strategy,
                 id=i[id_tag],
+                dtype=dtype,
             )
             # print ('ii',ii)
             if "extra_features" in i:
@@ -124,6 +129,7 @@ def get_torch_dataset(
     output_dir=".",
     tmp_name="dataset",
     sampler=None,
+    dtype="float32",
 ):
     """Get Torch Dataset."""
     df = pd.DataFrame(dataset)
@@ -147,6 +153,7 @@ def get_torch_dataset(
         cutoff_extra=cutoff_extra,
         max_neighbors=max_neighbors,
         id_tag=id_tag,
+        dtype=dtype,
     )
     data = StructureDataset(
         df,
@@ -160,5 +167,6 @@ def get_torch_dataset(
         id_tag=id_tag,
         classification=classification,
         sampler=sampler,
+        dtype=dtype,
     )
     return data
diff --git a/alignn/examples/sample_data/config_example.json b/alignn/examples/sample_data/config_example.json
@@ -5,6 +5,7 @@
     "atom_features": "cgcnn",
     "neighbor_strategy": "k-nearest",
     "id_tag": "jid",
+    "dtype": "float32",
     "random_seed": 123,
     "classification_threshold": null,
     "n_val": null,
@@ -18,7 +19,7 @@
     "batch_size": 2,
     "weight_decay": 1e-05,
     "learning_rate": 0.001,
-    "filename": "sample",
+    "filename": "A",
     "warmup_steps": 2000,
     "criterion": "mse",
     "optimizer": "adamw",
@@ -34,8 +35,15 @@
     "use_canonize": true,
     "num_workers": 0,
     "cutoff": 8.0,
+    "cutoff_extra": 3.0,
     "max_neighbors": 12,
     "keep_data_order": true,
+    "normalize_graph_level_loss": false,
+    "distributed": false,
+    "data_parallel": false,
+    "n_early_stopping": null,
+    "output_dir": "temp",
+    "use_lmdb": true,
     "model": {
         "name": "alignn_atomwise",
         "alignn_layers": 4,
@@ -46,8 +54,30 @@
         "embedding_features": 64,
         "hidden_features": 256,
         "output_features": 1,
+        "grad_multiplier": -1,
+        "calculate_gradient": false,
+        "atomwise_output_features": 0,
+        "graphwise_weight": 1.0,
+        "gradwise_weight": 1.0,
+        "stresswise_weight": 0.0,
+        "atomwise_weight": 0.0,
         "link": "identity",
         "zero_inflated": false,
-        "classification": false
+        "classification": false,
+        "force_mult_natoms": false,
+        "energy_mult_natoms": false,
+        "include_pos_deriv": false,
+        "use_cutoff_function": false,
+        "inner_cutoff": 3.0,
+        "stress_multiplier": 1.0,
+        "add_reverse_forces": true,
+        "lg_on_fly": true,
+        "batch_stress": true,
+        "multiply_cutoff": false,
+        "use_penalty": true,
+        "extra_features": 0,
+        "exponent": 5,
+        "penalty_factor": 0.1,
+        "penalty_threshold": 1.0
     }
-}
+}
diff --git a/alignn/examples/sample_data_ff/config_example_atomwise.json b/alignn/examples/sample_data_ff/config_example_atomwise.json
@@ -3,8 +3,9 @@
     "dataset": "user_data",
     "target": "target",
     "atom_features": "cgcnn",
-    "neighbor_strategy": "k-nearest",
+    "neighbor_strategy": "radius_graph",
     "id_tag": "jid",
+    "dtype": "float32",
     "random_seed": 123,
     "classification_threshold": null,
     "n_val": null,
@@ -18,7 +19,7 @@
     "batch_size": 2,
     "weight_decay": 1e-05,
     "learning_rate": 0.001,
-    "filename": "sample",
+    "filename": "B",
     "warmup_steps": 2000,
     "criterion": "l1",
     "optimizer": "adamw",
@@ -33,26 +34,25 @@
     "standard_scalar_and_pca": false,
     "use_canonize": true,
     "num_workers": 0,
-    "cutoff": 8.0,
+    "cutoff": 4.0,
     "max_neighbors": 12,
     "keep_data_order": true,
-    "distributed":false,
+    "distributed": false,
     "use_lmdb": true,
     "model": {
         "name": "alignn_atomwise",
         "atom_input_features": 92,
-        "calculate_gradient":true,
-        "atomwise_output_features":0,
-        "alignn_layers":4,
-        "gcn_layers":4,
+        "calculate_gradient": true,
+        "atomwise_output_features": 0,
+        "alignn_layers": 1,
+        "gcn_layers": 1,
+        "hidden_features": 64,
         "output_features": 1,
-        "graphwise_weight":0.85,
-        "gradwise_weight":0.05,
-        "atomwise_weight":0.0,
-        "stresswise_weight":0.05,
-        "add_reverse_forces":true,
-        "lg_on_fly":true
-
-
+        "graphwise_weight": 0.85,
+        "gradwise_weight": 0.05,
+        "atomwise_weight": 0.0,
+        "use_cutoff_function": false,
+        "stresswise_weight": 0.05,
+        "add_reverse_forces": true
     }
-}
+}
diff --git a/alignn/examples/sample_data_ff_additional/CMD b/alignn/examples/sample_data_ff_additional/CMD
@@ -0,0 +1,2 @@
+train_alignn.py --root_dir DataDir/ --config config.json --output_dir temp --target_key energy --additional_output_key dos --stresswise_key stresses
+
diff --git a/alignn/examples/sample_data_ff_additional/DataDir/id_prop.json b/alignn/examples/sample_data_ff_additional/DataDir/id_prop.json
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		train_alignn.py --root_dir DataDir/ --config config.json --output_dir temp --target_key energy --additional_output_key dos --stresswise_key stresses