Adds pH to charge balance options, and makes it default on blocks with exact speciation.

src/reaktoro_pse/core/reaktoro_inputs.py

@@ -94,9 +94,10 @@ def register_chemistry_modifiers(self, chemical_dict, index=None):
                 self.register_chemistry_modifier(chemical, obj)
 
     def register_chemistry_modifier(self, chemical, pyomo_var):
+        chemical = self.safe_modifier_name(chemical)
         if chemical not in self.chemical_to_elements:
             raise ValueError(
-                f"{chemical} is not avaialbe in chemical_to_element dict, please add"
+                f"{chemical} is not available in chemical_to_element dict, please add"
             )
         self.rkt_chemical_inputs[chemical] = RktInput(
             var_name=chemical, pyomo_var=pyomo_var
@@ -207,24 +208,30 @@ def add_specs(self, specs_object, assert_charge_neutrality, dissolveSpeciesInRkt
         pressure_not_set = True
         temperature_not_set = True
         for input_name, _ in self.state.inputs.items():
-            if input_name == "temperature":
+            if input_name == RktInputTypes.temperature:
                 specs_object.temperature()
                 temperature_not_set = False
-                self.rkt_inputs["temperature"] = self.state.inputs["temperature"]
-                self.rkt_inputs["temperature"].set_lower_bound(0)
-            elif input_name == "pressure":
+                self.rkt_inputs[RktInputTypes.temperature] = self.state.inputs[
+                    RktInputTypes.temperature
+                ]
+                self.rkt_inputs[RktInputTypes.temperature].set_lower_bound(0)
+            elif input_name == RktInputTypes.pressure:
                 specs_object.pressure()
                 pressure_not_set = False
-                self.rkt_inputs["pressure"] = self.state.inputs["pressure"]
-                self.rkt_inputs["pressure"].set_lower_bound(0)
-            elif input_name == "enthalpy":
+                self.rkt_inputs[RktInputTypes.pressure] = self.state.inputs[
+                    RktInputTypes.pressure
+                ]
+                self.rkt_inputs[RktInputTypes.pressure].set_lower_bound(0)
+            elif input_name == RktInputTypes.enthalpy:
                 specs_object.enthalpy()
-                self.rkt_inputs["enthalpy"] = self.state.inputs["enthalpy"]
-                self.rkt_inputs["enthalpy"].set_lower_bound(None)
-            elif input_name == "pH":
+                self.rkt_inputs[RktInputTypes.enthalpy] = self.state.inputs[
+                    RktInputTypes.enthalpy
+                ]
+                self.rkt_inputs[RktInputTypes.enthalpy].set_lower_bound(None)
+            elif input_name == RktInputTypes.pH:
                 specs_object.pH()
-                self.rkt_inputs["pH"] = self.state.inputs["pH"]
-                self.rkt_inputs["pH"].set_lower_bound(0)
+                self.rkt_inputs[RktInputTypes.pH] = self.state.inputs[RktInputTypes.pH]
+                self.rkt_inputs[RktInputTypes.pH].set_lower_bound(0)
             else:
                 pass
         if pressure_not_set:
@@ -236,11 +243,14 @@ def add_specs(self, specs_object, assert_charge_neutrality, dissolveSpeciesInRkt
         if assert_charge_neutrality:
             specs_object.charge()
             if self.neutrality_ion is not None:
-                self.ignore_elements_for_constraints.append(self.neutrality_ion)
+                if self.neutrality_ion == RktInputTypes.pH:
+                    specs_object.openTo("H+")
+                else:
+                    self.ignore_elements_for_constraints.append(self.neutrality_ion)
 
-                if self.neutrality_ion not in specs_object.namesInputs():
-                    # needs to be a species!
-                    specs_object.openTo(self.neutrality_ion)
+                    if self.neutrality_ion not in specs_object.namesInputs():
+                        # needs to be a species!
+                        specs_object.openTo(self.neutrality_ion)
 
         self._find_element_sums()
         # add/check if vars in rkt Inputs
@@ -287,7 +297,6 @@ def _find_element_sums(self):
         self.active_species = []
         rktState = self.state.state
         if self.exact_speciation == False:
-            # self.rktActiveSpecies.append(self.aqueousSolvent)
             for phase in self.state.inputs.registered_phases:
                 if phase in self.fixed_solvent_specie:
                     specie_elements = self.specie_to_elements[
@@ -377,8 +386,27 @@ def default_speciation(self):
             "NaOH": {"Na": 1, "O": 1, "H": 1},
             "H": {"H": 1},
             "OH": {"O": 1, "H": 1},
-            "H2O_evaporation": {"O": 1, "H": 2},
+            "H2O_evaporation": {"O": -1, "H": -2},
         }
+        self.ensure_safe_modifier_names()
+
+    def ensure_safe_modifier_names(self):
+        """we need to ensure any chemical has a safe subname added
+        so it does not override exact species (e.g. HCl can exist in database)"""
+        for key in list(self.chemical_to_elements.keys()):
+            self.chemical_to_elements[self.safe_modifier_name(key)] = (
+                self.chemical_to_elements.pop(key)
+            )
+
+    def safe_modifier_name(self, name):
+        """ensures we use safe modifiers that do not replicate
+        real species, e.g. exact species might contain HCl, but we might want
+        to also add HCl to system, internally we want to track it as
+        modifier_HCl it will be bound to provided pyomo var regardless"""
+        if "modifier_" not in name:
+            return f"modifier_{name}"
+        else:
+            return name
 
     def get_modifier_mw(self, elemental_composition):
         mw = 0
@@ -390,6 +418,7 @@ def get_modifier_mw(self, elemental_composition):
     def register_modifier(self, new_chemical):
         if new_chemical is not None:
             self.chemical_to_elements.update(new_chemical)
+        self.ensure_safe_modifier_names()
 
     def write_element_sum_constraint(self, spec_object, element):
         """writes a sum of elements constraint for reaktoro"""
@@ -458,7 +487,6 @@ def write_phase_volume_constraint(self, spec_object, phase, fixed_vlaue=1e-8):
         idx = spec_object.addInput(f"volume_{phase}")
         constraint = rkt.EquationConstraint()
         constraint.id = f"{phase}_volume_constraint"
-
         constraint.fn = lambda props, w: w[idx] - props.phaseProps(phase).volume()
         spec_object.addConstraint(constraint)
 

src/reaktoro_pse/core/tests/test_reaktoro_inputs.py

@@ -301,10 +301,10 @@ def test_with_chemical(build_rkt_state_with_species):
     expected_con_dict = {
         "C": [(1.0, "CO3-2"), (1.0, "CO2")],
         "Na": [(1, "Na+")],
-        "Ca": [(1, "Ca+2"), (1, "CaO")],
+        "Ca": [(1, "Ca+2"), (1, "modifier_CaO")],
         "Mg": [(1, "Mg+2")],
     }
-    expected_active_species = ["CaO", "CO3-2", "CO2", "Na+", "Mg+2", "Ca+2"]
+    expected_active_species = ["modifier_CaO", "CO3-2", "CO2", "Na+", "Mg+2", "Ca+2"]
     print(rkt_input.constraint_dict)
     for ion, ecd in expected_con_dict.items():
         rkt_ecd = rkt_input.constraint_dict[ion]
@@ -331,7 +331,7 @@ def test_input_with_pickle_copy(build_rkt_state_with_species):
     lime = Var(initialize=0.01, units=pyunits.mol / pyunits.s)
     lime.construct()
 
-    rkt_input.register_chemistry_modifier("CaO", lime)
+    rkt_input.register_chemistry_modifier("modifier_CaO", lime)
     rkt_input.configure_specs(dissolve_species_in_rkt=True)
     rkt_input.build_input_specs()
     export_object = rkt_input.export_config()
@@ -344,10 +344,10 @@ def test_input_with_pickle_copy(build_rkt_state_with_species):
     expected_con_dict = {
         "C": [(1.0, "CO3-2"), (1.0, "CO2")],
         "Na": [(1, "Na+")],
-        "Ca": [(1, "Ca+2"), (1, "CaO")],
+        "Ca": [(1, "Ca+2"), (1, "modifier_CaO")],
         "Mg": [(1, "Mg+2")],
     }
-    expected_active_species = ["CaO", "CO3-2", "CO2", "Na+", "Mg+2", "Ca+2"]
+    expected_active_species = ["modifier_CaO", "CO3-2", "CO2", "Na+", "Mg+2", "Ca+2"]
     print(new_rkt_input.constraint_dict)
     for ion, ecd in expected_con_dict.items():
         rkt_ecd = new_rkt_input.constraint_dict[ion]

src/reaktoro_pse/examples/reaktoro_pse_to_phreeqc_comparison/standard_model.py

@@ -62,7 +62,7 @@ def build_modification_example(water_comp):
     m.water_recovery.fix()
     m.eq_water_fow = Constraint(
         expr=m.water_recovery
-        == -m.modified_properties_water_removal / m.feed_composition["H2O"]
+        == m.modified_properties_water_removal / m.feed_composition["H2O"]
     )
     return m
 
@@ -94,11 +94,11 @@ def add_standard_properties(m):
             "H2O_evaporation": m.modified_properties_water_removal,
             "NaOH": m.base_addition,
         },
-        dissolve_species_in_reaktoro=False,
+        dissolve_species_in_reaktoro=True,
         # we can use default converter as its defined for default database (Phreeqc and pitzer)
         # we are modifying state and must speciate inputs before adding acid to find final prop state.
         build_speciation_block=True,
-        reaktoro_solve_options={"open_species_on_property_block": ["H+", "OH-"]},
+        # reaktoro_solve_options={"open_species_on_property_block": ["OH-", "H2O"]},
         jacobian_options={
             "user_scaling": {
                 ("saturationIndex", "Calcite"): 1,
@@ -116,7 +116,7 @@ def scale_model(m):
         iscale.set_scaling_factor(
             m.feed_composition[key], 1 / m.feed_composition[key].value
         )
-    iscale.set_scaling_factor(m.water_recovery, 1)
+    iscale.set_scaling_factor(m.water_recovery, 1 / 1)
     iscale.set_scaling_factor(m.acid_addition, 1 / 0.001)
     iscale.set_scaling_factor(m.base_addition, 1 / 0.001)
 

src/reaktoro_pse/examples/simple_desalination.py

@@ -140,7 +140,7 @@ def eq_desal_composition(fs, key):
 
         # However, this can result in incorrect speciation for some databases, so please use with caution.
 
-        species_to_open = ["H+", "OH-"]
+        species_to_open = ["OH-"]
     else:
         species_to_open = None
     m.eq_desal_properties = ReaktoroBlock(

src/reaktoro_pse/parallel_tools/reaktoro_block_manager.py

@@ -225,7 +225,7 @@ class ReaktoroBlockManagerData(ProcessBlockData):
     CONFIG.declare(
         "worker_timeout",
         ConfigValue(
-            default=20,
+            default=60,
             domain=int,
             description="Defines time in seconds for worker time out",
             doc="""This is time out for parallel workers to time out and shut down if they receive no 

src/reaktoro_pse/reaktoro_block.py

@@ -177,20 +177,21 @@ class ReaktoroBlockData(ProcessBlockData):
         ConfigValue(
             default="Cl",
             domain=str,
-            description="Ion to use for maintaining charge neutrality",
-            doc="""This will unfix specified ion during equilibrium calculations while enforcing charge==0 constraint in reaktoro""",
+            description="Ion to use for maintaining charge neutrality (pH, ion, or element)",
+            doc="""This will unfix specified ion during equilibrium calculations while enforcing charge==0 constraint
+              in reaktoro, if exact speciation is provided, pH will be used by default""",
         ),
     )
     CONFIG.declare(
         "assert_charge_neutrality_on_all_blocks",
         ConfigValue(
-            default=False,
+            default=True,
             domain=bool,
-            description="Defines if charge neutrality should be applied to both speciation and property block",
+            description="Defines if charge neutrality should be applied to property_block when build_speciation_block=True",
             doc="""
             When user provides chemical inputs, its assumed that user wants to modify an equilibrated state, 
-            as such a speciation and property block will be built. Charge neutrality would only be applied on speciation block if enabled
-            if this option is set to True then charge neutrality will also be applied on property block """,
+            as such a speciation and property block will be built. The speciation block will likely require charge balancing on ion basis (e.g Cl),
+            while property block that received exact speciation will balance on pH, to disable charge balance on property block set this option to False""",
         ),
     )
 
@@ -567,26 +568,41 @@ def build_rkt_inputs(
         block.rkt_inputs.register_free_elements(self.config.aqueous_phase.free_element)
         block.rkt_inputs.register_free_elements(self.config.liquid_phase.free_element)
         # register charge neutrality
-        if (
-            speciation_block_built == False
-            or self.config.assert_charge_neutrality_on_all_blocks
-        ):
-            # only ensure charge neutrality when doing first calculation
+        # only ensure charge neutrality when doing first calculation
+        if speciation_block_built == False:
+            # if exact speciation - then charge balance should be done on pH
+            assert_charge_neutrality = self.config.assert_charge_neutrality
+            if self.config.exact_speciation == True:
+                # only do so if we have 'H+' in species
+                ion_for_balancing = "pH"
+                if "H+" not in block.rkt_state.database_species:
+                    assert_charge_neutrality = False
+
+            else:
+                assert_charge_neutrality = self.config.assert_charge_neutrality
+                ion_for_balancing = self.config.charge_neutrality_ion
+
             block.rkt_inputs.register_charge_neutrality(
-                assert_neutrality=self.config.assert_charge_neutrality,
-                ion=self.config.charge_neutrality_ion,
+                assert_neutrality=assert_charge_neutrality,
+                ion=ion_for_balancing,
             )
             block.rkt_inputs.configure_specs(
                 dissolve_species_in_rkt=self.config.dissolve_species_in_reaktoro,
                 exact_speciation=self.config.exact_speciation,
             )
         else:
             # if we have built a speciation block, the feed should be charge neutral and
-            # exact speciation is provided
+            # exact speciation is provided, then we balance on pH only,
+            # user can disable this by setting self.assert_charge_neutrality_on_all_blocks to False
+            assert_charge_neutrality = False
+            if self.config.assert_charge_neutrality_on_all_blocks:
+                # only do so if we have 'H+' in species
+                if "H+" in block.rkt_state.database_species:
+                    assert_charge_neutrality = self.config.assert_charge_neutrality
             block.rkt_inputs.register_charge_neutrality(
-                assert_neutrality=False, ion=self.config.charge_neutrality_ion
+                assert_neutrality=assert_charge_neutrality,
+                ion="pH",
             )
-
             block.rkt_inputs.configure_specs(
                 dissolve_species_in_rkt=self.config.dissolve_species_in_reaktoro,
                 exact_speciation=True,