cpp/python/helpers.py

from __future__ import annotations

import pathlib
import re
import subprocess
from itertools import product


def prepare_rw_parameter(parameter: dict) -> list:
    """
    Converts a dictionary of iterables to list of dictionaries

    Example:
    If Input is {'a': [1, 2, 3], 'b' = [4, 5]}, output is cartesian product of dictionary values, i.e.,
    [{'a': 1, 'b': 4}, {'a': 1, 'b': 5}, {'a': 2, 'b': 4}, {'a': 2, 'b': 5}, {'a': 3, 'b': 4}, {'a': 3, 'b': 5}]

    :param parameter: dictionary of list values
    :return: list of dictionaries
    """
    output = {}
    for k, v in parameter.items():
        if isinstance(v, (float, int)):
            v = [v]

        output[k] = v

    output = list(dict(zip(parameter.keys(), step)) for step in product(*output.values()))

    return output


def run_sims(
        motion: str,
        distribution: str,
        ste: bool = True,
        spectrum: bool = False,
        exec_file: str = './rwsim',
        config_file: str = './config.txt',
        **kwargs
) -> None:

    # set positional arguments
    arguments = [exec_file, config_file, motion, distribution]

    if ste:
        arguments += ['--ste']
    if spectrum:
        arguments += ['--spectrum']

    # add optional parameters that overwrite those given by config file
    for k, v in kwargs.items():
        arguments += [f'-{k.upper()}', f'{v}']

    subprocess.run(arguments)


def find_config_file(motion: str, distribution: str, var_params: dict) -> pathlib.Path:
    # TODO handle situation if multiple files fit
    p_file = None
    if var_params:
        var_string = '|'.join(([f'{k}={v:1.6e}' for (k, v) in var_params.items()])).replace('.', '\.').replace('+', '\+')
        pattern = re.compile(var_string)
        for p_file in pathlib.Path('.').glob('*_parameter.txt'):
            if len(re.findall(pattern, str(p_file))) == len(var_params) and re.search(f'{motion}_{distribution}', str(p_file)):
                return p_file
        raise ValueError(f'No parameter file found for {motion}, {distribution}, {var_params}')
    else:
        for p_file in pathlib.Path('.').glob('*_parameter.txt'):
            if re.search(f'{motion}_{distribution}', str(p_file)):
                return p_file
        raise ValueError(f'No parameter file found for {motion}, {distribution}, {var_params}')




def read_parameter_file(path: str | pathlib.Path) -> dict[str, float]:
    path = pathlib.Path(path)
    if not path.exists():
        raise ValueError(f"No parameter file found at {path}")

    parameter_dict = {}
    with path.open('r') as f:
        for line in f.readlines():
            k, v = line.split('=')
            parameter_dict[k] = float(v)

            k, v = line.split('=')
    return parameter_dict