Repository: haddocking/pdb-tools (language: python)

Issue:
Implement `pdb_selmodel`
**Is your tool request related to a problem? Please describe.**
When we have a multi-model ensemble we need to use `pdb_splitmodel` and then (in disk) find the appropriate model we want.

**Describe the solution you'd like**
`pdb_selmodel` would select a given `MODEL` in a ensemble or a range of models:

```
$ pdb_selmodel -1 1ccm.pdb
$ pdb_selmodel -1:5 1ccm.pdb
```

In case the input is not an ensemble it could instead just return the same pdb instead of an error, then it would not break any pipelines.

**Does any of the existing tools or a combination of them do what you want?**
No, but `pdb_splitmodel` + scripting does the trick.



Discussion / hints:
We have talked about this in the past. I agree with having this tool, I
also find it very useful.

El mié., 19 ene. 2022 14:08, Rodrigo V Honorato ***@***.***>
escribió:

> *Is your tool request related to a problem? Please describe.*
> When we have a multi-model ensemble we need to use pdb_splitmodel and
> then (in disk) find the appropriate model we want.
>
> *Describe the solution you'd like*
> pdb_selmodel would select a given MODEL in a ensemble or a range of
> models:
>
> $ pdb_selmodel -1 1ccm.pdb
> $ pdb_selmodel -1:5 1ccm.pdb
>
> In case the input is not an ensemble it could instead just return the same
> pdb instead of an error, then it would not break any pipelines.
>
> *Does any of the existing tools or a combination of them do what you want?*
> No, but pdb_splitmodel + scripting does the trick.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/haddocking/pdb-tools/issues/122>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAWEOMWUJIHSHB6ORNQSLRDUW2ZUZANCNFSM5MJ4PHNA>
> .
> Triage notifications on the go with GitHub Mobile for iOS
> <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
> or Android
> <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.
>
> You are receiving this because you are subscribed to this thread.Message
> ID: ***@***.***>
>


Fix patch:
diff --git a/docs/index.md b/docs/index.md
index b876fcf0..b8552598 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -628,6 +628,21 @@ Example:
 </div>
 <div style="margin-bottom: 1em;">
 <details>
+<summary><b>pdb_selmodel</b><p>Extracts one or more models from a PDB file.</p></summary>
+<span style="font-family: monospace; white-space: pre;">
+If the PDB file has no MODEL records, returns the entire file.
+
+Usage:
+    python pdb_selmodel.py -&lt;model id&gt; &lt;pdb file&gt;
+
+Example:
+    python pdb_selmodel.py -1 1GGR.pdb  # selects model 1
+    python pdb_selmodel.py -1,3 1GGR.pdb  # selects models 1 and 3
+</span>
+</details>
+</div>
+<div style="margin-bottom: 1em;">
+<details>
 <summary><b>pdb_selres</b><p>Selects residues by their index, piecewise or in a range.</p></summary>
 <span style="font-family: monospace; white-space: pre;">
 The range option has three components: start, end, and step. Start and end
