AutoDiffBlock.hpp

Go to the documentation of this file.

/*
  Copyright 2013 SINTEF ICT, Applied Mathematics.

  This file is part of the Open Porous Media project (OPM).

  OPM 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 3 of the License, or
  (at your option) any later version.

  OPM 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 OPM.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef OPM_AUTODIFFBLOCK_HEADER_INCLUDED
#define OPM_AUTODIFFBLOCK_HEADER_INCLUDED

#include <opm/common/utility/platform_dependent/disable_warnings.h>

#include <Eigen/Eigen>
#include <Eigen/Sparse>
#include <opm/autodiff/fastSparseProduct.hpp>

#include <opm/common/utility/platform_dependent/reenable_warnings.h>

#include <opm/autodiff/AutoDiffMatrix.hpp>


#include <utility>
#include <vector>
#include <cassert>
#include <iostream>

namespace Opm
{

    template <typename Scalar>
    class AutoDiffBlock
    {
    public:
        typedef Eigen::Array<Scalar, Eigen::Dynamic, 1> V;
        typedef AutoDiffMatrix M;

        static AutoDiffBlock null()
        {
            return AutoDiffBlock(V(), {});
        }

        static AutoDiffBlock constant(V&& val)
        {
            return AutoDiffBlock(std::move(val));
        }

        static AutoDiffBlock constant(const V& val)
        {
            return AutoDiffBlock(val);
        }

        static AutoDiffBlock constant(const V& val, const std::vector<int>& blocksizes)
        {
            std::vector<M> jac;
            const int num_elem = val.size();
            const int num_blocks = blocksizes.size();
            // For constants, all jacobian blocks are zero.
            jac.resize(num_blocks);
            for (int i = 0; i < num_blocks; ++i) {
                jac[i] = M(num_elem, blocksizes[i]);
            }
            V val_copy(val);
            return AutoDiffBlock(std::move(val_copy), std::move(jac));
        }

        static AutoDiffBlock variable(const int index, V&& val, const std::vector<int>& blocksizes)
        {
            std::vector<M> jac;
            const int num_elem = val.size();
            const int num_blocks = blocksizes.size();
            // First, set all jacobian blocks to zero...
            jac.resize(num_blocks);
            for (int i = 0; i < num_blocks; ++i) {
                jac[i] = M(num_elem, blocksizes[i]);
            }
            // ... then set the one corrresponding to this variable to identity.
            assert(blocksizes[index] == num_elem);
            jac[index] = M::createIdentity(val.size());
            return AutoDiffBlock(std::move(val), std::move(jac));
        }

        static AutoDiffBlock variable(const int index, const V& val, const std::vector<int>& blocksizes)
        {
            V value = val;
            return variable(index, std::move(value), blocksizes);
        }

        static AutoDiffBlock function(V&& val, std::vector<M>&& jac)
        {
            return AutoDiffBlock(std::move(val), std::move(jac));
        }

        static AutoDiffBlock function(const V& val, const std::vector<M>& jac)
        {
            V val_copy(val);
            std::vector<M> jac_copy(jac);
            return AutoDiffBlock(std::move(val_copy), std::move(jac_copy));
        }

        static std::vector<AutoDiffBlock> variables(const std::vector<V>& initial_values)
        {
            const int num_vars = initial_values.size();
            std::vector<int> bpat;
            for (int v = 0; v < num_vars; ++v) {
                bpat.push_back(initial_values[v].size());
            }
            std::vector<AutoDiffBlock> vars;
            for (int v = 0; v < num_vars; ++v) {
                vars.emplace_back(variable(v, initial_values[v], bpat));
            }
            return vars;
        }

        AutoDiffBlock& operator+=(const AutoDiffBlock& rhs)
        {
            if (jac_.empty()) {
                jac_ = rhs.jac_;
            } else if (!rhs.jac_.empty()) {
                assert (numBlocks()    == rhs.numBlocks());
                assert (value().size() == rhs.value().size());

                const int num_blocks = numBlocks();
#pragma omp parallel for schedule(static)
                for (int block = 0; block < num_blocks; ++block) {
                    assert(jac_[block].rows() == rhs.jac_[block].rows());
                    assert(jac_[block].cols() == rhs.jac_[block].cols());
                    jac_[block] += rhs.jac_[block];
                }
            }

            val_ += rhs.val_;

            return *this;
        }

        AutoDiffBlock& operator-=(const AutoDiffBlock& rhs)
        {
            if (jac_.empty()) {
                const int num_blocks = rhs.numBlocks();
                jac_.resize(num_blocks);
#pragma omp parallel for schedule(static)
                for (int block = 0; block < num_blocks; ++block) {
                    jac_[block] = rhs.jac_[block] * (-1.0);
                }
            } else if (!rhs.jac_.empty()) {
                assert (numBlocks()    == rhs.numBlocks());
                assert (value().size() == rhs.value().size());

                const int num_blocks = numBlocks();
#pragma omp parallel for schedule(static)
                for (int block = 0; block < num_blocks; ++block) {
                    assert(jac_[block].rows() == rhs.jac_[block].rows());
                    assert(jac_[block].cols() == rhs.jac_[block].cols());
                    jac_[block] -= rhs.jac_[block];
                }
            }

            val_ -= rhs.val_;

            return *this;
        }

        AutoDiffBlock operator+(const AutoDiffBlock& rhs) const
        {
            if (jac_.empty() && rhs.jac_.empty()) {
                return constant(val_ + rhs.val_);
            }
            if (jac_.empty()) {
                return val_ + rhs;
            }
            if (rhs.jac_.empty()) {
                return *this + rhs.val_;
            }
            std::vector<M> jac = jac_;
            assert(numBlocks() == rhs.numBlocks());
            int num_blocks = numBlocks();
#pragma omp parallel for schedule(static)
            for (int block = 0; block < num_blocks; ++block) {
                assert(jac[block].rows() == rhs.jac_[block].rows());
                assert(jac[block].cols() == rhs.jac_[block].cols());
                jac[block] += rhs.jac_[block];
            }
            return function(val_ + rhs.val_, std::move(jac));
        }

        AutoDiffBlock operator-(const AutoDiffBlock& rhs) const
        {
            if (jac_.empty() && rhs.jac_.empty()) {
                return constant(val_ - rhs.val_);
            }
            if (jac_.empty()) {
                return val_ - rhs;
            }
            if (rhs.jac_.empty()) {
                return *this - rhs.val_;
            }
            std::vector<M> jac = jac_;
            assert(numBlocks() == rhs.numBlocks());
            int num_blocks = numBlocks();
#pragma omp parallel for schedule(static)
            for (int block = 0; block < num_blocks; ++block) {
                assert(jac[block].rows() == rhs.jac_[block].rows());
                assert(jac[block].cols() == rhs.jac_[block].cols());
                jac[block] -= rhs.jac_[block];
            }
            return function(val_ - rhs.val_, std::move(jac));
        }

        AutoDiffBlock operator*(const AutoDiffBlock& rhs) const
        {
            if (jac_.empty() && rhs.jac_.empty()) {
                return constant(val_ * rhs.val_);
            }
            if (jac_.empty()) {
                return val_ * rhs;
            }
            if (rhs.jac_.empty()) {
                return *this * rhs.val_;
            }
            int num_blocks = numBlocks();
            std::vector<M> jac(num_blocks);
            assert(numBlocks() == rhs.numBlocks());
            M D1(val_.matrix().asDiagonal());
            M D2(rhs.val_.matrix().asDiagonal());
#pragma omp parallel for schedule(dynamic)
            for (int block = 0; block < num_blocks; ++block) {
                assert(jac_[block].rows() == rhs.jac_[block].rows());
                assert(jac_[block].cols() == rhs.jac_[block].cols());
                if( jac_[block].nonZeros() == 0 && rhs.jac_[block].nonZeros() == 0 ) {
                    jac[block] = M( D2.rows(), jac_[block].cols() );
                }
                else if( jac_[block].nonZeros() == 0 )
                    jac[block] = D1*rhs.jac_[block];
                else if ( rhs.jac_[block].nonZeros() == 0 ) {
                    jac[block] = D2*jac_[block];
                }
                else {
                    jac[block] = D2*jac_[block] + D1*rhs.jac_[block];
                }
            }
            return function(val_ * rhs.val_, std::move(jac));
        }

        AutoDiffBlock operator/(const AutoDiffBlock& rhs) const
        {
            if (jac_.empty() && rhs.jac_.empty()) {
                return constant(val_ / rhs.val_);
            }
            if (jac_.empty()) {
                return val_ / rhs;
            }
            if (rhs.jac_.empty()) {
                return *this / rhs.val_;
            }
            int num_blocks = numBlocks();
            std::vector<M> jac(num_blocks);
            assert(numBlocks() == rhs.numBlocks());
            M D1(val_.matrix().asDiagonal());
            M D2(rhs.val_.matrix().asDiagonal());
            M D3((1.0/(rhs.val_*rhs.val_)).matrix().asDiagonal());
#pragma omp parallel for schedule(dynamic)
            for (int block = 0; block < num_blocks; ++block) {
                assert(jac_[block].rows() == rhs.jac_[block].rows());
                assert(jac_[block].cols() == rhs.jac_[block].cols());
                if( jac_[block].nonZeros() == 0 && rhs.jac_[block].nonZeros() == 0 ) {
                    jac[block] = M( D3.rows(), jac_[block].cols() );
                }
                else if( jac_[block].nonZeros() == 0 ) {
                    jac[block] =  D3 * ( D1*rhs.jac_[block]) * (-1.0);
                }
                else if ( rhs.jac_[block].nonZeros() == 0 )
                {
                    jac[block] = D3 * (D2*jac_[block]);
                }
                else {
                    jac[block] = D3 * (D2*jac_[block] + (D1*rhs.jac_[block]*(-1.0)));
                }
            }
            return function(val_ / rhs.val_, std::move(jac));
        }

        template <class Ostream>
        Ostream&
        print(Ostream& os) const
        {
            int num_blocks = jac_.size();
            os << "Value =\n" << val_ << "\n\nJacobian =\n";
            for (int i = 0; i < num_blocks; ++i) {
                Eigen::SparseMatrix<double> m;
                jac_[i].toSparse(m);
                os << "Sub Jacobian #" << i << '\n' << m << "\n";
            }
            return os;
        }

        void swap(AutoDiffBlock& other)
        {
            val_.swap(other.val_);
            jac_.swap(other.jac_);
        }

        int size() const
        {
            return val_.size();
        }

        int numBlocks() const
        {
            return jac_.size();
        }

        std::vector<int> blockPattern() const
        {
            const int nb = numBlocks();
            std::vector<int> bp(nb);
            for (int block = 0; block < nb; ++block) {
                bp[block] = jac_[block].cols();
            }
            return bp;
        }

        const V& value() const
        {
            return val_;
        }

        const std::vector<M>& derivative() const
        {
            return jac_;
        }

    private:
        AutoDiffBlock(const V& val)
            : val_(val)
        {
        }

        AutoDiffBlock(V&& val)
            : val_(std::move(val))
        {
        }

        AutoDiffBlock(V&& val, std::vector<M>&& jac)
            : val_(std::move(val)), jac_(std::move(jac))
        {
#ifndef NDEBUG
            const int num_elem = val_.size();
            const int num_blocks = jac_.size();
            for (int block = 0; block < num_blocks; ++block) {
                assert(num_elem == jac_[block].rows());
            }
#endif
        }

        V val_;
        std::vector<M> jac_;
    };


    // ---------  Free functions and operators for AutoDiffBlock  ---------

    template <class Ostream, typename Scalar>
    Ostream&
    operator<<(Ostream& os, const AutoDiffBlock<Scalar>& fw)
    {
        return fw.print(os);
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator*(const typename AutoDiffBlock<Scalar>::M& lhs,
                                    const AutoDiffBlock<Scalar>& rhs)
    {
        int num_blocks = rhs.numBlocks();
        std::vector<typename AutoDiffBlock<Scalar>::M> jac(num_blocks);
        assert(lhs.cols() == rhs.value().rows());
#pragma omp parallel for schedule(dynamic)
        for (int block = 0; block < num_blocks; ++block) {
            fastSparseProduct(lhs, rhs.derivative()[block], jac[block]);
        }
        typename AutoDiffBlock<Scalar>::V val = lhs*rhs.value().matrix();
        return AutoDiffBlock<Scalar>::function(std::move(val), std::move(jac));
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator*(const Eigen::SparseMatrix<Scalar>& lhs,
                                    const AutoDiffBlock<Scalar>& rhs)
    {
        int num_blocks = rhs.numBlocks();
        std::vector<typename AutoDiffBlock<Scalar>::M> jac(num_blocks);
        assert(lhs.cols() == rhs.value().rows());
        for (int block = 0; block < num_blocks; ++block) {
            fastSparseProduct(lhs, rhs.derivative()[block], jac[block]);
        }
        typename AutoDiffBlock<Scalar>::V val = lhs*rhs.value().matrix();
        return AutoDiffBlock<Scalar>::function(std::move(val), std::move(jac));
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator*(const typename AutoDiffBlock<Scalar>::V& lhs,
                                    const AutoDiffBlock<Scalar>& rhs)
    {
        return AutoDiffBlock<Scalar>::constant(lhs, rhs.blockPattern()) * rhs;
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator*(const AutoDiffBlock<Scalar>& lhs,
                                    const typename AutoDiffBlock<Scalar>::V& rhs)
    {
        return rhs * lhs; // Commutative operation.
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator+(const typename AutoDiffBlock<Scalar>::V& lhs,
                                    const AutoDiffBlock<Scalar>& rhs)
    {
        return AutoDiffBlock<Scalar>::constant(lhs, rhs.blockPattern()) + rhs;
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator+(const AutoDiffBlock<Scalar>& lhs,
                                    const typename AutoDiffBlock<Scalar>::V& rhs)
    {
        return rhs + lhs; // Commutative operation.
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator-(const typename AutoDiffBlock<Scalar>::V& lhs,
                                    const AutoDiffBlock<Scalar>& rhs)
    {
        return AutoDiffBlock<Scalar>::constant(lhs, rhs.blockPattern()) - rhs;
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator-(const AutoDiffBlock<Scalar>& lhs,
                                    const typename AutoDiffBlock<Scalar>::V& rhs)
    {
        return lhs - AutoDiffBlock<Scalar>::constant(rhs, lhs.blockPattern());
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator/(const typename AutoDiffBlock<Scalar>::V& lhs,
                                    const AutoDiffBlock<Scalar>& rhs)
    {
        return AutoDiffBlock<Scalar>::constant(lhs, rhs.blockPattern()) / rhs;
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator/(const AutoDiffBlock<Scalar>& lhs,
                                    const typename AutoDiffBlock<Scalar>::V& rhs)
    {
        return lhs / AutoDiffBlock<Scalar>::constant(rhs, lhs.blockPattern());
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator*(const AutoDiffBlock<Scalar>& lhs,
                                    const Scalar& rhs)
    {
        std::vector< typename AutoDiffBlock<Scalar>::M > jac;
        jac.reserve( lhs.numBlocks() );
        for (int block=0; block<lhs.numBlocks(); block++) {
            jac.emplace_back( lhs.derivative()[block] * rhs );
        }
        return AutoDiffBlock<Scalar>::function( lhs.value() * rhs, std::move(jac) );
    }


    template <typename Scalar>
    AutoDiffBlock<Scalar> operator*(const Scalar& lhs,
                                    const AutoDiffBlock<Scalar>& rhs)
    {
        return rhs * lhs; // Commutative operation.
    }


} // namespace Opm



#endif // OPM_AUTODIFFBLOCK_HEADER_INCLUDED