fluidxr source

import xml.etree.ElementTree as ET
from collections import namedtuple
from pathlib import Path

import pandas as pd

def fluidx(path: str, rsattrs: list = None, suattrs: list = None, filename: str = 'domain') -> None:
    # Variables = paramètres de la fonction
    # Invariables
    RS_file: str = 'RS'  # Description RS
    RSRS_file: str = 'RSRS'  # Connectivité RS à RS
    SU_file: str = 'SU'  # Description SU
    SURS_file: str = 'SURS'  # Connectivité SU à RS
    
    # DATA IMPORT ===============================================================================================
    path = Path(path)
    RS = pd.read_csv(path.joinpath(f'{RS_file}.csv'))[['ID', *rsattrs]]
    RSRS = pd.read_csv(path.joinpath(f'{RSRS_file}.csv'))
    SU = pd.read_csv(path.joinpath(f'{SU_file}.csv'))[['ID', *suattrs]]
    SURS = pd.read_csv(path.joinpath(f'{SURS_file}.csv'))
    # DATA PROCESSING ===========================================================================================
    # Ajout de la colonne PO (process order) en correspondance avec les RS cibles des SU et suppression des RS sans PO
    SURS = pd.merge(RSRS, SURS, left_on='FROM', right_on='TO', how='right', suffixes=('_RSRS', '')).dropna()[['FROM', 'TO', 'PO']]
    
    # Transformation de la connectivité en tuples
    def f(df: pd.DataFrame, src_col: str,  target_col: str) -> list:
        result = [] 
        for r in range(len(df)):
            src, target, po = df.iloc[r].values
            result.append([src_col, str(int(src)), str(int(po)), target_col, str(int(target))])
        return result
    
    Connections = namedtuple('Connections', 'src_class src_id src_po target_class target_id')
    CS = [Connections(*x) for x in f(SURS, 'SU', 'RS') + f(RSRS, 'RS', 'RS')]
    
    # Transformation des attributs en listes de str
    def g(df: pd.DataFrame):  # Returns tuple(str, list)
        cols = ';'.join(df.columns[1:])  # Remove id/ID column from the colorder
        result = []
        for i in range(len(df)):
            r = df.iloc[i].values
            r = [int(r[0]), *r[1:]]
            result.append(r" ".join(str(v) for v in r))
        return cols, result
    
    cols_RS, result_RS = g(RS)
    cols_SU, result_SU = g(SU)
    
    # EXPORT TO FLUIDX ==========================================================================================
    root = ET.Element("openfluid")
    domain = ET.SubElement(root, "domain")
    ET.SubElement(domain, "attributes", unitsclass="RS", colorder=cols_RS).text = f'\n{"\n".join(result_RS)}\n  '
    ET.SubElement(domain, "attributes", unitsclass="SU", colorder=cols_SU).text = f'\n{"\n".join(result_SU)}\n  '
    ET.SubElement(domain, "calendar").text = ''
    definition = ET.SubElement(domain, "definition")
    for con in CS:
        unit = ET.SubElement(definition, "unit", CLASS=con.src_class, ID=con.src_id, pcsorder=con.src_po)
        ET.SubElement(unit, "to", CLASS=con.target_class, ID=con.target_id)
    tree = ET.ElementTree(root)
    ET.indent(tree, space=' ', level=0)
    with open(f'{filename}.fluidx', 'wb') as f:
        f.write(b'<?xml version="1.0" standalone="yes"?>\n')
        tree.write(f, xml_declaration=False)
    # Fix python not accepting class as a kwarg
    with open(f'{filename}.fluidx', 'r') as f:
      data = f.read()
    data = data.replace('CLASS', 'class')
    with open(f'{filename}.fluidx', 'w') as f:
      f.write(data)
    
if __name__ == '__main__':
    import argparse
    # Declare cli
    parser = argparse.ArgumentParser(prog='fluidxr',
                                     description='Convert connectivity/attributes CSV files into an OpenFLUID (FLuidx) domain')
    parser.add_argument('path', help='Path to the data dir')
    parser.add_argument('--filename', help='Name of the output file', default='domain')
    parser.add_argument('--rsattrs', nargs='+', help='Names of the columns to keep in the RS attributes file')
    parser.add_argument('--suattrs', nargs='+', help='Names of the columns to keep in the Su attributes file')
    args = parser.parse_args()
    # If only 1 column is specified, type is not a list but a str, thus it must be cast as a list using [] (using list() would separate each letter)
    if isinstance((rsattrs := args.rsattrs), str):
        rsattrs = [args.rsattrs]
    if isinstance((suattrs := args.suattrs), str):
        suattrs = [args.suattrs]
    # Execute
    fluidx(args.path, rsattrs, suattrs, args.filename)