wmaee.codes.vasp.vasp_ml

  1"""
  2Collection of helper functions for handling ML functionality of VASP.
  3"""
  4
  5from typing import Callable
  6from typing import List, Union
  7
  8def ML_ABN_concat(files: List[str], output: str = 'ML_AB', overwrite: bool = False, verbose: bool = True) -> None:
  9    """
 10    Concatenates a list of `ML_AB` or `ML_ABN` files. This function merges headers and structures
 11    from multiple files into a single output file.
 12
 13    Parameters
 14    ----------
 15    files : List[str]
 16        List of paths to input files that should be concatenated.
 17    output : str, optional
 18        Path for the output file, by default 'ML_AB'.
 19    overwrite : bool, optional
 20        Whether to overwrite the output file if it already exists, by default False.
 21    verbose : bool, optional
 22        If True, prints details of the merging process, by default True.
 23
 24    Raises
 25    ------
 26    FileExistsError
 27        If the output file exists and overwrite is set to False.
 28
 29    Notes
 30    -----
 31    This function currently assumes that all input files contain the same chemistry.
 32    """
 33    
 34    from os.path import exists
 35    from os import remove, rename
 36
 37    if exists(output) and not overwrite:
 38        raise FileExistsError(f'Output file {output} exists in the current folder and overwrite is not permitted.')
 39    else:
 40        out = open(output+'.tmp_wmaee_tmp', 'w')
 41
 42    s = 0
 43    sp = list()
 44    max_types = 0
 45    max_atoms = 0
 46    max_species = 0
 47    masses = dict()
 48    for i, f in enumerate(files):
 49        if verbose:
 50            print(f'reading structures from {f}')
 51        with open(f, 'r') as src:
 52            # deal with header: collect merged information
 53            l = src.readline();
 54            while l.strip().split()[0] != 'Configuration':                
 55                if 'The atom types' in l.strip():
 56                    l = src.readline();
 57                    l = src.readline();
 58                    this_sp = l.strip().split()
 59                    sp += this_sp
 60                    sp = list(set(sp))
 61                    if verbose:
 62                        print(f'  found species: {sp}')
 63                if 'number of atom type' in l.strip():
 64                    l = src.readline();
 65                    l = src.readline();
 66                    max_at = int(l.strip())
 67                    if verbose:
 68                        print(f'  maximum number of atom types: {max_at}')
 69                    max_types = max(max_types, max_at)                        
 70                if 'atoms per system' in l.strip():
 71                    l = src.readline();
 72                    l = src.readline();
 73                    max_at = int(l.strip())
 74                    if verbose:
 75                        print(f'  maximum number of atoms: {max_at}')
 76                    max_atoms = max(max_atoms, max_at)
 77                if 'atoms per atom type' in l.strip():
 78                    l = src.readline();
 79                    l = src.readline();
 80                    max_at = int(l.strip())
 81                    if verbose:
 82                        print(f'  maximum number of atoms per type: {max_at}')
 83                    max_species = max(max_species, max_at)
 84                if 'Atomic mass' in l.strip():
 85                    l = src.readline();
 86                    m = ''
 87                    l = src.readline();
 88                    while '*********************' not in l.strip():
 89                        m += l
 90                        l = src.readline();
 91                    for X, M in zip(this_sp, m.strip().split()):
 92                        masses[X] = float(M)
 93                        if verbose:
 94                            print(f'  mass of {X}: {float(M)} a.u.')              
 95                l = src.readline();
 96            # done with headers, now go configuration by configuration            
 97            out.writelines("**************************************************\n")
 98            while l:
 99                s += 1
