/home/mark/model/software/ScrumPy/ScrumPy/LP/ScrumPyLP.py

"""
ScrumPyLP - (c) Mark Poolman 2008 onwards.
A module to allow LP optimisation and analysis of of metabolic models with ScrumPy.


This module is still experimental and is

  *************************   NOT TO BE REDISTRIBUTED  ***************************

without the prior consent of the author.

"""






import sys, types, exceptions

SysStderr = sys.stderr
NullF =  open("/dev/null", "w")


import glpk

from Util import Set, DynMatrix
from Data import DataSets
from Structural import StoMat


def SplitRev(sm,Conv=float): # copy of sm with reversible reacs split into forward and back irreversible pairs

    rv = sm.Copy(Conv)
    rv.Reversed = reversed = []
    rv.Backs = backs = []

    for r in sm.cnames:
        if sm.RevProps[r] == StoMat.t_Rever:
            sto = sm.InvolvedWith(r)
            rback = r+"_back"
            rv.NewReaction(rback,sto)
            rv.MulCol(rback,k=-1)
            backs.append(rback)

        elif sm.RevProps[r] == StoMat.t_BackIrrev:
           sm.MulCol(r,  k=-1)
           reversed.append(r)

    return rv

def SplitRevOld(sm,Conv=float): # copy of sm with reversible reacs split into forward and back irreversible pairs
    sm = sm.Copy(Conv)
    revs = Set.Complement(sm.cnames, sm.Irrevs)
    backs = []
    for r in revs:
        sto = sm.InvolvedWith(r)
        for met in sto:
            sto[met] *= -1
        rback = r+"_back"
        sm.NewReaction(rback,sto)
        backs.append(rback)
    sm.Irrevs = sm.cnames[:]
    sm.Backs = backs
    return sm



def SplitAll(sm,Conv=float):

    sm = sm.Copy(Conv)
    irbacks = []
    for r in sm.cnames[:]:
        sto = sm.InvolvedWith(r)
        for met in sto:
            sto[met] *= -1
        newreac = r+"_back"
        sm.NewReaction(newreac,sto)
        if r in sm.Irrevs:
            irbacks.append(newreac)

    sm.Irrevs = sm.cnames[:]
    sm.irbacks = irbacks
    return sm



