diff --git a/pycode/memilio-generation/memilio/generation/__init__.py b/pycode/memilio-generation/memilio/generation/__init__.py
index 65098b9693..3289ebf7be 100644
--- a/pycode/memilio-generation/memilio/generation/__init__.py
+++ b/pycode/memilio-generation/memilio/generation/__init__.py
@@ -25,3 +25,4 @@
 from .intermediate_representation import IntermediateRepresentation
 from .scanner import Scanner
 from .scanner_config import ScannerConfig
+from .ast import AST
diff --git a/pycode/memilio-generation/memilio/generation/ast.py b/pycode/memilio-generation/memilio/generation/ast.py
new file mode 100644
index 0000000000..77a0e8d837
--- /dev/null
+++ b/pycode/memilio-generation/memilio/generation/ast.py
@@ -0,0 +1,152 @@
+#############################################################################
+# Copyright (C) 2020-2024 MEmilio
+#
+# Authors: Maximilian Betz, Daniel Richter
+#
+# Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#############################################################################
+"""
+@file ast.py
+@brief Create the ast and assign ids. Get ids and nodes. 
+"""
+import subprocess
+import tempfile
+import logging
+from clang.cindex import Cursor, TranslationUnit, Index, CompilationDatabase
+from typing import TYPE_CHECKING
+from memilio.generation import utility
+
+
+if TYPE_CHECKING:
+    from memilio.generation import ScannerConfig
+
+from typing_extensions import Self
+
+
+class AST:
+    """! Create the ast and assign ids.
+    Functions for getting nodes and node ids.
+    """
+
+    def __init__(self: Self, conf: "ScannerConfig") -> None:
+        self.config = conf
+        self.cursor_id = -1
+        self.id_to_val = dict()
+        self.val_to_id = dict()
+        self.cursor = None
+        self.translation_unit = self.create_ast()
+
+    def create_ast(self: Self) -> TranslationUnit:
+        """! Create an abstract syntax tree for the main model.cpp file with a corresponding CompilationDatabase.
+        A compile_commands.json is required (automatically generated in the build process).
+        """
+        self.cursor_id = -1
+        self.id_to_val.clear()
+        self.val_to_id.clear()
+
+        idx = Index.create()
+
+        file_args = []
+
+        unwanted_arguments = [
+            '-Wno-unknown-warning', "--driver-mode=g++", "-O3", "-Werror", "-Wshadow"
+        ]
+
+        dirname = utility.try_get_compilation_database_path(
+            self.config.skbuild_path_to_database)
+        compdb = CompilationDatabase.fromDirectory(dirname)
+        commands = compdb.getCompileCommands(self.config.source_file)
+        for command in commands:
+            for argument in command.arguments:
+                if argument not in unwanted_arguments:
+                    file_args.append(argument)
+        file_args = file_args[1:-4]
+
+        clang_cmd = [
+            "clang-14", self.config.source_file,
+            "-std=c++17", '-emit-ast', '-o', '-']
+        clang_cmd.extend(file_args)
+
+        try:
+            clang_cmd_result = subprocess.run(
+                clang_cmd, stdout=subprocess.PIPE)
+            clang_cmd_result.check_returncode()
+        except subprocess.CalledProcessError as e:
+            # Capture standard error and output
+            logging.error(
+                f"Clang failed with return code {e.returncode}. Error: {clang_cmd_result.stderr.decode()}")
+            raise RuntimeError(
+                f"Clang AST generation failed. See error log for details.")
+
+        # Since `clang.Index.read` expects a file path, write generated abstract syntax tree to a
+        # temporary named file. This file will be automatically deleted when closed.
+        with tempfile.NamedTemporaryFile() as ast_file:
+            ast_file.write(clang_cmd_result.stdout)
+            translation_unit = idx.read(ast_file.name)
+
+        self._assing_ast_with_ids(translation_unit.cursor)
+
+        logging.info("AST generation completed successfully.")
+
+        return translation_unit
+
+    def _assing_ast_with_ids(self, cursor: Cursor) -> None:
+        """! Traverse the AST and assign a unique ID to each node during traversal.
+
+        @param cursor: The current node (Cursor) in the AST to traverse.
+        """
+        # assing_ids umschreiben -> mapping
+        self.cursor_id += 1
+        id = self.cursor_id
+        self.id_to_val[id] = cursor
+
+        if cursor.hash in self.val_to_id.keys():
+            self.val_to_id[cursor.hash].append(id)
+        else:
+            self.val_to_id[cursor.hash] = [id]
+
+        logging.info(
+            f"Node {cursor.spelling or cursor.kind} assigned ID {id}")
+
+        for child in cursor.get_children():
+            self._assing_ast_with_ids(child)
+
+    @property
+    def root_cursor(self):
+        return self.translation_unit.cursor
+
+    def get_node_id(self, cursor: Cursor) -> int:
+        """! Returns the id of the current node.
+
+        Extracs the key from the current cursor from the dictonary id_to_val
+
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        """
+        for cursor_id in self.val_to_id[cursor.hash]:
+
+            if self.id_to_val[cursor_id] == cursor:
+
+                return cursor_id
+        raise IndexError(f"Cursor {cursor} is out of bounds.")
+
+    def get_node_by_index(self, index: int) -> Cursor:
+        """! Returns the node at the specified index position.
+
+        @param index: Node_id from the ast.
+        """
+
+        if index < 0 or index >= len(self.id_to_val):
+            raise IndexError(f"Index {index} is out of bounds.")
+        return self.id_to_val[index]
diff --git a/pycode/memilio-generation/memilio/generation/graph_visualization.py b/pycode/memilio-generation/memilio/generation/graph_visualization.py
new file mode 100644
index 0000000000..9ef18c8964
--- /dev/null
+++ b/pycode/memilio-generation/memilio/generation/graph_visualization.py
@@ -0,0 +1,164 @@
+#############################################################################
+# Copyright (C) 2020-2024 MEmilio
+#
+# Authors: Maximilian Betz, Daniel Richter
+#
+# Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#############################################################################
+
+import os
+import logging
+from typing import Callable
+from graphviz import Digraph
+from clang.cindex import Cursor
+from memilio.generation.ast import AST
+
+
+class Visualization:
+    """! Class for plotting the abstract syntax tree in different formats.
+    """
+    @staticmethod
+    def output_ast_terminal(ast: AST, cursor: Cursor) -> None:
+        """! Output the abstract syntax tree to terminal.
+
+        @param ast: ast object from AST class.
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        """
+
+        def terminal_writer(level: int, cursor_label: str) -> None:
+            print(indent(level) + cursor_label)
+
+        _output_cursor_and_children(cursor, ast, terminal_writer)
+
+        logging.info("AST-Terminal written.")
+
+    @staticmethod
+    def output_ast_png(cursor: Cursor, max_depth: int, output_file_name: str = 'ast_graph') -> None:
+        """! Output the abstract syntax tree to a .png. Set the starting node and the max depth.
+
+        To save the abstract syntax tree as an png with a starting node and a depth u cann use the following command
+
+        Example command: aviz.output_ast_png(ast.get_node_by_index(1), 2)
+
+        aviz -> instance of the Visualization class.
+
+        ast -> instance of the AST class.
+
+        .get_node_by_index -> get a specific node by id (use .output_ast_formatted to see node ids)
+
+        The number 2 is a example for the depth the graph will show
+
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        @param max_depth: Maximal depth the graph displays.
+        """
+
+        graph = Digraph(format='png')
+
+        _output_cursor_and_children_graphviz_digraph(
+            cursor, graph, max_depth, 0)
+
+        graph.render(filename=output_file_name, view=False)
+
+        output_path = os.path.abspath(f"{output_file_name}.png")
+        logging.info(f"AST-png written to {output_path}")
+
+    @staticmethod
+    def output_ast_formatted(ast: AST, cursor: Cursor, output_file_name: str = 'ast_formated.txt') -> None:
+        """!Output the abstract syntax tree to a file.
+
+        @param ast: ast object from AST class.
+        @param cursor: The current node of the AST as a cursor object from libclang.
+        """
+
+        with open(output_file_name, 'w') as f:
+            def file_writer(level: int, cursor_label: str) -> None:
+                f.write(indent(level) + cursor_label + newline())
+            _output_cursor_and_children(cursor, ast,  file_writer)
+
+        output_path = os.path.abspath(f"{output_file_name}")
+        logging.info(f"AST-formated written to {output_path}")
+
+
+def indent(level: int) -> str:
+    """! Create an indentation based on the level.
+    """
+    return '│   ' * level + '├── '
+
+
+def newline() -> str:
+    """! Create a new line.
+    """
+    return '\n'
+
+
+def _output_cursor_and_children(cursor: Cursor, ast: AST, writer: Callable[[int, str], None], level: int = 0) -> None:
+    """!Generic function to output the cursor and its children with a specified writer.
+
+    @param cursor: The current node of the AST as a libclang cursor object.
+    @param ast: AST object from the AST class.
+    @param writer: Function that takes `level` and `cursor_label` and handles output.
+    @param level: The current depth in the AST for indentation purposes.
+    """
+
+    cursor_id = ast.get_node_id(cursor)
+
+    cursor_kind = f"<CursorKind.{cursor.kind.name}>"
+    file_path = cursor.location.file.name if cursor.location.file else ""
+
+    if cursor.spelling:
+        cursor_label = (f'ID:{cursor_id} {cursor.spelling} '
+                        f'{cursor_kind}   '
+                        f'{file_path}')
+    else:
+        cursor_label = f'ID:{cursor_id} {cursor_kind} {file_path}'
+
+    writer(level, cursor_label)
+
+    for child in cursor.get_children():
+        _output_cursor_and_children(
+            child, ast, writer, level + 1)
+
+
+def _output_cursor_and_children_graphviz_digraph(cursor: Cursor, graph: Digraph, max_d: int, current_d: int, parent_node: str = None) -> None:
+    """! Output the cursor and its children as a graph using Graphviz.
+
+    @param cursor: The current node of the AST as a Cursor object from libclang.
+    @param graph: Graphviz Digraph object where the nodes and edges will be added.
+    @param max_d: Maximal depth.
+    @param current_d: Current depth.
+    @param parent_node: Name of the parent node in the graph (None for the root node).
+    """
+
+    if current_d > max_d:
+        return
+
+    node_label = f"{cursor.kind.name}{newline()}({cursor.spelling})" if cursor.spelling else cursor.kind.name
+
+    current_node = f"{cursor.kind.name}_{cursor.hash}"
+
+    graph.node(current_node, label=node_label)
+
+    if parent_node:
+        graph.edge(parent_node, current_node)
+
+    if cursor.kind.is_reference():
+        referenced_label = f"ref_to_{cursor.referenced.kind.name}{newline()}({cursor.referenced.spelling})"
+        referenced_node = f"ref_{cursor.referenced.hash}"
+        graph.node(referenced_node, label=referenced_label)
+        graph.edge(current_node, referenced_node)
+
+    for child in cursor.get_children():
+        _output_cursor_and_children_graphviz_digraph(
+            child, graph, max_d, current_d + 1, current_node)
diff --git a/pycode/memilio-generation/memilio/generation/scanner.py b/pycode/memilio-generation/memilio/generation/scanner.py
index 2f14e646b3..317ce59889 100755
--- a/pycode/memilio-generation/memilio/generation/scanner.py
+++ b/pycode/memilio-generation/memilio/generation/scanner.py
@@ -24,24 +24,20 @@
 from __future__ import annotations
 
 import os
