Simple AMPL wrapper for Operating Research

amplix.py

#!/usr/bin/python3

import subprocess
import argparse
import re

SOLVERS = [
    "cplex",
    "minos",
    "snopt",
    "lgo",
    "gurobi",
    "baron",
    "knitro",
]

parser = argparse.ArgumentParser(description="AMPL with steroids")
parser.add_argument(
    "filename",
    help="the model file to parse/solve",
    type=str,
)
parser.add_argument(
    "--solver",
    metavar="SOLVER",
    help=(
        "the solver(s) you wish to use; "
        "you can specify multiple solvers "
        "(useful if combined with the --run option), "
        "but only the last one will be considered in the .run file."
    ),
    dest="solvers",
    action="append",
    required=True,
)
parser.add_argument(
    "-r",
    "--run",
    help="immediately execute in AMPL the output .run file",
    action="store_true",
    default=True,
)


def parse_file(filename: str) -> tuple[list[str], str | None]:
    try:
        with open(filename, "r+") as f:
            contents: list[str] = [line.strip() for line in f.readlines()]
    except FileNotFoundError:
        raise FileNotFoundError("File not found")

    variables: list[str] = []
    obj: str | None = None
    for line in contents:
        if line.startswith("var"):
            if res := re.findall(re.compile(r"(?<=var )(\w+)"), line):
                variables.append(res[0])

        if line.startswith("maximize") or line.startswith("minimize"):
            if res := re.findall(
                re.compile(r"(?<=(?<=maximize )|(?<=minimize ))(\w+)"), line
            ):
                obj = res[0]

        if line.startswith("data"):
            break

    return variables, obj


def generate_runfile(
    filename: str,
    solver: str,
    run: bool,
    variables: list[str],
    objective: str | None,
):
    contents = (
        "reset;"
        f'\nmodel "{filename}";'
        f'\noption solver "{solver}";'
        "\nsolve;"
        f"\ndisplay {', '.join(variables)};"
        "\n"
    )
    if objective:
        contents += f"\ndisplay: {objective};"
    contents += "\n"
    runfile = ".".join(filename.split(".")[:-1]) + ".run"
    with open(runfile, "w+") as f:
        f.write(contents)

    if run:
        subprocess.call(["ampl", runfile])


def main():
    args = parser.parse_args()
    filename: str = args.filename
    variables, objective = parse_file(filename)
    for solver in args.solvers:
        generate_runfile(filename, solver, args.run, variables, objective)


if __name__ == "__main__":
    main()