Source code for schrodinger.application.matsci.qexsd.utils

"""
Helper functions/classes for qexsd.

Copyright Schrodinger, LLC. All rights reserved.
"""

import xml.etree.ElementTree as Etree
from os.path import splitext

from schrodinger.application.matsci.qexsd import qespresso

NORUN_WARNING = """!   === WARNING ===
!   This file is not used by periodic_dft_driver.py or <jobname>.sh
!   If you want to run this file, use runner.py
!   ======
"""


def xml2in(xml_fn, warning=False):
    """
    Convert XML file to .in QE format.

    :type xml_fn: str
    :param xml_fn: XML filename

    :param bool warning: Print warning if True

    :rtype: str
    :return: .in filename
    """

    # Based on the qexsd/tests/test_converters.py MIT license: LEGAL-616
    tree = Etree.parse(xml_fn)
    root = tree.getroot()
    elementName = root.tag.split('}')[-1]
    if elementName == 'espresso':
        xml_conf = qespresso.PwDocument()
    elif elementName == 'nebRun':
        xml_conf = qespresso.NebDocument()
    elif elementName == 'espressoph':
        xml_conf = qespresso.PhononDocument()
    elif elementName == 'tddfpt':
        xml_conf = qespresso.TdDocument()
    elif elementName == 'spectrumIn':
        xml_conf = qespresso.SpectrumDocument()
    else:
        raise ValueError('Could not find correct XML in %s' % xml_fn)

    # Free memory
    root = None
    tree = None

    xml_conf.read(xml_fn)
    qe_in = xml_conf.get_qe_input()

    input_fn_name, input_fn_ext = splitext(xml_fn)
    outfile = input_fn_name + '.in'

    with open(outfile, 'w') as out_fh:
        if warning:
            out_fh.write(NORUN_WARNING)

        out_fh.write(qe_in)
        out_fh.write('\n')

    return outfile