LaTeX AutoImager

latex.py

import os
import math
import shutil
import argparse

def buildTEX(
    lc_files: list, 
    output_tex_file: str, 
    nr: int = 4, 
    nc: int = 4,
    nrows_first_page: int = 2):
        pages: int = math.ceil(len(lc_files)/(nr*nc))
        
        with open(output_tex_file,'w') as file:
            # Document class and usepackages
            file.write('\\documentclass[12pt, a4paper, twoside, openright]{book}\n')
            file.write('\\usepackage{graphicx}\n')
            file.write('\\usepackage{caption}\n')
            file.write('\\usepackage{subcaption}\n')
            file.write('\\usepackage[font=scriptsize]{caption}\n')
            file.write('\\begin{document}\n')

            # Define the appendix and the chapter
            file.write('\t\\appendix\n')
            file.write('\t\\chapter{\\textit{Swift}-XRT light curves fitted with JetFit}\n')
            file.write('\t\\label{appA}\n')

            for page in range(pages):
                if page:
                    file.write('\t\\clearpage\n')
                    file.write('\t\\newpage\n')
                    rows = nr
                else:
                    rows = nrows_first_page
                for row in range(rows):
                    file.write('\t\\begin{figure}[h!]\n')
                    file.write('\t\t\\centering\n')
                    for column in range(nc):
                        img_elm = page*nr*nc + row*nc + column
                        if img_elm > len(lc_files)-1:
                            continue
                        file.write('\t\t\\begin{subfigure}[b]{' + str((1*0.95)/nc) + '\\textwidth}\n')
                        file.write('\t\t\t\\centering\n')
                        file.write('\t\t\t\\includegraphics[width=\\textwidth]' + '{' + str(lc_files[img_elm]) + '}\n')
                        file.write('\t\t\t\\captionsetup{labelformat=empty}\n')
                        fig_description = f"GRB{lc_files[img_elm][lc_files[img_elm].rfind('/')+1:].split('_')[2]}"
                        file.write('\t\t\t\\caption{' + fig_description + '}\n')
                        file.write('\t\t\\end{subfigure}\n')
                        if column != nc-1:
                            file.write('\t\t\\hfill\n')
                    file.write('\t\\end{figure}\n')
            file.write('\\end{document}\n')
    

def get_images(input_dir: str) -> list:
    dirs = [os.path.join(input_dir, dir) for dir in os.listdir(input_dir) if os.path.isdir(os.path.join(input_dir, dir))]
    if len(dirs):
        lc_files = [os.path.join(dir, elm) for dir in sorted(dirs) for elm in os.listdir(dir) if elm.startswith('lc_') and elm.endswith('.png')]
    else:
        lc_files = [os.path.join(input_dir, elm) for elm in sorted(os.listdir(input_dir)) if elm.startswith('lc_') and elm.endswith('.png')]
    print(f'\n[INFO] Total number of images: {len(lc_files)}')
    return lc_files

def check_files(input_dir: str):
    dirs = [os.path.join(input_dir, dir) for dir in os.listdir(input_dir) if os.path.isdir(os.path.join(input_dir, dir))]
    lc_files = [os.path.join(dir, elm) for dir in sorted(dirs) for elm in os.listdir(dir) if elm.startswith('lc_') and elm.endswith('.png')]
    csv_files = [os.path.join(dir, elm) for dir in sorted(dirs) for elm in os.listdir(dir) if elm.startswith('new_sum_') and elm.endswith('.csv')]
    for lc_curve in lc_files:
        match = False
        for csv_file in csv_files:
            csv_name = csv_file[csv_file.rfind('/')+1:].split('_')[3][:-4]
            if csv_name in lc_curve:
                match = True
                break
        if not match:
            print(f'[WARNING] Non-matched GRB: {lc_curve}')

def export_files(images: list, output_dir: str = 'pngs', images_per_subfolder: int = 40):
    if os.path.isdir(output_dir):
        shutil.rmtree(output_dir)
    os.mkdir(output_dir)
    
    tmp_output_dir = output_dir
    for idx, image in enumerate(images):
        if images_per_subfolder:
            tmp_output_dir: str = os.path.join(output_dir, str((idx+1)//images_per_subfolder+1))
            if not os.path.isdir(tmp_output_dir):
                os.mkdir(tmp_output_dir)
        else:
            shutil.copy2(image, tmp_output_dir)
    print(f'[INFO] {len(images)} images have been exported to output folder: [{output_dir}]')

def main():
    parser = argparse.ArgumentParser(description='Utils to add light curves to latex document')
    parser.add_argument('--input', '-i', type=str, dest='input', help='input lightcurtve folder')
    parser.add_argument('--output', '-o', type=str, dest='output', help='output tex file')
    parser.add_argument('--verbose', '-v', default=False, dest='verbose', action='store_true', help='verbose mode')
    parser.add_argument('--check_files', '-c', default=False, dest='check_files', action='store_true', help='Check input files')
    parser.add_argument('--export_pngs', '-e', default=False, dest='export_files', action='store_true', help='Check input files')
    args = parser.parse_args()
    
    if os.path.isdir(args.input):
        files = get_images(args.input)
        if args.check_files:
            check_files(args.input)
        if args.export_files:
            export_files(files)
    else:
        print(f'[ERROR] Missing input dir... [{args.input}]')
        exit(100)

    buildTEX(files, args.output)
if __name__ == '__main__':
    main()