codegen.py

来自「CVXMOD is a Python-based tool for expres」· Python 代码 · 共 986 行 · 第 1/2 页

        print 'Generating reduced Newton system coefficient matrix.'
        (ASATnum, s, ss, p) = self.multstuff(A, AT, m, n)
        dt['CM_ASAT'] = s
        dt['CM_ASAT_size'] = len(ss)

        # Reverse the permutation. Ran into a weird bug when trying a fancier
        # scheme of inverting the permutation.
        pinv = [None,]*len(p)
        for i in range(len(p)):
            pinv[p[i]] = i

        print 'Generating code for b = A*gh.'
        s = ''
        # Take note of the (inverse) permutation p.
        for i in range(m):
            d = 0
            for j in range(n):
                if (p[i], j) in A:
                    d = d - A[p[i],j]*gh[j]
            s += 'CM_b[%d] = %s;\n' % (i, exprtoC(d))
        dt['CM_b'] = s


        print 'Generating Cholesky factorization and solution code:'
        (d, (L, ssL, KL)) = self.cholsolve(ASATnum, ss)

        dt.update(d)

        print 'Generating code for dx.'
        s = ''
        w = optvar('CM_w', m, 1)
        for i in range(n):
            d = 0
            for j in range(m):
                if (i, j) in AT:
                    # Take note of the (inverse) permutation p.
                    d = d + AT[i,j]*w[pinv[j]]
            s += 'CM_dx[%d] = -CM_h[%d]*(%s + CM_g[%d]);\n' % (i, i, exprtoC(d), i)
        dt['CM_dx'] = s

        writecode(self.outdir, self.skel, 'fastacent.c', dt)

    def fastbarr(self, As, bs, cs, m, n, p, suffix=''):
        dt = {}
        dt['header'] = header()
        dt['m'] = value(m)
        dt['n'] = value(n)
        dt['p'] = value(p)

        dt['suffix'] = suffix

        print 'Generating code for fast barrier method:'

        print 'Framing objective.'
        cs = value(cs)
        s = ''
        for i in range(n):
            s += 'CM_c[%d] = %s;\n' % (i, exprtoC(cs[i]))
        dt['CM_c'] = s

        print 'Creating method signature.'
        if getparams(self.p):
            dt['params'] = ', '.join(['double *%s' % x for x in \
                                      bylowerstr(getparams(self.p))])
            dt['params'] += ','
        else:
            dt['params'] = ''

        print 'Inspecting problem structure.'
        A = nzentries(As)
        AT = nzentries(tp(As))

        # Create a gh optvar simply to make it easy to convert the multiplication.
        gh = optvar('CM_gh', n, 1)
        
        # Now multiply and solve.
        print 'Generating reduced Newton system coefficient matrix.'
        (ASATnum, s, ss, p) = self.multstuff(A, AT, m, n)
        dt['CM_ASAT'] = s
        dt['CM_ASAT_size'] = len(ss)

        # Reverse the permutation. Ran into a weird bug when trying a fancier
        # scheme of inverting the permutation.
        pinv = [None,]*len(p)
        for i in range(len(p)):
            pinv[p[i]] = i

        print 'Generating code for b = A*gh.'
        s = ''
        # Take note of the (inverse) permutation p.
        for i in range(m):
            d = 0
            for j in range(n):
                if (p[i], j) in A:
                    d = d - A[p[i],j]*gh[j]
            s += 'CM_b[%d] = %s;\n' % (i, exprtoC(d))
        dt['CM_b'] = s


        print 'Generating Cholesky factorization and solution code:'
        (d, (L, ssL, KL)) = self.cholsolve(ASATnum, ss)

        dt.update(d)

        print 'Generating code for dx.'
        s = ''
        w = optvar('CM_w', m, 1)
        for i in range(n):
            d = 0
            for j in range(m):
                if (i, j) in AT:
                    # Take note of the (inverse) permutation p.
                    d = d + AT[i,j]*w[pinv[j]]
            s += 'CM_dx[%d] = -CM_h[%d]*(%s + CM_g[%d]);\n' % (i, i, exprtoC(d), i)
        dt['CM_dx'] = s

        writecode(self.outdir, self.skel, 'fastbarr.c', dt,
                  outname='fastbarr%s.c' % suffix)

    def testwithdata(self, dispvars=False):
        p = self.p

        d = {}
        d['header'] = header()

        d['dims'] = joinlist(bylowerstr(getdims(p)), True)
        d['params'] = joinlist(bylowerstr(getparams(p)), True)
        d['optvars'] = joinlist(bylowerstr(getoptvars(p)), True)
        d['numvars'] = str(len(getoptvars(p)))

        # jem: making the assumption that optvars order will be sorted like
        # this.
        s = ''
        i = 0
        for x in bylowerstr(getoptvars(p)):
            s += 'outvar %s;\n%s = vars[%d];\n' % (str(x), str(x), i)
            i += 1
        d['namevars'] = s

        s = joinlist(bylowerstr(getdims(p)) + bylowerstr(getparams(p)), True)
        d['tsf_args'] = s

        # jemjemjem.
        #if dispvars:
        #    d['tsf_args'] = s
        #// Display the variables.
        #pm(x->val, x->m, x->n);
        #pm(y->val, y->m, y->n);

        writecode(self.outdir, self.skel, 'testwithdata.c', d)

    def backsubs(self, L, n, ssL):
        # Include an implicit transpose here. ie, L is a lower triangular
        # version of an upper triangular matrix.
        s = ''
        for i in range(n - 1, -1, -1):
            #s += 'CM_w[%d] = CM_x[%d];\n' % (i, i)
            for j in range(i + 1, n):
                if not iszero(L[j,i]):
                    s += 'CM_w[%d] -= CM_L[%s]*CM_w[%d];\n' % (i, ssL[j,i], j)

        return s

    def forwardsubs(self, L, n, ssL):
        s = ''
        for i in range(n):
            s += 'CM_w[%d] = CM_b[%d];\n' % (i, i)
            for j in range(0, i):
                if not iszero(L[i,j]):
                    s += 'CM_w[%d] -= CM_L[%s]*CM_w[%d];\n' % (i, ssL[i,j], j)

        return s

    def cholsolve(self, Anum, ssA):
        n = rows(Anum)
        (L, ssL, KL, Aclean) = self.cholss(Anum, ssA)

        dt = {}
        dt['CM_L_size'] = KL

        print '- code for Cholesky factorization.'
        dt['chol'] = self.chol(Aclean, n, ssA, ssL)
        print '- code for forward substitution.'
        dt['forwardsubs'] = self.forwardsubs(L, n, ssL)

        print '- code for back substitution.'
        dt['backsubs'] = self.backsubs(L, n, ssL)
        return (dt, (L, ssL, ssA))

    def chol(self, A, n, ssA, ssL):
        v = [None,]*n
        d = [None,]*n
        L = {}
        d[0] = nze()
        s = 'CM_d[0] = CM_ASAT[%d];\n' % ssA[0,0]

        for i in range(1, n):
            L[i,0] = A[i,0]# / d[0]
            if (i,0) in ssA:
                s += 'CM_L[%d] = CM_ASAT[%d] / CM_d[0];\n' % (ssL[i,0], ssA[i,0])

        for j in range(1, n):
            for i in range(j):
                v[i] = L[j,i]*d[i]
                if isnonzero(L[j,i]) and isnonzero(d[i]):
                    s += 'CM_v[%d] = CM_L[%d]*CM_d[%d];\n' % (i, ssL[j,i], i)

            v[j] = A[j,j]
            s += 'CM_v[%d] = CM_ASAT[%d];\n' % (j, ssA[j,j])

            for i in range(j):
                v[j] = v[j] - L[j,i]*v[i]
                if isnonzero(L[j,i]*v[i]):
                    s += 'CM_v[%d] -= CM_L[%d]*CM_v[%d];\n' % (j, ssL[j,i], i)

            d[j] = v[j]
            s += 'CM_d[%d] = CM_v[%d];\n' % (j, j)

            if j < n - 1:
                for i in range(j + 1, n):
                    L[i,j] = A[i,j]
                    if (i,j) in ssL:
                        if isnonzero(A[i,j]):
                            s += 'CM_L[%d] = CM_ASAT[%d];\n' % (ssL[i,j], ssA[i,j])
                        else:
                            s += 'CM_L[%d] = 0;\n' % ssL[i,j]
                    else:
                        continue

                    for k in range(j):
                        L[i,j] = L[i,j] - L[i,k]*v[k]
                        if isnonzero(L[i,k]*v[k]):
                            s += 'CM_L[%d] -= CM_L[%d]*CM_v[%d];\n' % (ssL[i,j], ssL[i,k], k)

                    if isnonzero(L[i,j]):
                        s += 'CM_L[%d] /= CM_v[%d];\n' % (ssL[i,j], j)

        # Implicitly assuming CM_L_{ii} = 1 for all i.
        return s

    def cholss(self, Anum, ssA):
        # Returns (L, Lss) where Lss[i,j] returns the index to store L[i,j].
        # L[i,j] = nze() or 0.
        # Note: implicit about the fact that L_{ii} = 1 for all i.
        n = rows(Anum)

        # Work from a clean A.
        A = sparsedict()
        for i in range(n):
            for j in range(i + 1):
                if isnonzero(Anum[i,j]):
                    A[i,j] = nze()

        v = [None,]*n
        d = [None,]*n
        L = sparsedict()
        d[0] = A[0,0]

        for i in range(1, n):
            L[i,0] = A[i,0]

        for j in range(1, n):
            for i in range(j):
                v[i] = L[j,i]*d[i]

            v[j] = A[j,j]

            for i in range(j):
                v[j] = v[j] - L[j,i]*v[i]

            d[j] = v[j] # will probably be translated directly.

            if j < n - 1:
                for i in range(j + 1, n):
                    L[i,j] = A[i,j]
                    for k in range(j):
                        L[i,j] = L[i,j] - L[i,k]*v[k]

        # Compute (simple) storage scheme for L.
        k = 0
        ssL = {}
        Lclean = sparsedict()
        for i in range(n):
            for j in range(n):
                if isnonzero(L[i,j]):
                    ssL[i,j] = k
                    Lclean[i,j] = nze()
                    k += 1
        KL = k

        return (Lclean, ssL, KL, A)

    def cholmoddata(self, Anum):
        # jem: temporary function.
        # jem: later need to store only one half of A: pick which!

        d = {}
        d['header'] = header()

        # There must be a better way of doing this.
        n = rows(Anum)
        s = ''
        s += 'int *Ai = At->i;\n'
        s += 'int *Aj = At->j;\n'
        s += 'double *Ax = At->x;\n'

        k = 0
        for i in range(n):
            for j in range(i, n):
                if isnonzero(Anum[i,j]):

                    s += 'Ai[%d] = %d; ' % (k, i)
                    s += 'Aj[%d] = %d; ' % (k, j)
                    s += 'Ax[%d] = %.7g;\n' % (k, Anum[i,j])
                    k += 1

        s = 'At = cholmod_allocate_triplet(%d, %d, %d, 1, CHOLMOD_REAL, &common);\n' % \
                (n, n, k) + s
        s = 'cholmod_triplet *At;\n' + s
        s += 'At->nnz = %d;\n' % k

        d['At'] = s

        writecode(self.outdir, self.skel, 'cholmod_data.h', d)

    def choldata(self, Anum, bnum, n):
        # jem: later need to store only one half of A: pick which!

        d = {}
        d['header'] = header()

        ssA = self.cholss(Anum, True)

        s = ''
        for (i, j) in sorted(ssA.keys()):
            s += 'A[%d] = %.7g;\n' % (ssA[i, j], Anum[i,j])

        d['A'] = s

        s = ''
        for i in range(n):
            s += 'b[%d] = %.7g;\n' % (i, bnum[i])

        d['b'] = s

        writecode(self.outdir, self.skel, 'chol_data.h', d)

    def backsubsdata(self, L, ss, b):
        # jem: temporary function.

        d = {}
        d['header'] = header()

        (dss, K) = ss
        b = param('b', value=b)

        # Reverse the role of the keys and values.
        dss2 = {}
        for k in dss.keys():
            dss2[dss[k]] = k

        s = ''
        for k in range(K):
            s += 'L[%d] = %.7g;\n' % (k, L[dss2[k]])
        d['L'] = s

        d['b'] = self.declarematrix(b)

        writecode(self.outdir, self.skel, 'backsubs_data.h', d)

    def multstuff(self, A, AT, m, n):
        """Multiplies and stuffs matrices."""

        # Make a sparse mask of rn. Use A*A'.
        Am = zeros(m, n)
        for k in A.keys():
            Am[k] = 1
        Am = Am*tp(Am)

        p = order(Am)
        Ao = Am[p,p]

        # Design a simple storage scheme for Ao.
        k = 0
        ss = {}
        for (i, j) in zip(Ao.I, Ao.J):
            if i >= j:
                ss[i,j] = k
                k += 1

        # Find reduced Newton system matrix ASAT.
        d = {}
        ASAT = {}
        # Write some code.
        h = param('CM_h', m, 1)
        for i in range(m):
            ip = p[i]
            for j in range(i + 1):
                jp = p[j]
                for k in range(n):
                    if (ip,k) in A and (k,jp) in AT:
                        if ss[i,j] not in d:
                            d[ss[i,j]] = exprtoC(A[ip,k] * h[k] * AT[k,jp])
                        else:
                            d[ss[i,j]] += ' + ' + exprtoC(A[ip,k] * h[k] * AT[k,jp])

        s = ''
        for k in sorted(d.keys()):
            s += 'CM_ASAT[%d] = %s;\n' % (k, d[k])

        return (Ao, s, ss, p)

    def sparsematrixvec(self, ss, m, n, A, x, b):
        """Provides explicit C code for multiplying A*x, where ss is the
        storage scheme for a sparse matrix A, and x and b are both strings."""
        A = nzentries(A)

        x = optvar(x, n)

        for i in range(m):
            d = 0
            for j in range(n):
                if (i, j) in ss:
                    d = d + A[i,j]*x[j]
                    print A[i,j]
                    print d
            print '%s[%d] = %s;' % (b, i, exprtoC(d))

    def getss(self, A):
        # Design a simple storage scheme for Ao.
        m, n = value(size(A))
        A = nzentries(A)

        # Make a sparse mask of A.
        Am = zeros(m, n)
        for k in A.keys():
            Am[k] = 1

        k = 0
        ss = {}
        for (i, j) in zip(Am.I, Am.J):
            if i >= j:
                ss[i,j] = k
                k += 1

        return ss

