Metadata-Version: 2.1
Name: rush-py
Version: 1.2.1
Summary: Python SDK for interacting with the QDX Rush API and modules
Home-page: https://rush.qdx.co
Author: Ryan Swart
Author-email: ryan.swart@qdx.co
Requires-Python: >=3.10,<3.13
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Dist: httpx (>=0.25.2,<0.26.0)
Requires-Dist: pdb-tools (>=2.5.0,<3.0.0)
Requires-Dist: pydantic (>=2.5.2,<3.0.0)
Requires-Dist: websockets (>=11.0,<12.0)
Project-URL: Documentation, https://talo.github.io/rush-py
Description-Content-Type: text/markdown

# rush-py

<!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->

Below we’ll walk through the process of building and running a drug
discovery workflow, where we prepare a protein and ligand for molecular
dynamics simulation, run the molecular dynamics, perform a quantum
lattice interaction energy calculation.

First, install the following modules via pip - we require Python \> 3.10

    pip install rush-py pdb-tools

# 0) Setup

This is where we prepare the rush client, directories, and input data
we’ll be working with

## 0.0) Imports

``` python
import os
import tarfile
from datetime import datetime
from pathlib import Path

from pdbtools import pdb_fetch, pdb_delhetatm, pdb_selchain, pdb_rplresname, pdb_keepcoord, pdb_selresname
import requests
import py3Dmol

import rush
```

## 0.1) Credentials

``` python
# Set our token - ensure you have exported RUSH_TOKEN in your shell; or just replace the os.getenv with your token
TOKEN = os.getenv("RUSH_TOKEN")
# You might have a custom deployment url, by default it will use https://tengu.qdx.ai
URL = os.getenv("RUSH_URL") or "https://tengu.qdx.ai"
# These env variables will be read by default, so you can skip this step in future
```

## 0.2) Configuration

Lets set some global variables that define our project

``` python
# Define our project information
DESCRIPTION = "rush-py demo notebook"
TAGS = ["qdx", "rush-py-v2", "tutorial", "cdk2"]
WORK_DIR = Path.home() / "qdx" / "rush-py-demo"
# Set our inputs
SYSTEM_PDB_PATH = WORK_DIR / "test.pdb"
PROTEIN_PDB_PATH = WORK_DIR / "test_P.pdb"
LIGAND_SMILES_STR = (
    "c1nc(c2c(n1)n(cn2)[C@H]3[C@@H]([C@@H]([C@H](O3)CO[P@@](=O)(O)O[P@](=O)(O)OP(=O)(O)O)O)O)N"
)
LIGAND_PDB_PATH = WORK_DIR / "test_L.pdb"
# where want our jobs to run
TARGET = "NIX_SSH_3"
```

Ensure your workdir exists

``` python
os.makedirs(WORK_DIR)
```

## 0.2) Build your client

``` python
# Get our client, for calling modules and using the rush API
# Note, access_token and url are optional, if you have set the env variables RUSH_TOKEN and RUSH_URL
# Workspace sets the location where we will store our session history file and module lock file
# By using the `build_provider_with_functions` method, we will also build helper functions calling each module
client = await rush.build_provider_with_functions(
    access_token=TOKEN, url=URL, workspace=WORK_DIR, batch_tags=TAGS
)
```

## 0.3) Input selection

``` python
# fetch datafiles
complex = list(pdb_fetch.fetch_structure("1B39"))
protein = pdb_delhetatm.remove_hetatm(pdb_selchain.select_chain(complex, "A"))
# select the ATP residue
ligand = pdb_selresname.filter_residue_by_name(complex, "ATP")
# we require ligands to be labelled as UNL
ligand = pdb_rplresname.rename_residues(ligand, "ATP", "UNL")
# we don't want to repeat all of the remark / metadata that is already in the protein
ligand = pdb_keepcoord.keep_coordinates(ligand)
# write our files to the locations defined in the config block
with open(SYSTEM_PDB_PATH, "w") as f:
    for l in complex:
        f.write(str(l))
with open(PROTEIN_PDB_PATH, "w") as f:
    for l in protein:
        f.write(str(l))
with open(LIGAND_PDB_PATH, "w") as f:
    for l in ligand:
        f.write(str(l))
```

## 0.4) View rush modules

Rush modules are “functions” that perform various computational
chemistry tasks can be run on HPC infrastructure. We maintain multiple
versions of these functions so that your scripts will stay stable over
upgrades.

