src/pino/unifac_mapping.py

from __future__ import annotations

"""
Explicit group-translation mapping layer between ugropy's Dortmund-UNIFAC
fragmentation output and thermo's bundled DOUFSG parameter table.

This resolves nomenclature collisions (e.g. "CHO" belonging to both an
aldehyde main group and an ether linker main group) and maps ugropy's
idiosyncratic subgroup names to thermo's integer subgroup IDs.
"""

# -----------------------------------------------------------------------------
# Aliases for ugropy subgroup names that differ from thermo's labels.
# These are *name* translations only; they preserve the chemical meaning.
# -----------------------------------------------------------------------------
UGROPY_TO_THERMO_NAME: dict[str, str] = {
    # Aldehyde: ugropy calls it HCO, thermo calls it CHO (Main Group 10, id 20)
    "HCO": "CHO",
    "CH=O": "CHO",
    # Secondary/tertiary OH labels: remove spaces
    "OH (P)": "OH(P)",
    "OH (S)": "OH(S)",
    "OH (T)": "OH(T)",
}

# -----------------------------------------------------------------------------
# Explicit override for cases where the same thermo label maps to multiple
# subgroup IDs.  This pins the chemically correct ID.
# -----------------------------------------------------------------------------
# thermo.unifac.DOUFSG:
#   id 20 -> CHO, main_group 10 (aldehyde)
#   id 26 -> CHO, main_group 13 (ether CH-O)
# We want aldehydes to go to id 20, not 26.
EXPLICIT_THERMO_ID: dict[str, int] = {
    "HCO": 20,
    "CH=O": 20,
    "CHO__ALDEHYDE": 20,  # disambiguation handle
    "CHO__ETHER": 26,     # if we ever need the ether linker explicitly
}


def map_ugropy_to_thermo_id(
    name: str,
    name_to_id: dict[str, int],
) -> int | None:
    """
    Translate a ugropy subgroup name to the thermo DDB UNIFAC subgroup ID.

    Parameters
    ----------
    name : str
        ugropy subgroup name (e.g. "HCO", "OH (P)", "CH3").
    name_to_id : dict[str, int]
        Mapping from thermo subgroup label to thermo subgroup ID.

    Returns
    -------
    int | None
        Thermo subgroup ID, or None if the group is not representable.
    """
    # Hardcoded disambiguation overrides first.
    if name in EXPLICIT_THERMO_ID:
        return EXPLICIT_THERMO_ID[name]

    # Name-only aliases (e.g. "OH (P)" -> "OH(P)")
    normalised = UGROPY_TO_THERMO_NAME.get(name, name).replace(" ", "")
    if normalised in EXPLICIT_THERMO_ID:
        return EXPLICIT_THERMO_ID[normalised]

    return name_to_id.get(normalised)