74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
import os
|
|
|
|
|
|
def get_var_code(symbol, shape, type, definition, solver="gurobipy"):
|
|
if solver == "gurobipy":
|
|
if not shape:
|
|
return f'{symbol} = model.addVar(vtype=GRB.{type.upper()}, name="{symbol}")\n'
|
|
else:
|
|
return f"{symbol} = model.addVars(" + ", ".join([str(i) for i in shape]) + f', vtype=GRB.{type.upper()}, name="{symbol}")\n'
|
|
else:
|
|
raise NotImplementedError(f"Solver {solver} is not implemented")
|
|
|
|
|
|
def get_param_code(symbol, shape, definition):
|
|
return f'{symbol} = data["{symbol}"] # shape: {shape}, definition: {definition}\n'
|
|
|
|
|
|
def generate_code(state, dir):
|
|
code = []
|
|
code.append(
|
|
f"""
|
|
import os
|
|
import numpy as np
|
|
import json
|
|
from gurobipy import Model, GRB, quicksum
|
|
|
|
|
|
model = Model("OptimizationProblem")
|
|
|
|
with open("data.json", "r") as f:
|
|
data = json.load(f)
|
|
"""
|
|
)
|
|
|
|
code.append("\n\n### Define the parameters\n")
|
|
for symbol, v in state["parameters"].items():
|
|
print(v)
|
|
code.append(get_param_code(symbol, v["shape"], v["definition"]))
|
|
|
|
code.append("\n\n### Define the variables\n")
|
|
for symbol, v in state["variables"].items():
|
|
code.append(get_var_code(symbol, v["shape"], v["type"], v["definition"], solver="gurobipy"))
|
|
|
|
code.append("\n\n### Define the constraints\n")
|
|
for c in state["constraints"]:
|
|
code.append(c["code"])
|
|
|
|
code.append("\n\n### Define the objective\n")
|
|
code.append(state["objective"]["code"])
|
|
|
|
code.append("\n\n### Optimize the model\n")
|
|
code.append("model.optimize()\n")
|
|
|
|
code.append("\n\n### Output optimal objective value\n")
|
|
code.append(f'print("Optimal Objective Value: ", model.objVal)\n')
|
|
|
|
# code to save the optimal value if it exists
|
|
code.append(
|
|
"""
|
|
if model.status == GRB.OPTIMAL:
|
|
with open("output_solution.txt", "w") as f:
|
|
f.write(str(model.objVal))
|
|
print("Optimal Objective Value: ", model.objVal)
|
|
else:
|
|
with open("output_solution.txt", "w") as f:
|
|
f.write(model.status)
|
|
"""
|
|
)
|
|
|
|
code_str = "\n".join(code)
|
|
|
|
with open(os.path.join(dir, "code.py"), "w") as f:
|
|
f.write(code_str)
|