-import subprocess
-import sys
-import tempfile
 from typing import TYPE_CHECKING, Any, Callable
-from warnings import catch_warnings
 
 from clang.cindex import *
 from typing_extensions import Self
 
 from memilio.generation import IntermediateRepresentation, utility
 
+
 if TYPE_CHECKING:
     from memilio.generation import ScannerConfig
 
 
 class Scanner:
-    """
-    Analyze the model and extract the needed information.
+    """! Analyze the model and extract the needed information.
     """
 
     def __init__(self: Self, conf: ScannerConfig) -> None:
@@ -53,61 +49,22 @@ def __init__(self: Self, conf: ScannerConfig) -> None:
         self.config = conf
         utility.try_set_libclang_path(
             self.config.optional.get("libclang_library_path"))
-        self.ast = None
-        self.create_ast()
 
-    def create_ast(self: Self) -> None:
-        """
-        Create an abstract syntax tree for the main model.cpp file with a corresponding CompilationDatabase. 
-        A compile_commands.json is required (automatically generated in the build process).
-        """
-        idx = Index.create()
-
-        # Create the cmd arguments
-        file_args = []
-
-        dirname = utility.try_get_compilation_database_path(
-            self.config.skbuild_path_to_database)
-        compdb = CompilationDatabase.fromDirectory(dirname)
-        commands = compdb.getCompileCommands(self.config.source_file)
-        for command in commands:
-            for argument in command.arguments:
-                if (argument != '-Wno-unknown-warning' and
-                        argument != "--driver-mode=g++" and argument != "-O3"):
-                    file_args.append(argument)
-        file_args = file_args[1:-4]
-        clang_cmd = [
-            "clang-14", self.config.source_file,
-            "-std=c++17", '-emit-ast', '-o', '-']
-        clang_cmd.extend(file_args)
-
-        clang_cmd_result = subprocess.run(clang_cmd, stdout=subprocess.PIPE)
-        clang_cmd_result.check_returncode()
-
-        # Since `clang.Index.read` expects a file path, write generated abstract syntax tree to a
-        # temporary named file. This file will be automatically deleted when closed.
-        with tempfile.NamedTemporaryFile() as ast_file:
-            ast_file.write(clang_cmd_result.stdout)
-            self.ast = idx.read(ast_file.name)
-
-    def extract_results(self: Self) -> IntermediateRepresentation:
-        """
-        Extract the information of the abstract syntax tree and save them in the dataclass intermed_repr.
+    def extract_results(self: Self, root_cursor: Cursor) -> IntermediateRepresentation:
+        """! Extract the information of the abstract syntax tree and save them in the dataclass intermed_repr.
         Call find_node to visit all nodes of abstract syntax tree and finalize to finish the extraction.
 
+        @param root_cursor Represents the root node of the abstract syntax tree as a Cursor object from libclang.
         @return Information extracted from the model saved as an IntermediateRepresentation. 
         """
         intermed_repr = IntermediateRepresentation()