class lp(glpk.lp):

    def __init__(self, model):

        name = model.FileName.split("/")[-1].split(".")[0]
        sm = self.sm = SplitRev(model.sm)

        glpk.lp.__init__(self,name)
        self.AddRows(sm.rnames[:])
        self.AddCols(sm.cnames[:])
        for reac in sm.cnames:
            self.SetColBounds(reac,0)

        for row in sm.rnames:
            self.SetRowVals(row, map(float, sm[row]))
            self.SetRowBounds(row, 0,0)


    def _irrev_check(self,reac,rate):
        print "_irev_check disabled"
        return
        #
        #        if  rate <0 and reac in self.sm.Irrevs:
        #            sys.stderr.write("""
        #            ! Warning  -ve rate is inconsistent with irreverisibily
        #            ! assumptions. (doing it anyway) """)


    def _BackReac(self, reac):

        rback = reac + "_back"
        if rback in self.sm.Backs:
            return rback
        return None


    def SetObjective(self,reacs):

        if type(reacs)==types.DictType:
            reacs = dict(reacs)
            for reac in reacs.keys():
                back = self._BackReac(reac)
                if back != None:
                    reacs[back] = reacs[reac]
            self.SetObjCoefsFromDic(reacs)

        else:
            reacs = reacs[:]
            for reac in reacs:
                back = self._BackReac(reac)
                if back != None and back not in reacs:
                    reacs.append(back)
            self.SetObjCoefsFromLists(reacs,[1]*len(reacs))

    def SetObjCoef(self, c, coef):

        glpk.lp.SetObjCoef(self, c, coef)

        if not self.sm.IsIrrev(c):
            glpk.lp.SetObjCoef(self, c+"_back", coef)





    def SetFixedFlux(self, reacs):

        for reac in reacs:

            rate = reacs[reac]
            if reac in self.sm.Reversed:
                rate *= -1

            rback = reac + "_back"
            hasback = rback in self.sm.Backs

            if rate >= 0:
                self.SetColBounds(reac, rate, rate)
                if hasback:
                    self.SetColBounds(rback, 0,0)
            else:
                if not hasback:
                    sys.stderr.write("! Setting inconsitant flux in irreversible reaction %s\n" % reac)
                    self.SetColBounds(reac, rate, rate)
                else:
                    self.SetColBounds(reac, 0, 0)
                    self.SetColBounds(rback, -rate, -rate)


    def UnboundFlux(self, reac):

        self.SetColBounds(reac, 0, None)

        rback = self._BackReac(reac)
        if rback != None:
            self.SetColBounds(rback, 0, None)


    def FiniteBoundFlux(self, reac, lo, hi):

        if hi < lo:
            lo,hi = hi,lo

        rback = self._BackReac(reac)
        if rback == None:
            if lo < 0:
                raise exceptions.ValueError, "attempt to set -ve bound on irreversible "+reac
            self.SetColBounds(reac, lo,hi)
        else:
            if hi < 0:
                self.SetColBounds(reac,0,0)
                self.SetColBounds(rback, -hi, -lo)
            elif lo < 0:
                self.SetColBounds(reac,0,hi)
                self.SetColBounds(rback, 0, -lo)
            else:
                self.SetColBounds(reac,lo,hi)
                self.SetColBounds(rback, 0, 0)




    def SetFluxBounds(self, reacs):
        """ Pre: reacs = {name:(lo,hi)...}, (lo<=hi OR hi==None) """

        for reac in reacs:

            lo, hi = reacs[reac]

            if lo == hi == None:
                self.UnboundFlux(reac)

            elif lo == hi:
                self.SetFixedFlux({reac:lo})

            elif not None in (lo,hi):
                self.FiniteBoundFlux(reac, lo, hi)

            else:                                  # lo == None xor hi == None
                rback = self._BackReac(reac)

                if rback == None:                  # irreversible
                    self.SetColBounds(reac, lo, hi)
                else:                              # reversible reaction
                    if lo == None:                 # lo == None
                        if hi < 0.0:
                            self.SetColBounds(reac, 0,0)
                            self.SetColBounds(rback, -hi, None)
                        else:
                            self.SetColBounds(reac,  0, hi)
                            self.SetColBounds(rback, 0, None)
                    else:
                        if lo < 0.0:                # hi == None
                            self.SetColBounds(reac, 0,None)
                            self.SetColBounds(rback, 0, -lo)
                        else:
                            self.SetColBounds(reac, lo,None)
                            self.SetColBounds(rback, 0, 0)

    def SetSumFluxConstraint(self,  reacs, total,  name):

        self.AddRows([name])
        nc  =self.GetDims()[1]
        vals = [0.0] * nc  # TODO - need a GetRow/ColVec in glpk
        reacs = reacs[:]
        for r in reacs[:]:
            rback = self._BackReac(r)
            if rback != None:
                reacs.append(rback)
        idxs = glpk.SubOne(self.GetColIdxs(reacs))
        print idxs
        for idx in idxs:
            vals[idx] = 1.0
        self.SetRowVals(name,vals)
        self.SetRowBounds(name,  total, total)

        print "SetSumFluxCostraint testing !"





    def Solve(self,PrintStatus=True):
        glpk.lp.Solve(self)
        if PrintStatus:
            print self.GetStatusMsg()


    def GetPrimSol(self, IncZeroes=False,FixBack=True,AsMtx=False):

        sol = {}

        if self.GetStatusMsg()=="Optimal":
            for r in range(len(self.sm.cnames)):
                name = self.sm.cnames[r]
                val = self._glp_lpx_get_col_prim(r+1)
                if IncZeroes or not glpk.IsZero(val):
                    sol[name] = self._glp_lpx_get_col_prim(r+1)

        for rev in Set.Intersect(self.sm.Reversed,  sol.keys()):
            sol[rev] *= -1

        if FixBack:     #### all this lacks elegance - needs a re-think.svn
            for reac in sol.keys():
                if reac.endswith("_back"):
                    val = sol[reac]
                    if not glpk.IsZero(val):
                        r2 = reac[:-5]
                        if not sol.has_key(r2):
                            sol[r2] = 0.0
                        sol[r2] -= val
                    del sol[reac]

        if  AsMtx:
            rv = StoMat.StoMat(cnames=["lp_sol"],rnames=sol.keys(),Conv=float)
            for r in sol:
                rv[r][0] = sol[r]
        else:
            rv = sol

        return rv


    def GetSimplexAsMtx(self):

        rnames = self.GetRowNames()
        cnames = self.GetColNames()
        rows   = self.GetRowsAsLists()

        rv = DynMatrix.matrix(rnames=rnames, cnames=cnames, Conv=float)
        rv.rows =rows

        return rv








    def MatchFlux(self,fluxes, IncZeroes=False, AsMtx=False, SupressWarnings=True):

        if SupressWarnings:
            sys.stderr = NullF

        self.SetFixedFlux(fluxes)
        self.SetObjective(self.sm.cnames)
        self.Solve()

        if SupressWarnings:
            sys.stderr = SysStderr

        return self.GetPrimSol(IncZeroes=IncZeroes, AsMtx=AsMtx)




    def FixedFluxScan(self, reaction, lo, hi, n_p):

        rv = DataSets.DataSet(ItemNames=[reaction,"ObjVal"])
        lo = float(lo)
        hi = float(hi)

        inc = (hi - lo)/(n_p-1)
        cur = lo

        for n in range(n_p):
            self.SetFixedFlux({reaction:cur})
            self.Solve(False)

            sol = self.GetPrimSol()

            if len(sol)==0:
                obval = float("NaN") # indicate fail
            else:
                obval =  self.GetObjVal()

            sol["ObjVal"] = obval
            sol[reaction] = cur
            rv.UpdateFromDic(sol)

            cur += inc

        return rv