Source code for torx.equilibrium.io.xml_io_m

"""
Contains functionality to use equi XML files generated by pyequil [1].

[1] (https://gitlab.mpcdf.mpg.de/tal/pyequil)
"""
import numpy as np
import xml.etree.ElementTree as et
from pathlib import Path
from typing import Union



def read_xml_file(filepath: Union[str, Path]) -> dict:
    """Read equilibrium data from an XML file as generated by pyequil."""
    root = et.parse(filepath).getroot()
    equi = root.find("equilibrium")

    keys_paths = [
        ["r",   "equilibrium/Ri"],
        ["z",   "equilibrium/zj"],
        ["psi", "equilibrium/PFM"],
        ["op",  "equilibrium/singularities/op"],
        ["xp",  "equilibrium/singularities/xp"],
        ["xp2", "equilibrium/singularities/xp2"],
    ]
    data_dict = {
        key: _xml2np(root.find(path))
        for key, path in keys_paths
        if root.find(path) is not None
    }

    if root.find("pfcs") is not None:
        data_dict["divertor_polygon"] = _xml2np(root.find("pfcs"))

    data_dict["profiles"] = {
        data.tag: _xml2np(data) for data in root.find("equilibrium/profiles1D")
    }

    for key in ["command_line", "code_source"]:
        if root.find(key) is not None:
            data_dict[key] = root.find(key).text

    if "discharge" in root.attrib:
        data_dict["shot_number"] = int(root.attrib["discharge"])
    if "device" in root.attrib:
        data_dict["device"] = root.attrib["device"]    
    if "time" in root.attrib:
        data_dict["time"] = float(root.attrib["time"])

    equi_attrs = {key: float(equi.attrib[key]) for key in equi.attrib.keys()}
    data_dict.update(equi_attrs)

    return data_dict




def write_xml_file(filepath: Union[str, Path], content: dict):
    """Write magnetic equilibrium data to an XML file."""
    raise NotImplementedError("XML writer is not yet implemented")


def _xml2np(xml_input: et.Element, sep: str = ",") -> np.ndarray:
    """Interpret the numerical data in XML element as a numpy array."""
    try:
        keys = xml_input.attrib.keys()
        assert ("type" in keys) and ("dim" in keys)
        dtype = xml_input.attrib["type"]
        N = np.fromstring(xml_input.attrib["dim"], dtype=int, sep=sep)
    except AssertionError:
        raise Exception("xml-tag must contain type and dim information.")

    assert not np.all(N == 0), \
        f"scalar variable {xml_input.tag} cannot be converted to an array."

    if np.prod(N) == 0:
        return np.zeros(N)

    return np.fromstring(
        xml_input.text, sep=sep, dtype=dtype, count=np.prod(N)
    ).reshape(N)