DofVector.h

Go to the documentation of this file.

#pragma once

#include <map>
#include <vector>
#include <eigen3/Eigen/Core>
#include "nuto/mechanics/dofs/DofType.h"

namespace NuTo
{

template <typename T>
class DofVector
{
public:
    const Eigen::Matrix<T, Eigen::Dynamic, 1>& operator[](const DofType& dofType) const
    {
        return mData.at(dofType);
    }

    Eigen::Matrix<T, Eigen::Dynamic, 1>& operator[](const DofType& dofType)
    {
        return mData[dofType];
    }

    DofVector& operator+=(const DofVector& rhs)
    {
        for (auto& entry : rhs.mData)
        {
            if (mData.find(entry.first) != mData.end())
                mData[entry.first] += entry.second;
            else
                mData[entry.first] = entry.second;
        }
        return *this;
    }

    DofVector& operator-=(const DofVector& rhs)
    {
        for (auto& entry : rhs.mData)
        {
            if (mData.find(entry.first) != mData.end())
                mData[entry.first] -= entry.second;
            else
                mData[entry.first] = -entry.second;
        }
        return *this;
    }

    DofVector& operator*=(double scalar)
    {
        for (auto& entry : mData)
            entry.second *= scalar;
        return *this;
    }

    friend DofVector operator+(DofVector lhs, const DofVector& rhs)
    {
        lhs += rhs;
        return lhs;
    }

    friend DofVector operator-(DofVector lhs, const DofVector& rhs)
    {
        lhs -= rhs;
        return lhs;
    }

    friend DofVector operator*(DofVector lhs, double scalar)
    {
        lhs *= scalar;
        return lhs;
    }

    friend DofVector operator*(double scalar, DofVector rhs)
    {
        rhs *= scalar;
        return rhs;
    }

    double operator()(DofType dof, int globalDofNumber) const
    {
        return (*this)[dof][globalDofNumber];
    }

    void SetZero()
    {
        for (auto& entry : mData)
        {
            entry.second.setZero();
        }
    }

    std::vector<double> operator()(DofType dof, std::vector<int> globalDofNumbers) const
    {
        std::vector<double> v;
        v.reserve(globalDofNumbers.size());
        for (int globalDofNumber : globalDofNumbers)
            v.push_back((*this)(dof, globalDofNumber));
        return v;
    }

    friend std::ostream& operator<<(std::ostream& out, const DofVector& v)
    {
        for (auto& entry : v.mData)
        {
            out << "==== " << entry.first.GetName() << "====\n";
            out << entry.second << '\n';
        }
        return out;
    }

    void AddScaled(const DofVector& rhs, double scalar)
    {
        for (auto& entry : rhs.mData)
        {
            if (mData.find(entry.first) != mData.end())
                mData[entry.first] += entry.second * scalar;
            else
                mData[entry.first] = entry.second * scalar;
        }
    }

    std::vector<DofType> DofTypes() const
    {
        std::vector<DofType> dofTypes;
        for (const auto& data : mData)
            dofTypes.push_back(data.first);
        return dofTypes;
    }


protected:
    std::map<DofType, Eigen::Matrix<T, Eigen::Dynamic, 1>, CompareDofType> mData;
};

} /* NuTo */