From b96158fe779aeb0736b5b7827ff392216e78560a Mon Sep 17 00:00:00 2001
From: FNTwin <cristian@valencelabs.com>
Date: Fri, 1 Mar 2024 08:19:15 -0700
Subject: [PATCH] Upgraded to black>=24, fixed import tests, cleaning, gh
 action lint wip fix

---
 .github/workflows/code-check.yml |  2 +-
 docs/tutorials/usage.ipynb       | 15 ++---------
 env.yml                          |  2 +-
 openqdc/__init__.py              | 46 ++++++++++++++++----------------
 openqdc/utils/preprocess.py      | 28 ++-----------------
 tests/test_import.py             |  2 +-
 6 files changed, 30 insertions(+), 65 deletions(-)

diff --git a/.github/workflows/code-check.yml b/.github/workflows/code-check.yml
index 7c4be3b..767776e 100644
--- a/.github/workflows/code-check.yml
+++ b/.github/workflows/code-check.yml
@@ -24,7 +24,7 @@ jobs:
 
       - name: Install black
         run: |
-          pip install black>=23
+          pip install black>=24
 
       - name: Lint
         run: black --check .
diff --git a/docs/tutorials/usage.ipynb b/docs/tutorials/usage.ipynb
index cd71b6a..1ad74cd 100644
--- a/docs/tutorials/usage.ipynb
+++ b/docs/tutorials/usage.ipynb
@@ -748,7 +748,8 @@
    "source": [
     "### Chemical space from SOAP descriptors\n",
     "\n",
-    "openQDC offer a simple way to calculate the Smooth Overlaps of Atomic Positions (SOAP) descriptors for the molecules in the dataset. The method get_soap_descriptors returns the SOAP descriptors for the molecules in the dataset."
+    "openQDC offer a simple way to calculate the Smooth Overlaps of Atomic Positions (SOAP) descriptors for the molecules in the dataset.\n",
+    "The method get_soap_descriptors returns the SOAP descriptors for the molecules in the dataset."
    ]
   },
   {
@@ -813,18 +814,6 @@
    "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.2"
   }
  },
  "nbformat": 4,
diff --git a/env.yml b/env.yml
index 53663c9..d7a19db 100644
--- a/env.yml
+++ b/env.yml
@@ -30,7 +30,7 @@ dependencies:
   - pytest >=6.0
   - pytest-cov
   - nbconvert
-  - black >=23
+  - black >=24
   - jupyterlab
   - pre-commit
   - ruff
diff --git a/openqdc/__init__.py b/openqdc/__init__.py
index abf50ad..c757b23 100644
--- a/openqdc/__init__.py
+++ b/openqdc/__init__.py
@@ -9,29 +9,29 @@
 
 _lazy_imports_obj = {
     "__version__": "openqdc._version",
-    "BaseDataset": "openqdc.datasets",
-    "ANI1": "openqdc.datasets",
-    "ANI1CCX": "openqdc.datasets",
-    "ANI1X": "openqdc.datasets",
-    "Spice": "openqdc.datasets",
-    "GEOM": "openqdc.datasets",
-    "QMugs": "openqdc.datasets",
-    "ISO17": "openqdc.datasets",
-    "COMP6": "openqdc.datasets",
-    "GDML": "openqdc.datasets",
-    "Molecule3D": "openqdc.datasets",
-    "OrbnetDenali": "openqdc.datasets",
-    "SN2RXN": "openqdc.datasets",
-    "QM7X": "openqdc.datasets",
-    "DES": "openqdc.datasets",
-    "NablaDFT": "openqdc.datasets",
-    "SolvatedPeptides": "openqdc.datasets",
-    "WaterClusters": "openqdc.datasets",
-    "TMQM": "openqdc.datasets",
-    "Dummy": "openqdc.datasets",
-    "PCQM_B3LYP": "openqdc.datasets",
-    "PCQM_PM6": "openqdc.datasets",
-    "Transition1X": "openqdc.datasets",
+    "BaseDataset": "openqdc.datasets.base",
+    "ANI1": "openqdc.datasets.ani",
+    "ANI1CCX": "openqdc.datasets.ani",
+    "ANI1X": "openqdc.datasets.ani",
+    "Spice": "openqdc.datasets.spice",
+    "GEOM": "openqdc.datasets.geom",
+    "QMugs": "openqdc.datasets.qmugs",
+    "ISO17": "openqdc.datasets.iso_17",
+    "COMP6": "openqdc.datasets.comp6",
+    "GDML": "openqdc.datasets.gdml",
+    "Molecule3D": "openqdc.datasets.molecule3d",
+    "OrbnetDenali": "openqdc.datasets.orbnet_denali",
+    "SN2RXN": "openqdc.datasets.sn2_rxn",
+    "QM7X": "openqdc.datasets.qm7x",
+    "DES": "openqdc.datasets.des",
+    "NablaDFT": "openqdc.datasets.nabladft",
+    "SolvatedPeptides": "openqdc.datasets.solvated_peptides",
+    "WaterClusters": "openqdc.datasets.waterclusters3_30",
+    "TMQM": "openqdc.datasets.tmqm",
+    "Dummy": "openqdc.datasets.dummy",
+    "PCQM_B3LYP": "openqdc.datasets.pcqm",
+    "PCQM_PM6": "openqdc.datasets.pcqm",
+    "Transition1X": "openqdc.datasets.transition1x",
     "AVAILABLE_DATASETS": "openqdc.datasets",
 }
 
diff --git a/openqdc/utils/preprocess.py b/openqdc/utils/preprocess.py
index 9e55412..6ad102d 100644
--- a/openqdc/utils/preprocess.py
+++ b/openqdc/utils/preprocess.py
@@ -1,33 +1,9 @@
 import click
 import numpy as np
 from loguru import logger
+from openqdc import AVAILABLE_DATASETS
 
-from openqdc import datasets
-
-options = [
-    datasets.ANI1,
-    datasets.ANI1CCX,
-    datasets.ANI1X,
-    datasets.COMP6,
-    datasets.DES,
-    datasets.GDML,
-    datasets.GEOM,
-    datasets.ISO17,
-    datasets.Molecule3D,
-    datasets.NablaDFT,
-    datasets.OrbnetDenali,
-    datasets.PCQM_B3LYP,
-    datasets.PCQM_PM6,
-    datasets.QM7X,
-    datasets.QMugs,
-    datasets.SN2RXN,
-    datasets.SolvatedPeptides,
-    datasets.Spice,
-    datasets.TMQM,
-    datasets.Transition1X,
-    datasets.WaterClusters,
-]
-
+options = list(AVAILABLE_DATASETS.values())
 options_map = {d.__name__: d for d in options}
 
 
diff --git a/tests/test_import.py b/tests/test_import.py
index 0736e30..8c272b1 100644
--- a/tests/test_import.py
+++ b/tests/test_import.py
@@ -1,2 +1,2 @@
 def test_open_qdc():
-    import openQDC  # noqa
+    import openqdc  # noqa