File size: 2,407 Bytes
39c21b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | from cdvae.common.data_utils import preprocess
import pandas as pd
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice
from pymatgen.analysis.diffraction.xrd import XRDCalculator, WAVELENGTHS
import torch
from tqdm import tqdm
import argparse
import os
def split_data(args):
# list the directory
files = os.listdir(args.init_data_dir)
# filter the files
files = [f for f in files if f.endswith('.csv')]
print(f'Found {len(files)} files: {files}')
assert len(files) == 3
all_data = list()
for file in files:
print(f'Processing {file}')
# Load the data
curr_data = pd.read_pickle(os.path.join(args.init_data_dir, file))
print(curr_data.head())
all_data.append(curr_data)
print(len(curr_data))
all_data = pd.concat(all_data, axis=0, ignore_index=True)
print(all_data.head())
print(len(all_data))
print(all_data.columns)
new_train_len = int(args.train_percent * len(all_data))
new_test_len = int(args.test_percent * len(all_data))
train_data = all_data.iloc[0:new_train_len]
test_data = all_data.iloc[new_train_len:new_train_len+new_test_len]
val_data = all_data.iloc[new_train_len+new_test_len:]
print(train_data.head(), len(train_data))
print(val_data.head(), len(val_data))
print(test_data.head(), len(test_data))
train_data.to_pickle(os.path.join(args.new_data_dir, 'train.csv'))
val_data.to_pickle(os.path.join(args.new_data_dir, 'val.csv'))
test_data.to_pickle(os.path.join(args.new_data_dir, 'test.csv'))
return
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Split data')
parser.add_argument(
'--init_data_dir',
default = '/home/gabeguo/cdvae_xrd/data/mp_20_oldSplit',
type=str,
help='path to input CIF+XRD files'
)
parser.add_argument(
'--new_data_dir',
default = '/home/gabeguo/cdvae_xrd/data/mp_20',
type=str,
help='path to save re-split data'
)
parser.add_argument(
'--train_percent',
type=float,
default=0.9
)
parser.add_argument(
'--test_percent',
type=float,
default=0.025
)
args = parser.parse_args()
os.makedirs(args.new_data_dir, exist_ok=True)
split_data(args)
|