100                out.writelines(f"     Configuration num.{s:7d}\n")
101                if verbose:
102                    print(f"     Configuration num.{s:7d}")
103                l = src.readline()
104                while l and l.strip().split()[0] != 'Configuration':
105                    out.writelines(l)
106                    l = src.readline()
107    out.close()
108    
109    sp = list(set(sp))
110
111    # write header    
112    with open(output, 'w') as out:        
113        _SEP_star  = "**************************************************\n"
114        _SEP_minus = "--------------------------------------------------\n"
115        _SEP_equal = "==================================================\n"
116        out.writelines(' 1.0 Version\n'+_SEP_star)
117        out.writelines('     The number of configurations\n'+_SEP_minus+f'{s:11d}\n'+_SEP_star)
118        out.writelines('     The maximum number of atom type\n'+_SEP_minus+f'{len(sp):8d}\n'+_SEP_star)
119        out.writelines('     The atom types in the data file\n'+_SEP_minus+'     ')
120        for i, el in enumerate(sp):
121            out.writelines(el.ljust(3))
122            if (i%3 == 2) and i<len(sp)-1:
123                out.writelines('\n     ')
124        out.writelines('\n'+_SEP_star)
125        out.writelines('     The maximum number of atoms per system\n'+_SEP_minus+f'{max_atoms:8d}\n'+_SEP_star)
126        out.writelines('     The maximum number of atoms per atom type\n'+_SEP_minus+f'{max_species:8d}\n'+_SEP_star)
127        out.writelines('     Reference atomic energy (eV)\n'+_SEP_minus)
128        for i, _ in enumerate(sp):
129            out.writelines('   '+f'{0:18.16f}'+'     ') # dummy energy
130            if (i%3 == 2) and i<len(sp)-1:
131                out.writelines('\n')
132        out.writelines('\n'+_SEP_star)
133        out.writelines('     Atomic mass\n'+_SEP_minus)        
134        for i, X in enumerate(sp):
135            out.writelines('   '+f'{masses[X]:18.16f}'+'     ')
136            if (i%3 == 2) and i<len(sp)-1:
137                out.writelines('\n')
138        out.writelines('\n'+_SEP_star)
139        out.writelines('     The numbers of basis sets per atom type\n'+_SEP_minus+'    ')
140        for i, _ in enumerate(sp):
141            out.writelines('     1')
142            if (i%3 == 2) and i<len(sp)-1:
143                out.writelines('\n')
144        out.writelines('\n')
145        for i, el in enumerate(sp):
146            out.writelines(_SEP_star+f'     Basis set for {el}\n'+_SEP_minus+'          1      1\n')
147        
148        with open(output+'.tmp_wmaee_tmp', 'r') as src:
149            l = src.readline()
150            while l:
151                out.writelines(l)
152                l = src.readline()
153              
154    remove(output+'.tmp_wmaee_tmp')
155    
156    
157    
158def _read_table_pattern(filename: str, header_pattern: str, row_pattern: str,
159        footer_pattern: str, postprocess: Callable = str, first_one_only: bool = False,
160    ) -> list:
161    """
162    Parse table-like data from a file. A table comprises three parts: header,
163    main body, and footer. All the data matching the "row pattern" in the main body
164    will be returned.
165
166    Parameters
167    ----------
168    filename : str
169        The name of the file to read the table from.
170    header_pattern : str
171        The regular expression pattern that matches the table header.
172        This pattern should match all the text immediately before the main body of the table.
173        For multiple sections table, match the text until the section of interest.
174        MULTILINE and DOTALL options are enforced, so the "." meta-character will also match "\n" in this section.
175    row_pattern : str
176        The regular expression that matches a single line in the table. Capture interested fields using regular expression groups.
177    footer_pattern : str
178        The regular expression that matches the end of the table, e.g., a long dash line.
179    postprocess : Callable, optional
180        A post-processing function to convert all matches. Defaults to str, i.e., no change.
181    first_one_only : bool, optional
182        Only the first occurrence of the table will be parsed and the parsing procedure will stop. 
183        The enclosing list will be removed, i.e., only a single table will be returned. Incompatible with last_one_only.
184
185    Returns
186    -------
187    List[Union[List, Dict]]
188        A list of tables. A table is a list of rows. A row is either a list of attribute values (if capturing groups are defined without names in row_pattern), 
189        or a dictionary (if named capturing groups are defined by row_pattern).
190
191    Notes
192    -----
193    This function is adapted from pymatgen.io.vasp.outputs.py.
194    """
195    
196    from monty.io import zopen
197    import re
198    
199    with zopen(filename, mode='rt') as file:
200        text = file.read()
201    table_pattern_text = header_pattern + r"\s*^(?P<table_body>(?:\s+" + row_pattern + r")+)\s+" + footer_pattern
202    table_pattern = re.compile(table_pattern_text, re.MULTILINE | re.DOTALL)
203    rp = re.compile(row_pattern)
204    tables: list[list] = []
205    for mt in table_pattern.finditer(text):
206        table_body_text = mt.group("table_body")
207        table_contents = []
208        for line in table_body_text.split("\n"):
209            ml = rp.search(line)
210            # Skip empty lines
211            if not ml:
212                continue
213            d = ml.groupdict()
214            if len(d) > 0:
215                processed_line: dict | list = {k: postprocess(v) for k, v in d.items()}
216            else:
217                processed_line = [postprocess(v) for v in ml.groups()]
218            table_contents.append(processed_line)
219        tables.append(table_contents)
220        if first_one_only:
221            break
222    return tables
223    
224    
225
226def generate_ML_AB(input: str = 'OUTCAR', input_type: str = 'OUTCAR', output: str = 'ML_AB', 
227                   overwrite: bool = False, verbose: bool = True, system_name: str = 'parsed AIMD') -> None:   
228    """
229    Generate ML_AB formatted output from VASP calculation outputs.
230
231    Parameters
232    ----------
233    input : str, optional
234        Path to the input file, by default 'OUTCAR'
235    input_type : str, optional
236        Type of the input file, either 'OUTCAR' or 'vasprun', by default 'OUTCAR'
237    output : str, optional
238        Path to the output file, by default 'ML_AB'
239    overwrite : bool, optional
240        Flag to overwrite the existing output file, by default False
241    verbose : bool, optional
242        Flag to print detailed output, by default True
243    system_name : str, optional
244        Name of the system, by default 'parsed AIMD'
245
246    Raises
247    ------
248    Exception
249        If input_type is not 'OUTCAR' or 'vasprun'
250    FileExistsError
251        If the input file does not exist
252    """
253    
254    from os.path import exists
255    from monty.re import regrep
256    import numpy as np
257    
258     # Validate input_type
259    if input_type not in ['OUTCAR', 'vasprun']:
260        raise Exception(f'Unknown input file type, allowed values are `OUTCAR` or `vasprun`')
261
262    # Check if input file exists
263    if not exists(input):
264        raise FileExistsError(f'Input file {input} doesn\'t exist in the current folder.')
265
266    # Read species, numbers, and masses from the input file
267    if input_type=='vasprun':
268        chemistry = _read_table_pattern(
269            filename=input,            
270            header_pattern=r"\s*<field type=\"string\">pseudopotential<\/field>\s*<set>",
271            row_pattern=r"\s*<rc><c>\s*(\d*)<\/c><c>\s*(\w*)\s*<\/c><c>\s*(\d*.\d*)<\/c><c>.*",
272            footer_pattern=r"\s*<\/set>",
273            first_one_only=True
274        )[0]
275        sp = list()
276        num_sp = list()
277        masses = list()
278        for n, s, m in chemistry:
279            sp.append(s)
280            num_sp.append(int(n))
281            masses.append(float(m))
282    if input_type=='OUTCAR':
283        sp = regrep(
284            input, 
285            patterns=dict(species='TITEL\s+=\s\w+\s(\w+)')
286            )
287        sp = [x[0][0] for x in sp['species']]
288    if verbose:
289        print('found species: ', sp)
290    if input_type=='OUTCAR':
291        num_sp = regrep(
292            input,
293            patterns=dict(num_species='ions per type\s+=\s+'+''.join(['(\d+)\s*']*len(sp))),
294            terminate_on_match=True,
295            postprocess=lambda x: int(x)
296            )
297        # print('ions per type\s+=\s+'+''.join(['(\d+)\s*']*len(sp)))
298        num_sp = num_sp['num_species'][0][0]
299    if verbose:
300        print('num species: ', num_sp)
301    if input_type=='OUTCAR':
302        masses = regrep(
303            input,
304            patterns=dict(masses='POMASS\s+=\s+'+''.join(['(\d+.\d+)\s*']*len(sp))),
305            terminate_on_match=True,
306            postprocess=lambda x: float(x)
307            )
308        masses = masses['masses'][0][0]
309    if verbose:
310        print('masses: ', masses)
311        
312    # Read lattice vectors from the input file
313    if input_type == 'OUTCAR':
314        lattices = _read_table_pattern(
315            filename=input,
316            header_pattern=r"\sdirect lattice vectors\s+reciprocal lattice vectors",
317            row_pattern=r"\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+[+-]?\d+\.\d+\s+[+-]?\d+\.\d+\s+[+-]?\d+\.\d+",
318            footer_pattern=r"\n\s+length of vectors",
319            postprocess=lambda x: float(x)
320        )[2:]  # Skip initial lattice
321    else:
322        lattices = _read_table_pattern(
323            filename=input,
324            header_pattern=r"\s*<varray name=\"basis\"\s*>",
325            row_pattern=r"\s*<v>\s*([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s<\/v>\s*",
326            footer_pattern=r"\s*<\/varray>",
327            postprocess=lambda x: float(x)
328        )[2:]  # Skip initial lattice
329    if len(lattices[-1]) != 3:
330        lattices = lattices[:-1]  # Remove invalid lattice
331    if verbose:
332        print('found lattices: ', len(lattices))
333        
334    # Read energies and stresses from the input fi
335    if input_type=='OUTCAR':
336        energies_stresses = regrep(
337            input,
338            patterns=dict(
339                energies='free  energy\s+TOTEN\s+=\s+([+-]?\d+.\d+)\seV',
340                stresses='in kB\s+([+-]?\d+.\d+)\s+([+-]?\d+.\d+)\s+([+-]?\d+.\d+)\s+([+-]?\d+.\d+)\s+([+-]?\d+.\d+)\s+([+-]?\d+.\d+)'
341                ),
342            terminate_on_match=False,
343            postprocess=lambda x: float(x)
344            )
345        energies = [e[0][0] for e in energies_stresses['energies']]
346        stresses = [s[0] for s in energies_stresses['stresses']]
347    else:
348        energies = _read_table_pattern(
349            filename=input,
350            header_pattern=r"\s*<energy>",
351            row_pattern=r"\s*<i name=\"\w*\">\s*([+-]?\d+\.\d+)\s+<\/i>\s*",
352            footer_pattern=r"\s*<i name=\"kinetic\">",
353            postprocess=lambda x: float(x))        
354        for i, e in enumerate(energies):
355            energies[i] = e[0][0]
356        stresses = _read_table_pattern(
357            filename=input,
358            header_pattern=r"\s*<varray name=\"stress\"\s*>",
359            row_pattern=r"\s*<v>\s*([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s<\/v>\s*",
360            footer_pattern=r"\s*<\/varray>",
361            postprocess=lambda x: float(x))
362        for i, s in enumerate(stresses):
363            stresses[i] = [s[0][0], s[1][1], s[2][2], s[0][1], s[0][2], s[1][2]]
364    if verbose:
365        print('found energies: ', len(energies))
366        print('found stresses: ', len(stresses))
367        
368    # Read positions and forces from the input file
369    if input_type=='OUTCAR':
370        positions_and_forces = _read_table_pattern(
371            filename=input,
372            header_pattern=r"\sPOSITION\s+TOTAL-FORCE \(eV/Angst\)\n\s-+",
373            row_pattern=r"\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)",
374            footer_pattern=r"\s--+",
375            postprocess=lambda x: float(x))
376        if len(positions_and_forces[-1]) != sum(num_sp):
377            # invalid number of positions and forces of the last structure, remove it
378            positions_and_forces = positions_and_forces[:-1]
379        positions = [[x[:3] for x in ionic_step] for ionic_step in positions_and_forces]
380        forces = [[x[3:] for x in ionic_step] for ionic_step in positions_and_forces]
381    else:
382        positions = _read_table_pattern(
383            filename=input,
384            header_pattern=r"\s*<varray name=\"positions\"\s*>",
385            row_pattern=r"\s*<v>\s*([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s<\/v>\s*",
386            footer_pattern=r"\s*<\/varray>",
387            postprocess=lambda x: float(x))[2:] # the first lattices are printed during initialization, not after ionic steps
388        forces = _read_table_pattern(
389            filename=input,
390            header_pattern=r"\s*<varray name=\"forces\"\s*>",
391            row_pattern=r"\s*<v>\s*([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s+([+-]?\d+\.\d+)\s<\/v>\s*",
392            footer_pattern=r"\s*<\/varray>",
393            postprocess=lambda x: float(x))
394    if verbose:
395        print('found positions: ', len(positions))
396        print('found forces: ', len(forces))
397    
398    # Take the minimum number of data points across all lists    
399    take_only = min(len(lattices), len(positions), len(energies), len(stresses))
400    lattices = lattices[:take_only]
401    energies = energies[:take_only]
402    positions = positions[:take_only]
403    
404    # Write ML_AB file
405    out = open(output, 'w')
406    _SEP_star  = "**************************************************\n"
407    _SEP_minus = "--------------------------------------------------\n"
408    _SEP_equal = "==================================================\n"
409    out.writelines(' 1.0 Version\n'+_SEP_star)
410    out.writelines('     The number of configurations\n'+_SEP_minus+f'{len(lattices):11d}\n'+_SEP_star)
411    out.writelines('     The maximum number of atom type\n'+_SEP_minus+f'{len(num_sp):8d}\n'+_SEP_star)
412    out.writelines('     The atom types in the data file\n'+_SEP_minus+'     ')
413    for i, el in enumerate(sp):
414        out.writelines(el.ljust(3))
415        if (i%3 == 2) and i<len(sp)-1:
416            out.writelines('\n     ')
417    out.writelines('\n'+_SEP_star)
418    out.writelines('     The maximum number of atoms per system\n'+_SEP_minus+f'{sum(num_sp):8d}\n'+_SEP_star)
419    out.writelines('     The maximum number of atoms per atom type\n'+_SEP_minus+f'{max(num_sp):8d}\n'+_SEP_star)
420    out.writelines('     Reference atomic energy (eV)\n'+_SEP_minus)
421    for i, _ in enumerate(sp):
422        out.writelines('   '+f'{0:18.16f}'+'     ') # dummy energy
423        if (i%3 == 2) and i<len(sp)-1:
424            out.writelines('\n')
425    out.writelines('\n'+_SEP_star)
426    out.writelines('     Atomic mass\n'+_SEP_minus)
427    for i, m in enumerate(masses):
428        out.writelines('   '+f'{m:18.16f}'+'     ')
429        if (i%3 == 2) and i<len(sp)-1:
430            out.writelines('\n')
431    out.writelines('\n'+_SEP_star)
432    out.writelines('     The numbers of basis sets per atom type\n'+_SEP_minus+'    ')
433    for i, _ in enumerate(sp):
434        out.writelines('     1')
435        if (i%3 == 2) and i<len(sp)-1:
436            out.writelines('\n')
437    out.writelines('\n')
438    for i, el in enumerate(sp):
439        out.writelines(_SEP_star+f'     Basis set for {el}\n'+_SEP_minus+'          1      1\n')
440    for i in range(len(lattices)):
441        out.writelines(_SEP_star)
442        out.writelines(f'     Configuration num.{i+1:7d}\n'+_SEP_equal)
443        out.writelines('     System name\n'+_SEP_minus+'     '+system_name+'\n'+_SEP_equal)
444        out.writelines('     The number of atom types\n'+_SEP_minus+f'{len(num_sp):8d}\n'+_SEP_equal)
445        out.writelines('     The number of atoms\n'+_SEP_minus+f'{sum(num_sp):8d}\n'+_SEP_star)
446        out.writelines('     Atom types and atom numbers\n'+_SEP_minus)
447        for el, n in zip(sp, num_sp):
448            out.writelines('     '+el.ljust(3)+f'{n:6d}\n')
449        out.writelines(_SEP_equal)
450        out.writelines('     Primitive lattice vectors (ang.)\n'+_SEP_minus)
451        lattice = lattices[i]
452        for l in lattice:
453            out.writelines('  '+'  '.join([f'{x:18.16f}' for x in l])+'\n')
454        out.writelines(_SEP_equal)
455        out.writelines('     Atomic positions (ang.)\n'+_SEP_minus)
456        if input_type == 'OUTCAR':
457            for l in positions[i]:
458                out.writelines('  '+'  '.join([f'{x:18.16f}' for x in l])+'\n')
459        else:
460            # need to convert to Cartesian coordinates for vasprun.xml inputs
461            for l in positions[i]:
462                out.writelines('  '+'  '.join([f'{x:18.16f}' for x in np.matmul(l, np.array(lattice))])+'\n')
463        out.writelines(_SEP_equal)
464        out.writelines('     Total energy (eV)\n'+_SEP_minus+f'  {energies[i]}\n'+_SEP_equal)
465        out.writelines('     Forces (eV ang.^-1)\n'+_SEP_minus)
466        for l in forces[i]:
467            out.writelines('  '+'  '.join([f'{x:18.16f}' for x in l])+'\n')
468        out.writelines(_SEP_equal)
469        out.writelines('     Stress (kbar)\n'+_SEP_minus+'     XX YY ZZ\n'+_SEP_minus)
470        out.writelines('  '+'  '.join([f'{x:18.16f}' for x in stresses[i][:3]])+'\n'+_SEP_minus)
471        out.writelines('     XY YZ ZX\n'+_SEP_minus)
472        out.writelines('  '+'  '.join([f'{x:18.16f}' for x in stresses[i][3:]])+'\n')