PALEOS-API live tabulation (producer)

def make_grid_at_resolution(
    p_lo: float,
    p_hi: float,
    t_lo: float,
    t_hi: float,
    pts_per_decade: int = DEFAULT_PTS_PER_DECADE,
) -> GridSpec:
    """Build a ``GridSpec`` with a uniform log-log resolution.

    Parameters
    ----------
    p_lo, p_hi : float
        Pressure bounds [Pa], both > 0.
    t_lo, t_hi : float
        Temperature bounds [K], both > 0.
    pts_per_decade : int
        Nodes per decade on each axis. Default
        ``DEFAULT_PTS_PER_DECADE`` = 600 (4x the shipped Zenodo resolution).
    """
    return GridSpec(
        p_lo=p_lo,
        p_hi=p_hi,
        n_p=_n_points_per_decade(p_lo, p_hi, pts_per_decade),
        t_lo=t_lo,
        t_hi=t_hi,
        n_t=_n_points_per_decade(t_lo, t_hi, pts_per_decade),
    )

def make_default_grid_iron() -> GridSpec:
    """Default Fe grid: 1e5 to 1e14 Pa, 300 to 20000 K at 600 pts/decade."""
    return make_grid_at_resolution(1e5, 1e14, 300.0, 2.0e4)

def make_default_grid_mgsio3() -> GridSpec:
    """Default MgSiO3 grid: 1e5 to 1e14 Pa, 300 to 11500 K at 600 pts/decade."""
    return make_grid_at_resolution(1e5, 1e14, 300.0, 1.15e4)

def make_default_grid_h2o() -> GridSpec:
    """Default H2O grid: 0.1 Pa to 1e14 Pa, 150 to 1e5 K at 600 pts/decade.

    Wider P range than Fe/MgSiO3 to match AQUA table coverage (extends to
    1 micro-bar). H2O is table-backed in PALEOS (``Haldemann20``) so per-point
    cost is low.
    """
    return make_grid_at_resolution(1e-1, 1e14, 150.0, 1.0e5)

def check_paleos_commit(repo_dir: Path | None = None, tested: str | None = None) -> str:
    """Return ``'tested'``, ``'newer'`` or ``'dirty'`` for the PALEOS checkout in ``repo_dir``.

    A checkout at ``tested`` (default ``PALEOS_TESTED_SHA``) runs; one that descends
    from it runs with a warning. Package code that differs from its commit stops, unless
    ``ZALMOXIS_PALEOS_ALLOW_DIRTY=1`` lets it run with a warning.

    Returns
    -------
    str
        ``'tested'``, ``'newer'``, or ``'dirty'`` for an allowed checkout with changes.

    Raises
    ------
    PaleosVersionError
        If PALEOS is not installed, its package is not a tracked file of a git checkout,
        the checkout does not contain ``tested`` (an older or unrelated commit, or a
        clone that lacks it), its commit or files cannot be read, or its package code
        differs from the commit without the opt-in.
    """
    tested = PALEOS_TESTED_SHA if tested is None else tested
    repo_dir = _paleos_repo_dir() if repo_dir is None else Path(repo_dir)
    if repo_dir is None:
        raise PaleosVersionError('PALEOS is not installed; the PALEOS-API path needs it')
    head = _git(repo_dir, 'rev-parse', 'HEAD')
    if head.returncode != 0 or not head.stdout.strip():
        raise PaleosVersionError(
            f'cannot tell the PALEOS commit of {repo_dir} ({head.stderr.strip()[:200]}): the '
            f'PALEOS-API path needs PALEOS installed from a git checkout at {tested} or later'
        )
    if _git(repo_dir, 'ls-files', '--error-unmatch', 'paleos/__init__.py').returncode != 0:
        raise PaleosVersionError(
            f'the paleos package in {repo_dir} is not tracked by the git checkout that holds '
            'it, so its commit cannot be told; install PALEOS with pip install -e from a clone'
        )
    head_sha = head.stdout.strip()
    if head_sha != tested:
        if _git(repo_dir, 'merge-base', '--is-ancestor', tested, head_sha).returncode != 0:
            shallow = _git(repo_dir, 'rev-parse', '--is-shallow-repository').stdout.strip()
            fix = 'fetch --unshallow' if shallow == 'true' else 'pull'
            raise PaleosVersionError(
                f'PALEOS in {repo_dir} is at {head_sha}, which does not contain the tested '
                f'commit {tested}: update it, for example git -C {repo_dir} {fix}'
            )
        logger.warning(
            'PALEOS %s is newer than the commit the PALEOS-API path is tested with (%s)',
            head_sha,
            tested,
        )
    if _changes_hash(repo_dir, head_sha) is not None:
        if _os.environ.get(ALLOW_DIRTY_ENV) != '1':
            raise PaleosVersionError(
                f'the paleos package files in {repo_dir} differ from its commit {head_sha}: '
                f'commit or remove the changes, or set {ALLOW_DIRTY_ENV}=1 to run them '
                'with their own cache; in a clean checkout, a line-end conversion such as '
                'core.autocrlf is a likely cause'
            )
        logger.warning('PALEOS in %s has uncommitted changes to its package files', repo_dir)
        return 'dirty'
    return 'tested' if head_sha == tested else 'newer'

