{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tymworld/DROPPS/blob/v1.0.0/examples/colab/02_DROPPS_1_0_Simulation.ipynb)\n\n# DROPPS 1.0 Phase-Coexistence Simulation\n\nThis notebook implements the slab protocol in its required order:\n\n1. start from many protein monomers in a compact cubic box;\n2. condense/equilibrate the box in the NPT ensemble;\n3. expand the **z box axis tenfold**, recentering the condensed\n   configuration in the resulting elongated box;\n4. run the elongated box in the NVT ensemble;\n5. pass only `slab_nvt.tpr` and `slab_nvt.xtc` to trajectory processing and\n   phase-coexistence analysis.\n\nThe manuscript benchmark settings are 300 ns NPT at 0.01 ps followed by\n3 μs NVT at 0.02 ps. A short full-workflow mode validates all transitions\nwithout claiming a physically equilibrated phase-separated state.\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Install DROPPS 1.0\n#@markdown The default route installs the immutable wheel attached to the\n#@markdown public GitHub `v1.0.0` release. Upload remains available for offline use.\ninstallation_source = \"Install from the DROPPS v1.0.0 GitHub release\" #@param [\"Install from the DROPPS v1.0.0 GitHub release\", \"Upload the DROPPS 1.0 wheel\", \"Install from another wheel URL\"]\nwheel_url = \"https://github.com/tymworld/DROPPS/releases/download/v1.0.0/dropps-1.0-py3-none-any.whl\" #@param {type:\"string\"}\naccelerator_dependencies = \"auto\" #@param [\"auto\", \"cuda12\", \"cuda13\", \"latest\"]\n\nimport importlib.metadata\nimport re\nimport shutil\nimport subprocess\nimport sys\nfrom pathlib import Path\n\n\ndef detect_install_extra():\n    if accelerator_dependencies != \"auto\":\n        return accelerator_dependencies\n    if shutil.which(\"nvidia-smi\"):\n        probe = subprocess.run(\n            [\"nvidia-smi\"], text=True, capture_output=True, check=False\n        ).stdout\n        match = re.search(r\"CUDA Version:\\s*(\\d+)\", probe)\n        if match and int(match.group(1)) >= 13:\n            return \"cuda13\"\n        return \"cuda12\"\n    return \"latest\"\n\n\nif installation_source.startswith(\"Upload\"):\n    try:\n        from google.colab import files\n    except ImportError as exc:\n        raise RuntimeError(\n            \"This upload form is intended for Google Colab. Set installation_source \"\n            \"to the URL option when running elsewhere.\"\n        ) from exc\n    uploaded = files.upload()\n    wheel_candidates = [Path(name) for name in uploaded if name.endswith(\".whl\")]\n    if len(wheel_candidates) != 1:\n        raise ValueError(\"Upload exactly one DROPPS 1.0 .whl file.\")\n    wheel_target = str(wheel_candidates[0].resolve())\nelse:\n    if not wheel_url.strip():\n        raise ValueError(\"Provide the published DROPPS 1.0 wheel URL.\")\n    wheel_target = wheel_url.strip()\n\nextra = detect_install_extra()\nif wheel_target.startswith((\"https://\", \"http://\")):\n    package_spec = f\"dropps[{extra}] @ {wheel_target}\"\nelse:\n    package_spec = f\"{wheel_target}[{extra}]\"\nprint(f\"Installing {package_spec}\")\nsubprocess.run(\n    [\n        sys.executable,\n        \"-m\",\n        \"pip\",\n        \"install\",\n        \"--quiet\",\n        \"--upgrade\",\n        \"--upgrade-strategy\",\n        \"only-if-needed\",\n        package_spec,\n    ],\n    check=True,\n)\n\nversion = importlib.metadata.version(\"dropps\")\nif version != \"1.0\":\n    raise RuntimeError(f\"Expected DROPPS 1.0, but installed {version}.\")\nsubprocess.run([\"dps\", \"--version\"], check=True)\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Shared notebook helpers\nimport hashlib\nimport json\nimport os\nimport shlex\nimport subprocess\nimport zipfile\nfrom datetime import datetime, timezone\nfrom pathlib import Path\n\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nCOMMAND_LOG = []\n\n\ndef run_command(arguments, *, cwd=None, input_text=None, check=True):\n    command = [str(value) for value in arguments]\n    print(\"$\", shlex.join(command))\n    result = subprocess.run(\n        command,\n        cwd=None if cwd is None else str(cwd),\n        input=input_text,\n        text=True,\n        capture_output=True,\n        check=False,\n    )\n    if result.stdout:\n        print(result.stdout, end=\"\" if result.stdout.endswith(\"\\n\") else \"\\n\")\n    if result.stderr:\n        print(result.stderr, end=\"\" if result.stderr.endswith(\"\\n\") else \"\\n\")\n    COMMAND_LOG.append(\n        {\n            \"time_utc\": datetime.now(timezone.utc).isoformat(),\n            \"cwd\": str(Path(cwd or Path.cwd()).resolve()),\n            \"command\": command,\n            \"returncode\": result.returncode,\n        }\n    )\n    if check and result.returncode:\n        raise RuntimeError(\n            f\"Command failed with exit code {result.returncode}: {shlex.join(command)}\"\n        )\n    return result\n\n\ndef dps(*arguments, cwd=None, input_text=None, check=True):\n    return run_command(\n        [\"dps\", *arguments], cwd=cwd, input_text=input_text, check=check\n    )\n\n\ndef reset_task_directory(path):\n    path = Path(path).resolve()\n    if not path.name.startswith(\"dropps_\"):\n        raise ValueError(f\"Refusing to reset unexpected directory: {path}\")\n    if path.exists():\n        import shutil\n\n        shutil.rmtree(path)\n    path.mkdir(parents=True)\n    return path\n\n\ndef sha256(path):\n    digest = hashlib.sha256()\n    with Path(path).open(\"rb\") as stream:\n        for block in iter(lambda: stream.read(1024 * 1024), b\"\"):\n            digest.update(block)\n    return digest.hexdigest()\n\n\ndef safe_extract_zip(archive, destination):\n    destination = Path(destination).resolve()\n    destination.mkdir(parents=True, exist_ok=True)\n    with zipfile.ZipFile(archive) as bundle:\n        for member in bundle.infolist():\n            target = (destination / member.filename).resolve()\n            if destination not in target.parents and target != destination:\n                raise ValueError(f\"Unsafe ZIP member: {member.filename}\")\n        bundle.extractall(destination)\n\n\ndef find_unique(root, basename):\n    matches = [path for path in Path(root).rglob(basename) if path.is_file()]\n    if len(matches) != 1:\n        raise FileNotFoundError(\n            f\"Expected one file named {basename!r} below {root}, found {len(matches)}.\"\n        )\n    return matches[0]\n\n\ndef make_zip(paths, output_path, *, base=None):\n    output_path = Path(output_path)\n    base = Path(base or output_path.parent).resolve()\n    with zipfile.ZipFile(output_path, \"w\", compression=zipfile.ZIP_DEFLATED) as bundle:\n        for source in sorted({Path(path).resolve() for path in paths}):\n            if source == output_path.resolve() or not source.is_file():\n                continue\n            try:\n                arcname = source.relative_to(base)\n            except ValueError:\n                arcname = Path(source.name)\n            bundle.write(source, arcname)\n    print(f\"Wrote {output_path} ({output_path.stat().st_size / 1024:.1f} KiB)\")\n    return output_path\n\n\ndef download(path):\n    try:\n        from google.colab import files\n    except ImportError:\n        print(f\"Result available at {Path(path).resolve()}\")\n    else:\n        files.download(str(path))\n\n\ndef optional_time_arguments(start, end, interval):\n    arguments = []\n    for option, value in ((\"-b\", start), (\"-e\", end), (\"-dt\", interval)):\n        if str(value).strip():\n            arguments.extend([option, str(value).strip()])\n    return arguments\n\n\ndef read_xvg(path):\n    legends = {}\n    metadata = {}\n    rows = []\n    with Path(path).open(encoding=\"utf-8\") as stream:\n        for raw in stream:\n            line = raw.strip()\n            if not line or line.startswith(\"#\"):\n                continue\n            if line.startswith(\"@\"):\n                re_module = __import__(\"re\")\n                match = re_module.search(r's(\\d+)\\s+legend\\s+\"(.*)\"', line)\n                if match:\n                    legends[int(match.group(1)) + 1] = match.group(2)\n                for key, pattern in {\n                    \"title\": r'@\\s+title\\s+\"(.*)\"',\n                    \"xlabel\": r'@\\s+xaxis\\s+label\\s+\"(.*)\"',\n                    \"ylabel\": r'@\\s+yaxis\\s+label\\s+\"(.*)\"',\n                }.items():\n                    label_match = re_module.search(pattern, line)\n                    if label_match:\n                        metadata[key] = label_match.group(1)\n                continue\n            rows.append([float(value) for value in line.split()])\n    data = np.asarray(rows, dtype=float)\n    if data.ndim != 2 or data.shape[1] < 2:\n        raise ValueError(f\"No plottable data in {path}\")\n    return data, legends, metadata\n\n\ndef plot_xvg(path, *, title=None, output=None):\n    data, legends, metadata = read_xvg(path)\n    fig, ax = plt.subplots(figsize=(7.2, 4.2))\n    for column in range(1, data.shape[1]):\n        ax.plot(data[:, 0], data[:, column], label=legends.get(column, f\"series {column}\"))\n    ax.set_title(title or metadata.get(\"title\", Path(path).stem))\n    ax.set_xlabel(metadata.get(\"xlabel\", \"x / time\"))\n    ax.set_ylabel(metadata.get(\"ylabel\", \"value\"))\n    if data.shape[1] > 2:\n        ax.legend(frameon=False, fontsize=8)\n    ax.spines[[\"top\", \"right\"]].set_visible(False)\n    fig.tight_layout()\n    output_path = Path(output) if output is not None else Path(path).with_suffix(\".png\")\n    fig.savefig(output_path, dpi=220, bbox_inches=\"tight\")\n    print(f\"Saved figure: {output_path}\")\n    plt.show()\n    return fig, output_path\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Upload or reuse a workflow bundle\ninput_source = \"Upload ZIP or individual files\" #@param [\"Upload ZIP or individual files\", \"Reuse files from an earlier notebook in this runtime\"]\nreuse_directory = \"/content/dropps_model_builder\" #@param {type:\"string\"}\n\nimport shutil\n\nBASE = Path(\"/content\") if Path(\"/content\").is_dir() else Path.cwd()\nINPUT_ROOT = reset_task_directory(BASE / \"dropps_uploaded_inputs\")\n\nif input_source.startswith(\"Upload\"):\n    try:\n        from google.colab import files\n    except ImportError as exc:\n        raise RuntimeError(\"Upload this bundle in Colab or choose the reuse option.\") from exc\n    uploaded = files.upload()\n    for name in uploaded:\n        source = Path(name)\n        if source.suffix.lower() == \".zip\":\n            safe_extract_zip(source, INPUT_ROOT)\n        else:\n            shutil.copy2(source, INPUT_ROOT / source.name)\nelse:\n    source_root = Path(reuse_directory)\n    if not source_root.is_dir():\n        raise FileNotFoundError(source_root)\n    for source in source_root.rglob(\"*\"):\n        if source.is_file() and source.suffix.lower() in {\n            \".pdb\", \".itp\", \".top\", \".mdp\", \".tpr\", \".xtc\", \".ndx\", \".chk\", \".xml\", \".json\"\n        }:\n            target = INPUT_ROOT / source.name\n            if target.exists() and sha256(target) != sha256(source):\n                raise ValueError(f\"Duplicate filename with different content: {source.name}\")\n            shutil.copy2(source, target)\n\nprint(\"Available files:\")\nfor path in sorted(INPUT_ROOT.rglob(\"*\")):\n    if path.is_file():\n        print(\" \", path.relative_to(INPUT_ROOT), f\"({path.stat().st_size / 1024:.1f} KiB)\")\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 1 Resolve the compact cubic input system\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Input filenames\nstructure_filename = \"system.pdb\" #@param {type:\"string\"}\ntopology_filename = \"system.top\" #@param {type:\"string\"}\nrequire_cubic_input = True #@param {type:\"boolean\"}\ncubic_relative_tolerance = 0.001 #@param {type:\"number\"}\n\nBASE = Path(\"/content\") if Path(\"/content\").is_dir() else Path.cwd()\nWORKDIR = reset_task_directory(BASE / \"dropps_simulation\")\nINPUTS = WORKDIR / \"inputs\"\nINPUTS.mkdir()\n\nimport shutil\n\nfor source in INPUT_ROOT.rglob(\"*\"):\n    if source.is_file() and (\n        source.suffix.lower() in {\".pdb\", \".itp\", \".top\", \".ff\", \".chk\", \".xml\", \".json\", \".mdp\", \".tpr\"}\n        or source.name.endswith(\".state.xml\")\n        or source.name.endswith(\".restart.json\")\n    ):\n        target = INPUTS / source.name\n        if target.exists() and sha256(target) != sha256(source):\n            raise ValueError(f\"Conflicting duplicate input: {source.name}\")\n        shutil.copy2(source, target)\n\nSTRUCTURE = find_unique(INPUTS, structure_filename)\nTOPOLOGY = find_unique(INPUTS, topology_filename)\n\ndef pdb_box_lengths_nm(path):\n    for line in Path(path).read_text(encoding=\"utf-8\").splitlines():\n        if line.startswith(\"CRYST1\"):\n            return np.asarray(\n                [float(line[6:15]), float(line[15:24]), float(line[24:33])]\n            ) / 10.0\n    raise ValueError(f\"No CRYST1 box record found in {path}\")\n\ncompact_box_nm = pdb_box_lengths_nm(STRUCTURE)\nrelative_spread = np.ptp(compact_box_nm) / compact_box_nm.mean()\nif require_cubic_input and relative_spread > cubic_relative_tolerance:\n    raise ValueError(\n        \"The slab protocol must start from a compact cubic box; \"\n        f\"received {compact_box_nm.tolist()} nm.\"\n    )\nprint(\"Structure:\", STRUCTURE)\nprint(\"Topology:\", TOPOLOGY)\nprint(\"Compact box (nm):\", compact_box_nm)\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2 Define the NPT → z×10 → NVT protocol\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Stage durations, integration, and temperature\nworkflow_action = \"Short full-workflow validation\" #@param [\"Short full-workflow validation\", \"Run full configured workflow\", \"Prepare NPT TPR only\", \"Prepare NVT TPR from uploaded npt.pdb\", \"Run NVT from uploaded npt.pdb\"]\nproduction_temperature_K = 300.0 #@param {type:\"number\"}\nnpt_time_step_ps = 0.01 #@param {type:\"number\"}\nnpt_duration_ns = 300.0 #@param {type:\"number\"}\nnvt_time_step_ps = 0.02 #@param {type:\"number\"}\nnvt_duration_ns = 3000.0 #@param {type:\"number\"}\nvalidation_steps_per_stage = 100 #@param {type:\"integer\"}\nfriction_per_ps = 1.0 #@param {type:\"number\"}\nrandom_seed = 1215 #@param {type:\"integer\"}\nminimize_before_npt = True #@param {type:\"boolean\"}\nminimization_force_tolerance = 100.0 #@param {type:\"number\"}\nminimization_max_steps = 100000 #@param {type:\"integer\"}\nslab_axis = \"z\" #@param [\"z\", \"x\", \"y\"]\nslab_box_multiplier = 10.0 #@param {type:\"number\"}\n\nif slab_axis != \"z\":\n    raise ValueError(\"This Colab phase-coexistence workflow expands the z axis.\")\nif slab_box_multiplier != 10.0:\n    raise ValueError(\"This phase-coexistence workflow requires a 10-fold box expansion.\")\nif min(npt_time_step_ps, nvt_time_step_ps, npt_duration_ns, nvt_duration_ns) <= 0:\n    raise ValueError(\"Time steps and stage durations must be positive.\")\n\nconfigured_npt_steps = int(round(npt_duration_ns * 1000 / npt_time_step_ps))\nconfigured_nvt_steps = int(round(nvt_duration_ns * 1000 / nvt_time_step_ps))\nshort_workflow = workflow_action == \"Short full-workflow validation\"\nnpt_steps = validation_steps_per_stage if short_workflow else configured_npt_steps\nnvt_steps = validation_steps_per_stage if short_workflow else configured_nvt_steps\n\nprint(f\"NPT: {npt_steps:,} steps at {npt_time_step_ps:g} ps\")\nprint(f\"NVT: {nvt_steps:,} steps at {nvt_time_step_ps:g} ps\")\nprint(f\"Slab conversion: {slab_axis} axis × {slab_box_multiplier:g}\")\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Bond, COM, and nonbonded settings\nbond_treatment = \"bond\" #@param [\"bond\", \"constraint\"]\ncenter_of_mass_mode = \"Linear\" #@param [\"Linear\", \"none\"]\ncenter_of_mass_interval = 100 #@param {type:\"integer\"}\nvdw_type = \"pLJ\" #@param [\"pLJ\", \"MPiPi\"]\nlj_cutoff_scheme = \"static\" #@param [\"static\", \"dynamic\"]\nlj_cutoff_nm = 1.5 #@param {type:\"number\"}\nlj_sigma_multiplier = 3.0 #@param {type:\"number\"}\nshift_lj = True #@param {type:\"boolean\"}\ncoulomb_type = \"yukawa\" #@param [\"yukawa\", \"no\"]\ncoulomb_cutoff_nm = 1.5 #@param {type:\"number\"}\nshift_coulomb = True #@param {type:\"boolean\"}\nsalt_concentration_molar = 0.1 #@param {type:\"number\"}\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Pressure and output settings\nreference_pressure_bar = 1.0 #@param {type:\"number\"}\nbarostat_attempt_interval = 5 #@param {type:\"integer\"}\ntrajectory_interval_ns = 1.0 #@param {type:\"number\"}\nscreen_log_interval_ns = 10.0 #@param {type:\"number\"}\nfile_log_interval_ns = 1.0 #@param {type:\"number\"}\nenergy_interval_ns = 1.0 #@param {type:\"number\"}\nenergy_groups = \"potentialEnergy,kineticEnergy,totalEnergy,temperature,boxX,boxY,boxZ,volume,density\" #@param {type:\"string\"}\nstress_interval_steps = 0 #@param {type:\"integer\"}\nstress_strain = 0.0003 #@param {type:\"number\"}\nstress_platform = \"auto\" #@param [\"auto\", \"CUDA\", \"OpenCL\", \"CPU\", \"Reference\"]\nstress_precision = \"double\" #@param [\"single\", \"mixed\", \"double\"]\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title MDP renderer for both ensembles\ndef bool_text(value):\n    return \"True\" if value else \"False\"\n\ndef interval_steps(interval_ns, dt_ps, total_steps):\n    requested = max(1, int(round(interval_ns * 1000 / dt_ps)))\n    if short_workflow:\n        return min(requested, max(1, total_steps // 10))\n    return requested\n\ndef render_mdp(*, stage, dt_ps, nsteps, pressure_coupling, minimize):\n    effective_max_steps = (\n        min(minimization_max_steps, 100)\n        if short_workflow and minimize\n        else minimization_max_steps\n    )\n    return f\"\"\"# DROPPS 1.0 Colab: {stage} stage of NPT -> slab -> NVT\nintegrator = Langevin\ndt = {dt_ps}\nnsteps = {nsteps}\nfriction = {friction_per_ps}\nseed = {random_seed}\nminimize = {bool_text(minimize)}\nforcetol = {minimization_force_tolerance}\nmax-step = {effective_max_steps}\nproduction-temperature = {production_temperature_K}\ngen-vel = True\ninitial-temperature = {production_temperature_K}\nwarming-speed = 1.0\nkeep-warming-trajectory = False\nbondtype = {bond_treatment}\ncomm-mode = {center_of_mass_mode}\nnstcomm = {center_of_mass_interval}\nvdwtype = {vdw_type}\ncutoff-scheme-lj = {lj_cutoff_scheme}\ncutoff-lj = {lj_cutoff_nm}\ncutoff-lj-multi = {lj_sigma_multiplier}\nshift-lj = {bool_text(shift_lj)}\ncoulombtype = {coulomb_type}\ncutoff-coul = {coulomb_cutoff_nm}\nshift-coul = {bool_text(shift_coulomb)}\nsalt-conc = {salt_concentration_molar}\npcoulp = {bool_text(pressure_coupling)}\nref-P = {reference_pressure_bar}\ntau-P = {barostat_attempt_interval}\nnst-xout = {interval_steps(trajectory_interval_ns, dt_ps, nsteps)}\nnst-screenlog = {interval_steps(screen_log_interval_ns, dt_ps, nsteps)}\nnst-filelog = {interval_steps(file_log_interval_ns, dt_ps, nsteps)}\nnst-energy = {interval_steps(energy_interval_ns, dt_ps, nsteps)}\nenergy-grps = {energy_groups}\nnst-stress = {stress_interval_steps}\nstress-strain = {stress_strain}\nstress-output = auto\nstress-platform = {stress_platform}\nstress-precision = {stress_precision}\nstress-device = auto\nstress-threads = 0\n\"\"\"\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 3 Runtime, checkpoint, and restart controls\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Runtime controls shared by both stages\ncompute_platform = \"auto\" #@param [\"auto\", \"CUDA\", \"OpenCL\", \"CPU\", \"Reference\"]\ngpu_device_index = \"0\" #@param {type:\"string\"}\ngpu_precision = \"mixed\" #@param [\"single\", \"mixed\", \"double\"]\ncpu_threads = 0 #@param {type:\"integer\"}\ncheckpoint_minutes = 5.0 #@param {type:\"number\"}\nkeep_numbered_checkpoints = False #@param {type:\"boolean\"}\nmaximum_wall_hours = -1.0 #@param {type:\"number\"}\nrestart_npt_from_checkpoint = False #@param {type:\"boolean\"}\nnpt_checkpoint_filename = \"\" #@param {type:\"string\"}\nrestart_nvt_from_checkpoint = False #@param {type:\"boolean\"}\nnvt_checkpoint_filename = \"\" #@param {type:\"string\"}\nappend_restart_outputs = True #@param {type:\"boolean\"}\n\ndef run_stage(tpr, prefix, restart=False, checkpoint_filename=\"\"):\n    arguments = [\n        \"mdrun\", \"-s\", Path(tpr).name, \"-o\", prefix,\n        \"--platform\", compute_platform, \"-cpt\", checkpoint_minutes,\n        \"-maxh\", maximum_wall_hours,\n    ]\n    if compute_platform in {\"CUDA\", \"OpenCL\"}:\n        arguments.extend([\"--precision\", gpu_precision])\n        if gpu_device_index.strip():\n            arguments.extend([\"-gpu_id\", gpu_device_index.strip()])\n    elif compute_platform == \"CPU\" and cpu_threads > 0:\n        arguments.extend([\"-nt\", cpu_threads])\n    if keep_numbered_checkpoints:\n        arguments.append(\"-cpnum\")\n    if restart:\n        arguments.append(\"-cpi\")\n        if checkpoint_filename.strip():\n            arguments.append(checkpoint_filename.strip())\n    arguments.append(\"-append\" if append_restart_outputs else \"-noappend\")\n    dps(*arguments, cwd=INPUTS)\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 4 Stage 1 — compact cubic-box NPT\n\nPressure coupling is mandatory in this stage. The full defaults generate the\ncondensed configuration over 300 ns. `Prepare NPT TPR only` is intended for\ntransfer to a persistent GPU/HPC resource.\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Compile and optionally execute NPT\nNPT_MDP = INPUTS / \"npt.mdp\"\nNPT_TPR = INPUTS / \"npt.tpr\"\nrun_npt_here = workflow_action in {\n    \"Short full-workflow validation\", \"Run full configured workflow\"\n}\nuse_existing_npt = workflow_action in {\n    \"Prepare NVT TPR from uploaded npt.pdb\", \"Run NVT from uploaded npt.pdb\"\n}\n\nif not use_existing_npt:\n    NPT_MDP.write_text(\n        render_mdp(\n            stage=\"compact NPT\", dt_ps=npt_time_step_ps,\n            nsteps=npt_steps, pressure_coupling=True,\n            minimize=minimize_before_npt,\n        ),\n        encoding=\"utf-8\",\n    )\n    dps(\n        \"grompp\", \"-f\", STRUCTURE.name, \"-p\", TOPOLOGY.name,\n        \"-m\", NPT_MDP.name, \"-o\", NPT_TPR.name, cwd=INPUTS,\n    )\n    dps(\"check\", \"-s\", NPT_TPR.name, cwd=INPUTS)\n\nif run_npt_here:\n    run_stage(\n        NPT_TPR, \"npt\", restart_npt_from_checkpoint,\n        npt_checkpoint_filename,\n    )\n    dps(\"check\", \"-s\", NPT_TPR.name, \"-f\", \"npt.xtc\", cwd=INPUTS)\n    NPT_FINAL = INPUTS / \"npt.pdb\"\nelif use_existing_npt:\n    NPT_FINAL = find_unique(INPUTS, \"npt.pdb\")\nelse:\n    NPT_FINAL = None\n    print(\"NPT TPR prepared. Run it to completion before constructing the slab.\")\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 5 Stage 2 — tenfold box expansion and elongated-box NVT\n\n`dps editconf -mz 10` unwraps the expanded axis, increases only that box\nlength, and translates the NPT configuration to the center of the elongated\nbox. The generated `slab_nvt.mdp` has `pcoulp = False`. Because a PDB does\nnot store velocities, the NVT TPR generates new velocities at the same target\ntemperature.\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Build and optionally execute the NVT slab stage\nSLAB_PDB = INPUTS / \"slab.pdb\"\nNVT_MDP = INPUTS / \"slab_nvt.mdp\"\nNVT_TPR = INPUTS / \"slab_nvt.tpr\"\nrun_nvt_here = workflow_action in {\n    \"Short full-workflow validation\", \"Run full configured workflow\",\n    \"Run NVT from uploaded npt.pdb\",\n}\n\nif NPT_FINAL is not None:\n    npt_box_nm = pdb_box_lengths_nm(NPT_FINAL)\n    multiplier_option = {\"x\": \"-mx\", \"y\": \"-my\", \"z\": \"-mz\"}[slab_axis]\n    dps(\n        \"editconf\", \"-f\", NPT_FINAL.name, \"-o\", SLAB_PDB.name,\n        multiplier_option, slab_box_multiplier, cwd=INPUTS,\n    )\n    slab_box_nm = pdb_box_lengths_nm(SLAB_PDB)\n    axis_index = {\"x\": 0, \"y\": 1, \"z\": 2}[slab_axis]\n    ratios = slab_box_nm / npt_box_nm\n    expected = np.ones(3)\n    expected[axis_index] = slab_box_multiplier\n    if not np.allclose(ratios, expected, rtol=2e-3, atol=2e-3):\n        raise RuntimeError(\n            f\"Slab-box verification failed: ratios={ratios.tolist()}, \"\n            f\"expected={expected.tolist()}\"\n        )\n    print(\"NPT box (nm):\", npt_box_nm)\n    print(\"Slab box (nm):\", slab_box_nm)\n    print(\"Verified box-length ratios:\", ratios)\n\n    NVT_MDP.write_text(\n        render_mdp(\n            stage=\"elongated-box NVT\", dt_ps=nvt_time_step_ps,\n            nsteps=nvt_steps, pressure_coupling=False, minimize=False,\n        ),\n        encoding=\"utf-8\",\n    )\n    dps(\n        \"grompp\", \"-f\", SLAB_PDB.name, \"-p\", TOPOLOGY.name,\n        \"-m\", NVT_MDP.name, \"-o\", NVT_TPR.name, cwd=INPUTS,\n    )\n    dps(\"check\", \"-s\", NVT_TPR.name, cwd=INPUTS)\n\n    if run_nvt_here:\n        run_stage(\n            NVT_TPR, \"slab_nvt\", restart_nvt_from_checkpoint,\n            nvt_checkpoint_filename,\n        )\n        dps(\n            \"check\", \"-s\", NVT_TPR.name, \"-f\", \"slab_nvt.xtc\",\n            cwd=INPUTS,\n        )\n        print(\"Analysis input: slab_nvt.tpr + slab_nvt.xtc\")\n    else:\n        print(\"NVT TPR prepared: slab_nvt.tpr\")\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 6 Export the staged simulation and explicit analysis handoff\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Save manifest and download simulation bundle\nfiles_to_export = [path for path in INPUTS.iterdir() if path.is_file()]\nanalysis_ready = (INPUTS / \"slab_nvt.tpr\").is_file() and (INPUTS / \"slab_nvt.xtc\").is_file()\nmanifest = {\n    \"notebook\": \"02_DROPPS_1_0_Simulation.ipynb\",\n    \"dropps_version\": \"1.0\",\n    \"workflow_action\": workflow_action,\n    \"protocol\": [\n        \"compact cubic box\",\n        \"NPT condensation/equilibration\",\n        f\"{slab_axis}-axis x{slab_box_multiplier:g}\",\n        \"elongated-box NVT production\",\n        \"analyse NVT trajectory only\",\n    ],\n    \"scientific_parameters\": {\n        \"temperature_K\": production_temperature_K,\n        \"npt\": {\n            \"dt_ps\": npt_time_step_ps,\n            \"configured_duration_ns\": npt_duration_ns,\n            \"effective_steps\": npt_steps,\n            \"pressure_coupling\": True,\n        },\n        \"slab_expansion\": {\"axis\": slab_axis, \"multiplier\": slab_box_multiplier},\n        \"nvt\": {\n            \"dt_ps\": nvt_time_step_ps,\n            \"configured_duration_ns\": nvt_duration_ns,\n            \"effective_steps\": nvt_steps,\n            \"pressure_coupling\": False,\n        },\n        \"forcefield_vdw_type\": vdw_type,\n        \"salt_concentration_molar\": salt_concentration_molar,\n        \"seed\": random_seed,\n    },\n    \"operational_parameters\": {\n        \"platform\": compute_platform,\n        \"precision\": gpu_precision,\n        \"checkpoint_minutes\": checkpoint_minutes,\n        \"maximum_wall_hours\": maximum_wall_hours,\n    },\n    \"analysis_handoff\": {\n        \"ready\": analysis_ready,\n        \"run_input\": \"slab_nvt.tpr\",\n        \"trajectory\": \"slab_nvt.xtc\",\n        \"ensemble\": \"NVT\",\n    },\n    \"files\": {path.name: sha256(path) for path in files_to_export},\n    \"commands\": COMMAND_LOG,\n}\nmanifest_path = INPUTS / \"simulation_manifest.json\"\nmanifest_path.write_text(json.dumps(manifest, indent=2), encoding=\"utf-8\")\nfiles_to_export.append(manifest_path)\narchive = make_zip(\n    files_to_export,\n    WORKDIR / \"DROPPS_1_0_simulation_bundle.zip\",\n    base=INPUTS,\n)\ndownload(archive)\n"
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}