MRWatershedGraph.h

#pragma once
 
#include "MRGraph.h"
#include <cassert>
#include <cfloat>
 
namespace MR
{
 
class WatershedGraph
{
public:
    struct BasinInfo
    {
        VertId lowestVert; 
        float lowestLevel = FLT_MAX; 
        float area = 0;    
        float lowestBdLevel = FLT_MAX; 
        float maxVolume = 0; 
        float accVolume = 0; 
        float lastUpdateAmount = 0; 
        float lastMergeLevel = FLT_MAX; 
        float lastMergeVolume = 0; 
        Graph::EdgeId overflowVia; 
 
        BasinInfo() {} // Apparently I need this for `MR::Vector` to register default-constructibility inside the enclosing class.
 
        float amountTillOverflow() const
        {
            assert( !overflowVia );
            assert( maxVolume >= accVolume );
            return ( maxVolume - accVolume ) / area;
        }
 
        float approxLevel() const
        {
            assert( lastMergeLevel <= lowestBdLevel );
            assert( lastMergeVolume <= maxVolume );
            if ( maxVolume <= lastMergeVolume )
                return lowestBdLevel;
            const auto p = ( maxVolume - accVolume ) / ( maxVolume - lastMergeVolume );
            assert( p >= 0 && p <= 1 );
            return p * lastMergeLevel + ( 1 - p ) * lowestBdLevel;
        }
 
        void updateAccVolume( float amount )
        {
            assert( !overflowVia );
            assert( amount >= lastUpdateAmount );
            accVolume += ( amount - lastUpdateAmount ) * area;
            if ( accVolume > maxVolume ) // due to rounding errors
                accVolume = maxVolume;
            lastUpdateAmount = amount;
        }
    };
 
    struct BdInfo
    {
        VertId lowestVert; 
    };
 
public:
    MRMESH_API WatershedGraph( const Mesh & mesh, const Vector<int, FaceId> & face2basin, int numBasins );
 
    [[nodiscard]] MRMESH_API float getHeightAt( VertId v ) const;
 
    [[nodiscard]] const Graph & graph() const { return graph_; }
 
    [[nodiscard]] float totalArea() const { return totalArea_; }
 
    [[nodiscard]] int numBasins() const { return (int)graph_.validVerts().count() - 1; }
 
    [[nodiscard]] const BasinInfo & basinInfo( Graph::VertId v ) const { return basins_[v]; }
    [[nodiscard]] BasinInfo & basinInfo( Graph::VertId v ) { return basins_[v]; }
 
    [[nodiscard]] const BdInfo & bdInfo( Graph::EdgeId e ) const { return bds_[e]; }
    [[nodiscard]] BdInfo & bdInfo( Graph::EdgeId e ) { return bds_[e]; }
 
    [[nodiscard]] Graph::VertId outsideId() const { return outsideId_; }
 
    [[nodiscard]] MRMESH_API Graph::VertId getRootBasin( Graph::VertId v ) const;
 
    [[nodiscard]] MRMESH_API Graph::VertId flowsTo( Graph::VertId v ) const;
 
    [[nodiscard]] MRMESH_API Graph::VertId flowsFinallyTo( Graph::VertId v, bool exceptOutside = false ) const;
 
    MRMESH_API void setParentsToRoots();
 
    [[nodiscard]] MRMESH_API std::pair<Graph::EdgeId, float> findLowestBd() const;
 
    MRMESH_API Graph::VertId merge( Graph::VertId v0, Graph::VertId v1 );
 
    MRMESH_API Graph::VertId mergeViaBd( Graph::EdgeId bd );
 
    [[nodiscard]] MRMESH_API FaceBitSet getBasinFaces( Graph::VertId basin ) const;
 
    [[nodiscard]] MRMESH_API Vector<FaceBitSet, Graph::VertId> getAllBasinFaces( bool joinOverflowBasins = false ) const;
 
    [[nodiscard]] MRMESH_API FaceBitSet getBasinFacesBelowLevel( Graph::VertId basin, float waterLevel ) const;
 
    [[nodiscard]] MRMESH_API double computeBasinVolume( Graph::VertId basin, float waterLevel ) const;
 
    [[nodiscard]] MRMESH_API UndirectedEdgeBitSet getInterBasinEdges( bool joinOverflowBasins = false ) const;
 
    struct OverflowPoint
    {
        VertId v; // mesh vertex on the boundary of full basin and the other where it overflows
        Graph::VertId fullBasin;
        Graph::VertId overflowTo; // basin where the flow from v goes
    };
 
    [[nodiscard]] MRMESH_API std::vector<OverflowPoint> getOverflowPoints() const;
 
    [[nodiscard]] MRMESH_API Vector<Graph::VertId, Graph::VertId> iniBasin2Tgt( bool joinOverflowBasins = false ) const;
 
private:
    const Mesh & mesh_;
    const Vector<int, FaceId> & face2iniBasin_;
 
    Graph graph_;
    Vector<BasinInfo, Graph::VertId> basins_;
    Vector<BdInfo, Graph::EdgeId> bds_;
    float totalArea_ = 0;
 
    Graph::VertId outsideId_;
 
    Vector<Graph::VertId, Graph::VertId> parentBasin_;
};
 
} //namespace MR