Merge pull request #12 from avdudchenko/ph_charge_balance_fix

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

src/reaktoro_pse/core/reaktoro_block_builder.py

@@ -255,7 +255,7 @@ def calc_scale(value):
 
     def initialize_output_variables_and_constraints(self):
         for key, obj in self.solver.output_specs.user_outputs.items():
-            """update vars scaling in pyomo build ocnstraints
+            """update vars scaling in pyomo build constraints
             these are updated to actual value when we call solve_rektoro_block"""
             if PropTypes.pyomo_built_prop == obj.property_type:
                 for (
@@ -295,10 +295,10 @@ def initialize_output_variables_and_constraints(self):
                 )
 
         # update jacobian scaling
-        self.get_jacobian_scaling()
+        self.set_jacobian_scaling()
         self.set_user_jacobian_scaling()
 
-    def get_jacobian_scaling(self):
+    def set_jacobian_scaling(self):
         if self.jacobian_scaling_type == JacScalingTypes.no_scaling:
             for i, (key, obj) in enumerate(
                 self.solver.output_specs.rkt_outputs.items()
@@ -315,6 +315,9 @@ def get_jacobian_scaling(self):
                 np.sum(np.abs(self.solver.jacobian_matrix) ** 2, axis=1) ** 0.5
             )
 
+    def get_jacobian_scaling(self):
+        return self.solver.jacobian_scaling_values
+
     def set_user_jacobian_scaling(self, user_scaling=None):
 
         if user_scaling is None:

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
@@ -148,7 +149,7 @@ def configure_specs(
         """configures specification for the problem
 
         Keyword arguments:
-        dissolveSpeciesInRkt -- If true, species would be summed up to element amount in rkt, if false
+        dissolve_species_in_rkt-- If true, species would be summed up to element amount in rkt, if false
         mode will contain conditions to build pyomo constraints via raktorooutput class
         exact_speciation -- if True, will write exact element amount for all input species other wise
         will leave  H, and O open, while fixing aqueousSolvent to specified value (e.g. H2O)
@@ -201,30 +202,38 @@ def breakdown_species_to_elements(self):
                 ] = specie.elements().coefficients()[i]
         self.chemical_to_elements.update(self.specie_to_elements)
 
-    def add_specs(self, specs_object, assert_charge_neutrality, dissolveSpeciesInRkt):
+    def add_specs(
+        self, specs_object, assert_charge_neutrality, dissolve_species_in_rkt
+    ):
         # ignore elements for constraints
 
         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,15 +245,18 @@ 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
-        if dissolveSpeciesInRkt:
+        if dissolve_species_in_rkt:
             self.write_active_species(specs_object)
         else:
             for element in self.constraint_dict:
@@ -255,7 +267,7 @@ def add_specs(self, specs_object, assert_charge_neutrality, dissolveSpeciesInRkt
 
         # write reaktoro constraints to spec
         for element in self.constraint_dict:
-            if dissolveSpeciesInRkt:
+            if dissolve_species_in_rkt:
                 self.write_element_sum_constraint(specs_object, element)
             else:
                 self.write_elementAmount_constraint(specs_object, element)
@@ -287,7 +299,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[
@@ -354,6 +365,7 @@ def write_active_species(self, spec_object):
             input_name = f"input{specie}"
             idx = spec_object.addInput(input_name)
             if specie in self.state.inputs:
+
                 self.rkt_inputs[specie] = self.state.inputs[specie]
                 self.rkt_inputs[specie].set_rkt_index(idx)
                 self.rkt_inputs[specie].set_rkt_input_name(input_name)
@@ -371,14 +383,35 @@ def default_speciation(self):
         """defines species to element conversions"""
         self.chemical_to_elements = {
             "HCl": {"H": 1, "Cl": 1},
+            "H2SO4": {"H": 2, "S": 1, "O": 4},
             "CaO": {"Ca": 1, "O": 1},
+            "Ca(OH)2": {"Ca": 1, "O": 2, "H": 2},
             "Na2CO3": {"Na": 2, "C": 1, "O": 3},
             "CO2": {"C": 1, "O": 2},
             "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 +423,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 +492,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)
 
@@ -480,6 +513,7 @@ def export_config(self):
         export_object.neutrality_ion = self.neutrality_ion
         export_object.dissolve_species_in_rkt = self.dissolve_species_in_rkt
         export_object.exact_speciation = self.exact_speciation
+        export_object.chemical_to_elements = self.chemical_to_elements
         return export_object
 
     def load_from_export_object(self, export_object):
@@ -493,3 +527,4 @@ def load_from_export_object(self, export_object):
         self.neutrality_ion = export_object.neutrality_ion
         self.dissolve_species_in_rkt = export_object.dissolve_species_in_rkt
         self.exact_speciation = export_object.exact_speciation
+        self.chemical_to_elements = export_object.chemical_to_elements

src/reaktoro_pse/core/reaktoro_state.py

@@ -290,7 +290,7 @@ def copy_rkt_inputs(self, inputs):
             self.inputs[key].input_type = obj.input_type
             self.inputs[key].value = obj.value
             self.inputs[key].converted_value = obj.converted_value
-            self.inputs.registered_phases = inputs.registered_phases
+        self.inputs.registered_phases = inputs.registered_phases
         self.inputs.all_species = inputs.all_species
         self.inputs.species_list = inputs.species_list
         self.inputs.convert_to_rkt_species = inputs.convert_to_rkt_species
@@ -788,7 +788,6 @@ def set_rkt_state(self):
                                     ),
                                     self.inputs[species].main_unit,
                                 )
-        # _jac_phases = [phase.name() for phase in self.state.system().phases()]
 
     def equilibrate_state(self):
         self.set_rkt_state()

src/reaktoro_pse/core/tests/test_reaktoro_inputs.py

@@ -19,7 +19,7 @@
 
 import pickle
 
-__author__ = "Alexander V. Dudchenko (SLAC)"
+__author__ = "Alexander V. Dudchenko (NETL)"
 
 
 def test_with_rkt_sum(build_rkt_state_with_species):
@@ -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/core/tests/test_reaktoro_jacobian.py

@@ -22,7 +22,7 @@
 )
 import numpy as np
 
