<commit_before>'''
Created on May 24, 2012

@author: dstrauss

This set of routines implements the fast field splitting technique. 
Both parallel and serial methods are implemented
'''

import numpy as np
import scipy.sparse as sparse
import scipy.sparse.linalg as lin
from mpi4py import MPI
import scipy.io as spio
from optimize import optimizer
import time

class problem(optimizer):
    '''A function for implementing the field splitting methods'''
#    def __init__(self, freq, rho, xi):
#        self.f=freq
#        self.rho = rho
#        self.xi = xi
    
    def initOpt(self,uHat, D):
        ''' prepare for upcoming iterations '''
        self.rho = D['rho']
        self.xi = D['xi']
        self.uHat = uHat
        self.lmb = D['lmb']
        self.uBound = D['uBound']
        self.F = np.zeros(self.fwd.N,dtype='complex128')
        self.E = np.zeros(self.fwd.N,dtype='complex128')
        self.us = np.zeros(self.fwd.N,dtype='complex128')
        self.v = np.zeros(self.fwd.N,dtype='complex128')
        self.ub = self.fwd.sol[0] # -- update sol should be flat
        self.obj = np.zeros(D['maxIter'])
        
        # the update for the first step can be precomputed
        self.A = self.fwd.nabla2+self.fwd.getk(0)
        self.s = self.fwd.getS()