def withheading(s):
    return s.rstrip() + '\n' + '-'*len(s.rstrip()) + '\n'

def header():
    s = '// CVXMOD code generation.\n'
    s += '// Produced by CVXMOD, %s.\n' % strftime("%Y-%m-%d %H:%M:%S")
    s += '// CXVMOD is Copyright (C) 2006-2008 Jacob Mattingley.\n'
    return s

def writecode(outdir, skeleton, inname, d, outname=None):
    if outname is None:
        outname = inname

    infile = open(os.path.join(skeleton, inname))
    outname = os.path.join(outdir, outname)
    outfile = open(outname, 'w')
    print 'Writing to %s.' % outname
    ril = re.compile('(CVXMOD_IL_(\\w+))')
    rwl = re.compile('(\s*)CVXMOD_WL_(\\w+);?')
    ind = ''
    for l in infile:
        # Test for inline tags. Only one allowed max, per line.
        g = ril.search(l)
        # Test for whole line tags. These take on the same indent as the tag.
        h = rwl.search(l)

        if g:
            try:
                outfile.write(l[:g.start(1)] + str(d[g.group(2)]) + l[g.end(1):])
            except KeyError:
                Jwarn('could not match tag: ' + g.group(2))
                outfile.write(l[:-1] + ' // XXX failed to match tag.\n')
        elif h:
            try:
                for l2 in d[h.group(2)].splitlines():
                    outfile.write(h.group(1) + l2 + '\n')
            except KeyError:
                Jwarn('could not match tag: ' + h.group(2))
                outfile.write(l[:-1] + ' // XXX failed to match tag.\n')
        else:
            outfile.write(ind + l)

    infile.close()
    outfile.close()