rk-random/PACT-Net

1import argparse2 3# The runners module will contain the high-level logic for each model type4from runners import run_gnn_experiment, run_gp_experiment5 6 7def main():8    """9    Main entry point for running all experiments.10    Parses command-line arguments and calls the appropriate runner function.11    """12    parser = argparse.ArgumentParser(13        description="Run GNN or GP experiments for molecular property prediction."14    )15    parser.add_argument(16        "--model",17        choices=["gcn", "gin", "gat", "sage", "gp", "polyatomic"],18        required=True,19        help="The model to train and evaluate.",20    )21    parser.add_argument(22        "--rep",23        choices=["smiles", "selfies", "ecfp", "polyatomic"],24        required=True,25        help="The molecular representation to use.",26    )27    parser.add_argument(28        "--dataset",29        choices=[30            "esol",31            "freesolv",32            "lipophil",33            "boilingpoint",34            "qm9",35            "ic50",36            "bindingdb",37        ],38        required=True,39        help="The dataset to use for the experiment.",40    )41    parser.add_argument(42        "--n-trials",43        type=int,44        default=10,45        help="Number of Optuna trials to run for hyperparameter search in each fold.",46    )47 48    args = parser.parse_args()49 50    # --- Argument Validation ---51    if args.model == "gp" and args.rep == "polyatomic":52        raise ValueError(53            "The 'polyatomic' representation is not compatible with the 'gp' model."54        )55    if args.model == "polyatomic" and args.rep != "polyatomic":56        raise ValueError(57            "The 'polyatomic' model must be used with the 'polyatomic' representation."58        )59 60    # --- Delegate to the correct runner ---61    if args.model == "gp":62        run_gp_experiment(args)63    else:64        run_gnn_experiment(args)65 66 67if __name__ == "__main__":68    main()69