| """ | |
| The data splitting script for pretraining. | |
| """ | |
| import csv | |
| import os | |
| import shutil | |
| from argparse import ArgumentParser | |
| import grover.util.utils as fea_utils | |
| import numpy as np | |
| parser = ArgumentParser() | |
| parser.add_argument( | |
| "--data_path", default="../drug_data/grover_data/delaneyfreesolvlipo.csv" | |
| ) | |
| parser.add_argument( | |
| "--features_path", | |
| default="../drug_data/grover_data/delaneyfreesolvlipo_molbert.npz", | |
| ) | |
| parser.add_argument("--sample_per_file", type=int, default=1000) | |
| parser.add_argument( | |
| "--output_path", default="../drug_data/grover_data/delaneyfreesolvlipo" | |
| ) | |
| def load_smiles(data_path): | |
| with open(data_path) as f: | |
| reader = csv.reader(f) | |
| header = next(reader) | |
| res = [] | |
| for line in reader: | |
| res.append(line) | |
| return res, header | |
| def load_features(data_path): | |
| fea = fea_utils.load_features(data_path) | |
| return fea | |
| def save_smiles(data_path, index, data, header): | |
| fn = os.path.join(data_path, str(index) + ".csv") | |
| with open(fn, "w") as f: | |
| fw = csv.writer(f) | |
| fw.writerow(header) | |
| for d in data: | |
| fw.writerow(d) | |
| def save_features(data_path, index, data): | |
| fn = os.path.join(data_path, str(index) + ".npz") | |
| np.savez_compressed(fn, features=data) | |
| def run(): | |
| args = parser.parse_args() | |
| res, header = load_smiles(data_path=args.data_path) | |
| fea = load_features(data_path=args.features_path) | |
| assert len(res) == fea.shape[0] | |
| n_graphs = len(res) | |
| perm = np.random.permutation(n_graphs) | |
| nfold = int(n_graphs / args.sample_per_file + 1) | |
| print("Number of files: %d" % nfold) | |
| if os.path.exists(args.output_path): | |
| shutil.rmtree(args.output_path) | |
| os.makedirs(args.output_path, exist_ok=True) | |
| graph_path = os.path.join(args.output_path, "graph") | |
| fea_path = os.path.join(args.output_path, "feature") | |
| os.makedirs(graph_path, exist_ok=True) | |
| os.makedirs(fea_path, exist_ok=True) | |
| for i in range(nfold): | |
| sidx = i * args.sample_per_file | |
| eidx = min((i + 1) * args.sample_per_file, n_graphs) | |
| indexes = perm[sidx:eidx] | |
| sres = [res[j] for j in indexes] | |
| sfea = fea[indexes] | |
| save_smiles(graph_path, i, sres, header) | |
| save_features(fea_path, i, sfea) | |
| summary_path = os.path.join(args.output_path, "summary.txt") | |
| summary_fout = open(summary_path, "w") | |
| summary_fout.write("n_files:%d\n" % nfold) | |
| summary_fout.write("n_samples:%d\n" % n_graphs) | |
| summary_fout.write("sample_per_file:%d\n" % args.sample_per_file) | |
| summary_fout.close() | |
| if __name__ == "__main__": | |
| run() | |