Repository: ManimCommunity/manim
File: manim/mobject/graph.py (Python)

Code before:
        import networkx as nx

        class PartiteGraph(Scene):
            def construct(self):
                G = nx.Graph()
                G.add_nodes_from([0, 1, 2, 3])
                G.add_edges_from([(0, 2), (0,3), (1, 2)])
                graph = Graph(list(G.nodes), list(G.edges), layout="partite", partitions=[[0, 1]])
                self.play(Create(graph))

    The representation of a linear artificial neural network is facilitated
    by the use of the partite layout and defining partitions for each layer.

    .. manim:: LinearNN
        :save_last_frame:

        class LinearNN(Scene):

            def construct(self):
                edges = []
                partitions = []
                c=0
                layers = [2,3,3,2]         # the number of neurons on each layer

                for i in layers:
                    partitions.append([j for j in range(c+1, c+i+1)])
                    c += i
                for i, v in enumerate(layers[1:]):
                        last = int(np.sum(layers[:i+1]))
                        for j in range(v):
                            for k in range(last-layers[i], last):
                                edges.append((k+1, j+last+1))

                vertices = np.arange(1, int(np.sum(layers))+1)

                graph = Graph(vertices, edges, layout='partite', partitions=partitions, layout_scale=3, vertex_config={'radius':0.20})
                self.add(graph)

    The custom tree layout can be used to show the graph
    by distance from the root vertex. You must pass the root vertex
    of the tree.

    .. manim:: Tree

        import networkx as nx

        class Tree(Scene):
            def construct(self):
                G = nx.Graph()

                G.add_node("ROOT")

Review comment (suggestion):
The list comprehension here seems a bit superfluous. How about this?

```suggestion
                    partitions.append(list(range(c + 1, c + i + 1)))
```

Code after:
        import networkx as nx

        class PartiteGraph(Scene):
            def construct(self):
                G = nx.Graph()
                G.add_nodes_from([0, 1, 2, 3])
                G.add_edges_from([(0, 2), (0,3), (1, 2)])
                graph = Graph(list(G.nodes), list(G.edges), layout="partite", partitions=[[0, 1]])
                self.play(Create(graph))

    The representation of a linear artificial neural network is facilitated
    by the use of the partite layout and defining partitions for each layer.

    .. manim:: LinearNN
        :save_last_frame:

        class LinearNN(Scene):
            def construct(self):
                edges = []
                partitions = []
                c = 0
                layers = [2, 3, 3, 2]  # the number of neurons in each layer

                for i in layers:
                    partitions.append(list(range(c + 1, c + i + 1)))
                    c += i
                for i, v in enumerate(layers[1:]):
                        last = sum(layers[:i+1])
                        for j in range(v):
                            for k in range(last - layers[i], last):
                                edges.append((k + 1, j + last + 1))

                vertices = np.arange(1, sum(layers) + 1)

                graph = Graph(
                    vertices,
                    edges,
                    layout='partite',
                    partitions=partitions,
                    layout_scale=3,
                    vertex_config={'radius': 0.20},
                )
                self.add(graph)

    The custom tree layout can be used to show the graph
    by distance from the root vertex. You must pass the root vertex
    of the tree.

    .. manim:: Tree

        import networkx as nx