diff --git a/mdmix/cli/__init__.py b/mdmix/cli/__init__.py new file mode 100644 index 0000000..7d2632f --- /dev/null +++ b/mdmix/cli/__init__.py @@ -0,0 +1,12 @@ +from cli.base import app, DisplayUtils +from cli.project_plugin import ProjectPlugin +from cli.solvent_plugin import SolventPlugin + +# Register plugins +project_plugin = ProjectPlugin() +project_plugin.attach_to_main(app) + +solvent_plugin = SolventPlugin() +solvent_plugin.attach_to_main(app) + +__all__ = ['app', 'DisplayUtils', 'project_plugin', 'solvent_plugin'] \ No newline at end of file diff --git a/mdmix/cli/base.py b/mdmix/cli/base.py new file mode 100644 index 0000000..4b7290d --- /dev/null +++ b/mdmix/cli/base.py @@ -0,0 +1,130 @@ + +import typer +from typing import Optional, List +from pathlib import Path +from rich.console import Console +from rich.table import Table +from rich.panel import Panel +from rich import print as rprint +import logging +from abc import ABC, abstractmethod + +# Initialize console for rich output +console = Console() + +# Base CLI app +app = typer.Typer( + name="pymdmix3", + help="PyMDMix3 - Advanced Molecular Dynamics Mixed Solvent Simulations", + rich_markup_mode="rich" +) + + +class CLIPlugin(ABC): + """Base class for CLI plugins""" + + def __init__(self, name: str, help_text: str): + self.name = name + self.help_text = help_text + self.app = typer.Typer(help=help_text) + + @abstractmethod + def register_commands(self): + """Register plugin commands""" + pass + + def attach_to_main(self, main_app: typer.Typer): + """Attach plugin to main app""" + self.register_commands() + main_app.add_typer(self.app, name=self.name) + + +class DisplayUtils: + """Utilities for rich console output""" + + @staticmethod + def show_header(title: str, subtitle: str = ""): + """Display a styled header""" + if subtitle: + content = f"[bold cyan]{title}[/bold cyan]\n[dim]{subtitle}[/dim]" + else: + content = f"[bold cyan]{title}[/bold cyan]" + console.print(Panel(content, expand=False)) + + @staticmethod + def show_error(message: str): + """Display error message""" + console.print(f"[bold red]Error:[/bold red] {message}") + + @staticmethod + def show_success(message: str): + """Display success message""" + console.print(f"[bold green]Success:[/bold green] {message}") + + @staticmethod + def show_warning(message: str): + """Display warning message""" + console.print(f"[bold yellow]Warning:[/bold yellow] {message}") + + @staticmethod + def show_info(message: str): + """Display info message""" + console.print(f"[bold blue]Info:[/bold blue] {message}") + + @staticmethod + def create_table(title: str, columns: List[str]) -> Table: + """Create a styled table""" + table = Table(title=title, show_header=True, header_style="bold magenta") + for col in columns: + table.add_column(col) + return table + + +# Logging configuration +def setup_logging(verbose: bool = False): + """Setup logging configuration""" + level = logging.DEBUG if verbose else logging.INFO + logging.basicConfig( + level=level, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' + ) + + +# Main app callbacks +@app.callback() +def main_callback( + verbose: bool = typer.Option(False, "--verbose", "-v", help="Enable verbose output"), + version: bool = typer.Option(False, "--version", help="Show version") +): + """ + PyMDMix3 - Advanced Molecular Dynamics Mixed Solvent Simulations + + A modern Python implementation for mixed-solvent MD simulations. + """ + if version: + from .. import __version__ + console.print(f"PyMDMix3 version {__version__}") + raise typer.Exit() + + setup_logging(verbose) + + +# Main command +@app.command() +def info(): + """Show PyMDMix3 information and status""" + from .. import __version__, __description__ + + DisplayUtils.show_header("PyMDMix3", __description__) + + info_table = Table(show_header=False, box=None) + info_table.add_column("Property", style="cyan") + info_table.add_column("Value", style="white") + + info_table.add_row("Version", __version__) + info_table.add_row("Python", "3.8+") + info_table.add_row("License", "GPL-3.0") + + console.print(info_table) + console.print("\nUse [bold]pymdmix3 --help[/bold] to see available commands.") diff --git a/mdmix/cli/main.py b/mdmix/cli/main.py new file mode 100644 index 0000000..a559fa1 --- /dev/null +++ b/mdmix/cli/main.py @@ -0,0 +1,9 @@ + +from cli.base import app + +def main(): + """Main entry point for PyMDMix3 CLI""" + app() + +if __name__ == "__main__": + main() diff --git a/mdmix/cli/project_plugin.py b/mdmix/cli/project_plugin.py new file mode 100644 index 0000000..e69de29 diff --git a/mdmix/cli/solvent_plugin.py b/mdmix/cli/solvent_plugin.py new file mode 100644 index 0000000..3bc5403 --- /dev/null +++ b/mdmix/cli/solvent_plugin.py @@ -0,0 +1,275 @@ +import typer +from pathlib import Path +from typing import Optional, List +import yaml + +from cli.base import CLIPlugin, DisplayUtils, console +from core.models import ( + ProjectRequest, SystemRequest, ReplicaRequest, + SimulationRequest, RestraintMode, EnsembleType, + SimulationEngine, QueueSystem +) +from core.services import ProjectSetupService + + +class ProjectPlugin(CLIPlugin): + """Project management commands""" + + def __init__(self): + super().__init__("project", "Project setup and management commands") + self.service = ProjectSetupService() + + def register_commands(self): + """Register project commands""" + + @self.app.command("create") + def create_project( + config_file: Path = typer.Argument(..., + help="YAML configuration file for project", + exists=True, + file_okay=True, + dir_okay=False, + readable=True + ), + output_dir: Optional[Path] = typer.Option(None, "--output", "-o", + help="Output directory (default: current directory)" + ), + dry_run: bool = typer.Option(False, "--dry-run", + help="Show what would be created without actually creating files" + ) + ): + """Create a new PyMDMix project from YAML configuration""" + + DisplayUtils.show_header("PyMDMix3 Project Creation", + f"Configuration: {config_file.name}") + + try: + # Load configuration + with open(config_file, 'r') as f: + config_data = yaml.safe_load(f) + + # Create request from config + request = self._parse_project_config(config_data) + + if output_dir: + request.output_dir = output_dir + + # Show project summary + self._show_project_summary(request) + + if dry_run: + DisplayUtils.show_info("Dry run mode - no files created") + return + + # Confirm creation + if not typer.confirm("Create project with these settings?"): + raise typer.Abort() + + # Create project + with console.status("[bold green]Creating project..."): + response = self.service.setup_project(request) + + # Show results + self._show_creation_results(response) + + except Exception as e: + DisplayUtils.show_error(f"Project creation failed: {e}") + raise typer.Exit(1) + + @self.app.command("setup") + def setup_project( + name: str = typer.Argument(..., help="Project name"), + structure_file: Path = typer.Argument(..., + help="Structure file (OFF or PDB)", + exists=True + ), + solvents: List[str] = typer.Option(..., "--solvent", "-s", + help="Solvents to use (can specify multiple)" + ), + replicas: int = typer.Option(3, "--replicas", "-r", + help="Number of replicas per solvent" + ), + time: float = typer.Option(20.0, "--time", "-t", + help="Simulation time in nanoseconds" + ), + temperature: float = typer.Option(300.0, "--temp", + help="Temperature in Kelvin" + ), + restraints: RestraintMode = typer.Option(RestraintMode.FREE, "--restraints", + help="Restraint mode" + ), + output_dir: Optional[Path] = typer.Option(None, "--output", "-o", + help="Output directory" + ), + force_fields: List[str] = typer.Option( + ["leaprc.protein.ff14SB", "leaprc.water.tip3p"], + "--forcefield", "-ff", + help="Force field files" + ) + ): + """Quick project setup with command-line parameters""" + + DisplayUtils.show_header("PyMDMix3 Quick Project Setup", name) + + # Determine structure type + structure_type = "off" if structure_file.suffix.lower() == ".off" else "pdb" + + # Build request + request = ProjectRequest( + name=name, + description=f"PyMDMix3 project for {structure_file.stem}", + system=SystemRequest( + name=structure_file.stem, + structure_file=structure_file, + structure_type=structure_type, + force_fields=force_fields + ), + replicas=ReplicaRequest( + count=len(solvents) * replicas, + solvents=solvents * replicas # Replicate solvent list + ), + simulation=SimulationRequest( + total_time=time, + temperature=temperature, + restraint_mode=restraints + ), + output_dir=output_dir + ) + + # Show summary + self._show_project_summary(request) + + # Confirm + if not typer.confirm("Create project?"): + raise typer.Abort() + + # Create project + with console.status("[bold green]Creating project..."): + response = self.service.setup_project(request) + + self._show_creation_results(response) + + @self.app.command("template") + def generate_template( + output_file: Path = typer.Option("project_template.yaml", "--output", "-o", + help="Output file for template" + ) + ): + """Generate a template YAML configuration file""" + + template = { + "project": { + "name": "my_project", + "description": "Mixed-solvent MD simulation project" + }, + "system": { + "name": "protein_system", + "structure_file": "path/to/structure.off", + "structure_type": "off", # or "pdb" + "force_fields": [ + "leaprc.protein.ff14SB", + "leaprc.water.tip3p" + ], + "extra_files": [], # Additional parameter files + "box_buffer": 12.0 + }, + "replicas": { + "count": 3, + "solvents": ["WAT", "ETA", "WAT"], # One per replica + "naming_scheme": "{solvent}_{index}" + }, + "simulation": { + "total_time": 20.0, # nanoseconds + "timestep": 2.0, # femtoseconds + "temperature": 300.0, # Kelvin + "pressure": 1.0, # bar + "ensemble": "NPT", # or "NVT" + "restraint_mode": "FREE", # or "HA", "BB" + "restraint_force": 5.0, + "output_frequency": 5000, + "log_frequency": 1000, + "engine": "AMBER" # or "NAMD", "OPENMM" + }, + "queue_system": "SLURM" # or "SGE", "LOCAL", etc. + } + + with open(output_file, 'w') as f: + yaml.dump(template, f, default_flow_style=False, sort_keys=False) + + DisplayUtils.show_success(f"Template saved to: {output_file}") + DisplayUtils.show_info("Edit the template and use 'pymdmix3 project create' to create your project") + + def _parse_project_config(self, config: dict) -> ProjectRequest: + """Parse YAML configuration into ProjectRequest""" + + # Extract sections + project_cfg = config.get('project', {}) + system_cfg = config.get('system', {}) + replica_cfg = config.get('replicas', {}) + sim_cfg = config.get('simulation', {}) + + # Build request + return ProjectRequest( + name=project_cfg.get('name', 'pymdmix_project'), + description=project_cfg.get('description', ''), + system=SystemRequest(**system_cfg), + replicas=ReplicaRequest(**replica_cfg), + simulation=SimulationRequest(**sim_cfg), + queue_system=QueueSystem(config.get('queue_system', 'LOCAL')) + ) + + def _show_project_summary(self, request: ProjectRequest): + """Display project configuration summary""" + + summary_table = DisplayUtils.create_table("Project Configuration", + ["Setting", "Value"]) + + summary_table.add_row("Project Name", request.name) + summary_table.add_row("System", request.system.name) + summary_table.add_row("Structure", str(request.system.structure_file)) + summary_table.add_row("Total Replicas", str(request.replicas.count)) + summary_table.add_row("Solvents", ", ".join(set(request.replicas.solvents))) + summary_table.add_row("Simulation Time", f"{request.simulation.total_time} ns") + summary_table.add_row("Temperature", f"{request.simulation.temperature} K") + summary_table.add_row("Ensemble", request.simulation.ensemble.value) + summary_table.add_row("Restraints", request.simulation.restraint_mode.value) + + console.print(summary_table) + + def _show_creation_results(self, response): + """Display project creation results""" + + if response.status == "success": + DisplayUtils.show_success(f"Project created successfully!") + else: + DisplayUtils.show_warning(f"Project created with issues") + + # Summary + console.print(f"\nProject Directory: [bold]{response.project_dir}[/bold]") + console.print(f"Total Replicas: {response.total_replicas}") + console.print(f"Successful: [green]{response.successful_replicas}[/green]") + + if response.successful_replicas < response.total_replicas: + failed = response.total_replicas - response.successful_replicas + console.print(f"Failed: [red]{failed}[/red]") + + # Show any errors + if response.system.errors: + DisplayUtils.show_error("System setup errors:") + for error in response.system.errors: + console.print(f" - {error}") + + # Show failed replicas + failed_replicas = [r for r in response.replicas if r.status == "failed"] + if failed_replicas: + DisplayUtils.show_error("Failed replicas:") + for replica in failed_replicas: + console.print(f" - {replica.name}: {replica.error}") + + # Next steps + console.print("\n[bold]Next Steps:[/bold]") + console.print("1. Change to project directory:") + console.print(f" cd {response.project_dir}") + console.print("2. Run system preparation:") + console.print(" ./scripts/setup_project.sh") + console.print("3. Submit simulations to queue or run locally") \ No newline at end of file diff --git a/mdmix/core/models.py b/mdmix/core/models.py new file mode 100644 index 0000000..44ae3d8 --- /dev/null +++ b/mdmix/core/models.py @@ -0,0 +1,244 @@ +from pathlib import Path +from typing import List, Dict, Optional, Any, Union +from pydantic import BaseModel, Field, field_validator +from enum import Enum +from datetime import datetime + + +# Enums for validation +class RestraintMode(str, Enum): + FREE = "FREE" + HA = "HA" # Heavy atoms + BB = "BB" # Backbone atoms + + +class EnsembleType(str, Enum): + NVT = "NVT" + NPT = "NPT" + + +class SimulationEngine(str, Enum): + AMBER = "AMBER" + NAMD = "NAMD" + OPENMM = "OPENMM" + + +class QueueSystem(str, Enum): + SLURM = "SLURM" + SGE = "SGE" + LOCAL = "LOCAL" + CTE = "CTE" + DAINT = "DAINT" + + +# Request Models (for CLI input) +class SystemRequest(BaseModel): + """Request model for system setup""" + name: str = Field(..., description="System name") + structure_file: Path = Field(..., description="Path to OFF or PDB file") + structure_type: str = Field("off", description="Structure file type (off/pdb)") + force_fields: List[str] = Field( + default_factory=lambda: ["leaprc.protein.ff14SB", "leaprc.water.tip3p"], + description="Force field files" + ) + extra_files: List[Path] = Field(default_factory=list, description="Additional parameter files") + box_buffer: float = Field(12.0, description="Solvation buffer in Angstroms") + + @field_validator('structure_file') + @classmethod + def validate_structure_file(cls, v): + if not v.exists(): + raise ValueError(f"Structure file not found: {v}") + return v + + @field_validator('structure_type') + @classmethod + def validate_structure_type(cls, v): + if v.lower() not in ['off', 'pdb']: + raise ValueError("Structure type must be 'off' or 'pdb'") + return v.lower() + + +class ReplicaRequest(BaseModel): + """Request model for replica configuration""" + count: int = Field(3, ge=1, description="Number of replicas") + solvents: List[str] = Field(..., description="List of solvents for replicas") + naming_scheme: str = Field("{solvent}_{index}", description="Replica naming pattern") + + @field_validator('solvents') + @classmethod + def validate_solvents(cls, v, info): # Change parameter name + # Access other field values through info.data + if info.data and 'count' in info.data: + count = info.data['count'] + if len(v) != count: + # If single solvent provided, replicate for all replicas + if len(v) == 1: + return v * count + raise ValueError(f"Number of solvents ({len(v)}) must match replica count ({count})") + return v + + + +class SimulationRequest(BaseModel): + """Request model for simulation parameters""" + total_time: float = Field(20.0, gt=0, description="Total simulation time in ns") + timestep: float = Field(2.0, gt=0, description="Integration timestep in fs") + temperature: float = Field(300.0, gt=0, description="Temperature in K") + pressure: float = Field(1.0, description="Pressure in bar") + ensemble: EnsembleType = Field(EnsembleType.NPT, description="Thermodynamic ensemble") + restraint_mode: RestraintMode = Field(RestraintMode.FREE, description="Restraint mode") + restraint_force: float = Field(5.0, ge=0, description="Restraint force constant") + output_frequency: int = Field(5000, gt=0, description="Trajectory output frequency") + log_frequency: int = Field(1000, gt=0, description="Log output frequency") + engine: SimulationEngine = Field(SimulationEngine.AMBER, description="MD engine") + + +class ProjectRequest(BaseModel): + """Complete project request model""" + name: str = Field(..., description="Project name") + description: str = Field("", description="Project description") + system: SystemRequest + replicas: ReplicaRequest + simulation: SimulationRequest + output_dir: Optional[Path] = Field(None, description="Output directory") + queue_system: QueueSystem = Field(QueueSystem.LOCAL, description="Queue system") + + +# Response Models (for service output) +class SystemResponse(BaseModel): + """Response model for system setup results""" + name: str + status: str + files_created: Dict[str, Path] + warnings: List[str] = Field(default_factory=list) + errors: List[str] = Field(default_factory=list) + + +class ReplicaResponse(BaseModel): + """Response model for replica status""" + name: str + solvent: str + status: str + directory: Optional[Path] = None + files: Dict[str, Path] = Field(default_factory=dict) + error: Optional[str] = None + + +class ProjectResponse(BaseModel): + """Response model for complete project setup""" + project_name: str + status: str + created_at: datetime = Field(default_factory=datetime.now) + project_dir: Path + system: SystemResponse + replicas: List[ReplicaResponse] + total_replicas: int + successful_replicas: int + scripts_created: Dict[str, Path] = Field(default_factory=dict) + + def get_summary(self) -> Dict[str, Any]: + """Get project summary""" + return { + "project": self.project_name, + "status": self.status, + "total_replicas": self.total_replicas, + "successful": self.successful_replicas, + "failed": self.total_replicas - self.successful_replicas, + "directory": str(self.project_dir) + } + + +class SolventInfo(BaseModel): + """Solvent information response""" + name: str + display_name: str + description: str = "" + density: float + molecular_weight: float + water_model: str = "TIP3P" + ionic: bool = False + residue_counts: Dict[str, int] = Field(default_factory=dict) + probes: Dict[str, Dict[str, Any]] = Field(default_factory=dict) + + +class AnalysisRequest(BaseModel): + """Request model for analysis parameters""" + project_dir: Path + replicas: List[str] = Field(default_factory=list, description="Specific replicas to analyze") + probes: List[str] = Field(default_factory=list, description="Specific probes to analyze") + cutoff: float = Field(3.5, description="Probe cutoff distance") + grid_spacing: float = Field(0.5, description="Grid spacing for density calculation") + output_format: str = Field("dx", description="Output format for grids") + + +# Template Variables Models +class TemplateVariables(BaseModel): + """Base model for template variables""" + pass + + +class MinimizationTemplateVars(TemplateVariables): + """Variables for minimization templates""" + maxcyc: int = 5000 + ncyc: int = 2500 + ntpr: int = 100 + restraint_wt: float = 10.0 + restraintmask: str = "" + + +class EquilibrationTemplateVars(TemplateVariables): + """Variables for equilibration templates""" + nsteps: int = 25000 + timestep: float = 2.0 + temp0: float = 300.0 + tempi: float = 100.0 + ntpr: int = 1000 + ntwx: int = 5000 + restraint_wt: float = 5.0 + restraintmask: str = "" + stage: int = 1 + + +class ProductionTemplateVars(TemplateVariables): + """Variables for production templates""" + nsteps: int = 500000 + timestep: float = 2.0 + temp0: float = 300.0 + pres0: float = 1.0 + ntpr: int = 1000 + ntwx: int = 5000 + ioutfm: int = 1 # NetCDF format + iwrap: int = 1 # Wrap coordinates + restraint_wt: float = 0.0 + restraintmask: str = "" + + +# Configuration Models (for YAML) +class YAMLConfig(BaseModel): + """Base configuration model for YAML files""" + + class Config: + extra = "allow" + use_enum_values = True + + def to_yaml(self, path: Path): + """Save configuration to YAML file""" + import yaml + with open(path, 'w') as f: + yaml.dump(self.dict(), f, default_flow_style=False) + + @classmethod + def from_yaml(cls, path: Path): + """Load configuration from YAML file""" + import yaml + with open(path, 'r') as f: + data = yaml.safe_load(f) + return cls(**data) + + +class ProjectConfig(YAMLConfig): + """Project configuration for YAML storage""" + project: ProjectRequest + created_at: datetime = Field(default_factory=datetime.now) + pymdmix_version: str = "3.0.0" diff --git a/mdmix/core/services.py b/mdmix/core/services.py new file mode 100644 index 0000000..c109113 --- /dev/null +++ b/mdmix/core/services.py @@ -0,0 +1,610 @@ +from pathlib import Path +from typing import List, Dict, Optional, Any +import logging +from datetime import datetime +import shutil +import yaml + +from core.models import ( + ProjectRequest, SystemRequest, ReplicaRequest, SimulationRequest, + ProjectResponse, SystemResponse, ReplicaResponse, + SolventInfo, AnalysisRequest +) + +from core.template_engine import ScriptTemplateEngine +from solvents import SolventDatabase + + +class SystemSetupService: + """Service for setting up molecular systems""" + + def __init__(self, template_engine: ScriptTemplateEngine): + self.template_engine = template_engine + self.logger = logging.getLogger("SystemSetupService") + + def setup_system(self, request: SystemRequest, output_dir: Path) -> SystemResponse: + """Setup molecular system and create necessary scripts""" + + self.logger.info(f"Setting up system: {request.name}") + + # Create system directory + system_dir = output_dir / "system" + system_dir.mkdir(parents=True, exist_ok=True) + + files_created = {} + warnings = [] + errors = [] + + try: + # Copy structure file + dest_file = system_dir / f"system.{request.structure_file.suffix[1:]}" + shutil.copy2(request.structure_file, dest_file) + files_created['structure'] = dest_file + + # Create tLEaP script + leap_script = self._create_leap_script(request, system_dir) + files_created['leap_script'] = leap_script + + # Create execution script + exec_script = self._create_execution_script(request, system_dir) + files_created['execution_script'] = exec_script + + # Copy additional files + for extra_file in request.extra_files: + dest = system_dir / extra_file.name + shutil.copy2(extra_file, dest) + files_created[f'extra_{extra_file.stem}'] = dest + + status = "success" + + except Exception as e: + self.logger.error(f"System setup failed: {e}") + errors.append(str(e)) + status = "failed" + + return SystemResponse( + name=request.name, + status=status, + files_created=files_created, + warnings=warnings, + errors=errors + ) + + def _create_leap_script(self, request: SystemRequest, output_dir: Path) -> Path: + """Create tLEaP script for system preparation""" + + script_lines = [ + f"# PyMDMix3 System Preparation Script", + f"# System: {request.name}", + f"# Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", + "" + ] + + # Load force fields + for ff in request.force_fields: + if ff.startswith('leaprc'): + script_lines.append(f"source {ff}") + else: + script_lines.append(f"loadAmberParams {ff}") + + # Load additional parameter files + for extra_file in request.extra_files: + if extra_file.suffix == '.frcmod': + script_lines.append(f"loadAmberParams {extra_file.name}") + elif extra_file.suffix == '.off': + script_lines.append(f"loadOff {extra_file.name}") + + script_lines.append("") + + # Load structure + if request.structure_type == "off": + script_lines.extend([ + f"loadOff system.off", + "mol = first", # Get first unit + "check mol" + ]) + else: # PDB + script_lines.extend([ + f"mol = loadPdb system.pdb", + "check mol" + ]) + + script_lines.extend([ + "", + "# Save dry system", + f"saveAmberParm mol {request.name}_dry.prmtop {request.name}_dry.inpcrd", + f"savePdb mol {request.name}_dry.pdb", + f"saveOff mol {request.name}_dry.off", + "", + "quit" + ]) + + # Write script + script_path = output_dir / "prepare_system.leap" + with open(script_path, 'w') as f: + f.write('\n'.join(script_lines)) + + return script_path + + def _create_execution_script(self, request: SystemRequest, output_dir: Path) -> Path: + """Create bash execution script""" + + script_content = f"""#!/bin/bash +# PyMDMix3 System Preparation +# System: {request.name} + +echo "=== PyMDMix3 System Preparation ===" +echo "System: {request.name}" +echo "Date: $(date)" + +cd "$(dirname "$0")" + +if ! command -v tleap &> /dev/null; then + echo "ERROR: tLEaP not found. Please load AMBER." + exit 1 +fi + +echo "Running tLEaP..." +tleap -f prepare_system.leap + +if [ -f "{request.name}_dry.prmtop" ]; then + echo "SUCCESS: System prepared" + ls -la *.prmtop *.inpcrd *.pdb *.off 2>/dev/null +else + echo "ERROR: System preparation failed. Check leap.log" + exit 1 +fi +""" + + script_path = output_dir / "run_preparation.sh" + with open(script_path, 'w') as f: + f.write(script_content) + script_path.chmod(0o755) + + return script_path + + +class ReplicaSetupService: + """Service for setting up simulation replicas""" + + def __init__(self, template_engine: ScriptTemplateEngine, + solvent_db: SolventDatabase): + self.template_engine = template_engine + self.solvent_db = solvent_db + self.logger = logging.getLogger("ReplicaSetupService") + + def setup_replicas(self, request: ReplicaRequest, system_name: str, + simulation_request: SimulationRequest, + output_dir: Path) -> List[ReplicaResponse]: + """Setup all simulation replicas""" + + replicas_dir = output_dir / "replicas" + replicas_dir.mkdir(parents=True, exist_ok=True) + + responses = [] + + for i in range(request.count): + solvent_name = request.solvents[i] + replica_name = request.naming_scheme.format( + solvent=solvent_name, + index=i+1 + ) + + response = self._setup_single_replica( + replica_name, solvent_name, system_name, + simulation_request, replicas_dir + ) + responses.append(response) + + return responses + + def _setup_single_replica(self, name: str, solvent_name: str, + system_name: str, sim_request: SimulationRequest, + replicas_dir: Path) -> ReplicaResponse: + """Setup a single replica""" + + self.logger.info(f"Setting up replica: {name}") + + replica_dir = replicas_dir / name + files = {} + + try: + # Create directory structure + for subdir in ["system", "minimize", "equilibrate", "production", "logs"]: + (replica_dir / subdir).mkdir(parents=True, exist_ok=True) + + # Get solvent info + solvent = self.solvent_db.get_solvent(solvent_name) + if not solvent: + raise ValueError(f"Solvent not found: {solvent_name}") + + # Create solvation script + solv_script = self._create_solvation_script( + solvent, system_name, replica_dir / "system" + ) + files['solvation_script'] = solv_script + + # Create MD input files + md_files = self._create_md_inputs( + sim_request, replica_dir + ) + files.update(md_files) + + # Create run scripts + run_scripts = self._create_run_scripts( + name, replica_dir + ) + files.update(run_scripts) + + status = "success" + error = None + + except Exception as e: + self.logger.error(f"Replica setup failed: {e}") + status = "failed" + error = str(e) + + return ReplicaResponse( + name=name, + solvent=solvent_name, + status=status, + directory=replica_dir, + files=files, + error=error + ) + + def _create_solvation_script(self, solvent: SolventInfo, + system_name: str, output_dir: Path) -> Path: + """Create solvation tLEaP script""" + + script_lines = [ + f"# PyMDMix3 Solvation Script", + f"# Solvent: {solvent.display_name}", + f"# Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", + "", + "# Load force fields", + "source leaprc.protein.ff14SB", + "source leaprc.water.tip3p", + "" + ] + + # Load water model if different + if solvent.water_model and solvent.water_model != "TIP3P": + script_lines.append(f"source leaprc.water.{solvent.water_model.lower()}") + + # Load solvent parameters (simplified - in real implementation would handle OFF files) + script_lines.extend([ + "", + "# Load system", + f"mol = loadOff ../../system/{system_name}_dry.off", + "", + "# Solvate", + f"solvatebox mol {solvent.water_model}BOX 12.0", + "", + "# Add ions to neutralize", + "addions mol Na+ 0", + "addions mol Cl- 0", + "", + "# Save solvated system", + "saveAmberParm mol solvated.prmtop solvated.inpcrd", + "savePdb mol solvated.pdb", + "", + "quit" + ]) + + script_path = output_dir / "solvate.leap" + with open(script_path, 'w') as f: + f.write('\n'.join(script_lines)) + + return script_path + + def _create_md_inputs(self, sim_request: SimulationRequest, + replica_dir: Path) -> Dict[str, Path]: + """Create MD input files""" + + files = {} + + # Minimization + min_vars = { + 'minsteps': 5000, + 'ntpr': 100, + 'restraintmask': '@CA,C,N' if sim_request.restraint_mode != 'FREE' else '', + 'restraint_wt': 10.0 + } + + min_content = self.template_engine.create_amber_minimization(min_vars) + if min_content: + min_file = replica_dir / "minimize" / "min.in" + min_file.write_text(min_content) + files['minimize_input'] = min_file + + # Equilibration (5 stages) + for stage in range(1, 6): + eq_vars = { + 'final_temp': sim_request.temperature, + 'restraintmask': '@CA,C,N' if sim_request.restraint_mode != 'FREE' else '', + 'restraint_wt': 10.0 / stage, # Reduce restraints + 'nsteps': 25000 if stage < 5 else 50000 + } + + eq_content = self.template_engine.create_amber_equilibration(stage, eq_vars) + if eq_content: + eq_file = replica_dir / "equilibrate" / f"eq{stage}.in" + eq_file.write_text(eq_content) + files[f'equilibrate_{stage}_input'] = eq_file + + # Production + prod_vars = { + 'nsteps': int(sim_request.total_time * 500000), # 500k steps = 1ns + 'temp0': sim_request.temperature, + 'pres0': sim_request.pressure, + 'timestep': sim_request.timestep / 1000, # fs to ps + 'freq': sim_request.output_frequency, + 'ntpr': sim_request.log_frequency, + 'restraintmask': self._get_restraint_mask(sim_request.restraint_mode), + 'restraint_wt': sim_request.restraint_force if sim_request.restraint_mode != 'FREE' else 0 + } + + prod_content = self.template_engine.create_amber_production( + sim_request.ensemble.value, prod_vars + ) + if prod_content: + prod_file = replica_dir / "production" / "prod.in" + prod_file.write_text(prod_content) + files['production_input'] = prod_file + + return files + + def _get_restraint_mask(self, mode: str) -> str: + """Get restraint mask based on mode""" + if mode == 'HA': + return '!@H=' + elif mode == 'BB': + return '@CA,C,N' + else: + return '' + + def _create_run_scripts(self, replica_name: str, replica_dir: Path) -> Dict[str, Path]: + """Create execution scripts for the replica""" + + files = {} + + # Master run script + master_script = f"""#!/bin/bash +# PyMDMix3 Replica Execution Script +# Replica: {replica_name} + +echo "=== PyMDMix3 Replica: {replica_name} ===" +echo "Date: $(date)" + +cd "$(dirname "$0")" + +# Step 1: Solvation +echo "Step 1: Solvating system..." +cd system +tleap -f solvate.leap +if [ ! -f "solvated.prmtop" ]; then + echo "ERROR: Solvation failed" + exit 1 +fi +cd .. + +# Step 2: Minimization +echo "Step 2: Running minimization..." +cd minimize +pmemd.cuda -O -i min.in -p ../system/solvated.prmtop -c ../system/solvated.inpcrd \\ + -o min.out -r min.rst7 -inf min.mdinfo +cd .. + +# Step 3: Equilibration +echo "Step 3: Running equilibration..." +cd equilibrate +for i in {{1..5}}; do + echo " Stage $i..." + if [ $i -eq 1 ]; then + pmemd.cuda -O -i eq${{i}}.in -p ../system/solvated.prmtop \\ + -c ../minimize/min.rst7 -o eq${{i}}.out -r eq${{i}}.rst7 \\ + -x eq${{i}}.nc -inf eq${{i}}.mdinfo + else + prev=$((i-1)) + pmemd.cuda -O -i eq${{i}}.in -p ../system/solvated.prmtop \\ + -c eq${{prev}}.rst7 -o eq${{i}}.out -r eq${{i}}.rst7 \\ + -x eq${{i}}.nc -inf eq${{i}}.mdinfo + fi +done +cd .. + +# Step 4: Production +echo "Step 4: Running production MD..." +cd production +pmemd.cuda -O -i prod.in -p ../system/solvated.prmtop \\ + -c ../equilibrate/eq5.rst7 -o prod.out -r prod.rst7 \\ + -x prod.nc -inf prod.mdinfo +cd .. + +echo "Replica {replica_name} completed at: $(date)" +""" + + script_path = replica_dir / "run_replica.sh" + with open(script_path, 'w') as f: + f.write(master_script) + script_path.chmod(0o755) + files['run_script'] = script_path + + return files + + +class ProjectSetupService: + """Main service for project setup orchestration""" + + def __init__(self): + self.template_engine = ScriptTemplateEngine() + self.solvent_db = SolventDatabase() + self.system_service = SystemSetupService(self.template_engine) + self.replica_service = ReplicaSetupService(self.template_engine, self.solvent_db) + self.logger = logging.getLogger("ProjectSetupService") + + def setup_project(self, request: ProjectRequest) -> ProjectResponse: + """Setup complete PyMDMix project""" + + self.logger.info(f"Setting up project: {request.name}") + + # Determine output directory + if request.output_dir: + project_dir = request.output_dir / request.name + else: + project_dir = Path.cwd() / request.name + + project_dir.mkdir(parents=True, exist_ok=True) + + # Setup system + system_response = self.system_service.setup_system( + request.system, project_dir + ) + + # Setup replicas + replica_responses = self.replica_service.setup_replicas( + request.replicas, request.system.name, + request.simulation, project_dir + ) + + # Create project-level scripts + scripts = self._create_project_scripts(request, project_dir) + + # Save project configuration + self._save_project_config(request, project_dir) + + # Determine overall status + successful_replicas = sum(1 for r in replica_responses if r.status == "success") + status = "success" if successful_replicas == len(replica_responses) else "partial" + + return ProjectResponse( + project_name=request.name, + status=status, + project_dir=project_dir, + system=system_response, + replicas=replica_responses, + total_replicas=len(replica_responses), + successful_replicas=successful_replicas, + scripts_created=scripts + ) + + def _create_project_scripts(self, request: ProjectRequest, + project_dir: Path) -> Dict[str, Path]: + """Create project-level scripts""" + + scripts_dir = project_dir / "scripts" + scripts_dir.mkdir(exist_ok=True) + + scripts = {} + + # Master setup script + setup_script = f"""#!/bin/bash +# PyMDMix3 Project Setup +# Project: {request.name} + +echo "=== PyMDMix3 Project Setup ===" +echo "Project: {request.name}" +echo "Replicas: {request.replicas.count}" +echo "" + +cd "$(dirname "$0")/.." + +# System preparation +echo "Preparing system..." +cd system +./run_preparation.sh +cd .. + +# Setup all replicas +echo "Setting up replicas..." +for replica in replicas/*/; do + if [ -d "$replica" ]; then + echo "Processing $(basename "$replica")..." + cd "$replica/system" + tleap -f solvate.leap + cd ../../.. + fi +done + +echo "Setup completed!" +""" + + setup_path = scripts_dir / "setup_project.sh" + with open(setup_path, 'w') as f: + f.write(setup_script) + setup_path.chmod(0o755) + scripts['setup'] = setup_path + + return scripts + + def _save_project_config(self, request: ProjectRequest, project_dir: Path): + """Save project configuration to YAML""" + + config_file = project_dir / "project_config.yaml" + + config_data = { + 'project': request.dict(), + 'created_at': datetime.now().isoformat(), + 'pymdmix_version': '3.0.0' + } + + with open(config_file, 'w') as f: + yaml.dump(config_data, f, default_flow_style=False) + + +class SolventManagementService: + """Service for solvent database management""" + + def __init__(self): + self.solvent_db = SolventDatabase() + self.logger = logging.getLogger("SolventManagementService") + + def list_solvents(self) -> List[SolventInfo]: + """List all available solvents""" + solvent_names = self.solvent_db.list_solvents() + solvents = [] + + for name in solvent_names: + solvent = self.solvent_db.get_solvent(name) + if solvent: + solvents.append(SolventInfo( + name=solvent.name, + display_name=solvent.display_name, + description=solvent.description, + density=solvent.density, + molecular_weight=solvent.molecular_weight, + water_model=solvent.water_model, + ionic=solvent.ionic, + residue_counts=solvent.residue_counts, + probes=solvent.probes + )) + + return solvents + + def get_solvent_info(self, name: str) -> Optional[SolventInfo]: + """Get detailed information about a solvent""" + solvent = self.solvent_db.get_solvent(name) + + if solvent: + return SolventInfo( + name=solvent.name, + display_name=solvent.display_name, + description=solvent.description, + density=solvent.density, + molecular_weight=solvent.molecular_weight, + water_model=solvent.water_model, + ionic=solvent.ionic, + residue_counts=solvent.residue_counts, + probes=solvent.probes + ) + + return None + + def add_solvent_from_config(self, config_file: Path) -> bool: + """Add solvent from configuration file""" + try: + return self.solvent_db.load_from_config_file(config_file) + except Exception as e: + self.logger.error(f"Failed to add solvent: {e}") + return False \ No newline at end of file diff --git a/mdmix/core/template_engine.py b/mdmix/core/template_engine.py new file mode 100644 index 0000000..e69de29 diff --git a/mdmix/solvents.py b/mdmix/solvents.py new file mode 100644 index 0000000..e69de29