``` python
# Get our latest modules as a dict[module_name, module_path]
# If a lock file exists, load it so that the run is reproducable
# This will be done automatically if you use the `build_provider_with_functions` method
modules = await client.get_latest_module_paths()
```

``` python
module_name = "hermes_energy"
module_path = modules[module_name]
print(module_path)
```

    github:talo/tengu-prelude/efc6d8b3a8cc342cd9866d037abb77dac40a4d56#hermes_energy

- `module_name` is a descriptive string and indicates the “function” the
  module is calling;
- `module_path` is a versioned rush “endpoint” for a module accessible
  via the client.

Using the same `module_path` string across multiple runs provides
reproducibility.

## 0.5) Use module functions

Next, we’ll use helper functions for the modules that we’ve fetched

If we have built a provider with functions, we can use the python
`help()` function to describe their usage.

The QDX Type Description is a standard type definition across multiple
programing languages to assist in interoperablility. @ indicates that
the type is stored in a file, which will be synced to cloud storage

``` python
help(client.convert)
```

    Help on function convert in module rush.provider:

    async convert(*args: [list[typing.Union[str, ~T]], <class 'pathlib.Path'>], target: rush.graphql_client.enums.ModuleInstanceTarget | None = <ModuleInstanceTarget.NIX: 'NIX'>, resources: rush.graphql_client.input_types.ModuleInstanceResourcesInput | None = ModuleInstanceResourcesInput(gpus=0, gpu_mem=None, gpu_mem_units=None, cpus=None, nodes=None, mem=None, mem_units=None, storage=10, storage_units=<MemUnits.MB: 'MB'>, walltime=None, storage_mounts=None), tags: list[str] | None = None, restore: bool | None = None) -> [<class 'pathlib.Path'>]
        Convert biomolecular and chemical file formats to the QDX file format. Supports PDB and SDF
        
        Module version: github:talo/tengu-prelude/efc6d8b3a8cc342cd9866d037abb77dac40a4d56#convert
        
        QDX Type Description:
        
            format: PDB|SDF;
        
            input: @bytes 
        
        ->
        
            output: @[Conformer]
        
        
        
        :param format: the format of the input file
        :param input: the input file
        :return output: the output conformers

# 1) Running Rush Modules

Below we’ll call modules using the functions created on the client.

The parameters to a rush module function would look like the following

- `*args`: The values or ids passed to the :
  1.  For @Objects - A `pathlib.Path` or a file-like object like
      `BufferedReader`, `FileIO`, `StringIO` etc.: Loads the data in the
      file as an argument. **NOTE**: The uploaded value isn’t just the
      string of the file, so don’t pass the string directly; pass the
      path or wrap in StringIO.
  2.  An rush `Provider.Argument` or `ArgId` returned by a previous call
      to a rush module via `client.[some_module_name]()`: The `ArgId`
      type wraps data for use within rush. It may refer to an object
      already uploaded to rush storage, such as outputs of other run
      calls. See below for more details. It’s easier to understand when
      you see an example.
  3.  A parameter, i.e. a value of any other type, including `None`:
      Ensure the values match what is outlined in the \*args list
- \*\*kwargs
  - `target`: The machine we want to run on (eg. `NIX_SSH` for a
    cluster, `GADI` for a supercomputer).
  - `resources`: The resources to use on the target. The most commonly
    provided being
    `{"gpus": n, "storage": storage_in_units, "storage_units": "B" | "MB" | "GB", "walltime": mins}`.
  - `tags`: Tags to associate with our run, so we can easily look up our
    runs. They will be populated by the `batch_tags` passed to the
    cleint on constructionby default
  - `restore`: If this is set to True - the function will check if a
    single module_instance exists for the same version of the function
    with the same tags, and return that instead of re-running.

The return value is a list of `Provider.Arguments`. You can wait for
them to resolve by calling `await your_argument.get()`, or pass the
arguments directly to subsequent functions, which will cause Rush to do
the waiting for you.

You can see the status of all the the jobs submitted for your workspace
or session by going `client.status()`

We will now demonstrate how this works in action

## 1.1) Input Preparation

### 1.1.1) Prep the protein

First we will run the protein preparation routine (using pdbfixer
internally) to prepare the protein for molecular dynamics

