internal go contact map

bin/martinize2

@@ -22,6 +22,7 @@ import argparse
 import functools
 import logging
 import itertools
+import os.path
 from pathlib import Path
 import sys
 import networkx as nx
@@ -48,7 +49,7 @@ from vermouth.map_input import (
     generate_all_self_mappings,
     combine_mappings,
 )
-from vermouth.rcsu.contact_map import read_go_map
+from vermouth.rcsu.contact_map import GenerateContactMap
 from vermouth.rcsu.go_pipeline import GoPipeline
 from vermouth.gmx.topology import write_gmx_topology
 
@@ -546,11 +547,12 @@ def entry():
     go_group.add_argument(
         "-go",
         dest="go",
-        required=False,
+        nargs='?',
+        const=None,
         type=Path,
-        default=None,
-        help="Contact map to be used for the Martini Go model."
-             "Currently, only one format is supported. See docs.",
+        help="Use Martini Go model. Accepts either an input file from the server, "
+             "or just provide the flag to calculate as part of Martinize. Can be slow for large proteins "
+             "(> 1500 residues)"
     )
     go_group.add_argument(
         "-go-eps",
@@ -826,6 +828,22 @@ def entry():
             "be used together."
         )
 
+    """
+    Sort out the use of the go model:
+
+    go = True : do the go model
+    go_file = None : calculate the contact map in martinize2
+    go_file = some file : read in the contact map from elsewhere
+    """
+    go = False
+    go_file = None
+    if args.go is None:
+        go = True
+    else:
+        if Path(args.go).is_file():
+            go_file = args.go
+            go = True
+
     if args.to_ff.startswith("elnedyn"):
         # FIXME: This type of thing should be added to the FF itself.
         LOGGER.info(
@@ -967,6 +985,14 @@ def entry():
     elif args.cystein_bridge != "auto":
         vermouth.AddCysteinBridgesThreshold(args.cystein_bridge).run_system(system)
 
+    go_map = False
+    if go:
+        # need this here because have to get contact map at atomistic resolution
+        if go_file is None:
+            LOGGER.info("Generating Go model contact map.", type="step")
+            GenerateContactMap().run_system(system)
+            go_map = True
+
     # Run martinize on the system.
     system = martinize(
         system,
@@ -988,14 +1014,15 @@ def entry():
         vermouth.ApplyPosres(node_selector, args.posres_fc).run_system(system)
 
     # Generate the Go model if required
-    if args.go:
+
+    if go:
+        if not system.go_params["go_map"]:
+            LOGGER.info("Reading Go model contact map.", type="step")
+            GenerateContactMap(path=go_file).run_system(system)
         go_name_prefix = args.molname
-        LOGGER.info("Reading Go model contact map.", type="step")
-        go_map = read_go_map(args.go)
         LOGGER.info("Generating the Go model.", type="step")
         GoPipeline.run_system(system,
                               moltype=go_name_prefix,
-                              contact_map=go_map,
                               cutoff_short=args.go_low,
                               cutoff_long=args.go_up,
                               go_eps=args.go_eps,