/home/pclz-02/oschulz/Data/Science/Projects/Gerda/Software/blitzwave/src/WaveletDecomp.h

// Copyright (C) 2003-2013 Oliver Schulz <oliver.schulz@tu-dortmund.de>

// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.


#ifndef WAVELETDECOMP_H
#define WAVELETDECOMP_H

#include <cassert>
#include <iostream>
#include <blitz/array.h>

#include "Wavelet.h"

namespace bwave {



enum DecompType {
        STD_DECOMP=0,
        NONSTD_DECOMP=1
};



enum CoeffStorage {
        NESTED_COEFFS=0,
        SEPARATED_COEFFS=1
};



template<int tp_rank> class WaveletDecomp {
protected:
        Wavelet m_wavelet;
        DecompType m_decomp;
        CoeffStorage m_storageMode;
        ExtensionMode m_extMode;
        int m_maxLevel;
        blitz::TinyVector<bool, tp_rank> m_dimSelect;

        template<class tp_Type> inline void trafoStep
                (blitz::Array<tp_Type,tp_rank> &data, int targetDim, bool inverse) const;
                
        template<class tp_Type> inline blitz::TinyVector<int, tp_rank> waveletDecompose
                (blitz::Array<tp_Type,tp_rank> &data, int maxlevel=0) const;
        
        template<class tp_Type> inline blitz::TinyVector<int, tp_rank> waveletRecompose
                (blitz::Array<tp_Type,tp_rank> &data, int maxlevel=0) const;
        
public:
        const Wavelet& wavelet() const { return m_wavelet; }

        DecompType decompType() const { return m_decomp; }

        int maxLevel() const { return m_maxLevel; }

        CoeffStorage storageMode() const { return m_storageMode; }

        CoeffStorage storageMode(CoeffStorage newCS)
                { return m_storageMode=newCS; }

        ExtensionMode extensionMode() const { return m_extMode; }

        ExtensionMode extensionMode(ExtensionMode newEM)
                { return m_extMode=newEM; }

        const blitz::TinyVector<bool, tp_rank>& selectedDims() { return m_dimSelect; }

        bool dimSelected(int dim) const { return m_dimSelect(dim); }

        blitz::TinyVector<bool, tp_rank> dimSelection() const { return m_dimSelect; }

        blitz::TinyVector<bool, tp_rank> dimSelection(blitz::TinyVector<bool, tp_rank> selection)
                { return m_dimSelect = selection; }

        template<class tp_Type> inline blitz::TinyVector<int, tp_rank> apply
                (blitz::Array<tp_Type,tp_rank> &data) const
        {
                assert(storageMode()==NESTED_COEFFS);
                return waveletDecompose(data, m_maxLevel);
        }

        template<class tp_Type> inline blitz::TinyVector<int, tp_rank> applyInv
                (blitz::Array<tp_Type,tp_rank> &data) const
        {
                assert(storageMode()==NESTED_COEFFS);
                return waveletRecompose(data, m_maxLevel);
        }


        template<class tp_Type>
                blitz::Array< blitz::TinyVector<int, tp_rank>, 1> indices(
                        blitz::Array<tp_Type,tp_rank> &data) const;

        template<class tp_Type>
                blitz::Array<tp_Type,tp_rank> coeffs(blitz::Array<tp_Type,tp_rank> &data,
                        blitz::TinyVector<int, tp_rank> indices) const;
        
        double normFactor(blitz::TinyVector<int, tp_rank> indices) const {
                blitz::TinyVector<double, 2> baseNF(wavelet().normS(), wavelet().normD());
                double norm=1;
                for (int i=0; i<tp_rank; ++i) {
                        norm *= indices(i)>=0 ? pow(baseNF(0), indices(i))
                                              : pow(baseNF(0), -indices(i)-1) * baseNF(1);
                }
                return norm;
        }


        WaveletDecomp(WaveletDecomp<tp_rank> other, CoeffStorage cs)
                : m_wavelet(other.m_wavelet), m_decomp(other.m_decomp), m_storageMode(cs),
                  m_extMode(other.m_extMode), m_maxLevel(other.m_maxLevel), m_dimSelect(other.m_dimSelect)
        {}

        WaveletDecomp(Wavelet wl, DecompType decomp=NONSTD_DECOMP,
                      int maxlevel=0, CoeffStorage cs=NESTED_COEFFS, ExtensionMode em=CONSTANT_EXT)
                : m_wavelet(wl), m_decomp(decomp), m_storageMode(cs), m_extMode(em),
                  m_maxLevel(maxlevel)
          { m_dimSelect = true; }
};


template<int tp_rank> template<class tp_Type>
        inline void WaveletDecomp<tp_rank>::trafoStep
        (blitz::Array<tp_Type,tp_rank> &data, int targetDim, bool inverse) const
{
        using namespace blitz;
        assert( targetDim < tp_rank );
        TinyVector<int, tp_rank> stride=1, lboundS=data.lbound(), lboundD=data.lbound();
        stride(targetDim)=2;
        lboundD(targetDim)+=1;

        TinyVector<int, tp_rank> uboundS, uboundD;
        for (int dim=0; dim<tp_rank; ++dim) {
                uboundS(dim) = lboundS(dim) + (data.ubound()(dim)-lboundS(dim)) / stride(dim) * stride(dim);
                uboundD(dim) = lboundD(dim) + (data.ubound()(dim)-lboundD(dim)) / stride(dim) * stride(dim);
        }

        StridedDomain<tp_rank> subsetS(lboundS, uboundS, stride);
        StridedDomain<tp_rank> subsetD(lboundD, uboundD, stride);
        blitz::Array<tp_Type,tp_rank> s( data(subsetS) );
        blitz::Array<tp_Type,tp_rank> d( data(subsetD) );

        int n = data.size() / data.extent(targetDim);

        TinyVector<int, tp_rank> position = data.lbound();
        int lastDim = targetDim!=tp_rank-1 ? tp_rank-1 : tp_rank-2;

        for (int i=0; i<n; ++i) {
                Array<tp_Type, 1> sliceS( slice(s, targetDim, position) );
                Array<tp_Type, 1> sliceD( slice(d, targetDim, position) );
                if (!inverse) m_wavelet.lift(sliceS, sliceD, m_extMode);
                else m_wavelet.invLift(sliceS, sliceD, m_extMode);

                if (i < n-1) {
                        position(lastDim) += 1;
                        for (int d=lastDim; d>0; --d) {
                                if ( (d!=targetDim) && (position(d)>data.ubound(d)) ) {
                                        position(d)=data.lbound(d);
                                        ++position(targetDim!=d-1 ? d-1 : d-2);
                                }
                        }
                }
        }

}


template<int tp_rank> template<class tp_Type>
        inline blitz::TinyVector<int, tp_rank>
        WaveletDecomp<tp_rank>::waveletDecompose
        (blitz::Array<tp_Type,tp_rank> &data, int maxlevel) const
{
        using namespace blitz;

        TinyVector<int, tp_rank> depth=0;

        switch (decompType()) {
                case STD_DECOMP: assert(false); break;
                case NONSTD_DECOMP: {
                        for (int dim=0; dim<tp_rank; ++dim)
                                if ( (data.extent(dim)>1) && dimSelected(dim) ) {
                                        trafoStep(data, dim, false);
                                        depth(dim) += 1;
                                }
                                        
                        bool descend = false;
                        for (int dim=0; dim<tp_rank; ++dim)
                                if ((data.extent(dim)>2) && dimSelected(dim)) descend = true;

                        if ( descend && ((maxlevel==0) || (maxlevel>1)) ) {
                                TinyVector<int, tp_rank> stride, lbound, ubound;
                                for (int dim=0; dim<tp_rank; ++dim) {
                                        stride(dim) = dimSelected(dim) ? 2 : 1;
                                        lbound(dim) = data.lbound()(dim);
                                        ubound(dim) = lbound(dim) + (data.ubound()(dim)-lbound(dim)) / stride(dim) * stride(dim);
                                }
                                StridedDomain<tp_rank> subsetS(lbound, ubound, stride);
                                blitz::Array<tp_Type,tp_rank> sub(data(subsetS));
                
                                depth += waveletDecompose(sub, maxlevel>0 ? maxlevel-1 : 0);
                        }
                        
                        break;
                }
                default: assert(false);
        }
        
        return depth;
}


template<int tp_rank> template<class tp_Type>
        inline blitz::TinyVector<int, tp_rank>
        WaveletDecomp<tp_rank>::waveletRecompose
        (blitz::Array<tp_Type,tp_rank> &data, int maxlevel) const
{
        using namespace blitz;

        TinyVector<int, tp_rank> depth=0;

        switch (decompType()) {
                case STD_DECOMP: assert(false); break;
                case NONSTD_DECOMP: {
                        bool descend = false;
                        for (int dim=0; dim<tp_rank; ++dim)
                                if ((data.extent(dim)>2) && dimSelected(dim)) descend = true;

                        if ( descend && ((maxlevel==0) || (maxlevel>1)) ) {
                                TinyVector<int, tp_rank> stride, lbound, ubound;
                                for (int dim=0; dim<tp_rank; ++dim) {
                                        stride(dim) = dimSelected(dim) ? 2 : 1;
                                        lbound(dim) = data.lbound()(dim);
                                        ubound(dim) = lbound(dim) + (data.ubound()(dim)-lbound(dim)) / stride(dim) * stride(dim);
                                }
                                StridedDomain<tp_rank> subsetS(lbound, ubound, stride);
                                blitz::Array<tp_Type,tp_rank> sub(data(subsetS));
                
                                depth += waveletRecompose(sub, maxlevel>0 ? maxlevel-1 : 0);
                        }
                
                        for (int dim=tp_rank-1; dim>=0; --dim)
                                if ( (data.extent(dim)>1) && dimSelected(dim) ) {
                                        trafoStep(data, dim, true);
                                        depth(dim) -= 1;
                                }
                        
                        break;
                }
                default: assert(false);
        }       

        return depth;
}


template<int tp_rank> template<class tp_Type>
        blitz::Array< blitz::TinyVector<int, tp_rank>, 1>
        WaveletDecomp<tp_rank>::indices(
                blitz::Array<tp_Type,tp_rank> &data) const
{
        using namespace blitz;
        TinyVector<int, tp_rank> decompDepth;
        Array< TinyVector<int, tp_rank>, 1> idx;

        for (int dim=0; dim<tp_rank; ++dim) {
                // log(...-0.1) for numerical stability
                int val = (int) ceil(log((double)data.extent(dim)-0.1) / log(2.0));
                decompDepth(dim) = dimSelected(dim) ?
                        (m_maxLevel>0 ? min(m_maxLevel, val) : val) : 0;
        }

        switch (decompType()) {
                case STD_DECOMP: assert(false); break;
                case NONSTD_DECOMP: {
                        int maxLevel = max(decompDepth);
                        idx.resize(max(1, maxLevel*(1<<tp_rank)));
                        idx(0) = 0;
                        int begin=0, end=0;
                        for (int level=1; level<=maxLevel; ++level) {
                                for (int dim=0; dim<tp_rank; ++dim)
                                        if (decompDepth(dim)>=level)
                                {
                                        int offset = end-begin+1;
                                        for (int i=begin; i<=end; ++i) {
                                                idx(i+offset) = idx(i);
                                                idx(i)(dim) = -level;
                                                idx(i+offset)(dim) = level;
                                        }
                                        end += offset;
                                }
                                begin = end;
                        }
                        idx.resizeAndPreserve(end+1);
                        break;
                }
                default: assert(false);
        }       
        return idx;
}


template<int tp_rank> template<class tp_Type>
        blitz::Array<tp_Type,tp_rank> WaveletDecomp<tp_rank>::coeffs(
                blitz::Array<tp_Type,tp_rank> &data,
                blitz::TinyVector<int, tp_rank> indices) const
{
        using namespace blitz;

        Array<tp_Type,tp_rank> c;
        TinyVector<int, tp_rank> ubound, lbound, stride;
        
        for (int dim=0; dim<tp_rank; ++dim) {
                int i = indices(dim);
                // Remark: i>=0 specifies phi(i) coeff,
                //          i<0 specifies psi(-i) coeff;
                int n = data.extent(dim);
                
                if (storageMode()==NESTED_COEFFS) {
                        // Remark: coefficient location in data(l)
                        //         for phi(i):       0 <= l <= n-1, stride = 2^i
                        //         for psi(i): 2^(i-1) <= l <= n-1, stride = 2^i
                        stride(dim) = i>=0 ? 1 << i : 1 << -i;
                        lbound(dim) = data.lbound()(dim) + ( i>=0 ? 0 : 1 << (-i-1) );
                        ubound(dim) = lbound(dim) + (data.ubound()(dim)-lbound(dim)) / stride(dim) * stride(dim);
                } else if (storageMode()==SEPARATED_COEFFS) {
                        // Remark: coefficient location in data(l)
                        //         for phi(i):           0 <= l <= ceil(n/2^i)-1
                        //         for psi(i): ceil(n/2^i) <= l <= ceil(n/2^(i-1))-1
                        // Remark: ceil(n/2^i) = (n + (1<<i) - 1) >> i;
                        stride(dim) = 1;
                        lbound(dim) = data.lbound()(dim) +  ( i>=0 ? 0 : (n + (1<<-i) - 1) >> -i );
                        ubound(dim) = data.lbound()(dim) +  ( i>=0 ? ( (n + (1<<i) - 1) >> i) - 1
                                                                                                           : ( (n + (1<<(-i-1)) - 1) >> (-i-1) ) - 1 );
                } else assert(false);

                assert( lbound(dim) <= ubound(dim));
                assert( lbound(dim) >= data.lbound(dim));
                assert( ubound(dim) <= data.ubound(dim));
        }

        StridedDomain<tp_rank> subset(lbound, ubound, stride);
        c.reference(data(subset));

        return c;
}


} // namespace bwave

#endif