From 69437d20112aee3ef7c63d6fa9e0dc7df1a094aa Mon Sep 17 00:00:00 2001
From: Peter Kroon <pckroon@users.noreply.github.com>
Date: Wed, 8 May 2024 11:42:00 +0200
Subject: [PATCH] Change all molecule changing logic in bin/martinize2 to
 Processors

---
 bin/martinize2                         | 17 +++++----------
 vermouth/processors/__init__.py        |  1 +
 vermouth/processors/copy_node_attrs.py | 29 ++++++++++++++++++++++++++
 3 files changed, 35 insertions(+), 12 deletions(-)
 create mode 100644 vermouth/processors/copy_node_attrs.py

diff --git a/bin/martinize2 b/bin/martinize2
index cc279b16..7e3fe1bf 100755
--- a/bin/martinize2
+++ b/bin/martinize2
@@ -22,13 +22,12 @@ import argparse
 import functools
 import logging
 import itertools
-import textwrap
 from pathlib import Path
 import sys
 import networkx as nx
 import vermouth
 import vermouth.forcefield
-from vermouth.file_writer import deferred_open, DeferredFileWriter
+from vermouth.file_writer import DeferredFileWriter
 from vermouth import DATA_PATH
 from vermouth.processors.processor import ProcessorPipeline
 from vermouth.dssp import dssp
@@ -52,12 +51,8 @@ from vermouth.map_input import (
     combine_mappings,
 )
 from vermouth.rcsu.contact_map import read_go_map
-from vermouth.citation_parser import citation_formatter
 from vermouth.gmx.topology import write_gmx_topology
 
-# TODO Since vermouth's __init__.py does some logging (KDTree), this may or may
-# not work as intended. Investigation required.
-
 LOGGER = TypeAdapter(logging.getLogger("vermouth"))
 
 PRETTY_FORMATTER = logging.Formatter(
@@ -1045,20 +1040,18 @@ def entry():
         # if the go model hasn't been used we need to create virtual
         # sites for the biasing
         if not args.go:
-            vermouth.rcsu.go_vs_includes.VirtualSiteCreator().run_system(system)
+            postprocessing.add(VirtualSiteCreator())
             itp_paths = {"atomtypes": "virtual_sites_atomtypes.itp",
                          "nonbond_params": "virtual_sites_nonbond_params.itp"}
         # now we add a bias by defining specific virtual-site water interactions
-        vermouth.processors.ComputeWaterBias(args.water_bias,
+        postprocessing.add(vermouth.ComputeWaterBias(args.water_bias,
                                              { s:float(eps) for s, eps in args.water_bias_eps},
                                              [(int(start), int(stop)) for start, stop in args.water_idrs],
-                                            ).run_system(system)
+                                            ))
 
     # Here we need to add the resids from the PDB back if that is needed
     if args.resid_handling == "input":
-        for molecule in system.molecules:
-            old_resids = nx.get_node_attributes(molecule, "_old_resid")
-            nx.set_node_attributes(molecule, old_resids, "resid")
+        postprocessing.add(vermouth.CopyNodeAttrs("_old_resid", "resid"))
 
     # The Martini Go model assumes that we do not mess with the order of
     # particles in any way especially the virtual sites needed for the Go
diff --git a/vermouth/processors/__init__.py b/vermouth/processors/__init__.py
index 4cdcc58e..9a08273d 100644
--- a/vermouth/processors/__init__.py
+++ b/vermouth/processors/__init__.py
@@ -44,3 +44,4 @@
 from .merge_all_molecules import MergeAllMolecules
 from .annotate_mut_mod import AnnotateMutMod
 from .water_bias import ComputeWaterBias
+from .copy_node_attrs import CopyNodeAttrs
diff --git a/vermouth/processors/copy_node_attrs.py b/vermouth/processors/copy_node_attrs.py
new file mode 100644
index 00000000..549eb592
--- /dev/null
+++ b/vermouth/processors/copy_node_attrs.py
@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Copyright 2024 University of Groningen
+#
+# 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.
+
+import networkx as nx
+from .processor import Processor
+
+
+class CopyNodeAttrs(Processor):
+    def __init__(self, from_attr, to_attr):
+        self.from_attr = from_attr
+        self.to_attr = to_attr
+
+    def run_molecule(self, molecule):
+        values = nx.get_node_attributes(molecule, self.from_attr)
+        nx.set_node_attributes(molecule, values, self.to_attr)
+        return molecule
\ No newline at end of file