Fix get salt dict type

[ugognw]

Summary

Major changes:

closes #250

• Fixes type annotations in Solution.get_salt_dict
• minor refactoring to resolve mypy errors
• update docstring to represent changes
• changes return value of Solution.get_salt_dict to store Salt objects under "salt" keys
• modifies Solution.get_salt_dict to no longer return an entry for water
• modifies Solution.get_salt to return None if nothing returned from Solution.get_salt_dict
• fixes a few bugs in Solution.get_salt_dict
• fixes solution addition bugs
• fixes bugs in NativeEOS methods
• Introduces comprehensive unit test framework for testing

Todos

• update changelog
• decide on return type
• add unit tests
  • check that Salt object can be instantiated from dictionaries in Salt.from_dict
  • check that values corresponding to the 'mol' key are floats
• reduce test parametrization

Checklist

• Google format doc strings added.
• Code linted with ruff. (For guidance in fixing rule violates, see rule list)
• Type annotations included. Check with mypy.
• Tests added for new features/fixes.
• I have run the tests locally and they passed.

Note that I have changed the docstring in this patch to reflect the current function behaviour. The current patch still maps salt formulas to a dictionary that can be used to construct a representative Salt object. However, the previous docstring states that the return value maps salt formulas to Salt objects. Further, that this is the desired behaviour has been mentioned previously. So, I think this a good place to address this before the function's previous intent is disappeared into the ghost of git log's past. Hopefully, this can be resolved in this PR.

The current behaviour

To the output of Salt.as_dict is added to the 'mol' key which maps to the amount (in moles) of the salt. As a result, the dictionary cannot be directly passed to Salt.from_dict due to the fact that the MSONable methods are autogenerated from the constructor:

>>> from pyEQL import Solution
>>> from pyEQL.salt_ion_match import Salt
>>> s1 = Solution()
>>> salt_dict = s1.get_salt_dict()
>>> Salt.from_dict(salt_dict['HOH'])
...
TypeError: Salt.__init__() got an unexpected keyword argument 'mol'

Instead, to robustly obtain a Salt object, something awkward like this must be done:

>>> Salt.from_dict({k: v for k, v in salt_dict["HOH"] if k != 'mol'})
<pyEQL.salt_ion_match.Salt object at ...>

or:

>>> amount = salt_dict["HOH"].pop('mol')
>>> Salt.from_dict(salt_dict["HOH"])
<pyEQL.salt_ion_match.Salt object at ...>

Although so long as the Salt constructor requires only cation and anion as arguments, one can just do

>>> Salt.from_dict(salt_dict["HOH"]["cation"], salt_dict["HOH"]["anion"])
<pyEQL.salt_ion_match.Salt object at ...>

Proposed solutions for desired behaviour

Modify Salt to include property tracking amount

class Salt(MSONable):
    ...
    def __init__(self, cation, anion, *, mol: float | None = None) -> None:...
        ... 
        self.amount = mol
        ....

No changes are required to Salt.get_salt_dict.

I'm not in love with the choice of the parameter name 'mol' here, but it's chosen to minimize the amount of refactoring required since the amount is stored under the 'mol' key by Solution.get_salt_dict. Alternatively, one could do something like

class Salt(MSONable):
    ...
    def __init__(self, cation, anion, *, amount: float | None = None, **kwargs) -> None:...
        ... 
        self.amount = amount or kwargs.get('mol')
        ...

Note that this doesn't quite achieve the desired behaviour since the dictionary would still map salt formulas to dictionaries, but this is intended so as to minimize breaking things for anything that accesses the data as a dictionary before conversion to a Salt object. Overall, it simplifies the construction of Salt objects from the output of Solution.get_salt_dict(). The dictionaries can immediately be converted to Salt objects with Salt.from_dict:

