{
 "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/03_DROPPS_1_0_Trajectory_Preparation.ipynb)\n\n# DROPPS 1.0 Trajectory Preparation and Indexing\n\nPrepare the **elongated-box NVT trajectory** for analysis. The notebook rejects\nan NPT run input, creates reusable NDX groups through the public `make_ndx`\ninterface, and applies PBC reconstruction, centering, fitting, time selection,\nXTC/PDB conversion, and snapshot extraction with `trjconv`.\n\nTransform order in DROPPS 1.0 is: make whole/no-jump → center → pack → fit →\ntranslate. Time selections use physical units and choose the nearest saved frame.\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_simulation\" #@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 and validate the run files\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Run filenames\nrun_input_filename = \"slab_nvt.tpr\" #@param {type:\"string\"}\ntrajectory_filename = \"slab_nvt.xtc\" #@param {type:\"string\"}\n\nTPR = find_unique(INPUT_ROOT, run_input_filename)\nTRAJECTORY = find_unique(INPUT_ROOT, trajectory_filename)\n\nimport zipfile\n\nif not zipfile.is_zipfile(TPR):\n    raise ValueError(\"Expected a portable DROPPS 1.0 TPR v2 file.\")\nwith zipfile.ZipFile(TPR) as archive:\n    tpr_parameters = json.loads(archive.read(\"parameters.json\"))\nif tpr_parameters.get(\"pcoulp\"):\n    raise ValueError(\n        \"Trajectory preparation requires the elongated-box NVT TPR, not the NPT TPR.\"\n    )\n\nBASE = Path(\"/content\") if Path(\"/content\").is_dir() else Path.cwd()\nWORKDIR = reset_task_directory(BASE / \"dropps_trajectory\")\nimport shutil\nANALYSIS_TPR = WORKDIR / \"slab_nvt.tpr\"\nshutil.copy2(TPR, ANALYSIS_TPR)\ndps(\"check\", \"-s\", TPR, \"-f\", TRAJECTORY, cwd=WORKDIR)\nprint(\"Verified ensemble: NVT (pcoulp = False)\")\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 2 Create an index file\n\nCommands are separated by semicolons. Initial groups are `System` (group 0)\nfollowed by one group per molecule type. New selections append groups in order.\nExample: `mol SCAFFOLD; name 3 Scaffold_all; resid 1-4 & mol SCAFFOLD;\nname 4 Scaffold_N; q`. The notebook adds `q` if omitted.\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Non-interactive make_ndx command script\nindex_commands = \"q\" #@param {type:\"string\"}\nindex_filename = \"analysis.ndx\" #@param {type:\"string\"}\n\ncommands = [part.strip() for part in index_commands.split(\";\") if part.strip()]\nif not commands or commands[-1] != \"q\":\n    commands.append(\"q\")\nNDX = WORKDIR / index_filename\nresult = dps(\n    \"make_ndx\", \"-s\", TPR, \"-o\", NDX,\n    cwd=WORKDIR, input_text=\"\\n\".join(commands) + \"\\n\",\n    check=False,\n)\nif not NDX.is_file():\n    raise RuntimeError(\n        f\"make_ndx did not create {NDX}; subprocess status was {result.returncode}.\"\n    )\nif result.returncode:\n    print(\n        \"Note: DROPPS 1.0 make_ndx writes the requested file and then exits \"\n        \"through its legacy interactive quit path, which reports status 1.\"\n    )\ndps(\"check\", \"-s\", TPR, \"-f\", TRAJECTORY, \"-n\", NDX, cwd=WORKDIR)\nprint(NDX.read_text(encoding=\"utf-8\")[:4000])\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 3 Convert and transform the trajectory\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title trjconv settings\noutput_filename = \"processed.xtc\" #@param {type:\"string\"}\noutput_selection = \"group 0\" #@param {type:\"string\"}\nstart_time = \"\" #@param {type:\"string\"}\nend_time = \"\" #@param {type:\"string\"}\nsampling_interval = \"\" #@param {type:\"string\"}\ntime_unit = \"ns\" #@param [\"fs\", \"ps\", \"ns\", \"us\", \"ms\", \"s\"]\npbc_mode = \"whole\" #@param [\"none\", \"whole\", \"atom\", \"res\", \"mol\", \"nojump\"]\ncenter_mode = \"none\" #@param [\"none\", \"geometry\", \"mass\", \"dense\"]\ncenter_selection = \"\" #@param {type:\"string\"}\ncenter_axis = \"xyz\" #@param [\"x\", \"y\", \"z\", \"xyz\"]\ndense_phase_threshold = 0.5 #@param {type:\"number\"}\ndensity_bin_width_nm = 0.05 #@param {type:\"number\"}\nfit_mode = \"none\" #@param [\"none\", \"translation\", \"transxy\", \"rot+trans\", \"rotxy+transxy\", \"progressive\"]\nfit_selection = \"\" #@param {type:\"string\"}\nfit_reference = \"tpr\" #@param [\"tpr\", \"first\"]\nfit_weighting = \"mass\" #@param [\"mass\", \"uniform\"]\ntranslation_nm = \"\" #@param {type:\"string\"}\nper_frame_shift_nm = \"\" #@param {type:\"string\"}\nxtc_precision_decimals = 3 #@param {type:\"integer\"}\n\nOUTPUT_TRAJECTORY = WORKDIR / output_filename\narguments = [\n    \"trjconv\", \"-s\", TPR, \"-f\", TRAJECTORY, \"-n\", NDX,\n    \"-sel\", output_selection, \"-o\", OUTPUT_TRAJECTORY,\n    \"-tu\", time_unit, \"-pbc\", pbc_mode, \"--center\", center_mode,\n    \"--center-axis\", center_axis, \"-dpt\", dense_phase_threshold,\n    \"--density-bin-width\", density_bin_width_nm,\n    \"-fit\", fit_mode, \"--fit-reference\", fit_reference,\n    \"--fit-weighting\", fit_weighting, \"-ndec\", xtc_precision_decimals,\n    *optional_time_arguments(start_time, end_time, sampling_interval),\n]\nif center_selection.strip():\n    arguments.extend([\"--center-select\", center_selection.strip()])\nif fit_selection.strip():\n    arguments.extend([\"--fit-select\", fit_selection.strip()])\nif translation_nm.strip():\n    vector = translation_nm.replace(\",\", \" \").split()\n    if len(vector) != 3:\n        raise ValueError(\"translation_nm requires three values\")\n    arguments.extend([\"-trans\", *vector])\nif per_frame_shift_nm.strip():\n    vector = per_frame_shift_nm.replace(\",\", \" \").split()\n    if len(vector) != 3:\n        raise ValueError(\"per_frame_shift_nm requires three values\")\n    arguments.extend([\"-shift\", *vector])\ndps(*arguments, cwd=WORKDIR)\ndps(\"check\", \"-s\", TPR, \"-f\", OUTPUT_TRAJECTORY, \"-n\", NDX, cwd=WORKDIR)\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Extract a PDB snapshot\nextract_snapshot = True #@param {type:\"boolean\"}\nsnapshot_time = 0.0 #@param {type:\"number\"}\nsnapshot_time_unit = \"ns\" #@param [\"fs\", \"ps\", \"ns\", \"us\", \"ms\", \"s\"]\nsnapshot_selection = \"group 0\" #@param {type:\"string\"}\nadd_conect_records = True #@param {type:\"boolean\"}\n\nif extract_snapshot:\n    SNAPSHOT = WORKDIR / \"snapshot.pdb\"\n    arguments = [\n        \"trjconv\", \"-s\", TPR, \"-f\", OUTPUT_TRAJECTORY, \"-n\", NDX,\n        \"-sel\", snapshot_selection, \"-o\", SNAPSHOT,\n        \"-b\", snapshot_time, \"-e\", snapshot_time, \"-tu\", snapshot_time_unit,\n        \"-pbc\", \"whole\",\n    ]\n    if add_conect_records:\n        arguments.append(\"--conect\")\n    dps(*arguments, cwd=WORKDIR)\n    print(\"Snapshot:\", SNAPSHOT)\n"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 4 Export the prepared trajectory bundle\n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "#@title Save manifest and download\nmanifest = {\n    \"notebook\": \"03_DROPPS_1_0_Trajectory_Preparation.ipynb\",\n    \"dropps_version\": \"1.0\",\n    \"source_tpr_sha256\": sha256(TPR),\n    \"source_trajectory_sha256\": sha256(TRAJECTORY),\n    \"verified_ensemble\": \"NVT\",\n    \"index_commands\": commands,\n    \"commands\": COMMAND_LOG,\n}\nmanifest_path = WORKDIR / \"trajectory_manifest.json\"\nmanifest_path.write_text(json.dumps(manifest, indent=2), encoding=\"utf-8\")\nexport_files = [ANALYSIS_TPR, NDX, OUTPUT_TRAJECTORY, manifest_path]\nif extract_snapshot:\n    export_files.append(SNAPSHOT)\narchive = make_zip(\n    export_files,\n    WORKDIR / \"DROPPS_1_0_trajectory_bundle.zip\",\n    base=WORKDIR,\n)\ndownload(archive)\n"
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}