#        print self.xi
        self.Q = sparse.vstack([self.A, -self.xi*sparse.eye(self.fwd.N,self.fwd.N)])
        self.Moo = self.rho*(self.Q.conj().T*self.Q) + self.fwd.Ms.T*self.fwd.Ms
        
        tml = time.time()
        self.M = lin.factorized(self.Moo.tocsc())
        print 'factorization time (AtA + xiI) ' + repr(time.time()-tml)
        
        ''' in expansion to the 3d/TM, will need these operators...'''
        self.fwd.setCTRX()

    def runOpt(self, P):
        ''' method to run in each interaction of the ADMM routine - super parallel '''
        
        #update dual variables
        self.F = self.F + (self.A*self.us + self.s*(self.fwd.Md.T*P)*(self.ub+self.v))
        self.E = self.E + self.xi*(self.v-self.us)
        
        # reasses local uHat
        uHatLocal = self.uHat - (self.fwd.Ms*self.ub)
        
        #solve for us update
        b = self.fwd.Ms.T*uHatLocal - \
        self.rho*(self.Q.conj().T*np.append((self.s*(self.ub+self.v)*(self.fwd.Md.T*P) + self.F), (self.xi*self.v)+self.E))
        self.us = self.M(b)
        
        # I was curious about solver tolerance -- bottom line is that it is good
        # print np.linalg.norm(self.Moo*self.us - b)
        
        # solve for v update
        Q = sparse.vstack([sparse.spdiags((self.s*self.fwd.Md.T*P),0,self.fwd.N,self.fwd.N), self.xi*sparse.eye(self.fwd.N,self.fwd.N)])
        M = self.rho*(Q.conj().T*Q)
        b = -self.rho*(Q.conj().T*np.append(self.A*self.us + self.F + self.s*(self.fwd.Md.T*P)*self.ub, (self.E - self.xi*self.us)))
        
        self.v = lin.spsolve(M,b)
        
        #oh that's awkward. obj has to get passed out in order to be indexed.
        obj = np.linalg.norm(self.fwd.Ms*self.v-uHatLocal)
        return obj
        # gap = np.linalg.norm(self.v-self.us)
        # print 'obj = ' + repr(obj) + ' Split Gap ' + repr(gap)
        # return obj
    
    def writeOut(self, rank, ix=0):
        import os
        if not os.path.exists(self.outDir + 'Data'):
            os.mkdir(self.outDir + 'Data')
        #oh awkward again -- need to reshape and so on so that things get written out in the right column
        # form.
        sgm = self.fwd.parseFields(self.fwd.sigmap[0])
        ub = self.fwd.parseFields(self.fwd.sol[0])
        ut = self.fwd.parseFields(self.fwd.sol[1])
        
        us = self.fwd.parseFields(self.us)
        v = self.fwd.parseFields(self.v)
        
        D = {'f':self.fwd.f, 'angle':self.fwd.incAng, 'sigMat':sgm[0], 'ub':ub[0], \
             'us':us[0], 'uTrue':ut[0], \
             'v':v[0], 'rho':self.rho, 'xi':self.xi, 'obj':self.obj,\
             'flavor':self.fwd.flavor}
    
        spio.savemat(self.outDir + 'Data/splitField' + repr(rank) + '_' + repr(ix), D)
        
  
    
    def aggregatorSerial(self, S):
        '''Do the aggregate step updates '''
        # currently assumes that Md*Md' = I
        N = np.size(S)
        n = S[0].fwd.nRx*S[0].fwd.nRy
        # print N
        # print n
        
        U = np.zeros((n,N),dtype='complex128')
        Q = np.zeros((n,N),dtype='complex128')
        
        for ix in range(N):
            s = S[ix].s
            U[:,ix] = s*S[ix].fwd.Md*(S[ix].ub+S[ix].v)
            Q[:,ix] = S[ix].fwd.Md*((S[ix].A*S[ix].us) + S[ix].F)
            
        num = np.sum(U.real*Q.real + U.imag*Q.imag,1)
        print num.shape
        den = np.sum(U.conj()*U,1) + self.lmb/S[0].rho
        
        P = (-num/den).real
        P = np.maximum(P,0)
        P = np.minimum(P,self.uBound)
        
        gap = np.zeros(N)
        for ix in range(N):
            gap[ix] = np.linalg.norm(S[ix].fwd.Md.T*(U[:,ix]*P) + S[ix].A*S[ix].us)
            
        return P

    def aggregatorParallel(self,comm):
        ''' Do the aggregation step in parallel whoop! '''
        # currently assumes that Md*Md' is I
        N = np.size(self)
        print repr(N) + ' == better be 1!'
        
        U = self.s*self.fwd.Md*(self.ub+self.v)
        Q = self.fwd.Md*((self.A*self.us) + self.F)
        
        q = U.real*Q.real + U.imag*Q.imag
        num = np.zeros(q.shape)
        
        num = comm.allreduce(q,num,op=MPI.SUM)
        
        q = U.conj()*U + self.lmb/self.rho
        den = np.zeros(q.shape)
        den = comm.allreduce(q,den,op=MPI.SUM)
        
        P = (-num/den).real
        P = np.maximum(P,0)
        P = np.minimum(P,self.uBound)
        
        gap = np.linalg.norm(self.fwd.Md.T*(U*P) + self.A*self.us)
        print 'Proc ' + repr(comm.Get_rank()) + ' gap = ' + repr(gap)
        
        return P

    def aggregatorSemiParallel(self,S,comm):
        ''' Do the aggregation step in parallel whoop! '''
        N = np.size(S)
        n = S[0].fwd.nRx*S[0].fwd.nRy
        
        QL = sparse.lil_matrix( (n,n), dtype='complex128' )
        bL = np.zeros(n, dtype='complex128')
        # U = np.zeros((n,N),dtype='complex128')
        # Q = np.zeros((n,N),dtype='complex128')
        
        for L in S:
            # the flops are certainly not optimized here
            # looking at this it is seeking to do min || Av + s(v+ub)Md*theta + F ||
            D = sparse.spdiags(L.s*(L.ub+L.v),0, self.fwd.N,self.fwd.N)
            g = D*L.fwd.Md.T
            print 'size of F ' + repr(L.F.shape)
            print 'size of Aus ' + repr((L.A*L.us).shape)
            print 'size of g ' + repr(g.shape)
            bL += g.T.conj()*((L.A*L.us) + L.F)
            QL += g.T.conj()*g<commit_msg>has to do with nRx*nRy vs nRx*nRy*nRz<commit_after>'''
Created on May 24, 2012

@author: dstrauss

This set of routines implements the fast field splitting technique. 
Both parallel and serial methods are implemented
'''

