From 0e99a8ff6016d63a475eaf18ad4041cba4fb4e6c Mon Sep 17 00:00:00 2001
From: Chris Brasnett <35073246+csbrasnett@users.noreply.github.com>
Date: Mon, 6 Jan 2025 12:40:28 +0000
Subject: [PATCH] add CLI for backbone and go atomnames as per #585

---
 bin/martinize2                  | 18 +++++++++++++++++-
 vermouth/rcsu/go_pipeline.py    |  2 ++
 vermouth/rcsu/go_vs_includes.py | 14 +++++++++-----
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/bin/martinize2 b/bin/martinize2
index c48197a9f..0f61257eb 100755
--- a/bin/martinize2
+++ b/bin/martinize2
@@ -589,6 +589,20 @@ def entry():
         type=Path,
         help=("Write out contact map to file if calculating as part of Martinize2.")
     )
+    go_group.add_argument(
+        "-go-backbone",
+        dest="go_backbone",
+        default="BB",
+        type=str,
+        help=("Name of protein backbone bead on which to place a virtual interaction go site")
+    )
+    go_group.add_argument(
+        "-go-atomname",
+        dest="go_atomname",
+        default="CA",
+        type=str,
+        help=("Name of the virtual interaction go site atom")
+    )
 
     water_group = parser.add_argument_group("Apply water bias.")
     water_group.add_argument(
@@ -1022,7 +1036,9 @@ def entry():
                               cutoff_short=args.go_low,
                               cutoff_long=args.go_up,
                               go_eps=args.go_eps,
-                              res_dist=args.go_res_dist,)
+                              res_dist=args.go_res_dist,
+                              go_anchor_bead=args.go_backbone,
+                              go_atomname=args.go_atomname)
 
         defines = ("GO_VIRT",)
         itp_paths = {"atomtypes": "go_atomtypes.itp",
diff --git a/vermouth/rcsu/go_pipeline.py b/vermouth/rcsu/go_pipeline.py
index 1917cb178..4312bd566 100644
--- a/vermouth/rcsu/go_pipeline.py
+++ b/vermouth/rcsu/go_pipeline.py
@@ -40,6 +40,8 @@ def prepare_run(self, system, moltype):
         vermouth.MergeAllMolecules().run_system(system)
         molecule = system.molecules[0]
         molecule.meta['moltype'] = moltype
+        # add citations for the go model here
+        molecule.citations.add('M3_GO')
 
     def run_system(self, system, **kwargs):
         self.kwargs = kwargs
diff --git a/vermouth/rcsu/go_vs_includes.py b/vermouth/rcsu/go_vs_includes.py
index a65920ab1..0aa622c60 100644
--- a/vermouth/rcsu/go_vs_includes.py
+++ b/vermouth/rcsu/go_vs_includes.py
@@ -40,8 +40,10 @@ class VirtualSiteCreator(Processor):
         Assign molecule type names to the molecules in a system.
     :func:`add_virtual_sites`
     """
-    def __init__(self):
+    def __init__(self, go_anchor_bead, go_atomname):
         self.system = None
+        self.backbone = go_anchor_bead
+        self.atomname = go_atomname
 
     def run_molecule(self, molecule):
         moltype = molecule.meta.get('moltype')
@@ -50,9 +52,11 @@ def run_molecule(self, molecule):
 
         if not self.system:
             raise ValueError('This processor requires a system.')
-        molecule.citations.add('M3_GO')
-
-        self.add_virtual_sites(molecule, prefix=moltype)
+        print(self.backbone, self.atomname)
+        self.add_virtual_sites(molecule,
+                               prefix=moltype,
+                               backbone=self.backbone,
+                               atomname=self.atomname)
 
         return molecule
 
@@ -65,7 +69,7 @@ def add_virtual_sites(self, molecule, prefix, backbone='BB', atomname='CA', char
         """
         Add the virtual sites for GoMartini in the molecule.
 
-        One virtual site is added per backbone bead of the the Martini protein.
+        One virtual site is added per backbone bead of the Martini protein.
         Each virtual site copies the resid, resname, and chain of the backbone
         bead. It also copies the *reference* to the position array, so the virtual
         site position follows if the backbone bead is translated. The virtual sites