MultiSimu.py

Go to the documentation of this file.

from multiprocessing import Pool
import pandas as pd
from tqdm import tqdm
import itertools
from copy import deepcopy
from typing import Callable, Dict, List, Any, Union, Iterable
import inspect
 
_pool = None
 
def _pool_is_usable():
    global _pool
    if _pool is None:
        return False
 
    procs = getattr(_pool, "_pool", None)
    if not procs:
        return False
 
    alive = False
    for p in procs:
        try:
            if p is not None and p.is_alive():
                alive = True
                break
        except (ValueError, AssertionError, OSError):
            continue
 
    return alive
 
 
class MultiSimu():
    r"""
    Initializes the MultiSimu instance for running multiple simulations with various parameters, conditions, and replicates.
 
    Parameters
    ----------
    : **simuRunner** : callable
        The function to run the simulation. It should accept the parameters and conditions as input.
 
    : **params** : dict or list of dict
        The parameters for the simulation. If a single dictionary is provided, it is wrapped in a list. Each dictionary represents a distinct set of parameters for a simulation run.
 
    : **replicat** : int, optional, default=1
        The number of replicates to run for each parameter and condition combination.
 
    : **conditions** : list, optional, default=[0]
        A list of conditions to apply to each simulation run.
 
    : **batch_size_level** : int or str, optional, default=None
         Determines the size of the batches returned by the iterator. It can be an integer specifying the number of batches, or a string indicating the batch level ('param', 'replicat', or 'condition').
 
    : **cacheSize** : int or None, optional, default=200
        The size of the cache for the number of simulations to pre-fetch in parallel.
 
    : **withTqdm** : bool, optional, default=False
        If True, progress bars will be displayed using `tqdm` to show the progress of parameter, condition, and replicate iterations.
 
    : **parallel** : bool, optional, default=True
        If True, simulations will be run in parallel using multiprocessing. If False, simulations will be run sequentially.
 
    : **autoIndex** : bool, optional, default=True
        If True, adds index columns (ID_PARAMETER, ID_CONDITION, ID_REPLICAT) to the output data for identification.
 
    : **autoConcat** : bool, optional, default=True
        If True, concatenates the results from different simulations into a single pandas DataFrame.
 
    Raises
    ------
    : **KeyError**
        - If `params` is neither a list of dictionaries nor a single dictionary.
        - If `conditions` is not a list.
 
    Example
    -------
    ```python
    def dataCatcherOnStep(simu):
        return simu.cells.countState()
 
    def run_simu(params,condition,replicat):
        steps = list(range(0,2501,10))
        MOI, confluence = condition
        params['input']['Scenario']['infection'] = [[0,MOI]]
        params['input']['Scenario']['initConfluence'] = confluence
        data = Simu(params).compute(dataCatcherOnStep,steps)
        data.index *= params['input']['Scenario']['dt']
        return data
 
 
    paramsTest = [params1,params2]
    conditions = [(0,0.1),
                  (5,0.15),
                  (10,0.2)]
 
    def storeData(datas):
        pass
 
    for d in MultiSimu(run_simu,paramsTest,replicat = 5, conditions=conditions, batch_size_level='param', withTqdm=True):
        storeData(d)
 
    # or simple usage for replicate : 
 
    def dataCatcherOnStep(simu):
        return simu.cells.countState()
 
    def run_simu(params):
        steps = list(range(0,2501,10))
        data = Simu(params).compute(dataCatcherOnStep,steps)
        data.index *= params['input']['Scenario']['dt']
        return data
 
    data = MultiSimu(run_simu,params,replicat=5).get()
    ```
    """
    def __init__(self, 
                 simuRunner: Callable, 
                 params: dict | list[dict], 
                 replicat: int = 1, 
                 conditions: Iterable = None, 
                 batch_size_level: int | str | None = None, 
                 cacheSize: int | None = 200, 
                 withTqdm: bool = False, 
                 parallel: bool = True, 
                 autoIndex: bool = True, 
                 autoConcat: bool = True):
        if isinstance(params,dict): params = [params]
        elif not isinstance(params,Iterable): raise KeyError("params must a list of dict or dict")
        withCondition = conditions is not None
        if withCondition: 
            if not isinstance(conditions,Iterable): raise KeyError("conditions must be an iterable")
            conditions = list(conditions)
        else: conditions = []
        simuRunnerArity = len([p for p in inspect.signature(simuRunner).parameters.values() if p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD and p.default is inspect._empty])
        if simuRunnerArity == 3:
            self.matchRunner = lambda param,cond,rep: (param,cond,rep)
        elif simuRunnerArity == 2 and len(conditions)>0:
            self.matchRunner = lambda param,cond,rep: (param,cond)
        elif simuRunnerArity == 2 and len(conditions)==0:
            self.matchRunner = lambda param,cond,rep: (param,rep)
        elif simuRunnerArity == 1 and len(conditions)==0:
            self.matchRunner = lambda param,cond,rep: (param,)
        else:
            sig = f"{simuRunner.__name__}{inspect.signature(simuRunner)}"
 
            if withCondition:
                expected = f"{simuRunner.__name__}(params, condition) or {simuRunner.__name__}(params, condition, replicat)"
                note = ""
            else:
                expected = f"{simuRunner.__name__}(params) or {simuRunner.__name__}(params, replicat) or {simuRunner.__name__}(params, condition, replicat)"
                note = f"Note: when conditions is None, a 2-arg {simuRunner.__name__} is interpreted as (params, replicat)."
 
            raise TypeError(
                "simuRunner signature is not compatible with MultiSimu.\n"
                f"Found:    {sig}\n"
                f"Expected: {expected}\n"
                f"{note}")
        self.simuRunner = simuRunner
        self.parallel = parallel
        self.autoConcat = autoConcat
        self.autoIndex = autoIndex
        nbParams = len(params)
        nbCondition = len(conditions) if withCondition else 1
        self.nbParams = nbParams
        self.nbCondition = nbCondition
        self.nbReplicat = replicat
        self.WithParam = nbParams > 1
        self.WithCondition = withCondition
        self.WithReplicat = replicat > 1
        self.withTqdm = withTqdm
        if withTqdm: self._init_pbars(nbParams,nbCondition,replicat)
        if isinstance(batch_size_level,int):realBatchSize=replicat*nbCondition*batch_size_level
        elif (batch_size_level=='param'): realBatchSize=replicat*nbCondition
        elif batch_size_level=='condition': realBatchSize=replicat
        elif batch_size_level=='replicat': realBatchSize=1
        else: realBatchSize = replicat*nbCondition*nbParams
        self.realBatchSize = realBatchSize
        if self.parallel:
            self.argsIter = self._combine(params,conditions,range(replicat))
        else:
            self.argsIter = self._batched(self._combine(params,conditions,range(replicat)),realBatchSize)
        self.i = 0
        self.end = False
        self.cacheSize = cacheSize
 
    def _feedProcess(self,pool,futures,n=1):
        if not self.end:
            for _ in range(n):
                try:
                    ids, params = next(self.argsIter)
                    futures.append((ids,pool.apply_async(self.simuRunner,params)))
                except StopIteration:
                    self.end = True
                    break
 
 
    def __iter__(self):
        global _pool
        try:
            if self.parallel:
                if not _pool_is_usable():
                    _pool = Pool()
                try:
                    data = []
                    futures = []
                    self._feedProcess(_pool, futures, self.cacheSize)
 
                    while len(futures) > 0:
                        self._feedProcess(_pool, futures)
                        ids, simu = futures.pop(0)
                        data.append(self._postProcess(simu.get(), *ids))
 
                        if len(data) == self.realBatchSize:
                            yield self._autoConcat(data)
                            data = []
 
                    if len(data) > 0:
                        yield self._autoConcat(data)
 
                except KeyboardInterrupt:
                    _pool.terminate()
                    _pool.join()
                    _pool = None
                    raise
                except Exception:
                    _pool.terminate()
                    _pool.join()
                    _pool = None
                    raise
            else:
                for batch in self.argsIter:
                    yield self._autoConcat(
                        [self._postProcess(self.simuRunner(*params), *ids) for ids, params in batch]
                    )
 
        finally:
            self._close_pbars()
 
    def get(self) -> Union[pd.DataFrame, List[pd.DataFrame]]:
        """
        Retrieve all simulation results at once.
 
        Returns
        -------
        : Union[pd.DataFrame, List[pd.DataFrame]]
            Returns a single pandas DataFrame if there is only one batch of results.
            Returns a list of pandas DataFrames if there are multiple batches of results.
        """
        ret = list(self)
        if len(ret)==1: return ret[0]
        else: return ret
 
    def _autoConcat(self,datas):
        nbOutput = len(datas[0]) if (isinstance(datas[0],list) or isinstance(datas[0],tuple)) else 0
        if nbOutput==1: datas=datas[0]
        if nbOutput>1:
            datas = list(zip(*datas))
            if self.autoConcat:
                for i in range(len(datas)):
                    datas[i] = pd.concat(datas[i],copy=False)
        elif self.autoConcat:
            datas = pd.concat(datas,copy=False)
        return datas
 
    def _init_pbars(self,nbParam,nbCondition,nbReplicat):
        i= 0
        self.pbars={}
        if self.WithParam:
            self.pbars['param']=tqdm(total=nbParam, ncols=80, position=i        ,desc='Parameters')
            i+=1
        if self.WithCondition:
            self.pbars['condition']=tqdm(total=nbCondition, ncols=80, position=i,desc='Conditions')
            i+=1
        if self.WithReplicat: 
            self.pbars['replicat']=tqdm(total=nbReplicat, ncols=80, position=i  ,desc=' Replicats')
 
 
    def _set_pbar_value(self, pbar, value):
        if value < pbar.n:
            pbar.reset()
        if value > pbar.n:
            pbar.update(value - pbar.n)
        pbar.refresh()
 
    def _updateBar(self, id_param, id_cond, replicat):
        if self.withTqdm:
            is_replicat_complete = (not self.WithReplicat) or (replicat == self.nbReplicat - 1)
 
            if self.WithParam:
                is_param_complete = is_replicat_complete and ((not self.WithCondition) or (id_cond == self.nbCondition - 1))
                self._set_pbar_value(self.pbars['param'], id_param + int(is_param_complete))
 
            if self.WithCondition:
                completed_conditions = id_cond + int(is_replicat_complete)
                self._set_pbar_value(self.pbars['condition'], completed_conditions)
 
            if self.WithReplicat:
                self._set_pbar_value(self.pbars['replicat'], replicat + 1)
 
    def _postProcess(self,datas,id_param,id_cond,replicat):
        self.i+=1
        if self.autoIndex:
            if isinstance(datas,list) or isinstance(datas,tuple):
                for data in datas:
                    if isinstance(data,pd.DataFrame):
                        if self.WithParam: data['ID_PARAMETER'] = id_param
                        if self.WithCondition: data['ID_CONDITION'] = id_cond
                        if self.WithReplicat: data['ID_REPLICAT'] = replicat
            elif isinstance(datas,pd.DataFrame):
                if self.WithParam: datas['ID_PARAMETER'] = id_param
                if self.WithCondition: datas['ID_CONDITION'] = id_cond
                if self.WithReplicat: datas['ID_REPLICAT'] = replicat
 
        self._updateBar(id_param, id_cond, replicat)
        return datas
    
    def _combine(self,params,conditions,replicat):
        for idP,param in enumerate(params):
            if self.WithCondition:
                for idC,condition in enumerate(conditions):
                    if self.WithReplicat:
                        for idR in replicat:
                            yield ((idP,idC,idR),self.matchRunner(deepcopy(param),condition,idR))
                    else:
                        yield ((idP,idC,None),self.matchRunner(deepcopy(param),condition,None))
            else:
                if self.WithReplicat:
                    for idR in replicat:
                        yield ((idP,None,idR),self.matchRunner(deepcopy(param),None,idR))
                else:
                    yield ((idP,None,None),self.matchRunner(deepcopy(param),None,None))
 
    def _batched(self,iterable, n):
        if n < 1: raise ValueError('n must be at least one')
        iterator = iter(iterable)
        while batch := tuple(itertools.islice(iterator, n)):
            yield batch
 
    def _close_pbars(self):
        if self.withTqdm:
            for pbar in self.pbars.values():
                pbar.close()