import numpy as np
import scipy.sparse as sparse
import scipy.sparse.linalg as lin
from mpi4py import MPI
import scipy.io as spio
from optimize import optimizer
import time

class problem(optimizer):
    '''A function for implementing the field splitting methods'''
#    def __init__(self, freq, rho, xi):
#        self.f=freq
#        self.rho = rho
#        self.xi = xi
    
    def initOpt(self,uHat, D):
        ''' prepare for upcoming iterations '''
        self.rho = D['rho']
        self.xi = D['xi']
        self.uHat = uHat
        self.lmb = D['lmb']
        self.uBound = D['uBound']
        self.F = np.zeros(self.fwd.N,dtype='complex128')
        self.E = np.zeros(self.fwd.N,dtype='complex128')
        self.us = np.zeros(self.fwd.N,dtype='complex128')
        self.v = np.zeros(self.fwd.N,dtype='complex128')
        self.ub = self.fwd.sol[0] # -- update sol should be flat
        self.obj = np.zeros(D['maxIter'])
        
        # the update for the first step can be precomputed
        self.A = self.fwd.nabla2+self.fwd.getk(0)
        self.s = self.fwd.getS()
#        print self.xi
        self.Q = sparse.vstack([self.A, -self.xi*sparse.eye(self.fwd.N,self.fwd.N)])
        self.Moo = self.rho*(self.Q.conj().T*self.Q) + self.fwd.Ms.T*self.fwd.Ms
        
        tml = time.time()
        self.M = lin.factorized(self.Moo.tocsc())
        print 'factorization time (AtA + xiI) ' + repr(time.time()-tml)
        
        ''' in expansion to the 3d/TM, will need these operators...'''
        self.fwd.setCTRX()

    def runOpt(self, P):
        ''' method to run in each interaction of the ADMM routine - super parallel '''
        
        #update dual variables
        self.F = self.F + (self.A*self.us + self.s*(self.fwd.Md.T*P)*(self.ub+self.v))
        self.E = self.E + self.xi*(self.v-self.us)
        
        # reasses local uHat
        uHatLocal = self.uHat - (self.fwd.Ms*self.ub)
        
        #solve for us update
        b = self.fwd.Ms.T*uHatLocal - \
        self.rho*(self.Q.conj().T*np.append((self.s*(self.ub+self.v)*(self.fwd.Md.T*P) + self.F), (self.xi*self.v)+self.E))
        self.us = self.M(b)
        
        # I was curious about solver tolerance -- bottom line is that it is good
        # print np.linalg.norm(self.Moo*self.us - b)
        
        # solve for v update
        Q = sparse.vstack([sparse.spdiags((self.s*self.fwd.Md.T*P),0,self.fwd.N,self.fwd.N), self.xi*sparse.eye(self.fwd.N,self.fwd.N)])
        M = self.rho*(Q.conj().T*Q)
        b = -self.rho*(Q.conj().T*np.append(self.A*self.us + self.F + self.s*(self.fwd.Md.T*P)*self.ub, (self.E - self.xi*self.us)))
        
        self.v = lin.spsolve(M,b)
        
        #oh that's awkward. obj has to get passed out in order to be indexed.
        obj = np.linalg.norm(self.fwd.Ms*self.v-uHatLocal)
        return obj
        # gap = np.linalg.norm(self.v-self.us)
        # print 'obj = ' + repr(obj) + ' Split Gap ' + repr(gap)
        # return obj
    
    def writeOut(self, rank, ix=0):
        import os
        if not os.path.exists(self.outDir + 'Data'):
            os.mkdir(self.outDir + 'Data')
        #oh awkward again -- need to reshape and so on so that things get written out in the right column
        # form.
        sgm = self.fwd.parseFields(self.fwd.sigmap[0])
        ub = self.fwd.parseFields(self.fwd.sol[0])
        ut = self.fwd.parseFields(self.fwd.sol[1])
        
        us = self.fwd.parseFields(self.us)
        v = self.fwd.parseFields(self.v)
        
        D = {'f':self.fwd.f, 'angle':self.fwd.incAng, 'sigMat':sgm[0], 'ub':ub[0], \
             'us':us[0], 'uTrue':ut[0], \
             'v':v[0], 'rho':self.rho, 'xi':self.xi, 'obj':self.obj,\
             'flavor':self.fwd.flavor}
    
        spio.savemat(self.outDir + 'Data/splitField' + repr(rank) + '_' + repr(ix), D)
        
  
    
    def aggregatorSerial(self, S):
        '''Do the aggregate step updates '''
        # currently assumes that Md*Md' = I
        N = np.size(S)
        n = S[0].fwd.nRx*S[0].fwd.nRy
        # print N
        # print n
        
        U = np.zeros((n,N),dtype='complex128')
        Q = np.zeros((n,N),dtype='complex128')
        
        for ix in range(N):
            s = S[ix].s
            U[:,ix] = s*S[ix].fwd.Md*(S[ix].ub+S[ix].v)
            Q[:,ix] = S[ix].fwd.Md*((S[ix].A*S[ix].us) + S[ix].F)
            
        num = np.sum(U.real*Q.real + U.imag*Q.imag,1)
        print num.shape
        den = np.sum(U.conj()*U,1) + self.lmb/S[0].rho
        
        P = (-num/den).real
        P = np.maximum(P,0)
        P = np.minimum(P,self.uBound)
        
        gap = np.zeros(N)
        for ix in range(N):
            gap[ix] = np.linalg.norm(S[ix].fwd.Md.T*(U[:,ix]*P) + S[ix].A*S[ix].us)
            
        return P

    def aggregatorParallel(self,comm):
        ''' Do the aggregation step in parallel whoop! '''
        # currently assumes that Md*Md' is I
        N = np.size(self)
        print repr(N) + ' == better be 1!'
        
        U = self.s*self.fwd.Md*(self.ub+self.v)
        Q = self.fwd.Md*((self.A*self.us) + self.F)
        
        q = U.real*Q.real + U.imag*Q.imag
        num = np.zeros(q.shape)
        
        num = comm.allreduce(q,num,op=MPI.SUM)
        
        q = U.conj()*U + self.lmb/self.rho
        den = np.zeros(q.shape)
        den = comm.allreduce(q,den,op=MPI.SUM)
        
        P = (-num/den).real
        P = np.maximum(P,0)
        P = np.minimum(P,self.uBound)
        
        gap = np.linalg.norm(self.fwd.Md.T*(U*P) + self.A*self.us)
        print 'Proc ' + repr(comm.Get_rank()) + ' gap = ' + repr(gap)
        
        return P

    def aggregatorSemiParallel(self,S,comm):
        ''' Do the aggregation step in parallel whoop! '''
        N = np.size(S)
        n = self.fwd.getPSize()
        # S[0].fwd.nRx*S[0].fwd.nRy
        
        QL = sparse.lil_matrix( (n,n), dtype='complex128' )
        bL = np.zeros(n, dtype='complex128')
        # U = np.zeros((n,N),dtype='complex128')
        # Q = np.zeros((n,N),dtype='complex128')
        
        for L in S:
            # the flops are certainly not optimized here
            # looking at this it is seeking to do min || Av + s(v+ub)Md*theta + F ||
            D = sparse.spdiags(L.s*(L.ub+L.v),0, self.fwd.N,self.fwd.N)
            g = D*L.fwd.Md.T
            print 'size of F ' + repr(L.F.shape)
            print 'size of Aus ' + repr((L.A*L.us).shape)
            print 'size of g ' + repr(g.shape)
            bL += g.T.conj()*((L.A*L.us) + L.F)
            QL += g.T.conj()*g
{'fraction_non_alphanumeric': 0.16550908822225507, 'fraction_numerical': 0.009162110240874833, 'mean_word_length': 1.8500736997262581, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '318665', 'n_tokens_mistral': 5285, 'n_tokens_neox': 5136, 'n_words': 1002}