correct warp dose

fix_dose_tomostar.py

from pathlib import PureWindowsPath, Path
import mdocfile
import starfile
import pandas as pd
import typer

unbork = typer.Typer()


@unbork.command(no_args_is_help=True)
def correct_warp_dose(mdoc: Path = typer.Option(..., "-m", "--m", help="Mdoc file"),
                      tomostar: Path = typer.Option(..., "-t", "--t", help="Tomostar file"),
                      dose_per_tilt: float = typer.Option(..., "-d", "--d", help="dose per tilt file")):
    zero_tilt_dose = round((dose_per_tilt / 2), 2)

    mdoc_df = mdocfile.read(mdoc)
    tomostar_df = starfile.read(tomostar)
    mdoc_frames = mdoc_df.sort_values(by="DateTime")["SubFramePath"]
    mdoc_tilts = mdoc_df.sort_values(by="DateTime")["TiltAngle"]
    mdoc_frame_names = [PureWindowsPath(x).name for x in mdoc_frames]
    tilt_dose = [zero_tilt_dose]
    dose = 0
    count = len(mdoc_tilts) - 1
    i = 0
    while i < count:
        i += 1
        dose = dose + dose_per_tilt
        tilt_dose.append(dose)
    mdoc_frame_names = pd.Series(mdoc_frame_names, name="SubFramePath")
    tilt_dose = pd.Series(tilt_dose, name="AccumulatedDose")
    accumulated_dose = pd.concat((mdoc_frame_names, mdoc_tilts, tilt_dose), axis=1)
    merged_df = pd.merge(tomostar_df, accumulated_dose, left_on='wrpMovieName', right_on='SubFramePath', how='left')
    merged_df['wrpDose'] = merged_df['AccumulatedDose'].fillna(merged_df['wrpDose'])
    merged_df.drop(['SubFramePath', 'TiltAngle', 'AccumulatedDose'], axis=1, inplace=True)
    starfile.write(merged_df, tomostar, fmt="0.3f")


if __name__ == "__main__":
    unbork()