-__author__ = "Alexander V. Dudchenko (SLAC)"
+__author__ = "Alexander V. Dudchenko (NETL)"
 
 
 @pytest.fixture

src/reaktoro_pse/core/tests/test_reaktoro_outputs.py

@@ -20,7 +20,7 @@
 )
 import pickle
 
-__author__ = "Alexander V. Dudchenko (SLAC)"
+__author__ = "Alexander V. Dudchenko (NETL)"
 
 
 @pytest.fixture

src/reaktoro_pse/core/tests/test_reaktoro_solver.py

@@ -35,7 +35,7 @@
 )
 import pickle
 
-__author__ = "Alexander V. Dudchenko (SLAC)"
+__author__ = "Alexander V. Dudchenko (NETL)"
 
 
 @pytest.fixture

src/reaktoro_pse/core/tests/test_reaktoro_state.py

@@ -16,7 +16,7 @@
 from pyomo.environ import ConcreteModel, Var, units as pyunits
 import pickle
 
-__author__ = "Alexander V. Dudchenko (SLAC)"
+__author__ = "Alexander V. Dudchenko (NETL)"
 
 
 @pytest.fixture

src/reaktoro_pse/core/util_classes/rkt_inputs.py

@@ -282,17 +282,20 @@ def specie_to_rkt_species(species):
     # TODO: needs to be better automated
     name_dict = {
         "-2": ["SO4", "CO3"],
-        "-": ["Cl", "HCO3", "F"],
+        "-": ["Cl", "HCO3", "F", "NO3"],
         "+": ["Na", "K"],
         "+2": ["Mg", "Mn", "Ca", "Sr", "Ba"],
         "": ["H2O", "CO2"],
         "H4SiO4": ["Si", "SiO2"],
+        "SeO4-2": ["Se"],
     }
     for charge, species_list in name_dict.items():
         for spc in species_list:
             if spc in species:
                 if charge == "H4SiO4":
                     return charge
+                if charge == "SeO4-2":
+                    return charge
                 else:
                     return f"{spc}{charge}"
 

src/reaktoro_pse/core/util_classes/tests/test_rkt_inputs.py

@@ -13,7 +13,7 @@
 import reaktoro_pse.core.util_classes.rkt_inputs as RktInputs
 from pyomo.environ import Var, units as pyunits
 
-__author__ = "Alexander V. Dudchenko (SLAC)"
+__author__ = "Alexander V. Dudchenko (NETL)"
 
 
 @pytest.fixture

src/reaktoro_pse/examples/reaktoro_pse_to_phreeqc_comparison/standard_model.py

@@ -60,9 +60,9 @@ def build_modification_example(water_comp):
         units=pyunits.dimensionless,
     )
     m.water_recovery.fix()
-    m.eq_water_fow = Constraint(
+    m.eq_water_flow = 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,14 +116,14 @@ 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)
 
 
 def initialize(m):
     calculate_variable_from_constraint(
-        m.modified_properties_water_removal, m.eq_water_fow
+        m.modified_properties_water_removal, m.eq_water_flow
     )
     m.eq_modified_properties.initialize()
     solve(m)

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(