feat: start adding file linking

subpex/configs/subpex/main.py

@@ -5,13 +5,12 @@
 from .cluster import ClusteringConfig
 from .data import DataConfig
 from .pocket import PocketConfig
+from .sim import SimulationConfig
 
 
 class SubpexConfig(BaseSettings, YamlIO):
     restart: bool = Field(default=False)
     """Restart a subpex simulation if directly already exists."""
-    md_engine: str = Field(default="amber")
-    """Molecular dynamics engine to use. Both `amber` and `namd` are supported."""
     calculated_points: int = Field(default=-1)
     """Number of point to calculate per trajectory segment. If `-1`, it will
     calculate all.
@@ -20,3 +19,4 @@ class SubpexConfig(BaseSettings, YamlIO):
     pocket: PocketConfig = Field(default_factory=PocketConfig)
     data: DataConfig = Field(default_factory=DataConfig)
     clustering: ClusteringConfig = Field(default_factory=ClusteringConfig)
+    sim: SimulationConfig = Field(default_factory=SimulationConfig)

subpex/configs/subpex/sim.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel, Field
+
+
+class SimulationConfig(BaseModel):
+
+    md_engine: str = Field(default="amber")
+    """Molecular dynamics engine to use. Both `amber` and `namd` are supported."""
+    path_traj_relax: str | None = Field(default=None)
+    """Path to the final simulation trajectory file to use. The last frame of this
+    trajectory will be used as the starting frame for the enhanced simulations.
+    """

subpex/setup/westpa/main.py

@@ -1,22 +1,47 @@
+from typing import Any
+
 import argparse
 import os
 import sys
 
 from loguru import logger
 
 from ...configs import SubpexConfig, WestpaConfig
+from .utils import link_file
 
 
 def link_initial_sims(
     subpex_config: SubpexConfig,
     westpa_config: WestpaConfig,
     write_dir: str = "",
-    overwrite: bool = False,
+    *args: Any,
+    **kwargs: dict[str, Any],
 ) -> None:
     """WESTPA starts from a relaxed molecular simulation using the same MD engine.
     Instead of copying the preliminary simulation files, we create a soft link to
     the relevant files.
     """
+    logger.info("Linking simulation files")
+    link_file(
+        path_original=subpex_config.sim.path_traj_relax,
+        link_name="relax_traj",
+        write_dir=os.path.join(write_dir, kwargs.get("common_files", "")),  # type: ignore
+    )
+
+
+def create_westpa_dirs(write_dir: str = "", overwrite: bool = False) -> dict[str, str]:
+    paths = {}
+    logger.info("Creating all necessary WESTPA directories")
+    logger.debug("   - Creating common_files")
+    paths["common_files"] = os.path.join(write_dir, "common_files")
+    os.makedirs(paths["common_files"], exist_ok=overwrite)
+    logger.debug("   - Creating westpa_scripts")
+    paths["westpa_scripts"] = os.path.join(write_dir, "westpa_scripts")
+    os.makedirs(paths["westpa_scripts"], exist_ok=overwrite)
+    logger.debug("   - Creating bstates")
+    paths["bstates"] = os.path.join(write_dir, "bstates")
+    os.makedirs(paths["bstates"], exist_ok=overwrite)
+    return paths
 
 
 def run_westpa_setup(
@@ -39,16 +64,19 @@ def run_westpa_setup(
             logger.error("Aborting")
             sys.exit(1)
 
+    # Create all directories
     logger.info(f"Creating directory at {write_dir} if absent")
     os.makedirs(write_dir, exist_ok=overwrite)
 
-    logger.info("Creating all necessary WESTPA directories")
-    logger.debug("   - Creating common_files")
-    os.makedirs(os.path.join(write_dir, "common_files"), exist_ok=overwrite)
-    logger.debug("   - Creating westpa_scripts")
-    os.makedirs(os.path.join(write_dir, "westpa_scripts"), exist_ok=overwrite)
-    logger.debug("   - Creating bstates")
-    os.makedirs(os.path.join(write_dir, "bstates"), exist_ok=overwrite)
+    path_dirs_westpa = create_westpa_dirs(write_dir=write_dir, overwrite=overwrite)
+
+    # Link all necessary files
+    link_initial_sims(
+        subpex_config=subpex_config,
+        westpa_config=westpa_config,
+        write_dir=write_dir,
+        **path_dirs_westpa,  # type: ignore
+    )
 
 
 def cli_run_westpa_setup():

subpex/setup/westpa/utils.py

@@ -0,0 +1,14 @@
+import os
+
+
+def link_file(
+    path_original: str | None,
+    link_name: str,
+    write_dir: str = "",
+    f_ext: str | None = None,
+) -> None:
+    if path_original is None:
+        raise ValueError(f"Original final path for {link_name} cannot be None")
+    if f_ext is None:
+        f_ext = os.path.splitext(path_original)[1][1:]
+    os.symlink(path_original, os.path.join(write_dir, f"{link_name}.{f_ext}"))

subpex/setup/wizard.py

@@ -586,9 +586,6 @@ def run_wizard():
         # Exit
         sys.exit(0)
 
-    check_existing_params()
-    confirm_environment()
-    confirm_dependencies()
     engine, get_pcoord, runseg = amber_or_namd(get_pcoord, runseg)
     get_prelim_sim_files(engine)
     get_sim_last_frame()

tests/conftest.py

@@ -35,6 +35,9 @@ def m7g_config():
     subpex_config.pocket.radius = 6.5
     subpex_config.pocket.water_dist = 6.7
     subpex_config.pocket.selection_str = "resid 124 or resid 125 or resid 128 or resid 88 or resid 89 or resid 121 or resid 92 or resid 129 and (not name H*)"
+    subpex_config.sim.path_traj_relax = os.path.join(
+        TEST_DIR, "files/m7g/equil_frames.nc"
+    )
     for aux_data in subpex_config.data.aux.get_all():
         aux_data.active = True
     return subpex_config

tests/test_westpa_init.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 
 from subpex.configs import WestpaConfig
 from subpex.setup.westpa.main import run_westpa_setup
@@ -7,6 +8,8 @@
 def test_westpa_init_m7g(path_m7g_paths, m7g_config, tmp_dir):
     westpa_config = WestpaConfig()
     write_dir = os.path.join(tmp_dir, "westpa_init_m7g")
+    if os.path.exists(write_dir):
+        shutil.rmtree(write_dir)
     run_westpa_setup(
         subpex_config=m7g_config,
         westpa_config=westpa_config,