feat: custom step (#346)

* feat: custom topology mod as a recipe step

* chore: doc wording

* format

src/kimmdy/recipe.py

@@ -1,7 +1,7 @@
 """Contains the Reaction Recipe, RecipeStep and RecipeCollection.
 """
 from __future__ import annotations
-from typing import Optional
+from typing import TYPE_CHECKING, Callable, Optional
 
 from abc import ABC
 from dataclasses import dataclass, field
@@ -14,6 +14,9 @@
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from kimmdy.topology.topology import Topology
+
 
 class RecipeStep(ABC):
     """Base class for all RecipeSteps.
@@ -210,6 +213,19 @@ class Bind(BondOperation):
     """
 
 
+@dataclass
+class CustomTopMod(RecipeStep):
+    """A custom recipe step that can be used to define a custom topology modification.
+
+    Parameters
+    ----------
+    f : Callable[[Topology], Topology]
+        A function that takes a Topology object and modifies it in place.
+    """
+
+    f: Callable[[Topology], None]
+
+
 @dataclass
 class Recipe:
     """A reaction path defined by one series of RecipeSteps.
@@ -328,6 +344,9 @@ def get_recipe_name(self):
             elif isinstance(rs, Relax):
                 pass
 
+            elif isinstance(rs, CustomTopMod):
+                name += "🪛"
+
             else:
                 logger.warning(f"get_recipe_name got unknown step type: {type(rs)}")
                 name += "?".join(list(map(str, rs.__dict__.values())))

src/kimmdy/runmanager.py

@@ -22,7 +22,7 @@
     reaction_plugins,
     ReactionPlugin,
 )
-from kimmdy.recipe import RecipeCollection, Break, Bind, Place, Relax
+from kimmdy.recipe import CustomTopMod, RecipeCollection, Break, Bind, Place, Relax
 from kimmdy.utils import run_gmx, truncate_sim_files
 from kimmdy.coordinates import place_atom, break_bond_plumed, merge_top_slow_growth
 from kimmdy.tasks import Task, TaskFiles, get_plumed_out
@@ -624,6 +624,9 @@ def _apply_recipe(self, files: TaskFiles) -> TaskFiles:
                 logger.info(f"Finished task: {task.name}")
                 self._discover_output_files(task.name, md_files)
 
+            elif isinstance(step, CustomTopMod):
+                step.f(self.top)
+
         write_top(self.top.to_dict(), files.outputdir / self.config.top.name)
         files.output["top"] = files.outputdir / self.config.top.name