``` python
# we can check the arguments and outputs for prepare_protein with help()
help(client.prepare_protein)
```

    Help on function prepare_protein in module rush.provider:

    async prepare_protein(*args: [<class 'pathlib.Path'>], target: rush.graphql_client.enums.ModuleInstanceTarget | None = <ModuleInstanceTarget.NIX_SSH_2: 'NIX_SSH_2'>, resources: rush.graphql_client.input_types.ModuleInstanceResourcesInput | None = ModuleInstanceResourcesInput(gpus=1, gpu_mem=None, gpu_mem_units=None, cpus=None, nodes=None, mem=None, mem_units=None, storage=138, storage_units=<MemUnits.MB: 'MB'>, walltime=None, storage_mounts=None), tags: list[str] | None = None, restore: bool | None = None) -> [<class 'pathlib.Path'>, <class 'pathlib.Path'>]
        Prepare a PDB for downstream tasks: protonate, fill missing atoms, etc.
        
        Module version: github:talo/pdb2pqr/ff5abe87af13f31478ede490d37468a536621e9c#prepare_protein_tengu
        
        QDX Type Description:
        
            input_pdb: @bytes 
        
        ->
        
            output_qdxf: @[Conformer];
        
            output_pdb: @bytes
        
        
        
        :param input_pdb: An input protein as a file: one PDB file
        :return output_qdxf: An output protein a vec: one qdxf per model in pdb
        :return output_pdb: An output protein as a file: one PDB file

``` python
# Here we run the function, it will return a Provider.Arg which you can use to fetch the results
# We set restore = True so that we can restore a previous run to the same path with the same tags
(prepared_protein_qdxf, prepared_protein_pdb) = await client.prepare_protein(
    PROTEIN_PDB_PATH, target=TARGET, restore=True
)
print(f"{datetime.now().time()} | Running protein prep!")
prepared_protein_qdxf  # this initially only have the id of your result, we will show how to fetch the actual value later
```

    10:30:32.714308 | Running protein prep!

    Arg(id=06f12e87-6545-4092-92ef-f67a1e8cf1de, value=None)

### 1.1.2) Checking results

#### 1.1.2.1) Run statuses

This will show the status of all of your runs

``` python
await client.status()
```

    {'9e3b4032-e241-4675-8c16-fe56b139ce7e': (<ModuleInstanceStatus.RESOLVING: 'RESOLVING'>,
      'prepare_protein',
      1)}

#### 1.1.2.2) Run Logs

If any of our runs fail, we can check their logs with

``` python
for instance_id, (status, name, count) in (await client.status()).items():
    if status.value == "FAILED":
        async for log_page in client.logs(instance_id, "stderr"):
            for log in log_page:
                print(log)
        break
```

#### 1.1.2.3) Run Values

This will return the “value” of the output from the function - for files
you will recieve a url that you can download, otherwise you will recieve
them as python types

``` python
await prepared_protein_pdb.get()
```

    2024-01-24 10:30:32,940 - rush - INFO - Argument 22c9ff71-3b4e-4aef-8407-655b9575f830 is now ModuleInstanceStatus.RESOLVING
    2024-01-24 10:30:34,063 - rush - INFO - Argument 22c9ff71-3b4e-4aef-8407-655b9575f830 is now ModuleInstanceStatus.ADMITTED
    2024-01-24 10:30:48,298 - rush - INFO - Argument 22c9ff71-3b4e-4aef-8407-655b9575f830 is now ModuleInstanceStatus.DISPATCHED
    2024-01-24 10:30:53,778 - rush - INFO - Argument 22c9ff71-3b4e-4aef-8407-655b9575f830 is now ModuleInstanceStatus.RUNNING
    2024-01-24 10:31:04,999 - rush - INFO - Argument 22c9ff71-3b4e-4aef-8407-655b9575f830 is now ModuleInstanceStatus.AWAITING_UPLOAD

    {'url': 'https://storage.googleapis.com/qdx-store/a58eb07f-c59d-48d8-994b-d24873bf408d?x-goog-signature=81fad85408de4143b6e2ae2ddfd5cb2e0ee2f69020952714d0e63f12d71611cdc28fa53e9fc63f09dc8e336563c1db78c62f426be6bff87a656a90a82b4e0c3c52513d2cd9409ae1efffe3bef368a20638d4afa63d6872694b05f4b6fb896fc94634c066857e6ddb1cfe222bca6b2360d61f5770050169c9e46dbae116f984193a4678b33243564ed184a36ab538bb94dc9dbb3fe5233df6d99ab0060b0c3a719d9c4f195ec082c8ea178f2e711154587c2314e34b9854f33543c33912db27e1eced0ca68817865dd82181494088bd62410974a853ed4be437b3384d688a4b942388eb71c3f794a34b87836f60210c3ecbb01625e71b0535ff1b62d3b4946f7d&x-goog-algorithm=GOOG4-RSA-SHA256&x-goog-credential=qdx-store-user%40humming-bird-321603.iam.gserviceaccount.com%2F20240124%2Fasia-southeast1%2Fstorage%2Fgoog4_request&x-goog-date=20240124T023143Z&x-goog-expires=3600&x-goog-signedheaders=host'}

