tests/test_pino.py

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

from pino import evaporation, vle, verifier
from pino.models import Formula, Ingredient, IngredientResolutionError
from pino.structures import IngredientResolver, ingredient_from_smiles
from pino.verifier import FormulaVerifier, verify_formula_dict


# Reference data for the classic 5-component accord in an ethanol base.
# Vapor pressures are approximate room-temperature values (Pa) for 25 °C.
# Odor thresholds are rough literature estimates (µg/m³).
_TEST_FORMULA_DATA = {
    "temperature_k": 298.15,
    "ambient_pressure_pa": 101325.0,
    "headspace_volume_m3": 1e-3,
    "liquid_volume_m3": 1e-6,
    "metadata": {"density_g_ml": 0.9},
    "ingredients": [
        {
            "name": "Ethanol",
            "cas": "64-17-5",
            "smiles": "CCO",
            "molecular_weight": 46.07,
            "vapor_pressure_pa": 7900.0,
            "boiling_point_k": 351.45,
            "odor_threshold_ug_m3": 250000.0,
        },
        {
            "name": "Limonene",
            "cas": "138-86-3",
            "smiles": "CC1=CCC(C=C1)C(C)C",
            "molecular_weight": 136.23,
            "vapor_pressure_pa": 195.0,
            "boiling_point_k": 449.0,
            "odor_threshold_ug_m3": 30.0,
        },
        {
            "name": "Linalool",
            "cas": "78-70-6",
            "smiles": "CC=C(C)C(CC=C(C)C)O",
            "molecular_weight": 154.25,
            "vapor_pressure_pa": 22.0,
            "boiling_point_k": 471.0,
            "odor_threshold_ug_m3": 6.0,
        },
        {
            "name": "Hedione",
            "cas": "24851-98-7",
            "smiles": "CC/C(=C/C(=O)C(C)C)C(C)C(=O)O",
            "molecular_weight": 226.32,
            "vapor_pressure_pa": 1.2,
            "boiling_point_k": 548.0,
            "odor_threshold_ug_m3": 1.0,
        },
        {
            "name": "Galaxolide",
            "cas": "1222-05-5",
            "smiles": "CC(=O)C1=CC=C(C)C=C1",
            "molecular_weight": 258.40,
            "vapor_pressure_pa": 0.05,
            "boiling_point_k": 623.0,
            "odor_threshold_ug_m3": 0.1,
        },
        {
            "name": "Ambroxan",
            "cas": "6790-58-5",
            "smiles": "CC1CCCC2(C)C1C(C)CC(C)O2",
            "molecular_weight": 236.39,
            "vapor_pressure_pa": 0.01,
            "boiling_point_k": 573.0,
            "odor_threshold_ug_m3": 0.02,
        },
    ],
    "weight_fractions": [0.70, 0.10, 0.08, 0.06, 0.04, 0.02],
}


@pytest.fixture
def test_formula() -> Formula:
    return verify_formula_dict(_TEST_FORMULA_DATA).formula


def test_ingredient_from_smiles() -> None:
    ing = ingredient_from_smiles(
        name="Limonene",
        cas="138-86-3",
        smiles="CC1=CCC(C=C1)C(C)C",
        vapor_pressure_pa=195.0,
    )
    assert ing.name == "Limonene"
    assert ing.molecular_weight > 0
    assert ing.unifac_groups


def test_formula_validation() -> None:
    ing = Ingredient(
        name="Ethanol",
        cas="64-17-5",
        smiles="CCO",
        molecular_weight=46.07,
        vapor_pressure_pa=7900.0,
    )
    # Sum must be 1.0
    with pytest.raises(ValueError):
        Formula(
            ingredients=(ing,),
            weight_fractions=(0.5,),
            temperature_k=298.15,
            ambient_pressure_pa=101325.0,
        )


