added function to view rkt state

src/reaktoro_pse/core/reaktoro_block_builder.py

@@ -55,12 +55,18 @@ def __init__(self, block, reaktoro_solver, build_on_init=True):
             raise TypeError("Reaktoro block builder requires a ReaktoroSolver class")
         self.configure_jacobian_scaling()
         self.reaktoro_initialize_function = None  # used to provide external solve call
+        self.display_reaktoro_state_function = (
+            None  # used to specifying external function to display rkt state
+        )
         self.build_output_vars()
         if build_on_init:  # option to support legacy implementation
             self.build_reaktoro_block()
 
     def build_reaktoro_block(
-        self, gray_box_model=None, reaktoro_initialize_function=None
+        self,
+        gray_box_model=None,
+        reaktoro_initialize_function=None,
+        display_reaktoro_state_function=None,
     ):
         """build reaktoro model"""
         if gray_box_model is None:
@@ -73,6 +79,8 @@ def build_reaktoro_block(
             self.block.reaktoro_model = gray_box_model
         if reaktoro_initialize_function is not None:
             self.reaktoro_initialize_function = reaktoro_initialize_function
+        if display_reaktoro_state_function is not None:
+            self.display_reaktoro_state_function = display_reaktoro_state_function
         self.build_input_constraints()
         self.build_output_constraints()
 
@@ -350,3 +358,9 @@ def initialize_input_variables_and_constraints(self):
                 sf = self.get_sf(self.block.reaktoro_model.inputs[key])
             iscale.set_scaling_factor(self.block.reaktoro_model.inputs[key], sf)
             iscale.constraint_scaling_transform(self.block.input_constraints[key], sf)
+
+    def display_state(self):
+        if self.display_reaktoro_state_function is None:
+            print(self.solver.state)
+        else:
+            self.display_reaktoro_state_function()

src/reaktoro_pse/parallel_tools/parallel_manager.py

@@ -122,6 +122,9 @@ def get_params(self):
         for i, key in enumerate(self.inputs.rkt_inputs.keys()):
             self.params[key] = self.input_matrix[2][i]
 
+    def display_state(self):
+        print(self.state.state)
+
     def check_solve(self):
         if self.old_matrix is None:
             self.old_matrix = self.input_matrix[2].copy()
@@ -160,6 +163,7 @@ class WorkerMessages:
     CyIpoptEvaluationError = "CyIpoptEvaluationError"
     terminate = "terminate"
     failed = "failed"
+    display_state = "display_state"
 
 
 class LocalWorker:
@@ -260,6 +264,9 @@ def get_solution(self):
                     "The worker failed and did not return a solution terminated"
                 )
 
+    def display_state(self):
+        self.local_pipe.send(WorkerMessages.display_state)
+
     def terminate(self):
         self.local_pipe.send(WorkerMessages.terminate)
         _log.info("Worker terminated")
@@ -283,6 +290,9 @@ def get_solve_and_get_function(self, block_idx):
     def get_initialize_function(self, block_idx):
         return self.registered_workers[block_idx].initialize
 
+    def get_display_function(self, block_idx):
+        return self.registered_workers[block_idx].display_state
+
     def start_workers(self):
         for idx, local_worker in self.registered_workers.items():
             process = mp.Process(
@@ -330,6 +340,9 @@ def ReaktoroActor(
                 result = WorkerMessages.success
             if command == WorkerMessages.solve:
                 result = reaktoro_worker.solve()
+            if command == WorkerMessages.display_state:
+                result = reaktoro_worker.display_state()
+                result = WorkerMessages.success
             if command == WorkerMessages.terminate:
                 reaktoro_worker.close_shared_memory()
                 return

src/reaktoro_pse/parallel_tools/reaktoro_block_manager.py

@@ -316,11 +316,14 @@ def build_reaktoro_blocks(self):
             )
             if self.config.use_parallel_mode:
                 init_func = self.parallel_manager.get_initialize_function(block_idx)
+                disp_func = self.parallel_manager.get_display_function(block_idx)
             else:
                 init_func = self.aggregate_solver_state.solver_functions[block_idx]
+                disp_func = None
             block.builder.build_reaktoro_block(
                 gray_box_model=pseudo_gray_box_model,
                 reaktoro_initialize_function=init_func,
+                display_reaktoro_state_function=disp_func,
             )
             block.pseudo_gray_box = pseudo_gray_box_model
 

src/reaktoro_pse/parallel_tools/tests/test_manager.py

@@ -42,6 +42,11 @@ def test_blockBuild_with_speciation_block(build_rkt_state_with_species):
             "pressure": m.pressure,
             "pH": m.pH,
         },
+        jacobian_options={
+            "numerical_type": "average",
+            "numerical_order": 2,
+            "numerical_step": 1e-8,
+        },
         database="PhreeqcDatabase",
         database_file="pitzer.dat",
         chemistry_modifier=m.CaO,
@@ -51,6 +56,7 @@ def test_blockBuild_with_speciation_block(build_rkt_state_with_species):
     )
     m.reaktoro_manager.build_reaktoro_blocks()
     m.property_block.initialize()
+
     cy_solver = get_solver(solver="cyipopt-watertap")
     cy_solver.options["max_iter"] = 20
     m.pH.unfix()
@@ -115,6 +121,7 @@ def test_blockBuild_with_speciation_block(build_rkt_state_with_species):
     assert "property_block" in scaling_result
     for key in scaling_result["property_block"]:
         assert scaling_result["property_block"][key] == 1
+    m.property_block.display_reaktoro_state()
     m.reaktoro_manager.terminate_workers()
 
 

src/reaktoro_pse/reaktoro_block.py

@@ -759,9 +759,9 @@ def display_reaktoro_state(self):
         """Displays reaktoro state"""
         if self.config.build_speciation_block:
             _log.info("-----Displaying information for speciation block ------")
-            _log.info(self.speciation_block.rkt_state.state)
+            self.speciation_block.rkt_block_builder.display_state()
         _log.info("-----Displaying information for property block ------")
-        _log.info(self.rkt_state.state)
+        self.rkt_block_builder.display_state()
 
     def update_jacobian_scaling(
         self, user_scaling_dict=None, set_on_speciation_block=True