#### 1.1.2.3) Downloads

We provide a utility to download files into your workspace, you can
either provide a filename, which will be saved in
`workspace/objects/[filename]`, or you can provide your own filepath
which the client will use as-is

``` python
try:
    await prepared_protein_pdb.download(filename="01_prepared_protein.pdb")
except FileExistsError:
    # we will raise an error if you try to overwrite an existing file, you can force the file to overwrite
    # by passing an absolute filepath instead
    pass
```

``` python
# we can read our prepared protein pdb like this
with open(client.workspace / "objects" / "01_prepared_protein.pdb", "r") as f:
    print(f.readline(), "...")
```

    REMARK   1 PDBFIXER FROM: /home/ubuntu/.cache/tengu_store/run/9e3b4032-e241-4675-8c16-fe56b139ce7e/.tmp/m2_protein.pdb
     ...

You should visualize your prepared protein to spot check any incorrectly
transformed residues

``` python
view = py3Dmol.view()
with open(client.workspace / "objects" / "01_prepared_protein.pdb", "r") as f:
    view.addModel(f.read(), "pdb")
    view.setStyle({"cartoon": {"color": "spectrum"}})
    view.zoomTo()
    # view.show() # we can't have the widget in the readme
```

### 1.1.3) Prep the ligand

Next we will prepare the ligand (using gypsum_dl internally)

``` python
# we can check the inputs for prepare_ligand with help()
help(client.prepare_ligand)
```

    Help on function prepare_ligand in module rush.provider:

    async prepare_ligand(*args: [<class 'str'>, <class 'pathlib.Path'>, dict[str, ~T]], target: rush.graphql_client.enums.ModuleInstanceTarget | None = <ModuleInstanceTarget.NIX_SSH_3: 'NIX_SSH_3'>, resources: rush.graphql_client.input_types.ModuleInstanceResourcesInput | None = ModuleInstanceResourcesInput(gpus=1, gpu_mem=None, gpu_mem_units=None, cpus=None, nodes=None, mem=None, mem_units=None, storage=18, storage_units=<MemUnits.MB: 'MB'>, walltime=None, storage_mounts=None), tags: list[str] | None = None, restore: bool | None = None) -> [<class 'pathlib.Path'>, <class 'pathlib.Path'>]
        Prepare ligand for sim. or quantum energy calc. using gypsum_dl
        
        Module version: github:talo/gypsum_dl/04acd1852cb3e2c8d0347e15763926fdf9a93a5d#prepare_ligand_tengu
        
        QDX Type Description:
        
            in: string;
        
            in: @bytes;
        
            in: {
        
            job_manager:string,
        
            let_tautomers_change_chirality:bool,
        
            max_ph:f32,
        
            max_variants_per_compound:i32,
        
            min_ph:f32,
        
            num_processors:i32,
        
            output_folder:string,
        
            pka_precision:f32,
        
            separate_output_files:bool,
        
            skip_adding_hydrogen:bool,
        
            skip_alternate_ring_conformations:bool,
        
            skip_enumerate_chiral_mol:bool,
        
            skip_enumerate_double_bonds:bool,
        
            skip_making_tautomers:bool,
        
            skip_optimize_geometry:bool,
        
            source:string,
        
            thoroughness:i32,
        
            use_durrant_lab_filters:bool
        
            } 
        
        ->
        
            out: @bytes;
        
            out: @bytesPrepare ligand for sim. or quantum energy calc. using gypsum_dl.
        
        
        Inputs:
            - An input molecule as a SMILES string
            - The same input molecule as a PDB file
            - A json file representing any other options to pass; see gypsum_dl docs for details
        
        
        Outputs:
            - A pdb file containing the prepared version of the molecule, ready for downstream use

