diff --git a/tardis/io/logger/logger.py b/tardis/io/logger/logger.py
index 291878177be..552ef458240 100644
--- a/tardis/io/logger/logger.py
+++ b/tardis/io/logger/logger.py
@@ -11,6 +11,16 @@
     "ERROR": logging.ERROR,
     "CRITICAL": logging.CRITICAL,
 }
+
+LOG_LEVEL_COLORS = {
+    logging.INFO: '#D3D3D3',     
+    logging.WARNING: 'orange',
+    logging.ERROR: 'red', 
+    logging.CRITICAL: 'orange',
+    logging.DEBUG: 'blue', 
+    'default': 'black'  
+}
+
 DEFAULT_LOG_LEVEL = "INFO"
 DEFAULT_SPECIFIC_STATE = False
 
@@ -75,20 +85,45 @@ def logging_state(log_level, tardis_config, specific_log_level=None):
             for logger in tardis_loggers:
                 logger.removeFilter(filter)
 
+def create_output_widget(height='300px'):
+    """
+    Creates an Output widget with the specified layout.
+
+    Parameters
+    ----------
+    height : str, optional
+        Height of the widget in pixels, by default '300px'.
+
+    Returns
+    -------
+    ipywidgets.Output
+        Configured Output widget.
+    """
+    return Output(layout=Layout(height=height, overflow_y='auto'))
+
 log_outputs = {
-    "WARNING/ERROR": Output(layout=Layout(height='300px', overflow_y='auto')),
-    "INFO": Output(layout=Layout(height='300px', overflow_y='auto')),
-    "DEBUG": Output(layout=Layout(height='300px', overflow_y='auto')),
-    "ALL": Output(layout=Layout(height='300px', overflow_y='auto'))
+    "WARNING/ERROR": create_output_widget(),
+    "INFO": create_output_widget(),
+    "DEBUG": create_output_widget(),
+    "ALL": create_output_widget()
 }
 
-tab = Tab(children=[log_outputs["WARNING/ERROR"], log_outputs["INFO"], log_outputs["DEBUG"], log_outputs["ALL"]])
-tab.set_title(0, "WARNING/ERROR")
-tab.set_title(1, "INFO")
-tab.set_title(2, "DEBUG")
-tab.set_title(3, "ALL")
+def create_and_display_log_tab(log_outputs):
+    """
+    Creates a Tab widget for logging outputs and displays it.
 
-display(tab)
+    Parameters
+    ----------
+    log_outputs : dict
+        Dictionary containing Output widgets for different log levels.
+    """
+    tab = Tab(children=[log_outputs["WARNING/ERROR"], log_outputs["INFO"], log_outputs["DEBUG"], log_outputs["ALL"]])
+    tab.set_title(0, "WARNING/ERROR")
+    tab.set_title(1, "INFO")
+    tab.set_title(2, "DEBUG")
+    tab.set_title(3, "ALL")
+    
+    display(tab)
 
 def remove_ansi_escape_sequences(text):
     """
@@ -131,18 +166,7 @@ def emit(self, record):
         log_entry = self.format(record)
         clean_log_entry = remove_ansi_escape_sequences(log_entry)
 
-        if record.levelno == logging.INFO:
-            color = '#D3D3D3'
-        elif record.levelno == logging.WARNING:
-            color = 'orange'
-        elif record.levelno == logging.ERROR:
-            color = 'red'
-        elif record.levelno == logging.CRITICAL:
-            color = 'orange'
-        elif record.levelno == logging.DEBUG:
-            color = 'blue'
-        else:
-            color = 'black'
+        color = LOG_LEVEL_COLORS.get(record.levelno, LOG_LEVEL_COLORS['default'])
 
         parts = clean_log_entry.split(' ', 2)
         if len(parts) > 2:
@@ -182,6 +206,7 @@ def emit(self, record):
 
 logger.addHandler(widget_handler)
 logging.getLogger("py.warnings").addHandler(widget_handler)
+create_and_display_log_tab(log_outputs)
 
 class FilterLog(object):
     """
@@ -210,3 +235,4 @@ def filter(self, log_record):
             True if the log record's level is in the specified log_levels, False otherwise.
         """
         return log_record.levelno in self.log_levels
+    
\ No newline at end of file
diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py
index e717aff9a64..014bca84e8c 100644
--- a/tardis/simulation/base.py
+++ b/tardis/simulation/base.py
@@ -516,34 +516,16 @@ def log_plasma_state(
         plasma_state_log["w"] = dilution_factor
         plasma_state_log["next_w"] = next_dilution_factor
         plasma_state_log.columns.name = "Shell No."
-
-        if False and is_notebook(): #TODO: remove it with something better
-            logger.info("\n\tPlasma stratification:")
-
-            # Displaying the DataFrame only when the logging level is NOTSET, DEBUG or INFO
-            if logger.level <= logging.INFO:
-                if not logger.filters:
-                    display(
-                        plasma_state_log.iloc[::log_sampling].style.format(
-                            "{:.3g}"
-                        )
-                    )
-                elif logger.filters[0].log_level == 20:
-                    display(
-                        plasma_state_log.iloc[::log_sampling].style.format(
-                            "{:.3g}"
-                        )
-                    )
-        else:
-            output_df = ""
-            plasma_output = plasma_state_log.iloc[::log_sampling].to_string(
-                float_format=lambda x: f"{x:.3g}",
-                justify="center",
-            )
-            for value in plasma_output.split("\n"):
-                output_df = output_df + f"\t{value}\n"
-            logger.info("\n\tPlasma stratification:")
-            logger.info(f"\n{output_df}")
+        
+        output_df = ""
+        plasma_output = plasma_state_log.iloc[::log_sampling].to_string(
+            float_format=lambda x: f"{x:.3g}",
+            justify="center",
+        )
+        for value in plasma_output.split("\n"):
+            output_df = output_df + f"\t{value}\n"
+        logger.info("\n\tPlasma stratification:")
+        logger.info(f"\n{output_df}")
 
         logger.info(
             f"\n\tCurrent t_inner = {t_inner:.3f}\n\tExpected t_inner for next iteration = {next_t_inner:.3f}\n"