Python networkx 模块，balanced_tree() 实例源码

def simple_threestage_scenario_tree():
    from pyomo.pysp.scenariotree.tree_structure_model \
        import CreateConcreteTwoStageScenarioTreeModel
    import networkx
    G = networkx.balanced_tree(2,2,networkx.DiGraph())
    st_model = ScenarioTreeModelFromNetworkX(
        G,
        edge_probability_attribute=None)
    first_stage = st_model.Stages.first()
    second_stage = st_model.Stages.next(first_stage)
    third_stage = st_model.Stages.last()
    # First Stage
    st_model.StageCost[first_stage] = 'StageCost[1]'
    st_model.StageVariables[first_stage].add('x')
    # Second Stage
    st_model.StageCost[second_stage] = 'StageCost[2]'
    st_model.StageVariables[second_stage].add('y')
    # Third Stage
    st_model.StageCost[thrid_stage] = 'StageCost[3]'
    st_model.StageVariables[second_stage].add('z')
    return st_model

def setUp(self):
        integer_graph = nx.balanced_tree(2, 2, nx.DiGraph())

        package_mapping = {
            i: 'm.' + ('X' if i % 2 == 0 else 'Y') + '.' + letter
            for (i, letter) in enumerate(string.ascii_lowercase)
        }

        # Edges: [(X.a, Y.b), (X.a, X.c), (Y.b, Y.d), (Y.b, X.e), (X.c, Y.f), (X.c, X.g)]
        self.package_graph = nx.relabel_nodes(integer_graph, package_mapping)

        for node in self.package_graph:
            self.package_graph.node[node]['fqn'] = node.split('.')[1:]

        self.project = self.get_project()