``` python
ligand_prep_config = {
    "source": "",
    "output_folder": "./",
    "job_manager": "multiprocessing",
    "num_processors": -1,
    "max_variants_per_compound": 1,
    "thoroughness": 3,
    "separate_output_files": True,
    "min_ph": 6.4,
    "max_ph": 8.4,
    "pka_precision": 1.0,
    "skip_optimize_geometry": True,
    "skip_alternate_ring_conformations": True,
    "skip_adding_hydrogen": False,
    "skip_making_tautomers": True,
    "skip_enumerate_chiral_mol": True,
    "skip_enumerate_double_bonds": True,
    "let_tautomers_change_chirality": False,
    "use_durrant_lab_filters": True,
}
(prepared_ligand_pdb, prepared_ligand_sdf) = await client.prepare_ligand(
    LIGAND_SMILES_STR,
    LIGAND_PDB_PATH,
    ligand_prep_config,
    target=TARGET,
    restore=True,
)
print(f"{datetime.now().time()} | Running ligand prep!")
```

    10:31:49.690332 | Running ligand prep!

``` python
# we can check the status again
await client.status()
```

    {'98cf5555-0c4c-4432-bcfc-298a94f6f899': (<ModuleInstanceStatus.RESOLVING: 'RESOLVING'>,
      'prepare_ligand',
      1),
     '9e3b4032-e241-4675-8c16-fe56b139ce7e': (<ModuleInstanceStatus.COMPLETED: 'COMPLETED'>,
      'prepare_protein',
      1)}

``` python
# we can download our outputs
try:
    await prepared_ligand_pdb.download(filename="01_prepped_ligand.pdb")
    await prepared_ligand_sdf.download(filename="01_prepped_ligand.sdf")
except FileExistsError:
    pass

print(f"{datetime.now().time()} | Downloaded prepped ligand!")
```

    2024-01-24 10:31:49,926 - rush - INFO - Argument 3f04bb48-cc3e-4164-afce-632cdcd50abb is now ModuleInstanceStatus.RESOLVING
    2024-01-24 10:31:53,280 - rush - INFO - Argument 3f04bb48-cc3e-4164-afce-632cdcd50abb is now ModuleInstanceStatus.ADMITTED
    2024-01-24 10:32:05,305 - rush - INFO - Argument 3f04bb48-cc3e-4164-afce-632cdcd50abb is now ModuleInstanceStatus.DISPATCHED
    2024-01-24 10:32:11,870 - rush - INFO - Argument 3f04bb48-cc3e-4164-afce-632cdcd50abb is now ModuleInstanceStatus.AWAITING_UPLOAD
    10:32:48.809029 | Downloaded prepped ligand!

``` python
# we can read our outputs
with open(client.workspace / "objects" / "01_prepped_ligand.sdf", "r") as f:
    print(f.readline(), f.readline(), "...")
```

    untitled_0_molnum_0
          RDKit          3D
     ...

## 1.2) Run GROMACS (module: gmx_rush / gmx_rush_pdb)

Next we will run a molecular dynamics simulation on our protein and
ligand, using gromacs (gmx)

