get_indices_of_model_parameters_in_ode.py

from argparse import ArgumentParser
from pathlib import Path
import json
import logging
import re

from loggingUtility import set_log_level, add_log_level_to_argument_parser

from src.read_dataset import read_datasets_from_directory
from src.Dataset import ModelType, Dataset
from src.read_pcomp_model import read_pcomp_model


def find_whole_word(targetString: str, sourceString: str, ignoreCase: bool = True) -> bool:
    if ignoreCase:
        flag = re.IGNORECASE
    else:
        flag: re.RegexFlag = 0

    if re.search(r"\b" + re.escape(targetString) + r"\b", sourceString, flags=flag):
        return True
    
    return False


def any_function_contains_parameter(functions: list[str], parameterName: str) -> bool:
    return any([find_whole_word(parameterName, function, ignoreCase=True) for function in functions])


def get_indices_of_model_parameters_in_ode(modelFileName: Path, dataset: Dataset) -> list[tuple[int, str]]:
    pcompModelInfos = read_pcomp_model(modelFileName)

    numberOfTimeVariables = 1 
    numberOfConcentrationVariables = 1 if dataset.numberOfConcentrations > 0 else 0
    expectedNumberOfVariables = dataset.numberOfModelParameters + dataset.numberOfOdes + numberOfConcentrationVariables + numberOfTimeVariables

    expectedNumberOfFunctions = 2 * dataset.numberOfOdes + dataset.dimensionOfFittingFunction + dataset.totalNumberOfConstraints

    if len(pcompModelInfos.functions) != expectedNumberOfFunctions:
        logging.warning(f"{dataset.name}: wrong number of functions detected, expected: {expectedNumberOfFunctions}, observed: {len(pcompModelInfos.functions)}. Return no detected parameters.")
        return []
    elif len(pcompModelInfos.variables) != expectedNumberOfVariables:
        logging.warning(f"{dataset.name}: wrong number of variables detected, expected: {expectedNumberOfVariables}, observed: {len(pcompModelInfos.variables)}. Return no detected parameters")
        return []
    else:
        modelParameterNames = pcompModelInfos.variables[0:dataset.numberOfModelParameters]

        odeFunctions = pcompModelInfos.functions[0:dataset.numberOfOdes]
        modelParametersInOde = [(index, parameter) for index, parameter in enumerate(modelParameterNames) 
                            if any_function_contains_parameter(odeFunctions, parameter)]

    logging.info(f"Found parameters in {dataset.name}: {', '.join([f'({index}: {param})' for (index, param) in modelParametersInOde])}")

    return modelParametersInOde


def read_broken_models() -> list[str]:
    fileContent = Path("broken_models.txt").read_text()
    return [line for line in fileContent.splitlines() if line and not line.startswith("#")]


def main() -> int:
    parser = ArgumentParser(description='Read indices of model parameters in ODE from PCOMP file and write to JSON')
    parser.add_argument('datasetFolder', type=str, help='Path to dataset folder')
    parser.add_argument('modelFolder', type=str, help='Path to model folder')
    parser.add_argument('outputFilePath', type=str, help='Path to output file')
    add_log_level_to_argument_parser(parser)
    args = parser.parse_args()
    set_log_level(args.logLevel)

    datasets = read_datasets_from_directory(args.datasetFolder)
    brokenModels = read_broken_models() 

    indicesOfModelParametersInOde: dict[str, list[int]] = {}

    for dataset in datasets:
        if not dataset.modelType == ModelType.OrdinaryDifferentialEquations:
            logging.info(f"Dataset {dataset.name} is not an ODE model. Skip it.")
            continue
        
        if dataset.name in brokenModels:
            logging.info(f"Dataset {dataset.name} is broken. Skip it.")
            continue

        correspondingModeFileName = Path(args.modelFolder) / f"{dataset.name}.fun"
        modelParametersInOde = get_indices_of_model_parameters_in_ode(correspondingModeFileName, dataset)

        indicesOfModelParametersInOde[dataset.name] = [index for (index, _) in modelParametersInOde]

    outputFolder = Path(args.outputFilePath).parent
    outputFolder.mkdir(parents=True, exist_ok=True)
    
    with Path(args.outputFilePath).open("w") as outfile:
        outfile.write(json.dumps(indicesOfModelParametersInOde, indent=4))

    return 0

if __name__ == "__main__":
    exit(main())