Merge pull request #16 from Geode-solutions/feat_refactor_gibbs_energy

feat(GibbsEnergy): refactor energy terms

Author: BotellaA

include/geode/stochastic/sampling/mcmc/metropolis_hasting_sampler.hpp

@@ -188,7 +188,7 @@ namespace geode
         {
             const auto new_object = proposal.new_object();
             const auto delta_log_energy =
-                energy_.delta_log_energy_add( state, new_object );
+                energy_.delta_log_add( state, new_object );
             return accept_or_reject( proposal, state, engine, delta_log_energy,
                 []( auto& state, auto& proposal ) {
                     state.add_object(
@@ -203,7 +203,7 @@ namespace geode
         {
             const auto old_object_id = proposal.old_object_id();
             const auto delta_log_energy =
-                energy_.delta_log_energy_remove( state, old_object_id );
+                energy_.delta_log_remove( state, old_object_id );
             return accept_or_reject( proposal, state, engine, delta_log_energy,
                 []( auto& state, auto& proposal ) {
                     state.remove_object( proposal.old_object_id() );
@@ -216,8 +216,8 @@ namespace geode
         {
             const auto new_object = proposal.new_object();
             const auto old_object_id = proposal.old_object_id();
-            const auto delta_log_energy = energy_.delta_log_energy_change(
-                state, old_object_id, new_object );
+            const auto delta_log_energy =
+                energy_.delta_log_change( state, old_object_id, new_object );
             // should we test that objects are in the same group?
             // should be ensured by the dynamic
             return accept_or_reject( proposal, state, engine, delta_log_energy,

include/geode/stochastic/sampling/mcmc/models/components/density_term.hpp

@@ -33,14 +33,11 @@ namespace geode
     class DensityTerm : public EnergyTerm< ObjectType >
     {
     public:
-        explicit DensityTerm( std::string_view name, double lambda )
-            : EnergyTerm< ObjectType >( name, lambda )
-        {
-        }
-
-        explicit DensityTerm(
-            std::string_view name, double lambda, const uuid& subset_id )
-            : EnergyTerm< ObjectType >( name, lambda, subset_id )
+        explicit DensityTerm( std::string_view name,
+            double lambda,
+            absl::flat_hash_set< uuid > targeted_subset_ids )
+            : EnergyTerm< ObjectType >(
+                  name, lambda, std::move( targeted_subset_ids ) )
         {
         }
 
@@ -79,12 +76,12 @@ namespace geode
 
         double statistic( const ObjectSet< ObjectType >& state ) const override
         {
-            if( this->targeted_subset_id() )
+            index_t count{ 0 };
+            for( const auto& subset_uuid : this->targeted_subset_ids() )
             {
-                return static_cast< double >( state.nb_objects_in_subset(
-                    this->targeted_subset_id().value() ) );
+                count += state.nb_objects_in_subset( subset_uuid );
             }
-            return static_cast< double >( state.nb_objects() );
+            return static_cast< double >( count );
         }
     };
 } // namespace geode

include/geode/stochastic/sampling/mcmc/models/components/energy_term.hpp

@@ -22,8 +22,14 @@
  */
 #pragma once
 
+#include <absl/container/flat_hash_set.h>
+
+#include <geode/basic/identifier.hpp>
+#include <geode/basic/identifier_builder.hpp>
+
 #include <geode/stochastic/common.hpp>
 #include <geode/stochastic/spatial/object_set.hpp>
+
 #include <optional>
 
 namespace geode
@@ -93,43 +99,30 @@ namespace geode
     //    };
 
     template < typename ObjectType >
-    class EnergyTerm
+    class EnergyTerm : public Identifier
     {
     public:
-        explicit EnergyTerm( std::string_view name, double param )
-            : name_{ name }, energy_scale_{ param }
-        {
-        }
-
         explicit EnergyTerm( std::string_view name,
             double param,
-            const uuid& targeted_subset_id )
-            : name_{ name },
+            absl::flat_hash_set< uuid >&& targeted_subset_ids )
+            : Identifier{},
               energy_scale_{ param },
-              targeted_subset_id_{ targeted_subset_id }
+              targeted_subset_ids_{ std::move( targeted_subset_ids ) }
         {
+            IdentifierBuilder builder( *this );
+            builder.set_name( name );
         }
 
         virtual ~EnergyTerm() = default;
 
-        const uuid& id() const
-        {
-            return energy_term_id_;
-        }
-
-        std::string_view name() const
-        {
-            return name_;
-        }
-
         double parameter() const
         {
             return energy_scale_.parameter();
         }
 
-        std::optional< uuid > targeted_subset_id() const
+        const absl::flat_hash_set< uuid >& targeted_subset_ids() const
         {
-            return targeted_subset_id_;
+            return targeted_subset_ids_;
         }
 
         /// Energy contribution for a given statistic multiplier
@@ -158,45 +151,39 @@ namespace geode
         {
             auto message =
                 absl::StrCat( "Term : ", name(), "; uuid: ", id().string(),
-                    " parameter value: ", energy_scale_.parameter() );
-            if( targeted_subset_id_ )
+                    " parameter value: ", energy_scale_.parameter(),
+                    " applyied on ", targeted_subset_ids_.size(),
+                    " object subsets -->" );
+            for( const auto& subset_uuid : targeted_subset_ids_ )
             {
-                absl::StrAppend( &message, " targetted subset: ",
-                    targeted_subset_id_.value().string() );
+                absl::StrAppend( &message, "\t", subset_uuid.string() );
             }
             return message;
         }
 
     protected:
         bool is_targeted_subset( const uuid& subset_id ) const
         {
-            return !targeted_subset_id_ || subset_id == *targeted_subset_id_;
+            return targeted_subset_ids_.find( subset_id )
+                   != targeted_subset_ids_.end();
         }
 
         template < typename Func >
         void for_each_targeted_object(
             const ObjectSet< ObjectType >& state, Func&& do_apply ) const
         {
-            if( targeted_subset_id_ )
+            for( const auto& targeted_subset_id : targeted_subset_ids_ )
             {
                 for( const auto id : geode::Range{
-                         state.nb_objects_in_subset( *targeted_subset_id_ ) } )
+                         state.nb_objects_in_subset( targeted_subset_id ) } )
                 {
-                    do_apply( { id, *targeted_subset_id_ } );
+                    do_apply( ObjectId{ id, targeted_subset_id } );
                 }
-                return;
-            }
-            for( const auto& id : state.get_all_object() )
-            {
-                do_apply( id );
             }
         }
 
     private:
-        std::string name_;
         detail::EnergyScale energy_scale_;
-
-        std::optional< uuid > targeted_subset_id_{};
-        uuid energy_term_id_{};
+        absl::flat_hash_set< uuid > targeted_subset_ids_;
     };
 } // namespace geode

include/geode/stochastic/sampling/mcmc/models/components/energy_term_collection.hpp

@@ -0,0 +1,116 @@
+#pragma once
+
+#include <absl/container/flat_hash_map.h>
+#include <geode/stochastic/sampling/mcmc/models/components/energy_term.hpp>
+
+namespace geode
+{
+    template < typename ObjectType >
+    class EnergyTermCollection
+    {
+        OPENGEODE_DISABLE_COPY( EnergyTermCollection );
+
+    public:
+        EnergyTermCollection() = default;
+        EnergyTermCollection( EnergyTermCollection&& ) noexcept = default;
+        ~EnergyTermCollection() = default;
+
+        uuid add_energy_term( std::shared_ptr< EnergyTerm< ObjectType > > term )
+        {
+            const uuid id = term->id();
+            energy_terms_.emplace( id, term );
+            for( const uuid& subset_id : term->targeted_subset_ids() )
+            {
+                subset_to_terms_[subset_id].push_back( term );
+            }
+            return id;
+        }
+
+        bool remove_energy_term( const uuid& id )
+        {
+            auto it = energy_terms_.find( id );
+            if( it == energy_terms_.end() )
+                return false;
+
+            auto term = it->second;
+
+            for( const uuid& subset_id : term->targeted_subset_ids() )
+            {
+                auto vec_it = subset_to_terms_.find( subset_id );
+                if( vec_it == subset_to_terms_.end() )
+                    continue;
+
+                auto& vec = vec_it->second;
+                vec.erase(
+                    std::remove( vec.begin(), vec.end(), term ), vec.end() );
+
+                if( vec.empty() )
+                    subset_to_terms_.erase( vec_it );
+            }
+
+            energy_terms_.erase( it );
+            return true;
+        }
+
+        void clear()
+        {
+            energy_terms_.clear();
+            subset_to_terms_.clear();
+        }
+
+        [[nodiscard]] index_t size() const
+        {
+            return energy_terms_.size();
+        }
+
+        [[nodiscard]] std::shared_ptr< const EnergyTerm< ObjectType > > get(
+            const uuid& id ) const
+        {
+            auto it = energy_terms_.find( id );
+            OPENGEODE_EXCEPTION( it != energy_terms_.end(),
+                absl::StrCat( "[EnergyTermCollection] Unknown energy term: ",
+                    id.string() ) );
+            return it->second;
+        }
+
+        [[nodiscard]] const absl::flat_hash_map< uuid,
+            std::shared_ptr< EnergyTerm< ObjectType > > >&
+            all_terms() const
+        {
+            return energy_terms_;
+        }
+
+        [[nodiscard]] const std::vector<
+            std::shared_ptr< EnergyTerm< ObjectType > > >&
+            terms_for_subset( const uuid& subset_id ) const
+        {
+            const auto it = subset_to_terms_.find( subset_id );
+            OPENGEODE_EXCEPTION( it != subset_to_terms_.end(),
+                "[EnergyTermCollection] - Object Subset (", subset_id.string(),
+                ") does not have any energy term." );
+            return it->second;
+        }
+
+        std::string string() const
+        {
+            auto message = absl::StrCat(
+                "EnergyTermCollection: ", energy_terms_.size(), " terms:" );
+            for( const auto& [id, term] : energy_terms_ )
+            {
+                absl::StrAppend( &message, "\n\t --> ", term->string() );
+            }
+            return message;
+        }
+
+    private:
+        // strong ownership
+        absl::flat_hash_map< uuid, std::shared_ptr< EnergyTerm< ObjectType > > >
+            energy_terms_;
+
+        // subset index (shared ownership)
+        absl::flat_hash_map< uuid,
+            std::vector< std::shared_ptr< EnergyTerm< ObjectType > > > >
+            subset_to_terms_;
+    };
+
+} // namespace geode

include/geode/stochastic/sampling/mcmc/models/components/pairwise_term.hpp

@@ -37,17 +37,10 @@ namespace geode
     public:
         explicit PairwiseTerm( std::string_view name,
             double gamma,
-            std::unique_ptr< PairwiseInteraction< ObjectType > >&& interaction )
-            : EnergyTerm< ObjectType >( name, gamma ),
-              interaction_( std::move( interaction ) )
-        {
-        }
-
-        explicit PairwiseTerm( std::string_view name,
-            double gamma,
-            std::unique_ptr< PairwiseInteraction< ObjectType > >&& interaction,
-            const uuid& subset_id )
-            : EnergyTerm< ObjectType >( name, gamma, subset_id ),
+            absl::flat_hash_set< uuid > targeted_subset_ids,
+            std::unique_ptr< PairwiseInteraction< ObjectType > > interaction )
+            : EnergyTerm< ObjectType >(
+                  name, gamma, std::move( targeted_subset_ids ) ),
               interaction_( std::move( interaction ) )
         {
         }
@@ -157,6 +150,16 @@ namespace geode
                     state.get_all_object(); // state.neighbors( obj_id, 1.1 );
                 for( const auto& neigh_obj_id : neighbors )
                 {
+                    // if( neigh_obj_id.subset_id < obj_id.subset_id )
+                    //{
+                    //     continue;
+                    // }
+                    // if( neigh_obj_id.subset_id == obj_id.subset_id
+                    //     && neigh_obj_id.id <= obj_id.id )
+                    //{
+                    //     continue;
+                    // }
+
                     if( neigh_obj_id == obj_id )
                     {
                         continue;

include/geode/stochastic/sampling/mcmc/models/components/single_object_term.hpp

@@ -34,17 +34,9 @@ namespace geode
     public:
         explicit SingleObjectTerm( std::string_view name,
             double lambda,
+            absl::flat_hash_set< uuid > targeted_subset_ids,
             ObjectContributionFunc contribution_func )
-            : EnergyTerm< ObjectType >( name, lambda ),
-              contribution_func_( std::move( contribution_func ) )
-        {
-        }
-
-        explicit SingleObjectTerm( std::string_view name,
-            double lambda,
-            ObjectContributionFunc contribution_func,
-            std::optional< uuid > subset_id )
-            : EnergyTerm< ObjectType >( name, lambda, subset_id ),
+            : EnergyTerm< ObjectType >( name, lambda, targeted_subset_ids ),
               contribution_func_( std::move( contribution_func ) )
         {
         }