def paleos_installed_sha() -> str:
    """Return the git SHA of the installed PALEOS checkout, or a fallback string.

    A checkout whose package code differs from its commit gets ``-dirty-<8 hex>`` after
    the SHA, a hash of that code, so each set of changes has its own cache. Without git
    the result is ``version-<paleos.__version__>``, and without PALEOS
    ``paleos-not-installed``. The first result is kept for :func:`_same_code_as_import`;
    when paleos was imported before it, a warning says that changes before this call are
    not covered.
    """
    fresh = 'paleos' not in _sys.modules
    repo_dir = _paleos_repo_dir()
    if repo_dir is None:
        return 'paleos-not-installed'
    key = None
    head = _git(repo_dir, 'rev-parse', 'HEAD')
    if head.returncode == 0 and head.stdout.strip():
        sha = head.stdout.strip()
        try:
            changes = _changes_hash(repo_dir, sha)
            key = sha if changes is None else f'{sha}-dirty-{changes}'
        except PaleosVersionError:
            pass
    if key is None:
        import paleos as _paleos

        key = f'version-{getattr(_paleos, "__version__", "unknown")}'
    if 'key' not in _KEY_AT_IMPORT:
        if not fresh:
            logger.warning(
                'paleos was imported before its code key was taken; a change on disk before '
                'this call is not covered'
            )
        _KEY_AT_IMPORT['key'] = key
    return key

def build_inputs(material: str, h2o_table_path: str | None = None) -> dict:
    """Return what a PALEOS-API table depends on besides the PALEOS code and the grid: the
    RTpress install (the Wolf18 liquid), the md5 of the AQUA table of H2O, and the numpy,
    scipy and sympy versions."""
    aqua = None
    if material == 'h2o' and h2o_table_path is not None:
        aqua = _hashlib.md5(Path(h2o_table_path).read_bytes()).hexdigest()
    ids = {name: _dist_id(name) for name in ('RTpress', 'numpy', 'scipy', 'sympy')}
    return {**ids, 'aqua_md5': aqua}

def build_key(material: str, h2o_table_path: str | None = None) -> str:
    """Return an 8-character hash of :func:`build_inputs`, part of the cache key."""
    text = _json.dumps(build_inputs(material, h2o_table_path), sort_keys=True)
    return _hashlib.sha256(text.encode()).hexdigest()[:8]