def test_vle_non_ideal(test_formula: Formula) -> None:
    calc = vle.VLECalculator(test_formula)
    x = test_formula.mole_fractions
    gamma = calc.activity_coefficients(x)
    # Gamma should be finite and not all exactly 1.0 if UNIFAC produced groups
    assert np.isfinite(gamma).all()
    assert gamma.shape == (len(test_formula.ingredients),)
    # Partial pressures must be less than or equal to saturation pressures
    p_part = calc.partial_pressures(x)
    p_sat = np.array([i.vapor_pressure_pa for i in test_formula.ingredients])
    assert (p_part <= p_sat * 1.01).all()


def test_evaporation_stiff_solve() -> None:
    trajectory = verify_formula_dict(
        _TEST_FORMULA_DATA,
        t_span=(0.0, 28800.0),
        t_eval=np.linspace(0.0, 28800.0, 481),
        method="BDF",
    )
    # Shape checks
    assert trajectory.t.shape[0] == 481
    n_species = len(_TEST_FORMULA_DATA["ingredients"])
    assert trajectory.n_liquid.shape == (n_species, 481)
    assert trajectory.x_liquid.shape == (n_species, 481)
    assert trajectory.C_gas.shape == (n_species, 481)
    assert trajectory.y_gas.shape == (n_species, 481)
    assert trajectory.OAV.shape == (n_species, 481)

    # No negative moles (mass balance sanity)
    assert (trajectory.n_liquid >= 0).all()
    # Total liquid moles decrease monotonically over time
    total_moles = trajectory.n_liquid.sum(axis=0)
    assert np.all(np.diff(total_moles) <= 1e-12)

    # Compute true fractional liquid depletion per component
    depletion = (trajectory.n_liquid[:, 0] - trajectory.n_liquid[:, -1]) / trajectory.n_liquid[:, 0]

    # Top notes (ethanol, limonene) deplete faster than heart/base notes
    ethanol_idx, limonene_idx, linalool_idx, _, galax_idx, ambrox_idx = range(6)
    assert depletion[ethanol_idx] > depletion[limonene_idx]
    assert depletion[limonene_idx] > depletion[galax_idx]
    assert depletion[ethanol_idx] > depletion[ambrox_idx]

    # Gas concentrations should be non-negative
    assert (trajectory.C_gas >= 0).all()
    # All mole fractions should stay a valid probability vector
    assert np.allclose(trajectory.x_liquid.sum(axis=0), 1.0, atol=1e-8)
    assert (trajectory.x_liquid >= 0).all()

    # OAV should be non-negative
    assert (trajectory.OAV >= 0).all()

    # Distinct per-species dynamic skin sink rates are recorded
    assert "k_skin_s-1" in trajectory.notes
    assert len(trajectory.notes["k_skin_s-1"]) == n_species


def test_verify_and_write(tmp_path: Path) -> None:
    formula = verify_formula_dict(_TEST_FORMULA_DATA).formula
    v = verifier.FormulaVerifier(formula)
    out = v.verify_and_write(tmp_path / "trajectory_output.json", t_eval=np.linspace(0, 28800, 481))
    assert out.exists()
    data = json.loads(out.read_text())
    assert "t_s" in data
    assert "x_liquid" in data
    assert "C_gas_mg_m3" in data
    assert "OAV" in data
    assert "species" in data
    assert len(data["species"]) == len(_TEST_FORMULA_DATA["ingredients"])
    assert data["t_s"][-1] == 28800.0


def test_no_ideal_gas_simplification() -> None:
    """Ensure activity coefficients are not hard-coded to unity."""
    trajectory = verify_formula_dict(_TEST_FORMULA_DATA)
    calc = vle.VLECalculator(trajectory.formula)
    x = trajectory.formula.mole_fractions
    gamma = calc.activity_coefficients(x)
    # For a non-ideal mixture, at least one gamma should differ from 1.0 by >1%
    assert np.any(np.abs(gamma - 1.0) > 0.01)