create_starting_structures.py

import os
import argparse
import json

from UnitCell_Environment.unitcell_environment.env.utils import load_compositions, random_airss_structure

if __name__ == "__main__":

    cell_fixed = False

    parser = argparse.ArgumentParser(description='Optional app description')
    parser.add_argument('--output_dir', type=str, default="new_starting_structures",
                        help='The directory to save the structures')
    parser.add_argument("--n", type=int, default="10",
                        help="The number of structures to generate")
    parser.add_argument("--comp_config", type=str, default="comp_config.json",
                    help="the config file with compositions and structure sizes supported")
    parser.add_argument("--comp", type=str, default="SrTiO3x8",
                        help="the composition/structure size name")

    args = parser.parse_args()

    directory = args.output_dir
    
    if not os.path.exists(directory):
        os.makedirs(directory)

    with open(args.comp_config, 'r') as f:
        comp_config = json.load(f)

    load_compositions(comp_config)

    for i in range(args.n):

        atoms = random_airss_structure(args.comp, cell_fixed=cell_fixed)
        atoms.write(f"{directory}/{args.comp}_{i}.cif")

    quit()