code: def _draw_graph_with_custom_layout(
      graph: nx.Graph,
      colors: Dict[Any, str],
      edge_widths: Dict[Tuple[Any, Any], float],
      figure_size: Optional[List[int]] = None,
  ):
      # This layout tries to mimic the graphviz layout in a simpler form. The depth grows horizontally here instead of
      # vertically.
      if isinstance(graph, nx.DiGraph):
          graph = nx.DiGraph(graph)
      else:
          graph = nx.Graph(graph)
  
      layers = _custom_assign_layers(graph)
      nx.set_node_attributes(graph, layers, "layer")
      node_positions = nx.multipartite_layout(graph, subset_key="layer")
  
      if figure_size is None:
          # Set the figure size based on the number of nodes
          figure = plt.figure(
              figsize=(
                  max(np.max([v for v in layers.values()]) * 2, 5),
                  max(_custom_count_nodes_vertically(graph) * 1.25, 5),
              )
          )
      else:
          figure = plt.figure(figsize=figure_size)
  
      nx.draw_networkx_nodes(
          graph,
          node_positions,
          node_size=2000,
          node_color=[colors[node] for node in graph.nodes()],
          alpha=0.8,
      )
  
      vertical_neighbor_indicator = (
          _custom_create_vertical_neighbor_indicator(
              graph, node_positions
          )
      )
  
      # Nodes that are vertically connected, but not neighbors should be connected via a curved edge.
      edges_with_curved_line = [
          (u, v)
          for u, v in graph.edges()
          if graph.nodes[u]["layer"] == graph.nodes[v]["layer"]
          and not vertical_neighbor_indicator.loc[str(u), str(v)]
      ]
      nx.draw_networkx_edges(
          graph,
          node_positions,
          edgelist=edges_with_curved_line,
          width=[
              edge_widths[(s, t)] for (s, t) in edges_with_curved_line
          ],
          edge_color=[
              colors[(s, t)] for (s, t) in edges_with_curved_line
          ],
          arrowsize=20,
          connectionstyle="arc3,rad=0.5",
          alpha=0.8,
          min_source_margin=25,
          min_target_margin=21,
      )
  
      # All other nodes should be connected with a straight line.
      edges_with_straigth_line = [
          (u, v)
          for u, v in graph.edges()
          if (
              graph.nodes[u]["layer"] == graph.nodes[v]["layer"]
              and vertical_neighbor_indicator.loc[str(u), str(v)]
          )
          or graph.nodes[u]["layer"] != graph.nodes[v]["layer"]
      ]
      nx.draw_networkx_edges(
          graph,
          node_positions,
          edgelist=edges_with_straigth_line,
          width=[
              edge_widths[(s, t)] for (s, t) in edges_with_straigth_line
          ],
          edge_color=[
              colors[(s, t)] for (s, t) in edges_with_straigth_line
          ],
          arrowsize=20,
          alpha=0.8,
          min_source_margin=25,
          min_target_margin=21,
      )
  
      # Draw labels node labels
      for node, (x, y) in node_positions.items():
          plt.text(
              x,
              y,
              node,
              ha="center",
              va="center",
              color="black",
              fontweight="bold",
          )
  
      return figure
  
file_path: dowhy/dowhy/utils/networkx_plotting.py
imports: from copy import deepcopy
  from typing import Any, Dict, List, Optional, Tuple, Union
  import matplotlib.pyplot
  import networkx
  import numpy
  import pandas
  from networkx.drawing import nx_pydot
label_desc: Use image node_id 2 for calling a global function with example usage: _draw_graph_with_custom_layout(graph, colors, edge_widths, figure_size) and returns: figure