Source code for formulation_bench.reformulation

import json
import subprocess
from dataclasses import dataclass
from pathlib import Path

from ._codegen import generate_map
from .formulation import Formulation
from .models import ParameterMap




@dataclass(frozen=True)
class Reformulation:
    """A pair of MILP formulations with a reformulation label.

    Consists of two MILP formulations ``a`` and ``b`` and a boolean ``is_reformulation``
    label indicating whether ``b`` is a reformulation of ``a``. The formal definition
    of *reformulation* is given in :ref:`reformulation-definition`. Positive entries
    (``is_reformulation=True``) are accompanied by a Lean 4 proof whose path is
    accessible via the ``lean_proof_path`` attribute; negative entries have no proof
    and ``lean_proof_path`` resolves to ``None``.

    Attributes
    ----------
    a : Formulation
        The base formulation.
    b : Formulation
        The reformulation candidate.
    is_reformulation : bool
        ``True`` iff ``b`` is a *reformulation* of ``a``.
    path : pathlib.Path
        Resolved absolute path to this pair's directory.
    lean_proof_path : pathlib.Path or None
        For positive entries, the path to the accompanying Lean 4 proof file. For
        negative entries, ``None`` since no proof exists.
    parameter_map : ParameterMap
        The loaded ``map.json``, which states how each parameter of ``b`` is
        computed from the parameters of ``a``.

    Examples
    --------

    Formulation ``b`` of :doc:`/problems/p12` is a reformulation of formulation ``a``::

        >>> from formulation_bench import Dataset
        >>> ds = Dataset("dataset")
        >>> reform = ds.reformulations[73]  # corresponds to p12.a -> p12.b
        >>> reform.a.problem.name
        'Traveling Salesman Problem (TSP)'
        >>> reform.b.problem.name
        'Traveling Salesman Problem (TSP)'
        >>> reform.b.constraints[-1].description  # cutting plane added by p12.b
        'Depot-Exit Position Bound (EC1)...'
        >>> reform.is_reformulation
        True
    """

    a: Formulation
    b: Formulation
    is_reformulation: bool

    @property
    def path(self) -> Path:
        problem_dir = self.a.problem.path
        pair = f"{self.a.path.name}_{self.b.path.name}"
        return problem_dir.parent.parent / "reformulations" / problem_dir.name / pair

    @property
    def lean_proof_path(self) -> Path | None:
        if not self.is_reformulation:
            return None
        return self.path / "Reformulation.lean"

    @property
    def parameter_map(self) -> ParameterMap:
        return ParameterMap.from_dict(json.loads((self.path / "map.json").read_text()))



    def gen_map_py(self) -> str:
        """Generate a Python script computing ``b``'s parameters from ``a``'s.

        The script is generated from the ``python`` code snippets of this pair's
        ``map.json``. The resulting script takes the path to ``a``'s
        ``parameters.json`` and the path to write ``b``'s as positional arguments.

        Examples
        --------

        Generate the parameter-map script for the ``a`` to ``b`` pair of
        :doc:`/problems/p1`, whose parameters are a renaming::

            >>> from formulation_bench import Dataset
            >>> ds = Dataset("dataset")
            >>> reform = ds.reformulations[0]  # corresponds to p1.a -> p1.b
            >>> script = reform.gen_map_py()
            >>> print(script)
            import argparse
            import json
            <BLANKLINE>
            <BLANKLINE>
            def main(params_path: str, output_path: str) -> None:
                with open(params_path, "r") as f:
                    data = json.load(f)
            <BLANKLINE>
                # Source Parameters
                CashMachineProcessingRate = data["CashMachineProcessingRate"]
                ...
                # Parameter Map
                A = CashMachineProcessingRate
                ...

        """
        return generate_map(self)




    def run_map(
        self,
        input_path: str | Path | None = None,
        output_path: str | Path | None = None,
    ) -> None:
        """Write this pair's ``map.py`` and run it.

        The script generated by :meth:`gen_map_py` is written to ``map.py`` in
        the pair's directory, then applied to ``a``'s parameters to produce
        ``b``'s.

        Parameters
        ----------
        input_path : str or pathlib.Path, optional
            Path to a ``parameters.json`` holding formulation ``a``'s parameters.
            Defaults to ``parameters.json`` in ``a``'s formulation directory.
        output_path : str or pathlib.Path, optional
            Path to write the mapped parameters. Defaults to ``parameters.json``
            in this pair's directory.

        Examples
        --------

        Map formulation ``a``'s parameters to formulation ``b``'s for
        :doc:`/problems/p1`::

            >>> import json
            >>> from formulation_bench import Dataset
            >>> ds = Dataset("dataset")
            >>> reform = ds.reformulations[0]  # corresponds to p1.a -> p1.b

            >>> reform.a.run_gen_params()
            >>> params = json.load(open(reform.a.path / "parameters.json"))
            >>> params["CashMachineProcessingRate"]  # a's parameter
            20

            >>> reform.run_map()
            >>> params = json.load(open(reform.path / "parameters.json"))
            >>> params["A"]  # b's name for "CashMachineProcessingRate"
            20

        """
        script = self.path / "map.py"
        script.write_text(self.gen_map_py())
        if input_path is None:
            input_path = self.a.path / "parameters.json"
        if output_path is None:
            output_path = self.path / "parameters.json"
        subprocess.run(
            ["python", str(script), str(input_path), str(output_path)], check=True
        )