TCAMpy 0.1.0__py3-none-any.whl

TCAMpy.py

@@ -0,0 +1,1534 @@
+import io
+import time
+import random
+import hashlib
+import numpy as np
+import pandas as pd
+import altair as alt
+import streamlit as st
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+
+from sklearn.metrics import r2_score, mean_absolute_error
+from sklearn.model_selection import train_test_split
+from sklearn.ensemble import RandomForestRegressor
+from streamlit_javascript import st_javascript
+from scipy.ndimage import gaussian_filter
+from scipy.stats import skew, kurtosis
+from functools import wraps
+from tqdm import tqdm
+
+class TModel:
+    """
+    Class for a cellular automata, modeling tumor growth.
+
+    Parameters:
+        cycles (int): duration of the model given in hours
+        side (int): the length of the side of field (10um)
+        pmax (int): maximum proliferation potential of RTC
+        PA (int): chance for apoptosis of RTC (in percent)
+        CCT (int): cell cycle time of cells given in hours
+        Dt (float): time step of the model given in days
+        PS (int): STC-STC division chance (in percent)
+        mu (int): migration capacity of cancer cells
+        I (int): strength of the immune cells (1-5)
+        M (int): tumor mutation chance (in percent)
+    """
+    
+    def __init__(self, cycles, side, pmax, PA, CCT, Dt, PS, mu, I, M):
+        # Parameters
+        self.cycles = cycles
+        self.side = side
+        self.pmax = pmax
+        self.CCT = CCT
+        self.Dt = Dt
+        self.mu = mu
+        self.I = I
+        self.M = M
+        
+        # Single model data
+        self.stc_number = []
+        self.rtc_number = []
+        self.wbc_number = []
+        self.immune = []
+        self.mutate = []
+        self.mutmap = []
+        self.field  = []
+        self.images = []
+        
+        # Multiple models data
+        self.stats = []
+        self.runs  = []
+        
+        # Chances
+        self.PP = CCT*Dt/24*100
+        self.PM = 100*mu/24
+        self.PA = PA
+        self.PS = PS
+        
+        # Immune Data
+        self.it_ratio = []
+        self.kill_day = []
+    
+    # ---------------------------------------------------------------------
+    def init_state(self):
+        """
+        Creates the initial state with one STC in the middle.
+        Creates the field for immune cells and mutations too.
+        """
+
+        self.field  = np.zeros((self.side, self.side))
+        self.immune = np.zeros((self.side, self.side))
+        self.mutate = np.zeros((self.side, self.side))
+        self.mutmap = np.zeros((self.side, self.side))
+        
+        self.mod_cell(self.side//2, self.side//2, self.pmax+1)
+        
+    # ---------------------------------------------------------------------
+    def find_tumor_cells(self):
+        """
+        Saves the coordinates of tumor cells to self.tumor_cells.
+        """
+     
+        # Where are tumor cells?
+        coords = np.nonzero(self.field)
+        coords = np.transpose(coords)
+        
+        # Shuffle to randomize direction
+        np.random.shuffle(coords)
+        self.tumor_cells = coords
+
+    # ---------------------------------------------------------------------
+    def count_tumor_cells(self):
+        """
+        Saves the number of STCs/RTCs to self.stc_number/self.rtc_number.
+        """
+        
+        # Count RTC and STC
+        stc_count = np.count_nonzero(self.field == self.pmax + 1)
+        rtc_count = len(self.tumor_cells) - stc_count
+        
+        # Save the current number
+        self.stc_number.append(stc_count)
+        self.rtc_number.append(rtc_count)
+
+    # ---------------------------------------------------------------------
+    def get_neighbours(self, x, y, neighbour_type):
+        """
+        Returns the neighboring coordinates of a given cell in a 2D NumPy matrix.
+
+        Parameters:
+            x, y (int): representing the coordinates of the cell
+            neighbour_type (int): type of neighboring cells (1-5)
+            
+        Returns:
+            list: a list with the coords of the neighbouring cells
+        """
+
+        directions = [
+            (-1, 0), (1, 0),
+            (0, -1), (0, 1),
+            (-1,-1), (-1,1),
+            (1, -1), (1, 1)]
+        
+        coords = []
+        for dx, dy in directions:
+            nx, ny = x + dx, y + dy
+            if 0 < nx < self.side-1 and 0 < ny < self.side-1:
+                coords.append([nx, ny])
+                
+        neighbours = []
+        for n in coords:
+            match neighbour_type:
+                case 1:
+                    # Return list of empty cells
+                    if (self.field[n[0],n[1]] == 0 and self.immune[n[0],n[1]] == 0):
+                        neighbours.append(n)
+                case 2:
+                    # Return list of tumor cells
+                    if self.field[n[0],n[1]] != 0:
+                        neighbours.append(n)
+                case 3:
+                    # Return list of immune cells
+                    if self.immune[n[0],n[1]] != 0:
+                        neighbours.append(n)
+                case 4:
+                    # Return list of any cells
+                    if (self.field[n[0],n[1]] != 0 or self.immune[n[0],n[1]] != 0):
+                        neighbours.append(n)
+                case 5:
+                    # # Return list of free/immune cells
+                    if self.immune[n[0],n[1]] == 0:
+                        neighbours.append(n)
+                    
+        return neighbours   
+    
+    # ---------------------------------------------------------------------
+    def cell_step(self, x, y, step_type):
+        """
+        The function that makes a single cell do one of the following actions:
+        prolif STC - STC, prolif STC - RTC, prolif RTC - RTC, migration (1-4).
+        New mutations can appear every time a cell proliferates with M chance.
+
+        Parameters:
+            x, y (int): representing the coordinates of the cell
+            step_type (int): type of division or migration (1-4)
+        """
+        
+        # Choose random target position
+        free_nb = self.get_neighbours(x, y, 1)
+        nx, ny = free_nb[random.randint(1,len(free_nb)) - 1]
+        
+        match step_type:
+            case 1:
+                # Proliferation STC -> STC + STC
+                self.field[nx, ny] = self.pmax+1
+            case 2:
+                # Proliferation STC -> STC + RTC
+                self.field[nx, ny] = self.pmax
+            case 3:
+                # Proliferation RTC -> RTC + RTC
+                self.field[x,   y] -= 1
+                self.field[nx, ny] = self.field[x, y]
+            case 4:
+                # Migration
+                self.field[nx, ny] = self.field[x, y]
+                self.field[x,   y] = 0
+        
+        if step_type < 4 and self.field[x, y] == 0:
+            self.mutate[x, y] = 0
+            
+        elif step_type < 4:
+            # Inherit mother's mutation
+            self.mutate[nx, ny] = self.mutate[x, y]
+            
+            # Chance of a new mutation
+            if self.M >= random.randint(1, 100):
+                mut = random.choice([-1,1])
+                self.mutate[nx, ny] = np.clip(self.mutate[nx, ny]+mut, -3, 3)
+                
+                # Mutation influences pp value
+                if step_type != 1:
+                    self.field[nx, ny] = np.clip(self.field[nx, ny]+mut, 1, self.pmax)
+                    
+            self.mutmap[nx, ny] = self.mutate[nx, ny]
+        else:
+            self.mutate[nx, ny] = self.mutate[x, y]
+            self.mutmap[nx, ny] = self.mutate[x, y]
+            self.mutate[x,   y] = 0
+        
+    # ---------------------------------------------------------------------
+    def tumor_action(self):
+        """
+        This is the function that decides what action a cell will do.
+        Either kills the cell or calls the 'cell_step' function.
+        This function goes through every single cell in the field.
+        """
+        
+        for cell in self.tumor_cells:
+            x, y = cell
+            is_stc = (self.field[x, y] == self.pmax + 1)
+        
+            # Probabilities
+            probs = np.array([self.PA, self.PP, self.PM, 0], dtype=float)
+            if is_stc:
+                probs[0] = 0
+            if not self.get_neighbours(x, y, 1):
+                probs[1:3] = 0
+            probs = self.mutate_probs(probs, x, y)
+            probs /= probs.sum()
+        
+            # Choose action
+            choice = np.random.choice(4, p=probs)
+        
+            if choice == 0:  # apoptosis
+                self.field[x, y]  = 0
+                self.mutate[x, y] = 0
+                
+            elif choice == 1:  # proliferation
+                if is_stc and np.random.rand() < self.PS/100:
+                    self.cell_step(x, y, 1)   # STC-STC division
+                elif is_stc:
+                    self.cell_step(x, y, 2)   # STC-RTC division
+                else:
+                    self.cell_step(x, y, 3)   # RTC-RTC division
+                    
+            elif choice == 2:  # migration
+                self.cell_step(x, y, 4)
+
+    # ---------------------------------------------------------------------
+    def mutate_probs(self, chances, x, y):
+        """
+        The function that changes the cell action chances
+        based on the current mutation status of the cell.
+        
+        Parameters:
+            chances (list of float): the base action chances
+            x, y (int): representing coordinates of the cell
+        """
+        
+        mut_state = self.mutate[x, y]/2
+        
+        if mut_state > 0:
+            mut_state += 1
+            chances[0] = chances[0]/mut_state        # Decreased chance for apoptosis 
+            chances[1] = chances[1]*mut_state        # Increased proliferation chance
+        elif mut_state < 0:
+            mut_state -= 1
+            chances[0] = chances[0]*abs(mut_state)   # Increased chance for apoptosis 
+            chances[1] = chances[1]/abs(mut_state)   # Decreased proliferation chance
+            
+        if chances.sum() <= 100:
+            chances[3] = 100 - chances.sum()
+        return chances
+
+    # ---------------------------------------------------------------------
+    def immune_response(self, offset = 10, alpha = 0.002, it_targ = 0.1, infil = 0.3):
+        """
+        The function that simulates immune cells.
+        Spawns, moves and activates immune cells.
+        
+        Parameters:
+            offset (int): distance of spawnpoints ("frame") from the tumor
+            alpha (float): controls strength (slope) of immune exhaustion
+            it_targ (float): desired mean immune/tumor ratio during simulation
+            infil (float): "searching/infiltrating" threshold for wbcs (0-1)
+        """
+
+        # Current tumor cell locations
+        self.find_tumor_cells()
+        tumor_size = len(self.tumor_cells)
+        if tumor_size == 0:
+            # Count existing immune cells
+            coords = np.argwhere(self.immune > 0)
+            self.immune_cells = coords
+            immune_size = len(coords)
+            for x, y in coords:
+                self.immune[x, y] -= 1
+            self.wbc_number.append(immune_size)
+            return
+
+        tumor_x = [x for x, _ in self.tumor_cells]
+        tumor_y = [y for _, y in self.tumor_cells]
+    
+        # Find spawnpoints (a "frame" around tumor)
+        frame = []
+        left   = max(1, min(tumor_x) - offset)
+        right  = min(self.side - 2, max(tumor_x) + offset)
+        top    = max(1, min(tumor_y) - offset)
+        bottom = min(self.side - 2, max(tumor_y) + offset)
+        
+        for j in range(left, right + 1):
+            frame.append([top, j])
+            frame.append([bottom, j])
+        for i in range(top, bottom + 1):
+            frame.append([i, left])
+            frame.append([i, right])
+        self.spawnpoints = np.array(frame)
+        
+        # Immune exhaustion = time-dependent decline
+        IE = 1.0 / (1.0 + alpha * self.cycles)
+        IE = max(IE, 0.2)
+
+        # Saturating spawn (sigmoid-like), delayed onse
+        spawn = self.I * (tumor_size / (tumor_size + self.I * 100)) * IE
+        
+        coords = np.nonzero(self.immune)
+        it_ratio = len(np.transpose(coords)) / tumor_size
+        for cell in self.spawnpoints:
+            x, y = cell
+            if np.random.rand() < spawn / 50:
+                if self.immune[x, y] == 0 and self.field[x, y] == 0 and it_ratio <= it_targ:
+                    # Immune cell lifespan: I-1 weeks to I+1 weeks
+                    min_life = min(24, (self.I-1)*168)
+                    max_life = (self.I+1)*168
+                    self.immune[x, y] = np.random.randint(min_life, max_life)
+                  
+        # Chemoattractant map for tumor density
+        self.chemo = (self.field > 0).astype(float)
+        self.chemo = gaussian_filter(self.chemo, sigma=5)
+        self.chemo = self.chemo / np.max(self.chemo)
+                
+        # Immune action
+        coords = np.nonzero(self.immune)
+        kills_per_hour = 0
+        self.immune_cells = np.transpose(coords)
+        
+        # Temporary immune grid
+        new_immune = np.zeros_like(self.immune)
+
+        for (x, y) in self.immune_cells:
+            strength = self.immune[x, y]
+            if strength <= 0:
+                continue
+        
+            # Kill prob on contact: (0.15 - 0.3, if I=5, IE = 0)
+            tumor_nb = self.get_neighbours(x, y, 2)
+            if tumor_nb:
+                tx, ty = random.choice(tumor_nb)
+                kill = (0.05*self.I) * np.exp(-0.25*self.mutate[tx,ty]) * IE
+                kill = min(kill, 0.3)
+                if np.random.rand() < kill:
+                    self.field[tx, ty] = 0
+                    self.mutate[tx, ty] = 0
+                    kills_per_hour += 1
+
+            # # Multiple moves/cycle as immune cells are faster
+            moves = int(1 + self.I * (1 - self.chemo[x, y]))
+            for _ in range(moves):
+                free_nb = self.get_neighbours(x, y, 1)
+                if not free_nb:
+                    strength -= 1
+                    break
+        
+                # Biased movement towards tumor density (chemotaxis)
+                t_dens = [self.chemo[i, j] for (i, j) in free_nb]
+                if sum(t_dens) > 0:
+                    weights = np.array(t_dens) / sum(t_dens)
+                    tx, ty = free_nb[np.random.choice(len(free_nb), p=weights)]
+                else:
+                    tx, ty = random.choice(free_nb)
+                x, y = tx, ty
+                strength -= 1
+                
+            if strength > 0:
+                new_immune[x, y] = strength
+        self.immune = new_immune
+
+        # Save number of immune cells
+        immune_size = len(self.immune_cells)
+        self.wbc_number.append(immune_size)
+        self.it_ratio.append(immune_size / tumor_size)
+        
+        # Infiltrating immune cells
+        wbc_infil = sum(1 for (x,y) in self.immune_cells if self.chemo[x,y] >= infil)
+        if immune_size > 0:
+            self.kill_day.append(kills_per_hour / max(1, wbc_infil) * 24)
+
+    # ---------------------------------------------------------------------
+    def animate(self, mode):
+        """
+        Creates and returns animation of the growth.
+        
+        Parameters:
+            mode (int): create figure, save frame or display animation. (1-3)
+        
+        Returns:
+            ArtistAnimation: the animation of the growth (optional)
+        """
+        
+        if mode == 1:
+            # Create the figure
+            self.fig, self.ax = plt.subplots()
+            self.ax.imshow(self.field)
+            self.ax.set_title(str(self.cycles)+ " hour cell growth")
+            self.ax.set_xlabel(str(self.side*10) +   " micrometers")
+            self.ax.set_ylabel(str(self.side*10) +   " micrometers")
+        elif mode == 2:
+            # Save the current frame
+            growth = self.ax.imshow(self.field, animated=True)
+            immune_coords = np.argwhere(self.immune > 0)
+            immune = self.ax.scatter(immune_coords[:,1], immune_coords[:,0], c='blue', s=10)
+            self.images.append([growth, immune])
+        elif mode == 3:
+            # Display the animation
+            return animation.ArtistAnimation(self.fig, self.images, interval=50, blit=True)
+
+    # ---------------------------------------------------------------------
+    def save_field_to_excel(self, file_name):
+        """
+        Saves the current state of self.field to an excel file.
+        
+        Parameters:
+            file_name (str): name of the excel file
+        """
+
+        pd.DataFrame(self.field).to_excel(file_name, index=False)
+
+    # ---------------------------------------------------------------------
+    def mod_cell(self, x, y, value):
+        """
+        Modifies cell value. (Create initial state before this!)
+
+        Parameters:
+            x, y (int): representing coordinates of the cell
+            value (int): the new value at the given position
+        """
+        
+        self.field[y][x] = value
+
+    # ---------------------------------------------------------------------
+    def get_prolif_potentials(self):
+        """
+        Returns a dictionary of proliferation potential numbers.
+        
+        Returns:
+            dict: a dictionary of the proliferation potentials
+        """
+        
+        nonzero_field  = np.array(self.field)[np.array(self.field) > 0]
+        unique, counts = np.unique(nonzero_field, return_counts=True)
+        prolif_potents = {}
+        
+        for i in range(1, self.pmax + 2):
+            prolif_potents[i] = 0
+        for val, count in zip(unique, counts):
+            prolif_potents[int(val)] = count
+            
+        return prolif_potents
+
+    # ---------------------------------------------------------------------
+    def get_statistics(self):
+        """
+        Returns various statistical properties of the model.
+        
+        Returns:
+            dict: a dictionary of the statistical properties
+        """
+        
+        nonzero_field = self.field[self.field > 0]
+
+        # Statistics
+        if nonzero_field.size != 0:
+            stats = {
+                "Min":        nonzero_field.min(),
+                "Max":        nonzero_field.max(),
+                "Mean":       nonzero_field.mean(),
+                "Std":        nonzero_field.std(),
+                "Median":     np.median(nonzero_field),
+                "Skew":       skew(nonzero_field.ravel()),
+                "Kurtosis":   kurtosis(nonzero_field.ravel()),
+                "Final STC":  self.stc_number[self.cycles-1],
+                "Final RTC":  self.rtc_number[self.cycles-1],
+                "Final WBC":  self.wbc_number[self.cycles-1],
+                "Tumor Size": nonzero_field.size,
+                "Confluence": nonzero_field.size/self.field.size*100,
+            }
+            
+            if self.I > 0:
+                stats.update({
+                    "Mean I/T"  : sum(self.it_ratio)/len(self.it_ratio),
+                    "Mean k/d"  : sum(self.kill_day)/len(self.kill_day)
+                })
+        
+            # Proliferation potentials
+            stats.update(self.get_prolif_potentials())
+                
+            # Cell Numbers
+            checkpoints = np.linspace(0, self.cycles - 1, int(self.cycles/10) + 1, dtype=int)
+            for idx in checkpoints:
+                hour = (idx + 1)
+                stats[f"{hour}h_STC"] = self.stc_number[idx]
+                stats[f"{hour}h_RTC"] = self.rtc_number[idx]
+                stats[f"{hour}h_WBC"] = self.wbc_number[idx]
+            
+        else: stats = {"Status": "Extinct"}
+        return stats
+    
+    # ---------------------------------------------------------------------
+    def save_statistics(self, file_name):
+        """
+        Saves various statistical properties of the model to an excel file.
+        
+        Parameters:
+            file_name (str): name of the excel file
+        """
+        
+        stats_dict = self.get_statistics()
+        df = pd.DataFrame([stats_dict])
+        df.to_excel(file_name, index=False)
+
+    # ---------------------------------------------------------------------
+    def measure_runtime(func):
+        # Decorator to measure completion time
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            result     = func(*args, **kwargs)
+            end_time   = time.time()
+            runtime    = end_time - start_time
+            print("Model completion time (s): " + str(runtime))
+            return result
+        return wrapper
+
+    # ---------------------------------------------------------------------
+    @measure_runtime
+    def run_model(self, plot, animate, stats):
+        """
+        The function that runs a single entire simulation.
+        For animation matplotlib backend cannot be inline!
+
+        Parameters:
+            plot (bool): set to true to display the plots of the model
+            animate (bool): set to true to enable matplotlib animation
+            stats (bool): set to true to print statistics of the model
+        """
+
+        # Create initial state
+        if len(self.field) == 0: self.init_state()
+        self.find_tumor_cells()
+        if len(self.immune) == 0:
+            self.immune = np.zeros((self.side, self.side))
+        if len(self.mutate) == 0:
+            self.mutate = np.zeros((self.side, self.side))
+            self.mutmap = np.zeros((self.side, self.side))
+        
+        self.stc_number = []
+        self.rtc_number = []
+        self.wbc_number = []
+        
+        if animate: self.animate(1)
+
+        # Growth loop
+        for c in tqdm(range(self.cycles), desc="Running simulation..."):
+            self.tumor_action()
+            self.immune_response()
+            self.find_tumor_cells()
+            self.count_tumor_cells()
+            if animate: self.animate(2)
+        
+        # Store the results
+        self.store_model()
+        
+        # Output settings
+        if plot: self.plot_run(len(self.runs))
+        if animate: self.ani = self.animate(3)
+        if stats:
+            df = pd.DataFrame(self.stats)
+            base_cols = self.separate_columns(df)[0]
+            print(df[base_cols])
+        
+    # ---------------------------------------------------------------------
+    @measure_runtime
+    def run_multimodel(self, count, init_field, plot, stats):
+        """
+        Runs the model multiple times and returns a DataFrame of statistics.
+    
+        Parameters:
+            count (int): number of times to run the simulation
+            init_field (np.array): custom initial state of field/run
+            plot (bool): set to true to display the plots of the model
+            stats (bool): set to true to print statistics of the model
+    
+        Returns:
+            pd.DataFrame: collected statistics from each run
+        """
+        
+        stats = []
+        
+        for i in range(count):
+            self.field  = init_field.copy()
+            self.immune = []
+            self.mutate = []
+            self.mutmap = []
+            self.run_model(plot = False, animate = False, stats = False)
+            stats.append(self.get_statistics())
+        all_stats = pd.DataFrame(stats)
+        
+        if plot:
+            self.plot_averages(all_stats)
+        if stats:
+            df = pd.DataFrame(self.stats)
+            base_cols = self.separate_columns(df)[0]
+            print(df[base_cols])
+        
+        return all_stats
+
+    # ---------------------------------------------------------------------
+    def store_model(self):
+        """
+        Stores the results of the previous model executions.
+        """
+        
+        result = {}
+        
+        result["immune"] = self.immune
+        result["mutate"] = self.mutate
+        result["mutmap"] = self.mutmap
+        result["field"]  = self.field
+        result["stc"]    = self.stc_number
+        result["rtc"]    = self.rtc_number
+        result["wbc"]    = self.wbc_number
+        result["pp"]     = self.get_prolif_potentials().values()
+        
+        # Stores data for plotting
+        self.runs.append(result)     
+        
+        # Stores data for statistics
+        self.stats.append(self.get_statistics())
+        
+    # ---------------------------------------------------------------------
+    def separate_columns(self, data):
+        """
+        Separates the statistics DataFrame columns into logical groups:
+        base stats, STC, RTC, WBC counts, and proliferation potentials.
+
+        Parameters:
+            data (pd.DataFrame): Your data in a pandas dataframe format
+
+        Returns:
+            tuple of list[str]: A tuple containing 5 lists of column names:
+                - base: Columns with general statistical properties
+                - stc:  Columns with STC counts at each time point
+                - rtc:  Columns with RTC counts at each time point
+                - wbc:  Columns with WBC counts at each time point
+                - pp:   Columns for proliferation potential values
+        """
+        
+        base = [col for col in data.columns if not str(col).isdigit()
+                and "_STC" not in str(col)
+                and "_RTC" not in str(col)
+                and "_WBC" not in str(col)]
+        stc  = sorted([col for col in data.columns if "_STC" in str(col)],
+                      key=lambda x: int(str(x).split("h")[0]))
+        rtc  = sorted([col for col in data.columns if "_RTC" in str(col)],
+                      key=lambda x: int(str(x).split("h")[0]))
+        wbc  = sorted([col for col in data.columns if "_WBC" in str(col)],
+                      key=lambda x: int(str(x).split("h")[0]))
+        pp   = sorted([col for col in data.columns if isinstance(col, int)])
+        
+        return base, stc, rtc, wbc, pp
+
+    # ---------------------------------------------------------------------
+    def plot_run(self, run):
+        """
+        Creates growth and cell number plots, proliferation potential histograms.
+        
+        Paramteres:
+            run (int): which model execution to plot
+            
+        Returns:
+            matplotlib.figure.Figure: the generated plots of the specific run
+        """
+        
+        # Create the figue and axis
+        fig, axs = plt.subplots(2, 2, figsize=(14,14))
+
+        tumor = axs[0, 0].imshow(self.runs[run-1]["field"], vmin=0, vmax=self.pmax+1)
+        fig.colorbar(tumor, ax=axs[0, 0])
+        
+        immune_coords = np.argwhere(self.runs[run-1]["immune"] > 0)
+        axs[0, 0].scatter(immune_coords[:,1], immune_coords[:,0],
+                          c='blue', marker='v', s=10)
+        
+        axs[0, 1].plot(self.runs[run-1]["stc"], 'C1', label='STC')
+        axs[0, 1].plot(self.runs[run-1]["rtc"], 'C2', label='RTC')
+        axs[0, 1].plot(self.runs[run-1]["wbc"], 'C3', label='WBC')
+        axs[0, 1].legend()
+        
+        mutmap = axs[1, 0].imshow(self.runs[run-1]["mutmap"],
+                 cmap="RdBu_r", vmin=-3, vmax=3, interpolation="bicubic")
+        fig.colorbar(mutmap, ax=axs[1, 0])
+        
+        axs[1, 1].bar(range(1, self.pmax + 2), self.runs[run-1]["pp"], edgecolor='black')
+
+        # Titles/labels of the plots
+        titles = [str(self.cycles)+ "h cell growth", "Cell count",
+                  "Mutation history", "Final PP values"]
+        labs_x = [str(self.side*10) + " um", "Time (h)",
+                  str(self.side*10) + " um", "Proliferation potentials"]
+        labs_y = [str(self.side*10) + " um", "Cell numbers",
+                  str(self.side*10) + " um", "Number of appearance"]
+
+        fig.suptitle("Simulation " + str(run) + " Results", fontsize = 16)
+        for i, ax in enumerate(axs.flat):
+            ax.set_title(titles[i])
+            ax.set_xlabel(labs_x[i])
+            ax.set_ylabel(labs_y[i])
+
+    # ---------------------------------------------------------------------
+    def plot_averages(self, data):
+        """
+        The function that plots the averages of multiple model results.
+        Works with the results of the 'run_multimodel' function.
+        
+        Parameters:
+            data (pd.DataFrame): Your data in a pandas dataframe format
+            
+        Returns:
+            matplotlib.figure.Figure: The plots of the averages with SD values
+        """
+        
+        base_cols, stc_cols, rtc_cols, wbc_cols, pp_cols = self.separate_columns(data)
+        
+        avg_stc = data[stc_cols].mean()
+        std_stc = data[stc_cols].std()
+        avg_rtc = data[rtc_cols].mean()
+        std_rtc = data[rtc_cols].std()
+        avg_wbc = data[wbc_cols].mean()
+        std_wbc = data[wbc_cols].std()
+        avg_pp  = data[pp_cols].mean()
+        std_pp  = data[pp_cols].std()
+        
+        fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(14, 5))
+        timepoints      = np.linspace(0, self.cycles - 1, int(self.cycles/10) + 1)
+        
+        ax1.plot(timepoints, avg_stc, label='STC', color='C1')
+        ax1.fill_between(timepoints, avg_stc - std_stc, avg_stc + std_stc,
+                         color='C1', alpha=0.3)
+        ax1.plot(timepoints, avg_rtc, label='RTC', color='C2')
+        ax1.fill_between(timepoints, avg_rtc - std_rtc, avg_rtc + std_rtc,
+                         color='C2', alpha=0.3)
+        ax1.plot(timepoints, avg_wbc, label='WBC', color='C3')
+        ax1.fill_between(timepoints, avg_wbc - std_wbc, avg_wbc + std_wbc,
+                         color='C3', alpha=0.3)
+        
+        ax1.set_title("Average Tumor Cell Count")
+        ax1.set_xlabel("Model Time (hours)")
+        ax1.set_ylabel("Number of Cells")
+        ax1.legend()
+
+        ax2.bar(pp_cols, avg_pp, yerr=std_pp, capsize=5, edgecolor='black')
+        ax2.set_title("Average Proliferation Potential Distribution")
+        ax2.set_xlabel("Proliferation Potential")
+        ax2.set_ylabel("Average Count")
+        
+        fig.suptitle("Averages of " + str(len(self.stats)) + " Models", fontsize = 16)
+        plt.tight_layout()
+
+
+class TDashboard:
+    """
+    Class for a Streamlit dashboard providing a GUI for the model.
+    
+    Parameters:
+        model (TModel): The created model you want a dashboard for
+    """
+    
+    def __init__(self, model):
+        self.model = model
+
+    # ---------------------------------------------------------------------
+    def run_dashboard(self):
+        """
+        The function that creates the entire streamlit dashboard for the model.
+        """
+        
+        st.set_page_config(layout="wide")
+        st.markdown("<h1 style='text-align: center;'>TCAMpy</h1>", unsafe_allow_html=True)
+        self.screen_width = st_javascript("window.innerWidth", key="screen_width")
+
+        tab1, tab2 = st.tabs(["SIMULATION", "MACHINE LEARNING"])
+        with tab1:
+            self.columns = [4, 1, 12]
+            self.col1, _, self.col3 = st.columns(self.columns)
+
+            with self.col1:
+                self._initialize()
+                self._modify_cell()
+                self._execute_model()
+            with self.col3:
+                self._visualize_run("Last Simulation", len(self.model.runs))
+                self._show_statistics()
+                self._reset_save_stats()
+        
+        with tab2:
+            col1, col2 = st.columns(2)
+
+            with col1:
+                self._simdata_generator()
+            with col2:
+                self._train_and_predict()
+
+    # ---------------------------------------------------------------------
+    def print_title(self, title):
+        """
+        The function that prints text as a title on the dashboard.
+        
+        Parameters:
+            title (string): The text to print
+        """
+
+        st.markdown(
+            f"<h2 style='text-align: center;'>{title}</h2>",
+            unsafe_allow_html=True
+        )
+    
+    # ---------------------------------------------------------------------
+    def get_plot_height(self, col, scaler):
+        """
+        The function that calculates the height of plots
+        based on screen width, column width and a scaler.
+        
+        Parameters:
+            col (int): main column number
+            scalar (float): scaler for column width
+        """
+        
+        screen_width = st.session_state.get("screen_width")
+        col_width_px = screen_width * (self.columns[col-1] / sum(self.columns))
+        return int(col_width_px * scaler)
+
+    # ---------------------------------------------------------------------
+    def _initialize(self):
+        """
+        The function that sets the parameters and initializes the model.
+        """
+
+        self.print_title("Model Parameters")
+
+        self.model.cycles = st.slider("Model Duration (hours)", 50, 5000, value=self.model.cycles)
+        self.model.side   = st.slider("Field Side Length (10um)", 10, 200, value=self.model.side)
+        self.model.pmax   = st.slider("Max Proliferation Potential", 1, 20, value=self.model.pmax)
+        self.model.PA     = st.slider("Apoptosis Chance (RTC) (%)", 0, 100, value=self.model.PA)
+        self.model.CCT    = st.slider("Cell Cycle Time (hours)", 1, 48, value=self.model.CCT)
+        self.model.Dt     = st.slider("Time Step (days)", 0.01, 1.0, value=self.model.Dt, step=0.01)
+        self.model.PS     = st.slider("STC-STC Division Chance (%)", 0, 100, value=self.model.PS)
+        self.model.mu     = st.slider("Migration Capacity", 0, 10, value=self.model.mu)
+        self.model.I      = st.slider("Immune Strength", 0, 10, value=self.model.I)
+        self.model.M      = st.slider("Mutation Chance", 0, 50, value=self.model.M)
+
+        self.model.PP = int(self.model.CCT * self.model.Dt / 24 * 100)
+        self.model.PM = 100 * self.model.mu / 24
+
+        init_config = (
+            self.model.side, self.model.cycles, self.model.pmax,
+            self.model.PA, self.model.CCT, self.model.Dt, self.model.PS,
+            self.model.mu, self.model.I, self.model.M
+        )
+        config_hash = hashlib.md5(str(init_config).encode()).hexdigest()
+
+        # Storing data for model plotting
+        if "model_runs" in st.session_state:
+            self.model.runs = st.session_state.model_runs
+            
+        # Storing data for model statistics
+        if "model_stats" in st.session_state:
+            self.model.stats = st.session_state.model_stats
+
+        if (
+            "initialized" not in st.session_state
+            or "init_config_hash" not in st.session_state
+            or st.session_state.init_config_hash != config_hash
+        ):
+            self.model.init_state()
+            st.session_state.field  = self.model.field.copy()
+            st.session_state.immune = self.model.immune.copy()
+            st.session_state.mutate = self.model.mutate.copy()
+            st.session_state.mutmap = self.model.mutmap.copy()
+            st.session_state.initialized = True
+            st.session_state.init_config_hash = config_hash
+
+    # ---------------------------------------------------------------------
+    def _modify_cell(self):
+        """
+        The function for initial state modification logic.
+        """
+        
+        self.print_title("Initial State")
+
+        x_coord = st.number_input("X Coordinate", 0, self.model.side - 1, value=self.model.side // 2)
+        y_coord = st.number_input("Y Coordinate", 0, self.model.side - 1, value=self.model.side // 2)
+        cell_value = st.number_input("Cell Value", 0, self.model.pmax + 1, value=self.model.pmax + 1)
+        plots_height = self.get_plot_height(1, 0.9)
+
+        if st.button("Modify Cell"):
+            self.model.field = st.session_state.field.copy()
+            self.model.mod_cell(x_coord, y_coord, cell_value)
+            st.session_state.field = self.model.field.copy()
+            st.success(f"Cell modified at ({x_coord}, {y_coord}) to {cell_value}")
+
+        field   = st.session_state.field
+        heatmap = self._create_heatmap(
+            plots_height, "Initial state", "viridis",
+            "PP", 0, self.model.pmax+1, field
+            )
+        
+        st.altair_chart(heatmap, use_container_width=True)
+
+    # ---------------------------------------------------------------------
+    def _execute_model(self):
+        """
+        The function for model running logic.
+        """
+        
+        self.print_title("Execution")
+        
+        rep = st.number_input("How many simulations?", 1)
+        
+        if st.button("Run Model"):
+            with st.spinner("Running simulations..."):
+                for i in range(rep):
+                    self.model.field  = st.session_state.field.copy()
+                    self.model.immune = st.session_state.immune.copy()
+                    self.model.mutate = st.session_state.mutate.copy()
+                    self.model.mutmap = st.session_state.mutmap.copy()
+                    self.model.run_model(plot = False, animate=False, stats=False)
+    
+                st.session_state.model_runs = self.model.runs
+                st.session_state.model_stats = self.model.stats
+
+    # ---------------------------------------------------------------------
+    def _visualize_run(self, title, run):
+        """
+        The function for the result visualization logic.
+    
+        Parameters:
+            title (string): title of the visualization
+            run (int): which model execution to plot
+        """
+    
+        if "model_runs" not in st.session_state:
+            st.warning("Simulation results will appear here...")
+            return
+        self.print_title(title)
+    
+        # --- Get latest run ---
+        latest = self.model.runs[run - 1]
+        immune = latest["immune"]
+        mutmap = latest["mutmap"]
+        field  = latest["field"]
+        stc    = latest["stc"]
+        rtc    = latest["rtc"]
+        wbc    = latest["wbc"]
+        pp     = latest["pp"]
+    
+        # --- Create charts ---
+        plots_height = self.get_plot_height(3, 0.4)
+        
+        tumor_heatmap = self._create_heatmap(
+            plots_height, "Tumor growth", "viridis",
+            "PP", 0, self.model.pmax+1, field, immune
+        )
+        mutation_map = self._create_heatmap(
+            plots_height, "Mutation history", "redblue",
+            "M", -3, 3, mutmap
+        )
+        
+        bar_chart  = self._create_bar_chart(plots_height, list(pp))
+        line_chart = self._create_line_chart(plots_height, stc, rtc, wbc)
+
+        # --- Layout rules ---
+        col1, col2 = st.columns([4, 5])
+        with col1:
+            st.altair_chart(tumor_heatmap, use_container_width=True)
+            st.altair_chart(mutation_map, use_container_width=True)
+        with col2:
+            st.altair_chart(bar_chart, use_container_width=True)
+            st.altair_chart(line_chart, use_container_width=True)
+
+    # ---------------------------------------------------------------------
+    def _create_heatmap(
+            self, h, title, cmap, ctitle,
+            vmin, vmax, heatmap, scatter=None
+        ):
+        """
+        Creates an Altair heatmap with a scatter plot overlaid.
+        Used for tumor field with immune cells, and mutations.
+    
+        Parameters:
+            h (int): the height of the plot
+            title (string): title of the plot
+            cmap (stirng): colormap for the heatmap
+            vmin, vmax (int): domain for the colormap
+            heatmap (2D array-like): array for the heatmap
+            scatter (2D array-like): array for scatter plot
+    
+        Returns:
+            Altair.Chart: heatmap with scatter overlay
+        """
+        
+        # --- Heatmap data ---
+        heat_df = pd.DataFrame([
+            {"x": x, "y": y, "value": heatmap[y, x]}
+            for y in range(heatmap.shape[0])
+            for x in range(heatmap.shape[1])
+        ])
+    
+        heat_chart = alt.Chart(heat_df).mark_rect().encode(
+            x=alt.X("x:O", title="X"),
+            y=alt.Y("y:O", sort="descending", title="Y"),
+            color=alt.Color("value:Q", title=ctitle,
+                  scale=alt.Scale(scheme=cmap, domain=[vmin, vmax]))
+        ).properties(
+            title = title,
+            width='container',
+            height=h
+        )
+        
+        # --- Scatter plot data ---
+        if scatter is not None:
+            scatter_coords = np.argwhere(scatter > 0)
+            scatter_df = pd.DataFrame(scatter_coords, columns=["y", "x"])
+        
+            scatter = alt.Chart(scatter_df).mark_point(
+                color="blue", size=h/20, filled=True, shape="circle"
+            ).encode(
+                x=alt.X("x:O"),
+                y=alt.Y("y:O", sort="descending")
+            )
+        
+            # --- Combine layers ---
+            heat_chart = (heat_chart + scatter).properties(
+                title=title,
+                width='container',
+                height=h
+            )
+    
+        return heat_chart
+    
+    # ---------------------------------------------------------------------
+    def _create_line_chart(
+            self, h, stc, rtc, wbc, stc_l=None, stc_u=None,
+            rtc_l=None, rtc_u=None, wbc_l=None, wbc_u=None
+        ):
+        """
+        The function that creates an Altair line chart of the cell numbers.
+        
+        Parameters:
+            h (int): the height of the plot
+            stc, rtc, wbc (list): a list of the cell and immune numbers (mean or raw)
+            stc_l, rtc_l, wbc_l (list of float, optional): Lower bounds (e.g., mean - SD) for cell counts.
+            stc_u, rtc_u, wbc_u (list of float, optional): Upper bounds (e.g., mean + SD) for cell counts.
+            
+        Returns:
+            Altair.Chart: represents the line chart of the cell numbers
+        """
+
+        timepoints = list(range(len(stc)))
+        df = pd.DataFrame({
+            "Hour": timepoints * 3,
+            "Cell Type": ["STC"] * len(stc) + ["RTC"] * len(rtc) + ["WBC"] * len(wbc),
+            "Mean": stc + rtc + wbc
+        })
+    
+        if stc_l and rtc_l and wbc_l:
+            df["Lower"] = stc_l + rtc_l + wbc_l
+            df["Upper"] = stc_u + rtc_u + wbc_u
+    
+            area = alt.Chart(df).mark_area(opacity=0.3).encode(
+                x=alt.X("Hour:Q", title="Time (hours)"),
+                y=alt.Y("Lower:Q", title="Mean"),
+                y2="Upper:Q",
+                color="Cell Type:N"
+            )
+        else:
+            area = None
+    
+        line = alt.Chart(df).mark_line().encode(
+            x="Hour:Q",
+            y=alt.Y("Mean:Q", title="Mean"),
+            color="Cell Type:N"
+        )
+
+        chart = (area + line) if area else line
+        return chart.properties(title="Cell Counts Over Time", height=h)
+    
+    # ---------------------------------------------------------------------
+    def _create_bar_chart(self, h, pp, std=None):
+        """
+        Creates an Altair bar chart for proliferation potential distribution.
+    
+        Parameters:
+            h (int): the height of the plot
+            pp (list of float or int): Mean or raw counts of cells per proliferation potential class
+            std (list of float, optional): Standard deviation for each class
+    
+        Returns:
+            alt.Chart: An Altair chart representing the distribution of proliferation potentials
+        """
+        
+        pp_df = pd.DataFrame({
+            "Proliferation Potential": list(range(1, len(pp) + 1)),
+            "Mean": pp
+        })
+        chart = alt.Chart(pp_df).mark_bar().encode(
+            x="Proliferation Potential:O",
+            y="Mean:Q"
+        )
+    
+        if std is not None:
+            pp_df["Std"] = std
+            error = alt.Chart(pp_df).mark_errorbar(extent="stdev").encode(
+                x="Proliferation Potential:O",
+                y="Mean:Q",
+                yError="Std:Q"
+            )
+            chart = chart + error
+    
+        return chart.properties(title="Proliferation Potential Distribution", height=h)
+
+    # ---------------------------------------------------------------------
+    def _show_statistics(self):
+        """
+        The function for the statistics printing logic.
+        """
+        
+        if not self.model.stats: return
+        
+        self.print_title("All Simulations")
+        plots_height = self.get_plot_height(3, 0.4)
+        df = pd.DataFrame(self.model.stats)
+        base_cols, stc_cols, rtc_cols, wbc_cols, pp_cols = self.model.separate_columns(df)
+        df.index = df.index + 1
+
+        # Display Statistics
+        mean_row = df[base_cols].mean(numeric_only=True)
+        std_row  = df[base_cols].std(numeric_only=True)
+        mean_row.name = "Mean"
+        std_row.name  = "Std"
+        full_stats = pd.concat([df[base_cols], mean_row.to_frame().T, std_row.to_frame().T])
+        st.dataframe(full_stats)
+
+        # Create avg charts
+        stc_means = df[stc_cols].mean()
+        stc_stds  = df[stc_cols].std()
+        rtc_means = df[rtc_cols].mean()
+        rtc_stds  = df[rtc_cols].std()
+        wbc_means = df[wbc_cols].mean()
+        wbc_stds  = df[wbc_cols].std()
+        pp_means  = df[pp_cols].mean()
+        pp_stds   = df[pp_cols].std()
+        
+        line_chart = self._create_line_chart(plots_height,
+            list(stc_means.values), list(rtc_means.values), list(wbc_means.values),
+            list((stc_means - stc_stds).values), list((stc_means + stc_stds).values),
+            list((rtc_means - rtc_stds).values), list((rtc_means + rtc_stds).values),
+            list((wbc_means - wbc_stds).values), list((wbc_means + wbc_stds).values)
+        )
+        
+        bar_chart = self._create_bar_chart(plots_height, list(pp_means.values), list(pp_stds.values))
+        
+        col1, col2 = st.columns(2)
+        with col1:
+            st.altair_chart(line_chart, use_container_width=True)
+        with col2:
+            st.altair_chart(bar_chart, use_container_width=True)
+
+    # ---------------------------------------------------------------------
+    def _reset_save_stats(self):
+        """
+        The function for the reset/download statistics logic.
+        """
+        
+        if "model_stats" in st.session_state:
+            self.print_title("Simulation Options")
+            col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
+            selected_run = None
+            visualize = False
+            
+            with col1:
+                if st.button("Reset Model Executions and Data", use_container_width=True):
+                    del st.session_state.model_stats
+                    self.model.stats.clear()
+                    del st.session_state.model_runs
+                    self.model.runs.clear()
+                    
+                    st.success("Executions have been reset.")
+            with col2:
+                buffer = io.BytesIO()
+                pd.DataFrame(self.model.stats).to_excel(buffer, index=False)
+                buffer.seek(0)
+
+                st.download_button(
+                    label="Download Statistics (xlsx)",
+                    data=buffer,
+                    file_name="simulation_statistics.xlsx",
+                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+                    use_container_width=True
+                )
+            with col3:
+                run_select = st.selectbox(
+                    "", list(range(1, len(self.model.runs) + 1)),
+                    placeholder="Select simulation",
+                    label_visibility="collapsed",
+                    index=None
+                    )
+                selected_run = run_select
+            with col4:
+                if st.button("Visualize Selected Simulation", use_container_width=True):
+                    if selected_run: visualize = True
+                    else: st.warning('Please select a simulation!')
+            if visualize:
+                self._visualize_run("Selected Simulation", selected_run)
+
+    # ---------------------------------------------------------------------
+    def _simdata_generator(self):
+        """
+        The Machine Learning tab for dataset generation and download.
+        Uses the TML class to generate simulation data.
+        """
+
+        self.print_title("Simulation Data Generator")
+
+        # Initialize TML
+        tml = TML(self.model)
+
+        st.write("Select randomization ranges for each parameter:")
+
+        # Build parameter range inputs dynamically
+        param_ranges = {}
+        for param, default_val in tml.default_params.items():
+            col1, col2 = st.columns(2)
+            with col1:
+                low = st.number_input(
+                    f"{param} (min)", 
+                    value=float(default_val) * 0.8, 
+                    key=f"{param}_low"
+                )
+            with col2:
+                high = st.number_input(
+                    f"{param} (max)", 
+                    value=float(default_val) * 1.5, 
+                    key=f"{param}_high"
+                )
+            param_ranges[param] = (low, high)
+
+        n = st.number_input("Number of simulations", 5, 500, 50, step=5)
+
+        # Run simulation button
+        if st.button("Generate Dataset", use_container_width=True):
+            with st.spinner("Running simulations..."):
+                df = tml.generate_dataset(n=n, random_params=param_ranges)
+            st.success(f"Dataset generated successfully ({len(df)} rows).")
+
+            # Allow CSV download
+            csv = df.to_csv(index=False).encode("utf-8")
+            st.download_button(
+                label="Download Dataset (.csv)",
+                data=csv,
+                file_name="tumor_dataset.csv",
+                mime="text/csv",
+                use_container_width=True
+            )
+
+    # ---------------------------------------------------------------------
+    def _train_and_predict(self):
+        """
+        Streamlit UI for model training and prediction using the TML class.
+        """
+        
+        self.print_title("Model Trainer and Predictor")
+
+        tml = TML(self.model)
+    
+        uploaded_file = st.file_uploader("Upload CSV dataset", type=["csv"])
+        if uploaded_file is not None:
+            df = pd.read_csv(uploaded_file)
+            st.write("Uploaded dataset preview:")
+            st.dataframe(df.head())
+    
+            # Choose target column
+            target = st.selectbox("Select target attribute", df.columns, index=len(df.columns) - 1)
+    
+            # Three sliders for model parameters
+            test_size = st.slider("Test Size", 0.1, 0.5, 0.2, step=0.05)
+            random_state = st.slider("Random Seed", 0, 100, 42, step=1)
+            n_estimators = st.slider("Number of Trees (n_estimators)", 50, 500, 200, step=50)
+    
+            if st.button("Train Model", use_container_width=True):
+                with st.spinner("Training model..."):
+                    model, metrics = tml.train_predictor(
+                        file=df,
+                        target=target,
+                        test_size=test_size,
+                        random_state=random_state,
+                        n_estimators=n_estimators
+                    )
+    
+                st.success("Model trained successfully!")
+                st.write(f"**R^2:** {metrics['R2']:.3f}")
+                st.write(f"**MAE:** {metrics['MAE']:.3f}")
+    
+                # Store trained model for later prediction
+                st.session_state["trained_tml"] = tml
+                st.session_state["target"] = target
+        else:
+            st.info("Please upload a dataset to train a model.")
+
+        if "trained_tml" in st.session_state:
+            trained_tml = st.session_state["trained_tml"]
+            target = st.session_state.get("target", "Target")
+            feature_cols = trained_tml.feature_columns
+    
+            # Numeric inputs for each feature
+            new_params = []
+            self.print_title("Predict for new instance")
+            for col in feature_cols:
+                val = st.number_input(f"{col}", value=1.0, key=f"pred_{col}")
+                new_params.append(val)
+    
+            # Target selector (for user clarity, though single-target regression)
+            st.markdown("#### Select Target to Predict:")
+            st.text(f"Predicting: {target}")
+    
+            if st.button("🔮 Predict New", use_container_width=True):
+                try:
+                    prediction = trained_tml.predict_new(new_params)
+                    st.success(f"Predicted {target}: **{prediction:.3f}**")
+                except Exception as e:
+                    st.error(f"Prediction failed: {e}")
+        else:
+            st.info("Train a model first to enable prediction.")
+
+
+class TML:
+    """
+    Class for handling Machine Learning tasks related to the tumor model.
+    Allows dataset generation, parameter exploration, and result export.
+    Allows predicting the size/confluence for a new set of parameters.
+    
+    Parameters:
+        model (TModel): a created instance of the TModel class.
+    """
+
+    def __init__(self, model):
+        self.model = model
+
+        self.default_params = {
+            "cycles": self.model.cycles,
+            "side":   self.model.side,
+            "pmax":   self.model.pmax,
+            "PA":     self.model.PA,
+            "CCT":    self.model.CCT,
+            "Dt":     self.model.Dt,
+            "PS":     self.model.PS,
+            "mu":     self.model.mu,
+            "I":      self.model.I,
+            "M":      self.model.M,
+        }
+
+    # ---------------------------------------------------------------------
+    def generate_dataset(
+            self, n=50, random_params=None,
+            output_file="tumor_dataset.csv"
+        ):
+        """
+        Generate a dataset of tumor simulations by randomizing given parameters.
+
+        Parameters:
+            n_sims (int): Number of simulations to run.
+            randomize_params (dict): Parameters to randomize, e.g.
+                {
+                    "PA": (1, 20), "PS": (10, 40))
+                }
+            output_file (str): CSV filename to save dataset.
+
+        Returns:
+            pd.DataFrame: Combined DataFrame with all simulation results.
+        """
+
+        stats = []
+
+        # Randomize chosen parameters
+        for i in tqdm(range(n), desc="Generating simulations"):
+            params = self.default_params.copy()
+            for key, (low, high) in random_params.items():
+                if isinstance(params[key], int):
+                    params[key] = random.randint(int(low), int(high))
+                else:
+                    params[key] = random.uniform(float(low), float(high))
+
+            # Run simulation
+            model = TModel(**params)
+            model.run_model(plot = False, animate = False, stats = False)
+
+            run_stats = {}
+            for k, v in params.items():
+                run_stats[k] = v
+            run_stats["Tumor size"] = np.count_nonzero(model.field)
+            run_stats["Confluence"] = np.count_nonzero(model.field)/model.field.size*100
+            stats.append(run_stats)
+
+        if stats:
+            df = pd.DataFrame(stats)
+            df.to_csv(output_file, index=False)
+            print(f"Dataset saved to {output_file} ({len(df)} runs)")
+            return df
+            
+    # ---------------------------------------------------------------------
+    def train_predictor(
+            self, file, target, test_size=0.2, 
+            random_state=42, n_estimators=200
+        ):
+        """
+        Trains a regression model to predict final tumor size based on simulation parameters.
+    
+        Parameters:
+            file (str): CSV file containing the dataset
+            target (str): Column name of the target attribute
+            test_size (float): Fraction of dataset to use for testing
+            random_state (int): Random seed for reproducibility
+            n_estimators (int): Number of trees in the random forest
+    
+        Returns:
+            model (RandomForestRegressor): Trained model
+            metrics (dict): R^2 and MAE metrics on test set
+        """
+
+        if isinstance(file, pd.DataFrame):
+            df = file
+        else:
+            df = pd.read_csv(file)
+        x  = df[df.columns[0:10]]
+        y  = df[target]
+    
+        self.feature_columns = x.columns.tolist()
+
+        # Split into train/test
+        x_train, x_test, y_train, y_test = train_test_split(
+            x, y, test_size=test_size, random_state=random_state
+        )
+
+        # Train model
+        model = RandomForestRegressor(
+            n_estimators=n_estimators,
+            random_state=random_state,
+            n_jobs=-1
+        )
+        model.fit(x_train, y_train)
+    
+        # Predict & evaluate
+        y_pred = model.predict(x_test)
+        metrics = {
+            "R2": r2_score(y_test, y_pred),
+            "MAE": mean_absolute_error(y_test, y_pred)
+        }
+        self.trained_model = model
+    
+        print(f"Model trained on {len(x_train)} samples, tested on {len(x_test)}")
+        print(f"R^2: {metrics['R2']:.3f}, MAE: {metrics['MAE']:.3f}")
+    
+        # Feature importance summary
+        importance = pd.Series(model.feature_importances_, index=x.columns).sort_values(ascending=False)
+        print("\n Top influencing parameters:")
+        print(importance.head())
+    
+        return model, metrics
+    
+    # ---------------------------------------------------------------------
+    def predict_new(self, params):
+        """
+        Predicts an attribute value for a set of
+        parameters using a previously trained model.
+
+        Parameters:
+            params (list): List of parameters, e.g.
+                [500, 50, 10, 1, 24, 1/24, 15, 4, 4, 10]
+
+        Returns:
+            float: Predicted tumor size
+        """
+        if self.trained_model is None:
+            raise RuntimeError(
+                "No trained model found. Train one with train_predictor() first."
+                )
+        if self.feature_columns is None:
+            raise RuntimeError(
+                "Feature column list not found. Did you train the model?"
+                )
+
+        if len(params) != len(self.feature_columns):
+            raise ValueError(
+                f"Parameter list must have {len(self.feature_columns)} values "
+                f"(got {len(params)}). Expected order: {self.feature_columns}"
+            )
+        df = pd.DataFrame([params], columns=self.feature_columns)
+
+        # Ensure all expected features are present (fill missing ones with 0)
+        for col in self.feature_columns:
+            if col not in df.columns:
+                df[col] = 0
+        df = df[self.feature_columns]
+
+        return self.trained_model.predict(df)[0]