``` python
help(client.gmx)
```

    Help on function gmx in module rush.provider:

    async gmx(*args: [typing.Optional[~T], typing.Optional[~T], typing.Optional[~T], dict[str, ~T]], target: rush.graphql_client.enums.ModuleInstanceTarget | None = <ModuleInstanceTarget.NIX_SSH: 'NIX_SSH'>, resources: rush.graphql_client.input_types.ModuleInstanceResourcesInput | None = ModuleInstanceResourcesInput(gpus=4, gpu_mem=None, gpu_mem_units=None, cpus=None, nodes=None, mem=None, mem_units=None, storage=1034, storage_units=<MemUnits.MB: 'MB'>, walltime=None, storage_mounts=None), tags: list[str] | None = None, restore: bool | None = None) -> [<class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>]
        Runs a molecular dynamics simluation using GROMACS from either protein, ligand pdbs or conformers as inputs
        
        Module version: github:talo/gmx_tengu_support/6bd881c6bb32ac85ab9cb9c95d2b16676cc72b7a#gmx_tengu
        
        QDX Type Description:
        
            conformer: @Conformer?;
        
            protein: @bytes?;
        
            ligand: @bytes?;
        
            gmx-config: {
        
            frame_sel:{
        
            end_frame:u32,
        
            interval:u32,
        
            start_frame:u32
        
            }?,
        
            ligand_charge:i8?,
        
            num_gpus:u8,
        
            num_replicas:u8?,
        
            params_overrides:{
        
            em:{
        
            coulombtype:string?,
        
            cutoff_scheme:string?,
        
            emstep:f64?,
        
            emtol:f64?,
        
            integrator:string?,
        
            ns_type:string?,
        
            nsteps:i32?,
        
            nstlist:i32?,
        
            pbc:string?,
        
            rcoulomb:f64?,
        
            rlist:f64?,
        
            rvdw:f64?
        
            }?,
        
            ions:{
        
            coulombtype:string?,
        
            cutoff_scheme:string?,
        
            emstep:f64?,
        
            emtol:f64?,
        
            integrator:string?,
        
            nsteps:i32?,
        
            nstlist:i32?,
        
            pbc:string?,
        
            rcoulomb:f64?,
        
            rlist:f64?,
        
            rvdw:f64?
        
            }?,
        
            md:{
        
            compressibility:f64?,
        
            constraint_algorithm:string?,
        
            constraints:string?,
        
            continuation:string?,
        
            coulombtype:string?,
        
            cutoff_scheme:string?,
        
            disp_corr:string?,
        
            dt:f64?,
        
            fourierspacing:f64?,
        
            gen_vel:string?,
        
            integrator:string?,
        
            lincs_iter:i32?,
        
            lincs_order:i32?,
        
            ns_type:string?,
        
            nstenergy:i32?,
        
            nsteps:i32?,
        
            nstlist:i32?,
        
            nstlog:i32?,
        
            nstxout_compressed:i32?,
        
            pbc:string?,
        
            pcoupl:string?,
        
            pcoupltype:string?,
        
            pme_order:i32?,
        
            rcoulomb:f64?,
        
            ref_p:f64?,
        
            ref_t:[f64]?,
        
            rlist:f64?,
        
            rvdw:f64?,
        
            rvdw_switch:f64?,
        
            tau_p:f64?,
        
            tau_t:[f64]?,
        
            tc_grps:string?,
        
            tcoupl:string?,
        
            vdw_modifier:string?,
        
            vdwtype:string?
        
            }?,
        
            npt:{
        
            compressibility:f64?,
        
            constraint_algorithm:string?,
        
            constraints:string?,
        
            continuation:string?,
        
            coulombtype:string?,
        
            cutoff_scheme:string?,
        
            define:string?,
        
            disp_corr:string?,
        
            dt:f64?,
        
            fourierspacing:f64?,
        
            gen_vel:string?,
        
            integrator:string?,
        
            lincs_iter:i32?,
        
            lincs_order:i32?,
        
            ns_type:string?,
        
            nstenergy:i32?,
        
            nsteps:i32?,
        
            nstlist:i32?,
        
            nstlog:i32?,
        
            nstxout_compressed:i32?,
        
            pbc:string?,
        
            pcoupl:string?,
        
            pcoupltype:string?,
        
            pme_order:i32?,
        
            rcoulomb:f64?,
        
            ref_p:f64?,
        
            ref_t:[f64]?,
        
            refcoord_scaling:string?,
        
            rlist:f64?,
        
            rvdw:f64?,
        
            rvdw_switch:f64?,
        
            tau_p:f64?,
        
            tau_t:[f64]?,
        
            tc_grps:string?,
        
            tcoupl:string?,
        
            vdw_modifier:string?,
        
            vdwtype:string?
        
            }?,
        
            nvt:{
        
            constraint_algorithm:string?,
        
            constraints:string?,
        
            continuation:string?,
        
            coulombtype:string?,
        
            cutoff_scheme:string?,
        
            define:string?,
        
            disp_corr:string?,
        
            dt:f64?,
        
            fourierspacing:f64?,
        
            gen_seed:i32?,
        
            gen_temp:f64?,
        
            gen_vel:string?,
        
            integrator:string?,
        
            lincs_iter:i32?,
        
            lincs_order:i32?,
        
            ns_type:string?,
        
            nstenergy:i32?,
        
            nsteps:i32?,
        
            nstlist:i32?,
        
            nstlog:i32?,
        
            nstxout_compressed:i32?,
        
            pbc:string?,
        
            pcoupl:string?,
        
            pme_order:i32?,
        
            rcoulomb:f64?,
        
            ref_t:[f64]?,
        
            rlist:f64?,
        
            rvdw:f64?,
        
            rvdw_switch:f64?,
        
            tau_t:[f64]?,
        
            tc_grps:string?,
        
            tcoupl:string?,
        
            vdw_modifier:string?,
        
            vdwtype:string?
        
            }?
        
            }?,
        
            save_wets:bool
        
            } 
        
        ->
        
            output_gros: @bytes;
        
            outputs_tpr: @bytes;
        
            outputs_tops: @bytes;
        
            output_logs: @bytes;
        
            outputs_md.ligand_in.dry.xtc: @bytes;
        
            outputs_md.ligand_in.dry.pdb: @bytes;
        
            index.ligand_in.ndx: @bytes;
        
            outputs_md.ligand_in.xtc: @bytes
        
        
        
        :param conformer: Optional Conformer in QDXF format; must provide either this argument, or the Protein PDB argument .
        :param protein: Protein PDB file; provide this if no Conformer qdxf is provided.
        :param ligand: Ligand PDB file
        :param gmx-config: Configuration record
        :return output_gros: Intermediate and final structure files through the simulation
        :return outputs_tpr: .tpr files of the production MD runs
        :return outputs_tops: .top files, and relevant .itp files used in the simulation.
        :return output_logs: Logs of the production MD runs
        :return outputs_md.ligand_in.dry.xtc: Trajectories, i.e., without water molecules, from the production MD runs
        :return outputs_md.ligand_in.dry.pdb: Outputs of select_frame, pdb frames without water
        :return index.ligand_in.ndx: Index file
        :return outputs_md.ligand_in.xtc: Wet trajectories, i.e., with water molecules, from the production MD runs

