diff --git a/espaloma/__init__.py b/espaloma/__init__.py index 2855a327..87cda8b9 100644 --- a/espaloma/__init__.py +++ b/espaloma/__init__.py @@ -8,6 +8,7 @@ from .graphs.graph import Graph from .metrics import GraphMetric from .mm.geometry import * +from .utils.model_fetch import get_model, get_model_path # Add imports here # import espaloma diff --git a/espaloma/utils/tests/test_model_fetch.py b/espaloma/utils/tests/test_model_fetch.py index 1266d4d6..6e550c72 100644 --- a/espaloma/utils/tests/test_model_fetch.py +++ b/espaloma/utils/tests/test_model_fetch.py @@ -1,12 +1,11 @@ import espaloma as esp import torch -from espaloma.utils.model_fetch import get_model, get_model_path from openff.toolkit.topology import Molecule def test_get_model_path(tmp_path): model_dir = tmp_path / "latest" - model_path = get_model_path(model_dir=model_dir, disable_progress_bar=True) + model_path = esp.get_model_path(model_dir=model_dir, disable_progress_bar=True) molecule = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C") molecule_graph = esp.Graph(molecule) @@ -19,7 +18,7 @@ def test_get_model_path(tmp_path): def test_get_model(tmp_path): model_dir = tmp_path / "zoo" - espaloma_model = get_model() + espaloma_model = esp.get_model() molecule = Molecule.from_smiles("CN1C=NC2=C1C(=O)N(C(=O)N2C)C") molecule_graph = esp.Graph(molecule)