>>> Salt.from_dict(salt_dict['HOH"])
<pyEQL.salt_ion_match.Salt object at ...>. # no TypeError!

I like this, but this may be muddying the constructor for Salt, and expanding the Salt model with an amount attribute may be outside of the scope of what the class is meant to represent. Another alternative to this is to directly store the new Salt object in the salt dict; however, doing so induces a breaking change.

Store the Salt object under its own key

Otherwise, the Salt object could be nested under its own key:

class Solution(MSONable):
    ...
    def get_salt_dict(self) -> dict[str, dict[str, str] | dict[str, float]]:
        ...
        x = Salt(...)
        ...
        salt_dict.update({x.formula: {'salt': x}}). # remove call to Salt.as_dict()
        ...

So that the output of Salt.get_salt_dict is something like:

{'HOH': {'salt': {'@module': 'pyEQL.salt_ion_match',
        '@class': 'Salt',
        '@version': ...,
        'cation': 'H[+1]',
        'anion': 'OH[-1]'},
        'mol': 1e-07}}

Then, retrieving a given Salt is robust and simple:

>>> salt_dict['HOH']['salt']
<pyEQL.salt_ion_match.Salt object at ...>

The 'mol' key may be accessed as normal. But this is a major breaking change even for the library itself. Nearly everywhere Salt.get_salt_dict is used in pyEQL, the keys anion and cation are accessed prior to conversion into a Salt object, so this change would require refactoring everywhere.

Change the type of the values in the returned dictionary

Alternatively, the salt formulas could map to a tuple (salt, amount):

class Solution(MSONable):
    ...
    def get_salt_dict(self) -> dict[str, tuple[dict[str, str], float]]:
        ...
        salt_dict.update({x.formula: (x, amount)}).
        ...

Retrieving the Salt

>>> salt_dict['HOH'][0]
<pyEQL.salt_ion_match.Salt object at ...>

Again, this is a breaking change. But it has the benefit of keeping the amount and Salt entities programmatically distinct.

If none of these options are appealing, then this patch includes edits to the docstring compatible with the current behaviour.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.58%. Comparing base (0a36ec6) to head (a0cc4c1).
⚠️ Report is 9 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #258      +/-   ##
==========================================
+ Coverage   81.81%   83.58%   +1.76%     
==========================================
  Files           9        9              
  Lines        1496     1456      -40     
  Branches      258      248      -10     
==========================================
- Hits         1224     1217       -7     
+ Misses        226      210      -16     
+ Partials       46       29      -17     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[rkingsbury]

Awesome work @ugognw , thank you for digging into this.

After reviewing again how get_salt_dict() is used, I think the most straightforward way to address is to store the Salt under its own key, as you suggest in approach No. 2. This will only require refactoring in 2 places in engines.py (assuming we deprecated list_salts - see below) and has the benefit that the Salt object does not have to be serialized and re-built when get_activity_coefficient is called. I think all you need to do is change, e.g.,

v["cation"]

to

v["salt"].cation

where v is a Salt object.

Separately, since list_salts has been marked for deprecation for a long time now. please go ahead and remove it as part of this PR.

Finally, please update the CHANGELOG to include your changes here and add yourself to AUTHORS.md! (Before the next release I will also update the changelog with additional recent changes)

[ugognw]

Thanks! Okay, this should be nearly good to go pending your feedback on the changes in NativeEOS.get_activity_coefficient

[rkingsbury]

Looking great @ugognw . I made a comment to answer your HOH question and did some wordsmithing on the docstring.

[rkingsbury]

Summary of Most Recent Changes

In trying to resolve failing tests and writing additional unit tests, I found several apparent bugs and subtle changes due to this patch. The main takeaways are summarized below:

• cutoff is implemented in units of mol/kg solvent

• Concretize new NativeEOS method behaviour w/ tests

  • description of subtle differences in NativeEOS methods due to new get_salt_dict behaviour

    • activity coefficients for H+/OH- in pure water and dilute (<1e-3 m) salt solutions
    • osmotic coefficients for dilute (<1e-3 m) salt solutions
    • solute volumes for dilute (<1e-3 m) salt solutions

  • bug fix: Quantity is now returned by NativeEOS.get_osmotic_coefficient for empty solution

• bug fix: engine, solvent, and database now inherited by sum of Solution

• bug fix: activity coefficient for 1:1 (stoichiometric ratio) salts of divalent or more ions

• note: order of addition matters and affects Solution initialization

• suggestion: consider a cutoff parameter for Solution.get_salt as well as EOS methods

• more unit tests

Outstanding work. I am really glad you found the bug about engine not getting inherited on summation. I remember being very careful to implement that inheritance in from_dict / as_dict and I must have just overlooked in the summation stuff.

Interestingly, it seems like the test failure you're getting with MgCl2 is because get_salt is not correctly identifying the salt

[2025-07-11 13:43:42,149] [   ERROR] --- No salts found that contain solute Mg+2. Returning unit activity coefficient. (engines.py:327)
[2025-07-11 13:43:42,183] [   ERROR] --- No salts found that contain solute Cl-. Returning unit activity coefficient. (engines.py:327)

[rkingsbury]

Also, can you point me to a resource to better understand conftest.py? I see you use it to parameterize a bunch of fixtures for use in other test files. Is there any opportunity to consolidate some of the fixtures in the salt matching tests with the ones defined in conftest.py?

[ugognw]

Also, can you point me to a resource to better understand conftest.py?

👉🏾 conftest.py reference

Is there any opportunity to consolidate some of the fixtures in the salt matching tests with the ones defined in [conftest.py](http://conftest.py/)?

Yes, I think there is. There is also duplication in test_activity.py, test_osmotic_coeff.py, and test_solution.py that I'll address.

[ugognw]

Outstanding work. I am really glad you found the bug about engine not getting inherited on summation. I remember being very careful to implement that inheritance in from_dict / as_dict and I must have just overlooked in the summation stuff.

It's a sneaky one too. I only found it due to failing tests in test_mixing_functions.py–completely unrelated.

Interestingly, it seems like the test failure you're getting with MgCl2 is because get_salt is not correctly identifying the salt

Yes, after calling Solution.equilibrate, the concentrations of Mg[+2] and Cl[-1] decrease slightly due to speciation (into e.g., MgOH[+1], MgCl[+1], ClO[-1], etc.). The concentration of MgCl2 as determined by Solution.get_salt_dict falls just ~0.2 attomoles/kg below the cutoff.

[rkingsbury]

Yes, after calling Solution.equilibrate, the concentrations of Mg[+2] and Cl[-1] decrease slightly due to speciation (into e.g., MgOH[+1], MgCl[+1], ClO[-1], etc.). The concentration of MgCl2 as determined by Solution.get_salt_dict falls just ~0.2 attomoles/kg below the cutoff.

So speciation actually shouldn't matter, because get_salt_dict should be using the TOTAL element concentrations (use_totals=True).

[ugognw]

So speciation actually shouldn't matter, because get_salt_dict should be using the TOTAL element concentrations (use_totals=True).

Ah, yes! I thought this was related to oxidation states, but it appears to be a rounding error (?) introduced by speciation.

>>> s1 = Solution(solutes={'Mg[+2]': '1e-3 mol/kg', 'Cl[-1]': '2e-3 mol/kg'})
>>> before = s1.get_total_amount('Cl(-1.0)', 'mol').m
>>> s1.equilibrate()
>>> after = s1.get_total_amount('Cl(-1.0)', 'mol').m
>>> after - before
np.float64(–4.336808689942018e-19)

This corresponds to the ~0.2 attomoles/kg decrease in MgCl2 concentration. This seems to be a problem with NativeEOS.equilibrate. As a fix, I've introduced a small tolerance (~1e-16) around the cutoff in get_salt_dict. Tests are now passing.

On a side note:

get_salt_dict should be using the TOTAL element concentrations (use_totals=True).

But different oxidation states of the same element are to be treated different, yes? For example, in a solution of 1 M NaCl and 0.5 M NaClO4, should entries for both NaCl and NaClO4 be returned by get_salt_dict? I've just added tests assuming the affirmative.

[rkingsbury]

But different oxidation states of the same element are to be treated different, yes? For example, in a solution of 1 M NaCl and 0.5 M NaClO4, should entries for both NaCl and NaClO4 be returned by get_salt_dict? I've just added tests assuming the affirmative.

Correct! Although get_total_amount can operate both ways, the default behavior is to return totals by oxidation state, which is the appropriate behavior in the context of salt matching. See the docstring.

[rkingsbury]

This looks great @ugognw . I'm ready to merge pending a few small comments - let me know if you have questions. Thank you again for your diligent work on this!

[rkingsbury]

Super minor comment, for future maintainability - can you move the definition of _atol down close to where it is used (c. line 1646) and add a comment explaining what it does? As I understand it, this is a numerical tolerance to account for edge cases where calling equilibrate slightly changes the total amount of an element.

[rkingsbury]

If not too much trouble, can you please move this fix into a separate PR? Since it's a distinct issue that's being resolved, that keeps the history a bit cleaner and easier to track down issues.

[rkingsbury]

With my other comment, this concentration will need to be lowered to 1e-10

[rkingsbury]

Please lower the default cutoff to 1e-9 (1 part per billion). I know that seems extremely low, but the Debye-Huckel model will return an activity coefficient different from 1 (0.99996) even at this low level. I think 1 ppb makes sense as a "practical lower limit" below which it isn't really sensible to talk about an activity coefficient different from one.

If you encounter numerical or performance problems with that low a cutoff, we can use 1e-6 (1ppm ) instead, using the same logic.

>>> I = pyEQL.ureg('1e-9 mol/kg')
>>> pyEQL.activity_correction.get_activity_coefficient_debyehuckel(I)
<Quantity(0.9999628818205379, 'dimensionless')>

[rkingsbury]

Hi @ugognw just checking in. I'm ready to get this merged. I see from the "thumbs ups" that you were planning to take care of my remaining small changes; will you be able to do it soon? If not, let me know and we can possibly jump in.

You should know that your work has already had an impact. Last week our group got an inquiry from an engineer at a major water services company who is using pyEQL and noticed some discrepancies in osmotic coefficients for (NH4)2SO4. Your PR fixed it, and they were able to move on with their analysis! So thank you again for the great work spotting that stoichiometry bug.

Finally, if you're interested, I'd be happy to meet with you on Zoom to learn a little more about the work you're doing with pyEQL. Feel free to email me kingsbury princeton <.> edu

[rkingsbury]

Thank you again for the amazing work here @ugognw ! I'm going to go ahead and merge. There is a new test failure that was introduced by my changes in #201 , which I will address (along with the minor edits I requested before) in a separate PR shortly.

Closes #250

[ugognw]

Hi @ugognw just checking in. I'm ready to get this merged. I see from the "thumbs ups" that you were planning to take care of my remaining small changes; will you be able to do it soon? If not, let me know and we can possibly jump in.

Hi, sorry, for the delay! A few deadlines pulled me away from this in recent weeks. I do plan on finalizing the remaining edits. From what I can recall, the outstanding includes:

• consolidate solute parametrization in test fixtures (contest.py, test_activity.py, test_osmotic_coeff.py, and test_solution.py)
• ensure that concentration tolerances are consistent where changed (1e-10)
• redefine _atol closer to its use in Solution.get_salt_dict

I'll open a separate PR to finalize.

You should know that your work has already had an impact. Last week our group got an inquiry from an engineer at a major water services company who is using pyEQL and noticed some discrepancies in osmotic coefficients for (NH4)2SO4. Your PR fixed it, and they were able to move on with their analysis! So thank you again for the great work spotting that stoichiometry bug.

Oh, that's great to hear!

Finally, if you're interested, I'd be happy to meet with you on Zoom to learn a little more about the work you're doing with pyEQL. Feel free to email me kingsbury princeton <.> edu

Will do!