``` python
gmx_config = {
    "param_overrides": {
        "md": {"nsteps", "5000"},
        "em": {"nsteps", "1000"},
        "nvt": {"nsteps", "1000"},
        "npt": {"nsteps", "1000"},
        "ions": {},
    },
    "num_gpus": 0,
    "num_replicas": 1,
    "ligand_charge": None,
    "save_wets": False,
    "frame_sel": {
        "start_frame": 1,
        "end_frame": 10,
        "interval": 2,
    },
}
# we pass the outputs from our prior runs directly, instead of their values, to prevent them from being re-uploaded
gmx_results = await client.gmx(
    None,
    prepared_protein_pdb,
    prepared_ligand_pdb,
    gmx_config,
    resources={"gpus": 0, "storage": 1, "storage_units": "GB", "cpus": 48, "walltime": 60},
    target=TARGET,
)
print(f"{datetime.now().time()} | Running GROMACS simulation!")
```

    10:32:48.888017 | Running GROMACS simulation!

``` python
# we can check the status again
await client.status()
```

    {'18704b66-f74d-49f0-83ba-b1964c30339b': (<ModuleInstanceStatus.RESOLVING: 'RESOLVING'>,
      'gmx',
      1),
     '98cf5555-0c4c-4432-bcfc-298a94f6f899': (<ModuleInstanceStatus.COMPLETED: 'COMPLETED'>,
      'prepare_ligand',
      1),
     '9e3b4032-e241-4675-8c16-fe56b139ce7e': (<ModuleInstanceStatus.COMPLETED: 'COMPLETED'>,
      'prepare_protein',
      1)}

``` python
print("Fetching gmx results")
try:
    await gmx_results[5].download(filename="02_gmx_dry_frames.tar.gz")
    await gmx_results[0].download(filename="02_gmx_lig_gro.tar.gz")

except FileExistsError:
    pass

print(f"{datetime.now().time()} | Downloaded GROMACS output!")
```

    Fetching gmx results
    2024-01-24 10:32:49,109 - rush - INFO - Argument 7bc36817-f050-4c7d-8dac-7b70461d69b4 is now ModuleInstanceStatus.RESOLVING
    2024-01-24 10:32:50,229 - rush - INFO - Argument 7bc36817-f050-4c7d-8dac-7b70461d69b4 is now ModuleInstanceStatus.ADMITTED
    2024-01-24 10:33:01,203 - rush - INFO - Argument 7bc36817-f050-4c7d-8dac-7b70461d69b4 is now ModuleInstanceStatus.DISPATCHED
    2024-01-24 10:33:06,757 - rush - INFO - Argument 7bc36817-f050-4c7d-8dac-7b70461d69b4 is now ModuleInstanceStatus.RUNNING
    2024-01-24 10:50:47,692 - rush - INFO - Argument 7bc36817-f050-4c7d-8dac-7b70461d69b4 is now ModuleInstanceStatus.AWAITING_UPLOAD
    10:51:37.800711 | Downloaded GROMACS output!

``` python
# Extract the "dry" (i.e. non-solvated) pdb frames we asked for
with tarfile.open(client.workspace / "objects" / "02_gmx_dry_frames.tar.gz", "r") as tf:
    selected_frame_pdbs = [
        tf.extractfile(member).read() for member in tf if "pdb" in member.name and member.isfile()
    ]
    for i, frame in enumerate(selected_frame_pdbs):
        with open(client.workspace / "objects" / f"02_gmx_output_frame_{i}.pdb", "w") as pf:
            print(frame.decode("utf-8"), file=pf)
```