-        utility.output_cursor_print(self.ast.cursor, 1)
-        self.find_node(self.ast.cursor, intermed_repr)
-        # self.output_ast_file()
+        self.find_node(root_cursor, intermed_repr)
         self.finalize(intermed_repr)
         return intermed_repr
 
     def find_node(self: Self, node: Cursor,
                   intermed_repr: IntermediateRepresentation, namespace: str = "") -> None:
-        """
-        Recursively walk over every node of an abstract syntax tree. Save the namespace the node is in.
+        """! Recursively walk over every node of an abstract syntax tree. Save the namespace the node is in.
         Call check_node_kind for extracting information from the nodes.
 
         @param node Represents the current node of the abstract syntax tree as a Cursor object from libclang.
@@ -125,8 +82,7 @@ def find_node(self: Self, node: Cursor,
     def switch_node_kind(self: Self, kind: CursorKind) -> Callable[[Any,
                                                                    IntermediateRepresentation],
                                                                    None]:
-        """
-        Dictionary to map CursorKind to methods. Works like a switch.
+        """! Dictionary to map CursorKind to methods. Works like a switch.
 
         @param Underlying kind of the current node.
         @return Appropriate method for the given kind.
@@ -147,8 +103,7 @@ def switch_node_kind(self: Self, kind: CursorKind) -> Callable[[Any,
     def check_enum(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Inspect the nodes of kind ENUM_DECL and write needed information into intermed_repr.
+        """! Inspect the nodes of kind ENUM_DECL and write needed information into intermed_repr.
         Information: Name of Enum
 
         @param node Current node represented as a Cursor object.
@@ -160,8 +115,7 @@ def check_enum(
     def check_enum_const(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Inspect the nodes of kind ENUM_CONSTANT_DECL and write needed information into intermed_repr.
+        """! Inspect the nodes of kind ENUM_CONSTANT_DECL and write needed information into intermed_repr.
         Information: Keys of an Enum
 
         @param node Current node represented as a Cursor object.
@@ -174,8 +128,7 @@ def check_enum_const(
     def check_class(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Inspect the nodes of kind CLASS_DECL and write information 
+        """! Inspect the nodes of kind CLASS_DECL and write information 
         (model_class, model_base, simulation_class, parameterset_wrapper) into intermed_repr.
 
         @param node Current node represented as a Cursor object.
@@ -196,8 +149,7 @@ def check_class(
     def check_model_base(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Helper function to retreive the model base.
+        """! Helper function to retreive the model base.
 
         @param node Current node represented as a Cursor object.
         @param intermed_repr Dataclass used for saving the extracted model features.
@@ -211,8 +163,7 @@ def check_model_base(
     def check_base_specifier(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """ 
-        Not used yet.
+        """! Not used yet.
         Inspect nodes which represent base specifier.
         For now this is handled by the parent node, which represents the class.
         """
@@ -221,8 +172,7 @@ def check_base_specifier(
     def check_model_includes(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Helper function to retrieve the model specific includes needed for pybind.
+        """! Helper function to retrieve the model specific includes needed for pybind.
 
         @param node Current node represented as a Cursor object.
         @param intermed_repr Dataclass used for saving the extracted model features.
@@ -253,8 +203,7 @@ def check_model_includes(
     def check_age_group(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Inspect the nodes of kind CLASS_DECL with the name defined in 
+        """! Inspect the nodes of kind CLASS_DECL with the name defined in 
         config.age_group and write needed information into intermed_repr.
         Information: age_group
 
@@ -279,8 +228,7 @@ def check_age_group(
     def check_constructor(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Inspect the nodes of kind CONSTRUCTOR and write needed information into intermed_repr.
+        """! Inspect the nodes of kind CONSTRUCTOR and write needed information into intermed_repr.
         Information: intermed_repr.init
 
         @param node Current node represented as a Cursor object.
@@ -299,8 +247,7 @@ def check_constructor(
     def check_type_alias(
         self: Self, node: Cursor,
             intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Inspect the nodes of kind TYPE_ALIAS_DECL and write needed information into intermed_repr.
+        """! Inspect the nodes of kind TYPE_ALIAS_DECL and write needed information into intermed_repr.
         Information: intermed_repr.parameterset
 
         @param node Current node represented as a Cursor object.
@@ -316,8 +263,7 @@ def check_struct(
         pass
 
     def finalize(self: Self, intermed_repr: IntermediateRepresentation) -> None:
-        """
-        Finalize the IntermediateRepresenation as last step of the Scanner.
+        """! Finalize the IntermediateRepresenation as last step of the Scanner.
         Write needed information from config into intermed_repr,
         delet unnecesary enums and check for missing model features.
 
@@ -346,17 +292,3 @@ def finalize(self: Self, intermed_repr: IntermediateRepresentation) -> None:
 
         # check for missing data
         intermed_repr.check_complete_data(self.config.optional)
-
-    def output_ast(self: Self) -> None:
-        """
-        Output the abstract syntax tree to terminal.
-        """
-        utility.output_cursor_and_children(self.ast.cursor)
-
-    def output_ast_file(self: Self) -> None:
-        """
-        Output the abstract syntax tree to file.
-        """
-        with open('output_ast.txt', 'a') as f:
-            utility.output_cursor_and_children_file(self.ast.cursor, f)
-            print('AST written to ' + str(os.path.abspath(f.name)))
diff --git a/pycode/memilio-generation/memilio/generation/utility.py b/pycode/memilio-generation/memilio/generation/utility.py
index 72ace6a808..f5c67fbcfe 100644
--- a/pycode/memilio-generation/memilio/generation/utility.py
+++ b/pycode/memilio-generation/memilio/generation/utility.py
@@ -114,105 +114,3 @@ def get_base_class_string(base_type: Type) -> List[Any]:
         result.append(get_base_class_string(
             base_type.get_template_argument_type(i)))
     return result
-
-
-def indent(spaces: int) -> str:
-    """ 
-    Indentation string for pretty-printing.
-
-    @param spaces Number of spaces.
-    """
-    return '  '*spaces
-
-
-def output_cursor_print(cursor: Cursor, spaces: int) -> None:
-    """ 
-    Low level cursor output to the terminal.
-
-    @param cursor Represents the current node of the AST as an Cursor object from libclang.
-    @param spaces Number of spaces.
-    """
-    spelling = ''
-    displayname = ''
-
-    if cursor.spelling:
-        spelling = cursor.spelling
-    if cursor.displayname:
-        displayname = cursor.displayname
-    kind = cursor.kind
-
-    print(indent(spaces) + spelling, '<' + str(kind) + '>')
-    print(indent(spaces+1) + '"' + displayname + '"')
-
-
-def output_cursor_and_children_print(cursor: Cursor, spaces: int = 0) -> None:
-    """ 
-    Output this cursor and its children with minimal formatting to the terminal.
-
-    @param cursor Represents the current node of the AST as an Cursor object from libclang.
-    @param spaces [Default = 0] Number of spaces.
-    """
-    output_cursor_print(cursor, spaces)
-    if cursor.kind.is_reference():
-        print(indent(spaces) + 'reference to:')
-        output_cursor_print(cursor.referenced, spaces+1)
-
-    # Recurse for children of this cursor
-    has_children = False
-    for c in cursor.get_children():
-        if not has_children:
-            print(indent(spaces) + '{')
-            has_children = True
-        output_cursor_and_children_print(c, spaces+1)
-
-    if has_children:
-        print(indent(spaces) + '}')
-
-
-def output_cursor_file(cursor: Cursor, f: TextIO, spaces: int) -> None:
-    """ 
-    Low level cursor output to a file.
-
-    @param cursor Represents the current node of the AST as an Cursor object from libclang.
-    @param f Open file object for output.
-    @param spaces Number of spaces.
-    """
-    spelling = ''
-    displayname = ''
-
-    if cursor.spelling:
-        spelling = cursor.spelling
-    if cursor.displayname:
-        displayname = cursor.displayname
-    kind = cursor.kind
-
-    f.write(indent(spaces) + spelling + ' <' + str(kind) + '> ')
-    if cursor.location.file:
-        f.write(cursor.location.file.name + '\n')
-    f.write(indent(spaces+1) + '"' + displayname + '"\n')
-
-
-def output_cursor_and_children_file(
-        cursor: Cursor, f: TextIO, spaces: int = 0) -> None:
-    """ 
-    Output this cursor and its children with minimal formatting to a file.
-
-    @param cursor Represents the current node of the AST as an Cursor object from libclang.
-    @param f Open file object for output.
-    @param spaces Number of spaces.
-    """
-    output_cursor_file(cursor, f, spaces)
-    if cursor.kind.is_reference():
-        f.write(indent(spaces) + 'reference to:\n')
-        output_cursor_file(cursor.referenced, f, spaces+1)
-
-    # Recurse for children of this cursor
-    has_children = False
-    for c in cursor.get_children():
-        if not has_children:
-            f.write(indent(spaces) + '{\n')
-            has_children = True
-        output_cursor_and_children_file(c, f, spaces+1)
-
-    if has_children:
-        f.write(indent(spaces) + '}\n')
diff --git a/pycode/memilio-generation/memilio/generation_test/test_oseir_generation.py b/pycode/memilio-generation/memilio/generation_test/test_oseir_generation.py
index 5913acf8de..c08a218184 100644
--- a/pycode/memilio-generation/memilio/generation_test/test_oseir_generation.py
+++ b/pycode/memilio-generation/memilio/generation_test/test_oseir_generation.py
@@ -24,7 +24,7 @@
 import unittest
 from unittest.mock import patch
 
-from memilio.generation import Generator, Scanner, ScannerConfig
+from memilio.generation import Generator, Scanner, ScannerConfig, AST
 
 
 class TestOseirGeneration(unittest.TestCase):
@@ -71,9 +71,10 @@ def setUp(self, try_get_compilation_database_path_mock):
 
         conf = ScannerConfig.from_dict(config_json)
         self.scanner = Scanner(conf)
+        self.ast = AST(conf)
 
     def test_clean_oseir(self):
-        irdata = self.scanner.extract_results()
+        irdata = self.scanner.extract_results(self.ast.root_cursor)
 
         generator = Generator()
         generator.create_substitutions(irdata)
@@ -87,7 +88,7 @@ def test_clean_oseir(self):
     def test_wrong_model_name(self):
         self.scanner.config.model_class = "wrong_name"
         with self.assertRaises(AssertionError) as error:
-            self.scanner.extract_results()
+            self.scanner.extract_results(self.ast.root_cursor)
 
         error_message = "set a model name"
         self.assertEqual(str(error.exception), error_message)
diff --git a/pycode/memilio-generation/memilio/tools/example_oseir.py b/pycode/memilio-generation/memilio/tools/example_oseir.py
index a3226c5bd9..a2457d9825 100755
--- a/pycode/memilio-generation/memilio/tools/example_oseir.py
+++ b/pycode/memilio-generation/memilio/tools/example_oseir.py
@@ -31,7 +31,8 @@
     # For older python versions
     import importlib_resources
 
-from memilio.generation import Generator, Scanner, ScannerConfig
+from memilio.generation import Generator, Scanner, ScannerConfig, AST
+from memilio.generation.graph_visualization import Visualization
 
 
 def run_memilio_generation(print_ast=False):
@@ -41,9 +42,11 @@ def run_memilio_generation(print_ast=False):
         with open(path) as file:
             conf = ScannerConfig.schema().loads(file.read(), many=True)[0]
     scanner = Scanner(conf)
+    ast = AST(conf)
+    aviz = Visualization()
 
     # Extract results of Scanner into a intermed_repr
-    intermed_repr = scanner.extract_results()
+    intermed_repr = scanner.extract_results(ast.root_cursor)
 
     # Generate code
     generator = Generator()
@@ -52,7 +55,7 @@ def run_memilio_generation(print_ast=False):
 
     # Print the abstract syntax tree to a file
     if (print_ast):
-        scanner.output_ast_file()
+        aviz.output_ast_formatted(ast, ast.get_node_by_index(0))
 
 
 if __name__ == "__main__":
diff --git a/pycode/memilio-generation/setup.py b/pycode/memilio-generation/setup.py
index c415447a2e..1cc8cd362a 100644
--- a/pycode/memilio-generation/setup.py
+++ b/pycode/memilio-generation/setup.py
@@ -27,7 +27,7 @@
         where=os.path.dirname(os.path.abspath(__file__))),
     setup_requires=['cmake'],
     install_requires=['libclang==14.0.6',
-                      'dataclasses', 'dataclasses_json', 'importlib-resources>=1.1.0; python_version < \'3.9\''],
+                      'dataclasses', 'dataclasses_json', 'graphviz', 'importlib-resources>=1.1.0; python_version < \'3.9\''],
     extras_require={'dev': []},
     long_description='',
     test_suite='memilio.generation_test',