<s> from collections import defaultdict <EOL> import numpy as np <EOL> from pgmpy . base import UndirectedGraph <EOL> from pgmpy . factors import factor_product <EOL> from pgmpy . extern . six . moves import filter , range , zip <EOL> class ClusterGraph ( UndirectedGraph ) : <EOL> r"""<STR_LIT>""" <EOL> def __init__ ( self , ebunch = None ) : <EOL> super ( ClusterGraph , self ) . __init__ ( ) <EOL> if ebunch : <EOL> self . add_edges_from ( ebunch ) <EOL> self . factors = [ ] <EOL> def add_node ( self , node , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( node , ( list , set , tuple ) ) : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> node = tuple ( node ) <EOL> super ( ClusterGraph , self ) . add_node ( node , ** kwargs ) <EOL> def add_nodes_from ( self , nodes , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> for node in nodes : <EOL> self . add_node ( node , ** kwargs ) <EOL> def add_edge ( self , u , v , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> set_u = set ( u ) <EOL> set_v = set ( v ) <EOL> if set_u . isdisjoint ( set_v ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> super ( ClusterGraph , self ) . add_edge ( u , v ) <EOL> def add_factors ( self , * factors ) : <EOL> """<STR_LIT>""" <EOL> for factor in factors : <EOL> factor_scope = set ( factor . scope ( ) ) <EOL> nodes = [ set ( node ) for node in self . nodes ( ) ] <EOL> if factor_scope not in nodes : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> self . factors . append ( factor ) <EOL> def get_factors ( self , node = None ) : <EOL> """<STR_LIT>""" <EOL> if node is None : <EOL> return self . factors <EOL> else : <EOL> nodes = [ set ( n ) for n in self . nodes ( ) ] <EOL> if set ( node ) not in nodes : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> factors = filter ( lambda x : set ( x . scope ( ) ) == set ( node ) , self . factors ) <EOL> return next ( factors ) <EOL> def remove_factors ( self , * factors ) : <EOL> """<STR_LIT>""" <EOL> for factor in factors : <EOL> self . factors . remove ( factor ) <EOL> def get_cardinality ( self , check_cardinality = False ) : <EOL> """<STR_LIT>""" <EOL> cardinalities = defaultdict ( int ) <EOL> for factor in self . factors : <EOL> for variable , cardinality in zip ( factor . scope ( ) , factor . cardinality ) : <EOL> cardinalities [ variable ] = cardinality <EOL> if check_cardinality and len ( set ( ( x for clique in self . nodes ( ) for x in clique ) ) ) != len ( cardinalities ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> return cardinalities <EOL> def get_partition_function ( self ) : <EOL> r"""<STR_LIT>""" <EOL> if self . check_model ( ) : <EOL> factor = self . factors [ <NUM_LIT:0> ] <EOL> factor = factor_product ( factor , * [ self . factors [ i ] for i in range ( <NUM_LIT:1> , len ( self . factors ) ) ] ) <EOL> return np . sum ( factor . values ) <EOL> def check_model ( self ) : <EOL> """<STR_LIT>""" <EOL> for clique in self . nodes ( ) : <EOL> factors = filter ( lambda x : set ( x . scope ( ) ) == set ( clique ) , self . factors ) <EOL> if not any ( factors ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> cardinalities = self . get_cardinality ( ) <EOL> for factor in self . factors : <EOL> for variable , cardinality in zip ( factor . scope ( ) , factor . cardinality ) : <EOL> if ( cardinalities [ variable ] != cardinality ) : <EOL> raise ValueError ( <EOL> '<STR_LIT>' . format ( var = variable ) ) <EOL> return True <EOL> def copy ( self ) : <EOL> """<STR_LIT>""" <EOL> copy = ClusterGraph ( self . edges ( ) ) <EOL> if self . factors : <EOL> factors_copy = [ factor . copy ( ) for factor in self . factors ] <EOL> copy . add_factors ( * factors_copy ) <EOL> return copy </s>
