feat(solver): added callback functions to solver in C++ and python

src/python/solver/bindings.cpp

@@ -1,6 +1,10 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h> // Needed for vectors, maps, sets, strings
-#include <pybind11/stl_bind.h> // Needed for binding std::vector, std::map etc if needed directly
+#include <pybind11/stl_bind.h> // Needed for binding std::vector, std::map etc. if needed directly
+#include <pybind11/numpy.h>
+#include <pybind11/functional.h> // Needed for std::function
+
+#include <boost/numeric/ublas/vector.hpp>
 
 #include "bindings.h"
 
@@ -10,17 +14,60 @@
 namespace py = pybind11;
 
 
-void register_solver_bindings(py::module &m) {
+void register_solver_bindings(const py::module &m) {
     auto py_dynamic_network_solving_strategy = py::class_<gridfire::solver::DynamicNetworkSolverStrategy, PyDynamicNetworkSolverStrategy>(m, "DynamicNetworkSolverStrategy");
     auto py_direct_network_solver = py::class_<gridfire::solver::DirectNetworkSolver, gridfire::solver::DynamicNetworkSolverStrategy>(m, "DirectNetworkSolver");
 
     py_direct_network_solver.def(py::init<gridfire::DynamicEngine&>(),
         py::arg("engine"),
-        "Constructor for the DirectNetworkSolver. Takes a DynamicEngine instance to use for evaluating the network.");
+        "Constructor for the DirectNetworkSolver. Takes a DynamicEngine instance to use for evaluating the network."
+    );
 
     py_direct_network_solver.def("evaluate",
         &gridfire::solver::DirectNetworkSolver::evaluate,
         py::arg("netIn"),
-        "Evaluate the network for a given timestep. Returns the output conditions after the timestep.");
+        "Evaluate the network for a given timestep. Returns the output conditions after the timestep."
+    );
+
+    py_direct_network_solver.def("set_callback",
+        [](gridfire::solver::DirectNetworkSolver &self, gridfire::solver::DirectNetworkSolver::TimestepCallback cb) {
+            self.set_callback(cb);
+        },
+        py::arg("callback"),
+        "Sets a callback function to be called at each timestep."
+    );
+
+    py::class_<gridfire::solver::DirectNetworkSolver::TimestepContext>(py_direct_network_solver, "TimestepContext")
+        .def_readonly("t", &gridfire::solver::DirectNetworkSolver::TimestepContext::t, "Current time in the simulation.")
+        .def_property_readonly(
+            "state", [](const gridfire::solver::DirectNetworkSolver::TimestepContext& ctx) {
+                std::vector<double> state(ctx.state.size());
+                std::ranges::copy(ctx.state, state.begin());
+                return py::array_t<double>(static_cast<ssize_t>(state.size()), state.data());
+            })
+        .def_readonly("dt", &gridfire::solver::DirectNetworkSolver::TimestepContext::dt, "Current timestep size.")
+        .def_readonly("cached_time", &gridfire::solver::DirectNetworkSolver::TimestepContext::cached_time, "Cached time for the last computed result.")
+        .def_readonly("last_observed_time", &gridfire::solver::DirectNetworkSolver::TimestepContext::last_observed_time, "Last time the state was observed.")
+        .def_readonly("last_step_time", &gridfire::solver::DirectNetworkSolver::TimestepContext::last_step_time, "Last step time taken for the integration.")
+        .def_readonly("T9", &gridfire::solver::DirectNetworkSolver::TimestepContext::T9, "Temperature in units of 10^9 K.")
+        .def_readonly("rho", &gridfire::solver::DirectNetworkSolver::TimestepContext::rho, "Temperature in units of 10^9 K.")
+        .def_property_readonly("cached_result", [](const gridfire::solver::DirectNetworkSolver::TimestepContext& ctx) -> py::object {
+            if (ctx.cached_result.has_value()) {
+                const auto&[dydt, nuclearEnergyGenerationRate] = ctx.cached_result.value();
+                return py::make_tuple(
+                    py::array_t<double>(static_cast<ssize_t>(dydt.size()), dydt.data()),
+                    nuclearEnergyGenerationRate
+                );
+            }
+            return py::none();
+        }, "Cached result of the step derivatives.")
+        .def_readonly("num_steps", &gridfire::solver::DirectNetworkSolver::TimestepContext::num_steps, "Total number of steps taken in the simulation.")
+        .def_property_readonly("engine", [](const gridfire::solver::DirectNetworkSolver::TimestepContext &ctx) -> const gridfire::DynamicEngine & {
+            return ctx.engine;
+        }, py::return_value_policy::reference)
+
+        .def_property_readonly("network_species", [](const gridfire::solver::DirectNetworkSolver::TimestepContext &ctx) -> const std::vector<fourdst::atomic::Species> & {
+            return ctx.networkSpecies;
+        }, py::return_value_policy::reference);
 }
 

src/python/solver/bindings.h

@@ -2,4 +2,4 @@
 
 #include <pybind11/pybind11.h>
 
-void register_solver_bindings(pybind11::module &m);
+void register_solver_bindings(const pybind11::module &m);

src/python/solver/trampoline/py_solver.cpp

@@ -5,6 +5,9 @@
 #include <pybind11/functional.h> // Needed for std::function
 
 #include <vector>
+#include <tuple>
+#include <string>
+#include <any>
 
 #include "py_solver.h"
 
@@ -19,3 +22,21 @@ gridfire::NetOut PyDynamicNetworkSolverStrategy::evaluate(const gridfire::NetIn
         netIn  // Arguments
     );
 }
+
+void PyDynamicNetworkSolverStrategy::set_callback(const std::any &callback) {
+    PYBIND11_OVERRIDE_PURE(
+        void,
+        gridfire::solver::DynamicNetworkSolverStrategy,  // Base class
+        set_callback,  // Method name
+        callback  // Arguments
+    );
+}
+
+std::vector<std::tuple<std::string, std::string>> PyDynamicNetworkSolverStrategy::describe_callback_context() const {
+    using DescriptionVector = std::vector<std::tuple<std::string, std::string>>;
+    PYBIND11_OVERRIDE_PURE(
+        DescriptionVector,  // Return type
+        gridfire::solver::DynamicNetworkSolverStrategy,  // Base class
+        describe_callback_context  // Method name
+    );
+}

src/python/solver/trampoline/py_solver.h

@@ -3,8 +3,13 @@
 #include "gridfire/solver/solver.h"
 
 #include <vector>
+#include <tuple>
+#include <string>
+#include <any>
 
 class PyDynamicNetworkSolverStrategy final : public gridfire::solver::DynamicNetworkSolverStrategy {
     explicit PyDynamicNetworkSolverStrategy(gridfire::DynamicEngine &engine) : gridfire::solver::DynamicNetworkSolverStrategy(engine) {}
     gridfire::NetOut evaluate(const gridfire::NetIn &netIn) override;
+    void set_callback(const std::any &callback) override;
+    std::vector<std::tuple<std::string, std::string>> describe_callback_context() const override;
 };