def generate_paleos_api_unified_table(
    material: str,
    out_path: str | Path,
    grid: GridSpec,
    *,
    h2o_table_path: str | None = None,
    log_every: int = 50,
    n_workers: int = 1,
) -> dict:
    """Generate a unified PALEOS ``.dat`` for one material.

    Stable-phase dispatch: for each (P, T) the PALEOS top-level class
    (``IronEoS`` / ``MgSiO3EoS`` / ``WaterEoS``) chooses the thermodynamically
    stable phase and evaluates that phase's EoS. Phase label is PALEOS's
    intrinsic string (``'solid-hpcen'``, ``'liquid'``, ``'hcp-Fe'``, etc).

    Parameters
    ----------
    material : {'iron', 'mgsio3', 'h2o'}
        Material selector.
    out_path : path-like
        Destination ``.dat`` path. Parent directories are created.
    grid : GridSpec
        (P, T) grid. ``p_lo`` must be > 0.
    h2o_table_path : str or None
        AQUA table path. Passed to ``WaterEoS`` when ``material == 'h2o'``.
    log_every : int
        Emit a progress line every ``log_every`` pressure rows.
    n_workers : int
        Parallelism. ``1`` = serial (default, backward-compatible). ``-1`` =
        ``os.cpu_count()``. Any other positive int = that many workers.
        Workers parallelize over P-rows; each worker instantiates its own
        PALEOS EoS at init (Wolf18 sympy compile cost paid once per worker,
        not per task).

    Returns
    -------
    dict
        ``{'n_valid', 'n_skipped', 'sha'}``.
    """
    if grid.p_lo <= 0:
        raise ValueError('GridSpec.p_lo must be > 0 (log spacing)')
    if material not in _HUMAN_MATERIAL:
        raise ValueError(f"material must be one of 'iron', 'mgsio3', 'h2o'; got {material!r}")

    human_material = _HUMAN_MATERIAL[material]
    n_workers = _resolve_n_workers(n_workers)
    sha = paleos_installed_sha()
    _check_paleos_once(sha)
    build = build_key(material, h2o_table_path)
    P_axis, T_axis = grid.axes()
    T_axis_list = T_axis.tolist()  # avoid repeated numpy scalar conversions in workers

    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = out_path.with_suffix(out_path.suffix + '.tmp')

    tasks = [(i, float(P), T_axis_list) for i, P in enumerate(P_axis)]

    rows_by_i: dict[int, str] = {}
    n_valid = 0
    n_skipped = 0
    n_done = 0

    if n_workers == 1:
        # Serial path: no Pool, no worker init — tabulate in-process.
        _same_code_as_import(sha)
        _worker_init(material, h2o_table_path=h2o_table_path)
        iterator = (_worker_unified_row(t) for t in tasks)
    else:
        ctx = _mp.get_context('spawn')
        pool = ctx.Pool(
            processes=n_workers,
            initializer=_worker_init,
            initargs=(material, h2o_table_path),
        )
        # chunksize kept small so progress logs stay meaningful at high cost/cell.
        chunksize = max(1, len(tasks) // (n_workers * 16))
        iterator = pool.imap_unordered(_worker_unified_row, tasks, chunksize=chunksize)

    try:
        for i_P, rows_text, nv, ns in iterator:
            rows_by_i[i_P] = rows_text
            n_valid += nv
            n_skipped += ns
            n_done += 1
            if log_every and (n_done % log_every == 0):
                logger.info(
                    'paleos_api %s unified: %d / %d P rows done (valid=%d skipped=%d)',
                    human_material,
                    n_done,
                    grid.n_p,
                    n_valid,
                    n_skipped,
                )
    finally:
        if n_workers > 1:
            pool.close()
            pool.join()

    with open(tmp_path, 'w') as f:
        _write_header(
            f,
            human_material,
            'unified (stable-phase)',
            grid,
            sha,
            build=build,
            n_valid=n_valid,
            n_skipped=n_skipped,
        )
        for i in range(grid.n_p):
            f.write(rows_by_i[i])
    _os.replace(tmp_path, out_path)

    logger.info(
        'paleos_api %s unified: wrote %s (valid=%d skipped=%d sha=%s n_workers=%d)',
        human_material,
        out_path,
        n_valid,
        n_skipped,
        sha[:10],
        n_workers,
    )
    return {'n_valid': n_valid, 'n_skipped': n_skipped, 'sha': sha}

def generate_paleos_api_2phase_mgsio3_tables(
    out_solid: str | Path,
    out_liquid: str | Path,
    grid: GridSpec,
    *,
    log_every: int = 50,
    n_workers: int = 1,
) -> dict:
    """Generate solid and liquid MgSiO3 P-T tables with metastable extensions.

    Replaces the shipped ``paleos_mgsio3_tables_pt_proteus_solid.dat`` and
    ``paleos_mgsio3_tables_pt_proteus_liquid.dat`` — bit-format-compatible,
    consumed unchanged by ``eos_export.generate_spider_eos_tables`` via
    ``solid_eos_file`` / ``liquid_eos_file``.

    Parameters
    ----------
    out_solid, out_liquid : path-like
        Destinations for the two files. Parent directories are created.
    grid : GridSpec
        (P, T) grid; ``p_lo`` must be > 0.
    log_every : int
        Progress log cadence in pressure rows.
    n_workers : int
        Parallelism. ``1`` = serial (default, backward-compatible). ``-1`` =
        ``os.cpu_count()``. Workers parallelize over P-rows; each worker
        instantiates ``MgSiO3EoS`` + ``Wolf18`` once (Wolf18's sympy compile
        cost is paid once per worker, not per task).

    Returns
    -------
    dict
        Summary with per-side ``n_valid`` / ``n_skipped`` and the installed
        PALEOS ``sha``.

    Notes
    -----
    The liquid-side file evaluates ``paleos.mgsio3_eos.Wolf18`` on every
    grid node, including below the solidus (metastable liquid). The
    solid-side file dispatches through ``MgSiO3EoS()._phase_eos_map`` using
    ``_get_mgsio3_solid_phase`` (melting-curve test removed) so every node
    gets a solid-polymorph EoS evaluated metastably above the liquidus if
    needed.
    """
    if grid.p_lo <= 0:
        raise ValueError('GridSpec.p_lo must be > 0 (log spacing)')

    n_workers = _resolve_n_workers(n_workers)
    sha = paleos_installed_sha()
    _check_paleos_once(sha)
    build = build_key('mgsio3')
    P_axis, T_axis = grid.axes()
    T_axis_list = T_axis.tolist()

    out_solid = Path(out_solid)
    out_liquid = Path(out_liquid)
    out_solid.parent.mkdir(parents=True, exist_ok=True)
    out_liquid.parent.mkdir(parents=True, exist_ok=True)
    tmp_solid = out_solid.with_suffix(out_solid.suffix + '.tmp')
    tmp_liquid = out_liquid.with_suffix(out_liquid.suffix + '.tmp')

    tasks = [(i, float(P), T_axis_list) for i, P in enumerate(P_axis)]

    solid_by_i: dict[int, str] = {}
    liquid_by_i: dict[int, str] = {}
    n_valid_s = n_skip_s = 0
    n_valid_l = n_skip_l = 0
    n_done = 0

    if n_workers == 1:
        _same_code_as_import(sha)
        _worker_init('mgsio3')
        iterator = (_worker_2phase_row(t) for t in tasks)
    else:
        ctx = _mp.get_context('spawn')
        pool = ctx.Pool(
            processes=n_workers,
            initializer=_worker_init,
            initargs=('mgsio3',),
        )
        chunksize = max(1, len(tasks) // (n_workers * 16))
        iterator = pool.imap_unordered(_worker_2phase_row, tasks, chunksize=chunksize)

    try:
        for i_P, s_text, l_text, nv_s, ns_s, nv_l, ns_l in iterator:
            solid_by_i[i_P] = s_text
            liquid_by_i[i_P] = l_text
            n_valid_s += nv_s
            n_skip_s += ns_s
            n_valid_l += nv_l
            n_skip_l += ns_l
            n_done += 1
            if log_every and (n_done % log_every == 0):
                logger.info(
                    'paleos_api MgSiO3 2-phase: %d / %d P rows done '
                    '(solid valid=%d skipped=%d | liquid valid=%d skipped=%d)',
                    n_done,
                    grid.n_p,
                    n_valid_s,
                    n_skip_s,
                    n_valid_l,
                    n_skip_l,
                )
    finally:
        if n_workers > 1:
            pool.close()
            pool.join()

    with open(tmp_solid, 'w') as f:
        _write_header(
            f,
            'MgSiO3',
            'solid (metastable extension above liquidus)',
            grid,
            sha,
            build=build,
            n_valid=n_valid_s,
            n_skipped=n_skip_s,
            extra_lines=[
                'Solid polymorph picked by _get_mgsio3_solid_phase; melting-curve test suppressed.',
                'Phase column is the polymorph label even if (P, T) is above the liquidus.',
            ],
        )
        for i in range(grid.n_p):
            f.write(solid_by_i[i])
    _os.replace(tmp_solid, out_solid)

    with open(tmp_liquid, 'w') as f:
        _write_header(
            f,
            'MgSiO3',
            'liquid (Wolf18, metastable extension below solidus)',
            grid,
            sha,
            build=build,
            n_valid=n_valid_l,
            n_skipped=n_skip_l,
            extra_lines=[
                'Wolf18 RTpress evaluated everywhere; phase column is always "liquid".',
            ],
        )
        for i in range(grid.n_p):
            f.write(liquid_by_i[i])
    _os.replace(tmp_liquid, out_liquid)

    logger.info(
        'paleos_api MgSiO3 2-phase: wrote %s (valid=%d/%d) and %s (valid=%d/%d) '
        'sha=%s n_workers=%d',
        out_solid,
        n_valid_s,
        grid.n_p * grid.n_t,
        out_liquid,
        n_valid_l,
        grid.n_p * grid.n_t,
        sha[:10],
        n_workers,
    )
    return {
        'solid': {'n_valid': n_valid_s, 'n_skipped': n_skip_s, 'path': str(out_solid)},
        'liquid': {'n_valid': n_valid_l, 'n_skipped': n_skip_l, 'path': str(out_liquid)},
        'sha': sha,
    }