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
 
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: list[int] = [0], 
                 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")
        if not isinstance(conditions,list): raise KeyError("conditions must be a list")
        self.simuRunner = simuRunner
        self.parallel = parallel
        self.autoConcat = autoConcat
        self.autoIndex = autoIndex
        nbParams = len(params)
        nbCondition = len(conditions)
        self.WithParam = nbParams > 1
        self.WithCondition = nbCondition > 1
        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):
        if self.parallel:
            with Pool() as pool:
                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)
                    data = []
        else:
            for batch in self.argsIter:
                yield self._autoConcat([self._postProcess(self.simuRunner(*params),*ids) for ids,params in batch])
        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')
            self.pbars['param'].updateEach = nbReplicat*nbCondition
            i+=1
        if self.WithCondition:
            self.pbars['condition']=tqdm(total=nbCondition, ncols=80, position=i,desc='Conditions')
            self.pbars['condition'].updateEach = nbReplicat
            i+=1
        if self.WithReplicat: 
            self.pbars['replicat']=tqdm(total=nbReplicat, ncols=80, position=i  ,desc=' Replicats')
            self.pbars['replicat'].updateEach = 1
        self.total = nbReplicat*nbParam*nbCondition
 
 
    def _updateBar(self):
        if self.withTqdm:
            for pbar in self.pbars.values():
                if pbar.n >= pbar.total:
                    pbar.n=0
                    pbar.refresh()
                elif (self.i%pbar.updateEach)==0:
                    pbar.update()
                    pbar.refresh()
 
    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()
        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),(deepcopy(param),condition))
                    else:
                        yield ((idP,idC,None),(deepcopy(param),condition))
            else:
                if self.WithReplicat:
                    for idR in replicat:
                        yield ((idP,None,idR),(deepcopy(param),))
                else:
                    yield ((idP,None,None),(deepcopy(param),))
 
    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()