GI_SR/make_crys.py at master · angus-gentles/GI_SR · GitHub

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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
#!/usr/bin/env python3

import numpy as np
import os
import sys
import argparse
import yaml

def replace_atom_labels(input_file, output_file, a_list, b_list):
    with open(input_file, 'r') as f:
        lines = f.readlines()

    new_lines = lines[:2]  # Keep header lines (atom count and comment)
    a_index = 0
    b_index = 0

    for line in lines[2:]:
        parts = line.strip().split()
        label, coords = parts[0], parts[1:]

        if label == 'A':
            new_label = a_list[a_index]
            a_index += 1
        elif label == 'B':
            new_label = b_list[b_index]
            b_index += 1
        else:
            new_label = label  # in case of unexpected label

        new_line = f"{new_label} {' '.join(coords)}\n"
        new_lines.append(new_line)

    with open(output_file, 'w') as f:
        f.writelines(new_lines)

def extract_configuration(yaml_file):
    with open(yaml_file, 'r') as f:
        data = yaml.safe_load(f)

    # Access the first (and possibly only) configuration ID
    config_id = next(iter(data['configurations']))
    configuration = data['configurations'][config_id]['configuration']

    return configuration

def extract_all_configurations(yaml_file):
    with open(yaml_file, 'r') as f:
        data = yaml.safe_load(f)

    all_configs = {}
    for config_id, config_data in data['configurations'].items():
        all_configs[int(config_id)] = config_data['configuration']

    return all_configs
"""
# Example usage
configs = extract_all_configurations('In0.0Ga1.0As0.5Sb0.5_cube_2x2x2.result.yaml')

for config_id, configuration in configs.items():
    print(f"ID: {config_id}, Configuration: {configuration}")
"""

if __name__=="__main__":
    parser = argparse.ArgumentParser(description="Generate a random SQS structure")
    parser.add_argument("-ba", "--base_a", type=str, required=True, help="result yaml file")
    parser.add_argument("-bc", "--base_crys", type=str, required=True, help="base_configuration_file")
    #parser.add_argument("-o", "--output_file", type=str, required=True, help="Output file name for the generated SQS structure")
    parser.add_argument("-sl", "--sublattice", type=str, required=True, help="Output file name for the generated SQS structure")
    parser.add_argument("-ol", "--other_lattice", type=str, required=True, help="other lattice element")

    args = parser.parse_args()

    base = args.base_a
    base_crys= args.base_crys
    #output_file = args.output_file
    sublattice = args.sublattice
    ol=args.other_lattice

    # Generate a random SQS structure
    if sublattice=='A':
        A_configs=extract_all_configurations(f"{base}.result.yaml")
        first_key = next(iter(A_configs))
        B_config=[ol for j in range(len(A_configs[str(first_key)]))]
        for i, config in enumerate(A_configs):
            replace_atom_labels(base_crys, f"crys_files/{base}.{i}.crys", config, B_config)
    elif sublattice=='B':
        B_configs=extract_all_configurations(f"{base}.result.yaml")
        first_key = list(B_configs.keys())[0]
        print(first_key)
        A_config=[ol for j in range(len(B_configs[first_key]))]
        for i, config_id in enumerate(B_configs):

            replace_atom_labels(base_crys, f"crys_files/{base}.{i}.crys", A_config, B_configs[config_id])
    else:
        exit('not correct lattice')

    #replace_atom_labels(base_file, output_file, A_config, B_config)