SwitchingFCalongAtask/findInitailCond.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jun 29 13:20:23 2020

@author: p277634
"""
import scipy.io as sio
import pygmo as po
import numpy as np
from itertools import combinations
import copy

class InitalProblem():
    def __init__(self, mIx=0):
        self.mIx        = mIx # model index
        self.nroi       = 68
        self.pathData   = './data/'
        self.efcData    = 'FCstagesDTIrois.mat'
        self.iniData    = 'partialInitialCondit.mat'
        self.piVar      = np.pi/8
        self.pairs      = np.array(list(combinations(range(self.nroi),2)), dtype=np.int32)
        self.withNoise  = True
        self.noise      = 0.01
        self.nModel     = 50
        self.sizePop    = 300
        self.genera     = 1000  
        self.pathSave   = './dataSave/'
        self.nameSave   = 'OptiNoisy0p01IniState.npz'
        self.info       = """intial amplitude and cosistent (ITPC) pahses defined from 10 samples backwards of the first bump\n
                            with file boundIniState_2.m. ROIS are ordered as in DTI\n
                            edgeSt has the edges and onsets of simulated FC stages in seconds \n
                            edgeSt(n stages, [low edge upper edge, onset,]) \n'
                            phase values from significant inter trial phase coherence were given and the resta were zeros \n.
                            the file findInitalCong uses CMAES optimzer to find 50 possible inital states of the model"""
    
        self._importFCstage()
        self._getEmpiricalInitialState()

    def _importFCstage(self):        
        filepath    = self.pathData + self.efcData
        self.eFC    = np.float32(sio.loadmat(filepath)['fc1d'][1,:])
        self.eFCs   = np.sign(self.eFC)
        self.neFCs  = np.sum(np.abs(self.eFCs), axis=-1)
        
        
    def _getEmpiricalInitialState(self):
        # load initial conditions
        loaddata    = self.pathData + self.iniData
        iniSta      = sio.loadmat(loaddata) 
        self.iniPhi0    = iniSta['iniState'][:,0]
        self.iniAmp1    = iniSta['iniState'][:,1]
        self.iniAmp2    = iniSta['iniState'][:,2]
        self.edgeSt     = iniSta['edgeSt']
        self.nIphi      = np.sum(self.iniPhi0 != 0)
        self.missPhi    = self.nroi - self.nIphi
        self.noIniPhi   = self.iniPhi0 == 0        
        self.hasIniPhi  = self.iniPhi0 != 0 
        if self.withNoise:
            self.noiseIniPhase()
        else:
            self.mPhi      = np.mean(self.iniPhi[self.iniPhi.nonzero()])
            self.iniPhi     = copy.deepcopy(self.iniPhi0)

        
    def get_bounds(self):
        lowboud = [self.mPhi - self.piVar]*self.missPhi 
        upbound = [self.mPhi + self.piVar]*self.missPhi 
        return (lowboud,upbound)
    
    def costFunction(self, fci):
        return  -1 * np.sum(fci * self.eFCs) / self.neFCs
    
    def noiseIniPhase(self):
        iniPhi     = copy.deepcopy(self.iniPhi0)
        self.iniPhi     = np.repeat(iniPhi[:,np.newaxis], self.nModel, axis=1)
        self.iniPhi[self.hasIniPhi,:] += self.noise * np.random.randn(self.nIphi, self.nModel)
        self.mPhi      = np.mean(self.iniPhi[self.iniPhi[:,0].nonzero(),:], axis=1)

    
    def fullfilConstrin(self, fci):
        # unequal signs are converted to positive values that violate contraints
        # constraint fulfiltmnet is associatied with a negative value (pygmo requirement)
        match = -1 * fci * self.eFCs 
        return match[self.eFCs != 0]
     
    def fitness(self, x):
        phis = copy.deepcopy(self.iniPhi[:, self.mIx])
        phis[self.noIniPhi] = x
        # fci = np.sign(phis[self.pairs[:,0]] - phis[self.pairs[:,1]])
        fci = np.sign( np.imag( np.exp(1j*phis[self.pairs[:,0]]) * np.conj(np.exp(1j*phis[self.pairs[:,1]]) )))

        # cnstr = self.fullfilConstrin(fci)
        fit = self.costFunction(fci)
        # return np.insert(cnstr, 0, fit)
        return np.array([fit])
    
    def gradient(self, x):
        return po.estimate_gradient_h(lambda x: self.fitness(x), x)
    
    def putNoiseAmpli(self):
        self.iniAmp1  = np.repeat(self.iniAmp1[:,np.newaxis], self.nModel, axis=1)
        self.iniAmp2  = np.repeat(self.iniAmp2[:,np.newaxis], self.nModel, axis=1)
        self.iniAmp1 += self.noise * np.random.randn(self.nroi, self.nModel)
        self.iniAmp2 += self.noise * np.random.randn(self.nroi, self.nModel)  
        self.iniAmp1[self.iniAmp1<0.02] = 0.02 
        self.iniAmp2[self.iniAmp2<0.02] = 0.02 
        
    
        
    def nInitialStates(self):
        self.phi    = copy.deepcopy(self.iniPhi)
        self.fitpli = np.zeros(self.nModel) 
        for iXm in range(self.nModel):
            # prob = po.problem(udp=self.__class__(mIx=iXm))
            pop = po.population(prob = FindIniState(self, iXm), size = self.sizePop)
            algo = po.algorithm(uda=po.cmaes(gen=self.genera,force_bounds=True))
            pop = algo.evolve(pop)
            self.fitpli[iXm] = pop.champion_f[0]
            self.phi[self.noIniPhi,iXm] = pop.champion_x
        self.putNoiseAmpli()
            
  
    def saveIniState(self):
        ''' info: inicond are computed at 10 sampels before onset '''
        data = np.zeros((self.nroi,3,self.nModel), dtype=np.float32)
        data[:,0,:] = self.phi
        data[:,1,:] = self.iniAmp1
        data[:,2,:] = self.iniAmp2
        saveto = self.pathSave + self.nameSave
        np.savez(saveto, info=self.info,iniState=data,  edgeSt=self.edgeSt, piVar=self.piVar, noise=self.noise, fit=self.fitpli)
        


class FindIniState():
    def __init__(self, ini, mIx):
        self.ini = ini
        self.mIx = mIx
        # print('Here')
        
    def get_bounds(self):
        lowboud = [self.ini.mPhi[0][self.mIx] - self.ini.piVar]*self.ini.missPhi 
        upbound = [self.ini.mPhi[0][self.mIx] + self.ini.piVar]*self.ini.missPhi
        return (lowboud, upbound)
        
    def fitness(self, x):
        phis = copy.deepcopy(self.ini.iniPhi[:, self.mIx])
        phis[self.ini.noIniPhi] = x
        fci = np.sign( np.imag( np.exp(1j*phis[self.ini.pairs[:,0]]) * np.conj(np.exp(1j*phis[self.ini.pairs[:,1]]) )))
        fit = self.costFunction(fci)
        # return np.insert(cnstr, 0, fit)
        return np.array([fit])
    
    def costFunction(self, fci):
        return  -1 * np.sum(fci * self.ini.eFCs) / self.ini.neFCs
    
if __name__ == "__main__":
    
    ini = InitalProblem()
    
    ini.nInitialStates( )
    ini.saveIniState()