diff --git a/pdbtools/pdb_selmodel.py b/pdbtools/pdb_selmodel.py
new file mode 100644
index 00000000..c51a3537
--- /dev/null
+++ b/pdbtools/pdb_selmodel.py
@@ -0,0 +1,194 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2022 João Pedro Rodrigues
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Extracts one or more models from a PDB file.
+
+If the PDB file has no MODEL records, returns the entire file.
+
+Usage:
+    python pdb_selmodel.py -<model id> <pdb file>
+
+Example:
+    python pdb_selmodel.py -1 1GGR.pdb  # selects model 1
+    python pdb_selmodel.py -1,3 1GGR.pdb  # selects models 1 and 3
+
+This program is part of the `pdb-tools` suite of utilities and should not be
+distributed isolatedly. The `pdb-tools` were created to quickly manipulate PDB
+files using the terminal, and can be used sequentially, with one tool streaming
+data to another. They are based on old FORTRAN77 code that was taking too much
+effort to maintain and compile. RIP.
+"""
+
+import os
+import sys
+
+__author__ = "Joao Rodrigues"
+__email__ = "j.p.g.l.m.rodrigues@gmail.com"
+
+
+def check_input(args):
+    """Checks whether to read from stdin/file and validates user input/options.
+    """
+
+    # Defaults
+    option = ''
+    fh = sys.stdin  # file handle
+
+    if not len(args):
+        # Reading from pipe with default option
+        if sys.stdin.isatty():
+            sys.stderr.write(__doc__)
+            sys.exit(1)
+
+    elif len(args) == 1:
+        # One of two options: option & Pipe OR file & default option
+        if args[0].startswith('-'):
+            option = args[0][1:]
+            if sys.stdin.isatty():  # ensure the PDB data is streamed in
+                emsg = 'ERROR!! No data to process!\n'
+                sys.stderr.write(emsg)
+                sys.stderr.write(__doc__)
+                sys.exit(1)
+
+        else:
+            if not os.path.isfile(args[0]):
+                emsg = 'ERROR!! File not found or not readable: \'{}\'\n'
+                sys.stderr.write(emsg.format(args[0]))
+                sys.stderr.write(__doc__)
+                sys.exit(1)
+
+            fh = open(args[0], 'r')
+
+    elif len(args) == 2:
+        # Two options: option & File
+        if not args[0].startswith('-'):
+            emsg = 'ERROR! First argument is not an option: \'{}\'\n'
+            sys.stderr.write(emsg.format(args[0]))
+            sys.stderr.write(__doc__)
+            sys.exit(1)
+
+        if not os.path.isfile(args[1]):
+            emsg = 'ERROR!! File not found or not readable: \'{}\'\n'
+            sys.stderr.write(emsg.format(args[1]))
+            sys.stderr.write(__doc__)
+            sys.exit(1)
+
+        option = args[0][1:]
+        fh = open(args[1], 'r')
+
+    else:  # Whatever ...
+        sys.stderr.write(__doc__)
+        sys.exit(1)
+
+    # Validate option
+    option_set = set()
+    for opt in option.split(','):
+        opt = opt.strip()
+        if not opt:
+            continue
+        try:
+            opt = int(opt)
+        except ValueError:
+            emsg = 'ERROR!! Model identifier is invalid: \'{}\'\n'
+            sys.stderr.write(emsg.format(opt))
+            sys.stderr.write(__doc__)
+            sys.exit(1)
+        option_set.add(opt)
+
+    if not option_set:
+        emsg = 'ERROR!! You must provide at least one model identifier\n'
+        sys.stderr.write(emsg)
+        sys.stderr.write(__doc__)
+        sys.exit(1)
+
+    return (fh, option_set)
+
+
+def run(fhandle, model_set):
+    """
+    Filter the PDB file for specific model identifiers.
+
+    This function is a generator.
+
+    Parameters
+    ----------
+    fhandle : a line-by-line iterator of the original PDB file.
+
+    model_set : set, or list, or tuple
+        The group of models to kepp.
+        Example: (1, 3), keeps only atoms from models 1 and 3
+
+    Yields
+    ------
+    str (line-by-line)
+        The PDB lines for those matching the selected chains.
+    """
+    records = ('ATOM', 'HETATM', 'TER', 'ANISOU')
+    ignore_model = False
+    for line in fhandle:
+        if line.startswith('MODEL'):
+            model_id = int(line[10:14])
+            if model_id not in model_set:
+                ignore_model = True
+                continue
+
+        elif line.startswith('ENDMDL'):
+            if ignore_model:
+                ignore_model = False
+                continue
+
+        elif ignore_model and line.startswith(records):
+            continue
+        yield line
+
+
+select_model = run
+
+
+def main():
+    # Check Input
+    pdbfh, models = check_input(sys.argv[1:])
+
+    # Do the job
+    new_pdb = run(pdbfh, models)
+
+    try:
+        _buffer = []
+        _buffer_size = 5000  # write N lines at a time
+        for lineno, line in enumerate(new_pdb):
+            if not (lineno % _buffer_size):
+                sys.stdout.write(''.join(_buffer))
+                _buffer = []
+            _buffer.append(line)
+
+        sys.stdout.write(''.join(_buffer))
+        sys.stdout.flush()
+    except IOError:
+        # This is here to catch Broken Pipes
+        # for example to use 'head' or 'tail' without
+        # the error message showing up
+        pass
+
+    # last line of the script
+    # We can close it even if it is sys.stdin
+    pdbfh.close()
+    sys.exit(0)
+
+
+if __name__ == '__main__':
+    main()


Test patch:
diff --git a/tests/test_pdb_selmodel.py b/tests/test_pdb_selmodel.py
new file mode 100644
index 00000000..49a3cdf7
--- /dev/null
+++ b/tests/test_pdb_selmodel.py
@@ -0,0 +1,165 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2022 João Pedro Rodrigues
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Unit Tests for `pdb_selmodel`.
+"""
+
+import os
+import sys
+import unittest
+
+from config import data_dir
+from utils import OutputCapture
+
+
+class TestTool(unittest.TestCase):
+    """
+    Generic class for testing tools.
+    """
+
+    def setUp(self):
+        # Dynamically import the module
+        name = 'pdbtools.pdb_selmodel'
+        self.module = __import__(name, fromlist=[''])
+
+    def exec_module(self):
+        """
+        Execs module.
+        """
+
+        with OutputCapture() as output:
+            try:
+                self.module.main()
+            except SystemExit as e:
+                self.retcode = e.code
+
+        self.stdout = output.stdout
+        self.stderr = output.stderr
+
+        return
+
+    def test_one_option(self):
+        """$ pdb_selmodel -1 data/ensemble_OK.pdb"""
+
+        # Simulate input
+        # pdb_selmodel ensemble_OK.pdb
+        sys.argv = ['', '-1', os.path.join(data_dir, 'ensemble_OK.pdb')]
+
+        # Execute the script
+        self.exec_module()
+
+        # Validate results
+        self.assertEqual(self.retcode, 0)  # ensure the program exited OK.
+        self.assertEqual(len(self.stdout), 7)  # selected model 1
+        self.assertEqual(len(self.stderr), 0)  # no errors
+
+    def test_multiple(self):
+        """
+        $ pdb_selmodel -1,2 data/ensemble_OK.pdb
+        """
+
+        sys.argv = ['', '-1,2', os.path.join(data_dir, 'ensemble_OK.pdb')]
+
+        self.exec_module()
+
+        self.assertEqual(self.retcode, 0)
+        self.assertEqual(len(self.stdout), 11)  # all models
+        self.assertEqual(len(self.stderr), 0)
+
+    def test_file_not_found(self):
+        """$ pdb_selmodel not_existing.pdb"""
+
+        afile = os.path.join(data_dir, 'not_existing.pdb')
+        sys.argv = ['', afile]
+
+        self.exec_module()
+
+        self.assertEqual(self.retcode, 1)  # exit code is 1 (error)
+        self.assertEqual(len(self.stdout), 0)  # nothing written to stdout
+        self.assertEqual(self.stderr[0][:22],
+                         "ERROR!! File not found")  # proper error message
+
+    @unittest.skipIf(os.getenv('SKIP_TTY_TESTS'), 'skip on GHA - no TTY')
+    def test_file_missing(self):
+        """$ pdb_selmodel -1"""
+
+        sys.argv = ['', '-1']
+
+        self.exec_module()
+
+        self.assertEqual(self.retcode, 1)
+        self.assertEqual(len(self.stdout), 0)  # no output
+        self.assertEqual(self.stderr[0],
+                         "ERROR!! No data to process!")
+
+    @unittest.skipIf(os.getenv('SKIP_TTY_TESTS'), 'skip on GHA - no TTY')
+    def test_helptext(self):
+        """$ pdb_selmodel"""
+
+        sys.argv = ['']
+
+        self.exec_module()
+
+        self.assertEqual(self.retcode, 1)  # ensure the program exited gracefully.
+        self.assertEqual(len(self.stdout), 0)  # no output
+        self.assertEqual(self.stderr, self.module.__doc__.split("\n")[:-1])
+
+    def test_invalid_option(self):
+        """$ pdb_selmodel data/ensemble_OK.pdb"""
+
+        sys.argv = ['', os.path.join(data_dir, 'ensemble_OK.pdb')]
+
+        self.exec_module()
+
+        self.assertEqual(self.retcode, 1)
+        self.assertEqual(len(self.stdout), 0)
+        self.assertEqual(self.stderr[0][:47],
+                         "ERROR!! You must provide at least one model ide")
+
+    def test_invalid_option_2(self):
+        """$ pdb_selmodel -AB data/ensemble_OK.pdb"""
+
+        sys.argv = ['', '-AB', os.path.join(data_dir, 'ensemble_OK.pdb')]
+
+        self.exec_module()
+
+        self.assertEqual(self.retcode, 1)
+        self.assertEqual(len(self.stdout), 0)
+        self.assertEqual(self.stderr[0][:40],
+                "ERROR!! Model identifier is invalid: 'AB")
+
+    def test_not_an_option(self):
+        """$ pdb_selmodel 20 data/ensemble_OK.pdb"""
+
+        sys.argv = ['', '20', os.path.join(data_dir, 'ensemble_OK.pdb')]
+
+        self.exec_module()
+
+        self.assertEqual(self.retcode, 1)
+        self.assertEqual(len(self.stdout), 0)
+        self.assertEqual(self.stderr[0],
+                         "ERROR! First argument is not an option: '20'")
+
+
+if __name__ == '__main__':
+    from config import test_dir
+
+    mpath = os.path.abspath(os.path.join(test_dir, '..'))
+    sys.path.insert(0, mpath)  # so we load dev files before  any installation
+
+    unittest.main()