wls.h

// -*- c++ -*-

// Copyright (C) 2005,2009 Tom Drummond (twd20@cam.ac.uk)

//All rights reserved.
//
//Redistribution and use in source and binary forms, with or without
//modification, are permitted provided that the following conditions
//are met:
//1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//2. Redistributions in binary form must reproduce the above copyright
//   notice, this list of conditions and the following disclaimer in the
//   documentation and/or other materials provided with the distribution.
//
//THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND OTHER CONTRIBUTORS ``AS IS''
//AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
//IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
//ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR OTHER CONTRIBUTORS BE
//LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
//CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
//SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
//INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
//CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
//ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
//POSSIBILITY OF SUCH DAMAGE.

#ifndef TOON_INCLUDE_WLS_H
#define TOON_INCLUDE_WLS_H

#include <TooN/TooN.h>
#include <TooN/Cholesky.h>
#include <TooN/helpers.h>

#include <cmath>

namespace TooN {

template <int Size=Dynamic, class Precision=DefaultPrecision,
          template<int DecompSize, class DecompPrecision> class Decomposition = Cholesky>
class WLS {
public:

    WLS(int size=0) :
        my_C_inv(size,size),
        my_vector(size),
        my_decomposition(size),
        my_mu(size)
    {
        clear();
    }

    void clear(){
        my_C_inv = Zeros;
        my_vector = Zeros;
    }

    void add_prior(Precision val){
        for(int i=0; i<my_C_inv.num_rows(); i++){
            my_C_inv(i,i)+=val;
        }
    }
  
    template<class B2>
    void add_prior(const Vector<Size,Precision,B2>& v){
        SizeMismatch<Size,Size>::test(my_C_inv.num_rows(), v.size());
        for(int i=0; i<my_C_inv.num_rows(); i++){
            my_C_inv(i,i)+=v[i];
        }
    }

    template<class B2>
    void add_prior(const Matrix<Size,Size,Precision,B2>& m){
        my_C_inv+=m;
    }

    template<class B2, class P2>
    inline void add_mJ(Precision m, const Vector<Size, P2, B2>& J, Precision weight = 1) {
        
        //Upper right triangle only, for speed
        for(int r=0; r < my_C_inv.num_rows(); r++)
        {
            double Jw = weight * J[r];
            my_vector[r] += m * Jw;
            for(int c=r; c < my_C_inv.num_rows(); c++)
                my_C_inv[r][c] += Jw * J[c];
        }
    }

    template<class VB, int S2, class B2>
    inline void add_mJ(const Vector<S2, Precision, VB>& m, const Matrix<S2,Size, Precision, B2>& J, Precision weight = 1) {

        SizeMismatch<S2,S2>::test(m.size(), J.num_rows());
        for(int r=0; r < J.num_rows(); r++)
            add_mJ(m[r], J[r], weight);
            
    }

    template<int N, class B1, class B2, class B3>
    inline void add_mJ(const Vector<N,Precision,B1>& m,
                       const Matrix<Size,N,Precision,B2>& J,
                       const Matrix<N,N,Precision,B3>& invcov){
        const Matrix<Size,N,Precision> temp =  J * invcov;
        my_C_inv += temp * J.T();
        my_vector += temp * m;
    }

    template<int N, class B1, class B2, class B3>
    inline void add_mJ_rows(const Vector<N,Precision,B1>& m,
                       const Matrix<N,Size,Precision,B2>& J,
                       const Matrix<N,N,Precision,B3>& invcov){
        const Matrix<Size,N,Precision> temp =  J.T() * invcov;
        my_C_inv += temp * J;
        my_vector += temp * m;
    }

    template<int N, typename B1>
    inline void add_sparse_mJ(const Precision m,
                       const Vector<N,Precision,B1>& J1, const int index1,
                       const Precision weight = 1){
        //Upper right triangle only, for speed
        for(int r=0; r < J1.size(); r++)
        {
            double Jw = weight * J1[r];
            my_vector[r+index1] += m * Jw;
            for(int c = r; c < J1.size(); c++)
                my_C_inv[r+index1][c+index1] += Jw * J1[c];
        }
    }

    template<int N, int S1, class P1, class P2, class P3, class B1, class B2, class B3>
    inline void add_sparse_mJ_rows(const Vector<N,P1,B1>& m,
                       const Matrix<N,S1,P2,B2>& J1, const int index1,
                       const Matrix<N,N,P3,B3>& invcov){
        const Matrix<S1,N,Precision> temp1 = J1.T() * invcov;
        const int size1 = J1.num_cols();
        my_C_inv.slice(index1, index1, size1, size1) += temp1 * J1;
        my_vector.slice(index1, size1) += temp1 * m;
    }

    template<int N, int S1, int S2, class B1, class B2, class B3, class B4>
    inline void add_sparse_mJ_rows(const Vector<N,Precision,B1>& m,
                       const Matrix<N,S1,Precision,B2>& J1, const int index1,
                       const Matrix<N,S2,Precision,B3>& J2, const int index2,
                       const Matrix<N,N,Precision,B4>& invcov){
        const Matrix<S1,N,Precision> temp1 = J1.T() * invcov;
        const Matrix<S2,N,Precision> temp2 = J2.T() * invcov;
        const Matrix<S1,S2,Precision> mixed = temp1 * J2;
        const int size1 = J1.num_cols();
        const int size2 = J2.num_cols();
        my_C_inv.slice(index1, index1, size1, size1) += temp1 * J1;
        my_C_inv.slice(index2, index2, size2, size2) += temp2 * J2;
        my_C_inv.slice(index1, index2, size1, size2) += mixed;
        my_C_inv.slice(index2, index1, size2, size1) += mixed.T();
        my_vector.slice(index1, size1) += temp1 * m;
        my_vector.slice(index2, size2) += temp2 * m;
    }

    void compute(){
    
        //Copy the upper right triangle to the empty lower-left.
        for(int r=1; r < my_C_inv.num_rows(); r++)
            for(int c=0; c < r; c++)
                my_C_inv[r][c] = my_C_inv[c][r];

        my_decomposition.compute(my_C_inv);
        my_mu=my_decomposition.backsub(my_vector);
    }

    void operator += (const WLS& meas){
        my_vector+=meas.my_vector;
        my_C_inv += meas.my_C_inv;
    }

    Matrix<Size,Size,Precision>& get_C_inv() {return my_C_inv;}
    const Matrix<Size,Size,Precision>& get_C_inv() const {return my_C_inv;}
    Vector<Size,Precision>& get_mu(){return my_mu;}  
    const Vector<Size,Precision>& get_mu() const {return my_mu;} 
    Vector<Size,Precision>& get_vector(){return my_vector;} 
    const Vector<Size,Precision>& get_vector() const {return my_vector;} 
    Decomposition<Size,Precision>& get_decomposition(){return my_decomposition;} 
    const Decomposition<Size,Precision>& get_decomposition() const {return my_decomposition;} 


private:
    Matrix<Size,Size,Precision> my_C_inv;
    Vector<Size,Precision> my_vector;
    Decomposition<Size,Precision> my_decomposition;
    Vector<Size,Precision> my_mu;

    // comment out to allow bitwise copying
    // WLS( WLS& copyof );
    // int operator = ( WLS& copyof );
};

}

#endif