Additional changes to fix bugs and suggestions from erdem

ORNL · Sep 3, 2024 · 86886cc · 86886cc
1 parent 6a1b917
commit 86886cc
Show file tree

Hide file tree

Showing 10 changed files with 14 additions and 10 deletions.
diff --git a/examples/ising_model/train_ising.py b/examples/ising_model/train_ising.py
@@ -27,7 +27,7 @@
 import numpy as np
 
 try:
-    from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+    from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 except ImportError:
     pass
 

diff --git a/examples/multidataset/energy_linear_regression.py b/examples/multidataset/energy_linear_regression.py
@@ -7,10 +7,15 @@
 import numpy as np
 
 import hydragnn
-from hydragnn.utils import nsplit
-from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+from hydragnn.utils.distributed import nsplit
+from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 from tqdm import tqdm
-from mpi_list import Context, DFM
+
+# This import requires having installed the package mpi_list
+try:
+    from mpi_list import Context, DFM
+except ImportError:
+    print("mpi_list requires having installed: https://github.com/frobnitzem/mpi_list")
 
 
 def subset(i):

diff --git a/examples/ogb/train_gap.py b/examples/ogb/train_gap.py
@@ -32,7 +32,7 @@
 import numpy as np
 
 try:
-    from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+    from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 except ImportError:
     pass
 

diff --git a/examples/qm7x/train.py b/examples/qm7x/train.py
@@ -61,7 +61,7 @@ def info(*args, logtype="info", sep=" "):
     getattr(logging, logtype)(sep.join(map(str, args)))
 
 
-from hydragnn.utils.datasets import AbstractBaseDataset
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
 
 # FIXME: this radis cutoff overwrites the radius cutoff currently written in the JSON file
 create_graph_fromXYZ = RadiusGraph(r=5.0)  # radius cutoff in angstrom

diff --git a/hydragnn/preprocess/stratified_sampling.py b/hydragnn/preprocess/stratified_sampling.py
@@ -39,7 +39,7 @@ def stratified_sampling(dataset: [Data], subsample_percentage: float, verbosity=
         n_splits=1, train_size=subsample_percentage, random_state=0
     )
 
-    for subsample_index, rest_of_data_index in sss.split(dataset, dataset_categories):
+    for subsample_index, _ in sss.split(dataset, dataset_categories):
         subsample_indices = subsample_index.tolist()
 
     for index in subsample_indices:

diff --git a/utils/lsms/__init__.py → hydragnn/utils/lsms/__init__.py b/utils/lsms/__init__.py → hydragnn/utils/lsms/__init__.py
diff --git a/utils/lsms/compositional_histogram_cutoff.py → ...ls/lsms/compositional_histogram_cutoff.py b/utils/lsms/compositional_histogram_cutoff.py → ...ls/lsms/compositional_histogram_cutoff.py
diff --git a/...onvert_total_energy_to_formation_gibbs.py → ...onvert_total_energy_to_formation_gibbs.py b/...onvert_total_energy_to_formation_gibbs.py → ...onvert_total_energy_to_formation_gibbs.py
diff --git a/tests/test_enthalpy.py b/tests/test_enthalpy.py
@@ -11,9 +11,9 @@
 
 import os
 import numpy as np
-import hydragnn, tests
+import tests
 import pytest
-from utils.lsms import (
+from hydragnn.utils.lsms import (
     convert_raw_data_energy_to_gibbs,
 )
 

diff --git a/utils/__init__.py b/utils/__init__.py