Warp XML to SG wedgelist

warp_xml2_sg_wedgelist.py

                                                                                                                                                                                                                                                                                                                                                                                                                                                                                               1,1           Top
from pathlib import Path
import starfile
import pandas as pd
from lxml import etree
import typer


def flatten(t):
    return [item for sublist in t for item in sublist]


app = typer.Typer(add_completion=False)


@app.command(no_args_is_help=True)
def main(px:float = typer.Option(1.0, '-p', '--pixel_size', help="tomo unbinned pixel size"),
         sx: int = typer.Option( ..., '-x', '--size_x', help="tomo unbinned dimX"),
         sy: int = typer.Option( ..., '-y', '--size_y', help="tomo unbinned dimY"),
         sz: int = typer.Option( ..., '-z', '--size_z', help="tomo unbinned dimZ"),
         xmld: str = typer.Option(None, '-d', '--xml_dir', help="Path to folder containing xml files"),
         xmlf: str = typer.Option( None, '-f', '--xml_file', help="Path to specific TS xml file"),
         outwedge: str = typer.Option(..., '-o', '--out', help="Path to where wedgelist file should be stored")):

    # Validate input: either xmlf or xmld must be provided, but not both
    if bool(xmlf) == bool(xmld):
        raise typer.BadParameter("You must provide either --xml_file or --xml_dir, but not both.")

    if xmlf:
        xmlfiles = [Path(xmlf)]
    else:
        xmlfiles = sorted(list(Path(f"{xmld}/").glob("*.xml")))

    if not xmlfiles or not xmlfiles[0].is_file():
        print(f"No XML files found in {xmld}")
        return

    print("Found XML files. Working...")
    tilt_names = []
    tilt_angles = []
    tilt_defocus = []
    tilt_dose = []
    tomo_num = []
    pixel_size = []
    tomo_x = []
    tomo_y = []
    tomo_z = []
    tomo_z_shift = []
    voltage = []
    amp_contrast = []
    cs = []

    for index, xmlfile in enumerate(xmlfiles):
        tree = etree.parse(xmlfile)
        tilt_name_nodes = tree.findall(".//MoviePath")
        tilt_angle_nodes = tree.findall(".//Angles")
        tilt_defocus_nodes = tree.findall(".//GridCTF/Node")
        tilt_dose_nodes = tree.findall(".//Dose")

        for ts_name in tilt_name_nodes:
            text = ts_name.text
            names = text.split('\n')
            tilt_names.append(names)

        for ts_angle in tilt_angle_nodes:
            text = ts_angle.text
            angles = text.split('\n')
            tilt_angles.append(angles)

        ts_defocus_values = []
        for node in tilt_defocus_nodes:
            ts_defocus_values.append(node.attrib['Value'])
        tilt_defocus.append(ts_defocus_values)

        for ts_dose in tilt_dose_nodes:
            text = ts_dose.text
            ts_cum_dose = text.split('\n')
            tilt_dose.append(ts_cum_dose)

        # This section sets up the tomo number counter
        ts_length = len(angles)
        counter = index + 1
        num = [counter] * ts_length
        tomo_num.append(num)

        # This section takes in some constants and just tiles out to be the same
        # length as the tilt series
        ts_pixel_size = [px] * ts_length
        pixel_size.append(ts_pixel_size)

        ts_tomo_x = [sx] * ts_length
        tomo_x.append(ts_tomo_x)

        ts_tomo_y = [sy] * ts_length
        tomo_y.append(ts_tomo_y)

        ts_tomo_z = [sz] * ts_length
        tomo_z.append(ts_tomo_z)

        ts_z_shift = [0] * ts_length
        tomo_z_shift.append(ts_z_shift)

        ts_voltage = [300] * ts_length
        voltage.append(ts_voltage)

        ts_amp_contrast = [0.1] * ts_length
        amp_contrast.append(ts_amp_contrast)

        ts_cs = [2.7] * ts_length
        cs.append(ts_cs)

    tomo_num = flatten(tomo_num)
    tilt_names = flatten(tilt_names)
    tilt_angles = flatten(tilt_angles)
    tilt_defocus = flatten(tilt_defocus)
    tilt_defocus = [float(i) for i in tilt_defocus]
    tilt_dose = flatten(tilt_dose)
    pixel_size = flatten(pixel_size)
    tomo_x = flatten(tomo_x)
    tomo_y = flatten(tomo_y)
    tomo_z = flatten(tomo_z)
    tomo_z_shift = flatten(tomo_z_shift)
    voltage = flatten(voltage)
    amp_contrast = flatten(amp_contrast)
    cs = flatten(cs)

    # Assemble a dataframe
    df = pd.DataFrame({
        'tomo_num': tomo_num,
        'pixelsize': pixel_size,
        'tomo_x': tomo_x,
        'tomo_y': tomo_y,
        'tomo_z': tomo_z,
        'z_shift': tomo_z_shift,
        'tilt_angle': tilt_angles,
        'defocus': tilt_defocus,
        'exposure': tilt_dose,
        'voltage': voltage,
        'amp_contrast': amp_contrast,
        'cs': cs,
        'tilt_name': tilt_names
    })

    # Save the dataframe to a star file
    df1 = df.drop(['tilt_name'], axis=1)
    starfile.write({"stopgap_wedgelist": df1}, f"{Path(outwedge)}/wedgelist.star", overwrite='True')


if __name__ == '__main__':
    app()(main)