``` python
# Extract the ligand.gro file
with tarfile.open(client.workspace / "objects" / "02_gmx_lig_gro.tar.gz", "r") as tf:
    gro = [tf.extractfile(member).read() for member in tf if "md.ligand" in member.name][0]
    with open(client.workspace / "objects" / f"02_gmx_lig.gro", "w") as pf:
        print(gro.decode("utf-8"), file=pf)
```

## 1.3) Run MM-PBSA

``` python
help(client.gmx_mmpbsa)
```

    Help on function gmx_mmpbsa in module rush.provider:

    async gmx_mmpbsa(*args: [<class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>, <class 'pathlib.Path'>, dict[str, ~T]], target: rush.graphql_client.enums.ModuleInstanceTarget | None = <ModuleInstanceTarget.NIX_SSH_3: 'NIX_SSH_3'>, resources: rush.graphql_client.input_types.ModuleInstanceResourcesInput | None = ModuleInstanceResourcesInput(gpus=0, gpu_mem=None, gpu_mem_units=None, cpus=None, nodes=None, mem=None, mem_units=None, storage=10, storage_units=<MemUnits.MB: 'MB'>, walltime=None, storage_mounts=None), tags: list[str] | None = None, restore: bool | None = None) -> [<class 'pathlib.Path'>]
        Run gmx mmpbsa on the outputs of a gmx simulation 
        
        Module version: github:talo/gmx_tengu_support/6bd881c6bb32ac85ab9cb9c95d2b16676cc72b7a#gmx_mmpbsa_tengu
        
        QDX Type Description:
        
            tpr_tar_gz: @bytes;
        
            dry_xtc_tar_gz: @bytes;
        
            index_file: @bytes;
        
            topol_file: @bytes;
        
            mmpbsa_config: {
        
            end_frame:u64,
        
            ie_segment:u32?,
        
            interaction_entropy:bool?,
        
            interval:u32?,
        
            num_cpus:u32,
        
            start_frame:u64,
        
            with_gb:bool?
        
            } 
        
        ->
        
            output: @bytes
        
        
        
        :param tpr_tar_gz: Compressed GROMACS output folder
        :param dry_xtc_tar_gz: Compressed GROMACS output folder
        :param index_file: Compressed GROMACS output folder
        :param topol_file: Compressed GROMACS output folder
        :param mmpbsa_config: Configuration record for mmpbsa:
            start_frame: Frame to start with
            end_frame: Frame to end with
            num_cpus: Number of CPUs to use - cannot be larger than the number of frames
            interaction_entropy: Calculate interaction entropy
            
        :return output: Compressed mmpbsa output folder

``` python
mmpbsa_config = {
    "start_frame": 1,
    "end_frame": 2,
    "num_cpus": 1,  # cannot be greater than number of frames
}
(mmpbsa_result_tar,) = await client.gmx_mmpbsa(
    gmx_results[1],
    gmx_results[4],
    gmx_results[6],
    gmx_results[2],
    mmpbsa_config,
    resources=rush.Resources(storage=100, storage_units="MB", cpus=42, gpus=0, walltime=60, mem=1024 * 42),
    target="GADI",
)
print(f"{datetime.now().time()} | Running GROMACS MM-PBSA calculation!")
```

    10:51:37.932948 | Running GROMACS MM-PBSA calculation!

``` python
print("Fetching gmx_mmpbsa results")
try:
    await mmpbsa_result_tar.download(filename="04_gmx_mmpbsa_run_folder.tar.gz")
except FileExistsError:
    pass
print(f"{datetime.now().time()} | Downloaded MM-PBSA results!")
```

    Fetching gmx_mmpbsa results
    2024-01-24 10:51:38,107 - rush - INFO - Argument f4f59056-8bf7-490c-b905-d0c4fac1b254 is now ModuleInstanceStatus.RESOLVING
    2024-01-24 10:55:22,967 - rush - INFO - Argument f4f59056-8bf7-490c-b905-d0c4fac1b254 is now ModuleInstanceStatus.ADMITTED
    2024-01-24 10:55:53,845 - rush - INFO - Argument f4f59056-8bf7-490c-b905-d0c4fac1b254 is now ModuleInstanceStatus.DISPATCHED
    2024-01-24 10:55:58,254 - rush - INFO - Argument f4f59056-8bf7-490c-b905-d0c4fac1b254 is now ModuleInstanceStatus.QUEUED

