From 5a2c08006f275a479a86485fafc7e07889355ff6 Mon Sep 17 00:00:00 2001
From: Greg Bronevetsky <bronevet@google.com>
Date: Tue, 28 Oct 2025 15:27:52 +0000
Subject: [PATCH] Experiment with adding pest plugin.

Prompt:
Add a model of pest infestation of corn plants and the impact of these infestations on crop plant health. Allow a new configuration parameter that controls the severity of the infestation
---
 pythia/plugin.py                            |  9 +++--
 pythia/plugins/pest_infestation/__init__.py | 37 +++++++++++++++++++++
 pythia/plugins/test_plugin/__init__.py      |  3 +-
 pythia/tests/plugin_test.py                 |  4 +--
 sample.json                                 |  6 ++++
 5 files changed, 54 insertions(+), 5 deletions(-)
 create mode 100644 pythia/plugins/pest_infestation/__init__.py

diff --git a/pythia/plugin.py b/pythia/plugin.py
index 004bc730..e4eb7274 100644
--- a/pythia/plugin.py
+++ b/pythia/plugin.py
@@ -1,5 +1,6 @@
 import logging
 from enum import Enum, unique
+import copy
 
 
 @unique
@@ -98,14 +99,18 @@ def load_plugins(config, plugins={}, module_prefix="pythia.plugins"):
     return _imported
 
 
+import copy
+
 def run_plugin_functions(hook, plugins, **kwargs):
     _return = {**kwargs}
     if hook in plugins:
         for plugin_fun in plugins[hook]:
-            plugin_fun_return = plugin_fun["fun"](plugin_fun.get("config", {}), _return, **kwargs)
+            plugin_fun_return = plugin_fun["fun"](
+                plugin_fun.get("config", {}), **copy.deepcopy(_return)
+            )
             _return = {
                 **_return,
-                **({} if plugin_fun_return is None else plugin_fun_return)
+                **({} if plugin_fun_return is None else plugin_fun_return),
             }
 
     return _return
diff --git a/pythia/plugins/pest_infestation/__init__.py b/pythia/plugins/pest_infestation/__init__.py
new file mode 100644
index 00000000..9731f4b7
--- /dev/null
+++ b/pythia/plugins/pest_infestation/__init__.py
@@ -0,0 +1,37 @@
+import pythia.plugin
+import numpy as np
+import pandas as pd
+
+def pest_infestation(config, input, **kwargs):
+    """
+    This is a simple pest infestation model that reduces yield based on a severity parameter.
+    """
+    if "err" in input and len(input["err"].decode()) > 0:
+        return None
+    if "params" not in config or "pest_infestation_severity" not in config["params"]:
+        return None
+
+    severity = config["params"]["pest_infestation_severity"]
+    biomass_effect = config["params"].get("pest_infestation_biomass_effect", 0.5)
+    summary_file = f"{input['loc']}/summary.csv"
+    try:
+        summary = pd.read_csv(summary_file)
+        summary["HWAM"] = summary["HWAM"] * (1 - severity)
+        summary["CWAM"] = summary["CWAM"] * (1 - severity * biomass_effect)
+        summary.to_csv(summary_file, index=False)
+    except FileNotFoundError:
+        pass
+    return input
+
+
+def initialize(plugin_config, plugins, config):
+    """
+    This function is called by the plugin manager to initialize the plugin.
+    """
+    plugins = pythia.plugin.register_plugin_function(
+        pythia.plugin.PluginHook.post_run_pixel_success,
+        pest_infestation,
+        plugin_config,
+        plugins,
+    )
+    return plugins
diff --git a/pythia/plugins/test_plugin/__init__.py b/pythia/plugins/test_plugin/__init__.py
index c5c21508..cc2f0ebf 100644
--- a/pythia/plugins/test_plugin/__init__.py
+++ b/pythia/plugins/test_plugin/__init__.py
@@ -19,8 +19,9 @@ def sample_function(config={}, **kwargs):
     return {**kwargs, "config": config, "retval": retval}
 
 
-def contexted_function(context={}, **kwargs):
+def contexted_function(config, **kwargs):
     logging.info("[TEST PLUGIN] Running the contexted_function()")
+    context = kwargs.get("context", {})
     context["context_value"] = context.get("context_value", 2) + 1
     return {**kwargs, "context": context}
 
diff --git a/pythia/tests/plugin_test.py b/pythia/tests/plugin_test.py
index fa197a57..342d3aed 100644
--- a/pythia/tests/plugin_test.py
+++ b/pythia/tests/plugin_test.py
@@ -68,7 +68,7 @@ def test_load_plugin():
     plugins1 = load_plugins(config, plugins)
     assert PluginHook.post_config in plugins1
     assert plugins1[PluginHook.post_config] == [
-        {"fun": pythia.plugins.test_plugin.sample_function, "config": {}}
+        {"fun": pythia.plugins.test_plugin.sample_function, "config": config["plugins"][0]}
     ]
     assert plugins1 != {
         PluginHook.post_config: [{"fun": sample_function, "config": {}}]
@@ -82,7 +82,7 @@ def test_plugin_manual_execution():
     plugins1 = load_plugins(config, plugins)
 
     for plugin in plugins1[PluginHook.post_config]:
-        assert 1 == plugin["fun"]()
+        assert 1 == plugin["fun"](config={})["retval"]
 
 
 def test_plugin_auto_execution():
diff --git a/sample.json b/sample.json
index da7ca445..319205d1 100644
--- a/sample.json
+++ b/sample.json
@@ -15,6 +15,12 @@
                 "end_date": "2018-04-30:",
                 "wsta": "SSUF"
             }
+        },
+        {
+            "plugin": "pest_infestation",
+            "params": {
+                "pest_infestation_severity": 0.2
+            }
         }
     ],
     "default_setup": {