DooABLe / scripts /sweep_preferences.py
pranamanam's picture
Upload 309 files
81ae663 verified
Raw
History Blame Contribute Delete
2.42 kB
"""Evaluate new public-property weights on one stored reaction graph."""
import argparse, json
from pathlib import Path
import pandas as pd
from dooable.graph import Graph
from dooable.properties import property_rewards
from dooable.exact import solve, endpoint_distribution, expected_cost
from dooable.learning import train
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--graph", required=True)
parser.add_argument("--models", required=True)
parser.add_argument("--output", default="results/preferences")
parser.add_argument("--neural-steps", type=int, default=0)
args = parser.parse_args()
g = Graph.load(args.graph)
out = Path(args.output)
out.mkdir(parents=True, exist_ok=True)
rows = []
previous = None
for beta in [1.0, 2.0, 5.0]:
for w in [0.0, 0.25, 0.5, 0.75, 1.0]:
rewards, scores = property_rewards(
g, args.models, weights=(w, 1 - w), concentration=beta
)
directory = out / f"beta{beta:g}_weight{w:g}"
directory.mkdir(exist_ok=True)
(directory / "rewards.json").write_text(json.dumps(rewards, indent=2))
policies = {"exact": solve(g, rewards, 0.7).forward}
if args.neural_steps:
model, _ = train(
g,
rewards,
0.7,
steps=args.neural_steps,
output=directory,
initialize=previous,
)
policies["dooable"] = model.probabilities()
previous = directory
for name, p in policies.items():
masses = endpoint_distribution(g, p)
s = scores.set_index("smiles")
a = pd.Series(masses).reindex(s.index)
rows.append(
{
"method": name,
"bace_weight": w,
"concentration": beta,
"mean_bace_utility": float(a @ s.bace_utility),
"mean_caco2_utility": float(a @ s.caco2_utility),
"mean_cost": expected_cost(g, p),
}
)
pd.DataFrame(rows).to_csv(out / "measurements.csv", index=False)
print(f"Saved {len(rows)} preference measurements to {out}")
if __name__ == "__main__":
main()