spacr 0.0.20__tar.gz → 0.0.35__tar.gz

{spacr-0.0.20/spacr.egg-info → spacr-0.0.35}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: spacr
-Version: 0.0.20
+Version: 0.0.35
 Summary: Spatial phenotype analysis of crisp screens (SpaCr)
 Home-page: https://github.com/EinarOlafsson/spacr
 Author: Einar Birnir Olafsson
@@ -19,6 +19,7 @@ Requires-Dist: scikit-image>=0.22.0
 Requires-Dist: scikit-learn>=1.4.1
 Requires-Dist: seaborn>=0.13.2
 Requires-Dist: matplotlib>=3.8.3
+Requires-Dist: shap>=0.45.0
 Requires-Dist: pillow>=10.2.0
 Requires-Dist: imageio>=2.34.0
 Requires-Dist: scipy>=1.12.0
@@ -31,6 +32,8 @@ Requires-Dist: IPython>=8.18.1
 Requires-Dist: opencv-python-headless>=4.9.0.80
 Requires-Dist: umap>=0.1.1
 Requires-Dist: ttkthemes>=3.2.2
+Requires-Dist: xgboost>=2.0.3
+Requires-Dist: PyWavelets>=1.6.0
 Requires-Dist: lxml>=5.1.0
 Provides-Extra: dev
 Requires-Dist: pytest>=3.9; extra == "dev"
@@ -49,7 +52,7 @@ Requires-Dist: opencv-python; extra == "full"
 <tr>
 <td>
   
-Spatial phenotype analysis of crisp screens (SpaCr). A collection of functions for generating cellpose masks -> single object images and measurements -> annotation and classification of single object images. Spacr uses batch normalization to facilitate accurate segmentation of objects with low foreground representation.
+Spatial phenotype analysis of CRISPR-Cas9 screens (SpaCr). The spatial organization of organelles and proteins within cells constitutes a key level of functional regulation. In the context of infectious disease, the spatial relationships between host cell structures and intracellular pathogens are critical to understand host clearance mechanisms and how pathogens evade them. Spacr is a Python-based software package for generating single cell image data for deep-learning sub-cellular/cellular phenotypic classification from pooled genetic CRISPR-Cas9 screens. Spacr provides a flexible toolset to extract single cell images and measurements from high content cell painting experiments, train deep-learning models to classify cellular/ subcellular phenotypes, simulate and analyze pooled CRISPR-Cas9 imaging screens.
 
 </td>
 <td>

{spacr-0.0.20 → spacr-0.0.35}/README.md

@@ -8,7 +8,7 @@
 <tr>
 <td>
   
-Spatial phenotype analysis of crisp screens (SpaCr). A collection of functions for generating cellpose masks -> single object images and measurements -> annotation and classification of single object images. Spacr uses batch normalization to facilitate accurate segmentation of objects with low foreground representation.
+Spatial phenotype analysis of CRISPR-Cas9 screens (SpaCr). The spatial organization of organelles and proteins within cells constitutes a key level of functional regulation. In the context of infectious disease, the spatial relationships between host cell structures and intracellular pathogens are critical to understand host clearance mechanisms and how pathogens evade them. Spacr is a Python-based software package for generating single cell image data for deep-learning sub-cellular/cellular phenotypic classification from pooled genetic CRISPR-Cas9 screens. Spacr provides a flexible toolset to extract single cell images and measurements from high content cell painting experiments, train deep-learning models to classify cellular/ subcellular phenotypes, simulate and analyze pooled CRISPR-Cas9 imaging screens.
 
 </td>
 <td>
@@ -86,4 +86,4 @@ gui_classify
 Simulate spatial phenotype screens.
 ```
 gui_sim
-```
+```

{spacr-0.0.20 → spacr-0.0.35}/setup.py

@@ -15,6 +15,7 @@ dependencies = [
     'scikit-learn>=1.4.1',
     'seaborn>=0.13.2',
     'matplotlib>=3.8.3',
+    'shap>=0.45.0',
     'pillow>=10.2.0',
     'imageio>=2.34.0',
     'scipy>=1.12.0',
@@ -27,12 +28,14 @@ dependencies = [
     'opencv-python-headless>=4.9.0.80',
     'umap>=0.1.1',
     'ttkthemes>=3.2.2',
+    'xgboost>=2.0.3',
+    'PyWavelets>=1.6.0',
     'lxml>=5.1.0'
 ]
 
 setup(
     name="spacr",
-    version="0.0.20",
+    version="0.0.35",
     author="Einar Birnir Olafsson",
     author_email="olafsson@med.umich.com",
     description="Spatial phenotype analysis of crisp screens (SpaCr)",
@@ -42,12 +45,13 @@ setup(
     install_requires=dependencies,
     entry_points={
         'console_scripts': [
-            'gui_mask=spacr.gui_mask_app:gui_mask',
-            'gui_measure=spacr.gui_measure_app:gui_measure',
-            'gui_make_masks=spacr.mask_app:gui_make_masks',
-            'gui_annotation=spacr.annotate_app:gui_annotation',
-            'gui_classify=spacr.gui_classify_app:gui_classify',
-            'gui_sim=spacr.gui_sim_app:gui_sim',
+            'mask=spacr.gui_mask_app:gui_mask',
+            'measure=spacr.gui_measure_app:gui_measure',
+            'make_masks=spacr.mask_app:gui_make_masks',
+            'annotate=spacr.annotate_app:gui_annotation',
+            'classify=spacr.gui_classify_app:gui_classify',
+            'sim=spacr.gui_sim_app:gui_sim',
+            'gui=spacr.gui:gui_app',
         ],
     },
     extras_require={

spacr-0.0.35/spacr/alpha.py

@@ -0,0 +1,295 @@
+from skimage import measure, feature
+from skimage.filters import gabor
+from skimage.color import rgb2gray
+from skimage.feature.texture import greycomatrix, greycoprops, local_binary_pattern
+from skimage.util import img_as_ubyte
+import numpy as np
+import pandas as pd
+from scipy.stats import skew, kurtosis, entropy, hmean, gmean, mode
+import pywt
+
+import os
+import pandas as pd
+from PIL import Image
+from torchvision import transforms
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.data import Data, DataLoader
+from torch_geometric.nn import GCNConv, global_mean_pool
+from torch.optim import Adam
+
+# Step 1: Data Preparation
+
+# Load images
+def load_images(image_dir):
+    images = {}
+    for filename in os.listdir(image_dir):
+        if filename.endswith(".png"):
+            img = Image.open(os.path.join(image_dir, filename))
+            images[filename] = img
+    return images
+
+# Load sequencing data
+def load_sequencing_data(seq_file):
+    seq_data = pd.read_csv(seq_file)
+    return seq_data
+
+# Step 2: Data Representation
+
+# Image Representation (Using a simple CNN for feature extraction)
+class CNNFeatureExtractor(nn.Module):
+    def __init__(self):
+        super(CNNFeatureExtractor, self).__init__()
+        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
+        self.fc = nn.Linear(32 * 8 * 8, 128)  # Assuming input images are 64x64
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.max_pool2d(x, 2)
+        x = F.relu(self.conv2(x))
+        x = F.max_pool2d(x, 2)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+# Graph Representation
+def create_graph(wells, sequencing_data):
+    nodes = []
+    edges = []
+    node_features = []
+    
+    for well in wells:
+        # Add node for each well
+        nodes.append(well)
+        
+        # Get sequencing data for the well
+        seq_info = sequencing_data[sequencing_data['well'] == well]
+        
+        # Create node features based on gene knockouts and abundances
+        features = torch.tensor(seq_info['abundance'].values, dtype=torch.float)
+        node_features.append(features)
+        
+        # Define edges (for simplicity, assume fully connected graph)
+        for other_well in wells:
+            if other_well != well:
+                edges.append((wells.index(well), wells.index(other_well)))
+    
+    edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
+    x = torch.stack(node_features)
+    
+    data = Data(x=x, edge_index=edge_index)
+    return data
+
+# Step 3: Model Architecture
+
+class GraphTransformer(nn.Module):
+    def __init__(self, in_channels, hidden_channels, out_channels):
+        super(GraphTransformer, self).__init__()
+        self.conv1 = GCNConv(in_channels, hidden_channels)
+        self.conv2 = GCNConv(hidden_channels, hidden_channels)
+        self.fc = nn.Linear(hidden_channels, out_channels)
+        self.attention = nn.MultiheadAttention(hidden_channels, num_heads=8)
+
+    def forward(self, x, edge_index, batch):
+        x = F.relu(self.conv1(x, edge_index))
+        x = F.relu(self.conv2(x, edge_index))
+        
+        # Apply attention mechanism
+        x, _ = self.attention(x.unsqueeze(1), x.unsqueeze(1), x.unsqueeze(1))
+        x = x.squeeze(1)
+        
+        x = global_mean_pool(x, batch)
+        x = self.fc(x)
+        return x
+
+# Step 4: Training
+
+# Training Loop
+def train(model, data_loader, criterion, optimizer, epochs=10):
+    model.train()
+    for epoch in range(epochs):
+        for data in data_loader:
+            optimizer.zero_grad()
+            out = model(data.x, data.edge_index, data.batch)
+            loss = criterion(out, data.y)
+            loss.backward()
+            optimizer.step()
+        print(f'Epoch {epoch+1}, Loss: {loss.item()}')
+        
+def evaluate(model, data_loader):
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for data in data_loader:
+            out = model(data.x, data.edge_index, data.batch)
+            _, predicted = torch.max(out, 1)
+            total += data.y.size(0)
+            correct += (predicted == data.y).sum().item()
+    accuracy = correct / total
+    print(f'Accuracy: {accuracy * 100:.2f}%')
+
+def spacr_transformer(image_dir, seq_file, nr_grnas=1350, lr=0.001, mode='train'):
+    images = load_images(image_dir)
+    
+    sequencing_data = load_sequencing_data(seq_file)
+    wells = sequencing_data['well'].unique()
+    graph_data = create_graph(wells, sequencing_data)
+    model = GraphTransformer(in_channels=nr_grnas, hidden_channels=128, out_channels=nr_grnas)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = Adam(model.parameters(), lr=lr)
+    data_list = [graph_data]
+    loader = DataLoader(data_list, batch_size=1, shuffle=True)
+    if mode == 'train':
+        train(model, loader, criterion, optimizer)
+    elif mode == 'eval':
+        evaluate(model, loader)
+    else:
+        raise ValueError('Invalid mode. Use "train" or "eval".')
+
+def _calculate_glcm_features(intensity_image):
+    glcm = greycomatrix(img_as_ubyte(intensity_image), distances=[1, 2, 3, 4], angles=[0, np.pi/4, np.pi/2, 3*np.pi/4], symmetric=True, normed=True)
+    features = {}
+    for prop in ['contrast', 'dissimilarity', 'homogeneity', 'energy', 'correlation', 'ASM']:
+        for i, distance in enumerate([1, 2, 3, 4]):
+            for j, angle in enumerate([0, np.pi/4, np.pi/2, 3*np.pi/4]):
+                features[f'glcm_{prop}_d{distance}_a{angle}'] = greycoprops(glcm, prop)[i, j]
+    return features
+
+def _calculate_lbp_features(intensity_image, P=8, R=1):
+    lbp = local_binary_pattern(intensity_image, P, R, method='uniform')
+    lbp_hist, _ = np.histogram(lbp, density=True, bins=np.arange(0, P + 3), range=(0, P + 2))
+    return {f'lbp_{i}': val for i, val in enumerate(lbp_hist)}
+
+def _calculate_wavelet_features(intensity_image, wavelet='db1'):
+    coeffs = pywt.wavedec2(intensity_image, wavelet=wavelet, level=3)
+    features = {}
+    for i, coeff in enumerate(coeffs):
+        if isinstance(coeff, tuple):
+            for j, subband in enumerate(coeff):
+                features[f'wavelet_coeff_{i}_{j}_mean'] = np.mean(subband)
+                features[f'wavelet_coeff_{i}_{j}_std'] = np.std(subband)
+                features[f'wavelet_coeff_{i}_{j}_energy'] = np.sum(subband**2)
+        else:
+            features[f'wavelet_coeff_{i}_mean'] = np.mean(coeff)
+            features[f'wavelet_coeff_{i}_std'] = np.std(coeff)
+            features[f'wavelet_coeff_{i}_energy'] = np.sum(coeff**2)
+    return features
+
+
+from .measure import _estimate_blur, _calculate_correlation_object_level, _calculate_homogeneity, _periphery_intensity, _outside_intensity, _calculate_radial_distribution, _create_dataframe, _extended_regionprops_table, _calculate_correlation_object_level
+
+def _intensity_measurements(cell_mask, nucleus_mask, pathogen_mask, cytoplasm_mask, channel_arrays, settings, sizes=[3, 6, 12, 24], periphery=True, outside=True):
+    radial_dist = settings['radial_dist']
+    calculate_correlation = settings['calculate_correlation']
+    homogeneity = settings['homogeneity']
+    distances = settings['homogeneity_distances']
+    
+    intensity_props = ["label", "centroid_weighted", "centroid_weighted_local", "max_intensity", "mean_intensity", "min_intensity", "integrated_intensity"]
+    additional_props = ["standard_deviation_intensity", "median_intensity", "sum_intensity", "intensity_range", "mean_absolute_deviation_intensity", "skewness_intensity", "kurtosis_intensity", "variance_intensity", "mode_intensity", "energy_intensity", "entropy_intensity", "harmonic_mean_intensity", "geometric_mean_intensity"]
+    col_lables = ['region_label', 'mean', '5_percentile', '10_percentile', '25_percentile', '50_percentile', '75_percentile', '85_percentile', '95_percentile']
+    cell_dfs, nucleus_dfs, pathogen_dfs, cytoplasm_dfs = [], [], [], []
+    ls = ['cell','nucleus','pathogen','cytoplasm']
+    labels = [cell_mask, nucleus_mask, pathogen_mask, cytoplasm_mask]
+    dfs = [cell_dfs, nucleus_dfs, pathogen_dfs, cytoplasm_dfs]
+    
+    for i in range(0,channel_arrays.shape[-1]):
+        channel = channel_arrays[:, :, i]
+        for j, (label, df) in enumerate(zip(labels, dfs)):
+            
+            if np.max(label) == 0:
+                empty_df = pd.DataFrame()
+                df.append(empty_df)
+                continue
+                
+            mask_intensity_df = _extended_regionprops_table(label, channel, intensity_props)
+            
+            # Additional intensity properties
+            region_props = measure.regionprops_table(label, intensity_image=channel, properties=['label'])
+            intensity_values = [channel[region.coords[:, 0], region.coords[:, 1]] for region in measure.regionprops(label)]
+            additional_data = {prop: [] for prop in additional_props}
+            
+            for values in intensity_values:
+                if len(values) == 0:
+                    continue
+                additional_data["standard_deviation_intensity"].append(np.std(values))
+                additional_data["median_intensity"].append(np.median(values))
+                additional_data["sum_intensity"].append(np.sum(values))
+                additional_data["intensity_range"].append(np.max(values) - np.min(values))
+                additional_data["mean_absolute_deviation_intensity"].append(np.mean(np.abs(values - np.mean(values))))
+                additional_data["skewness_intensity"].append(skew(values))
+                additional_data["kurtosis_intensity"].append(kurtosis(values))
+                additional_data["variance_intensity"].append(np.var(values))
+                additional_data["mode_intensity"].append(mode(values)[0][0])
+                additional_data["energy_intensity"].append(np.sum(values**2))
+                additional_data["entropy_intensity"].append(entropy(values))
+                additional_data["harmonic_mean_intensity"].append(hmean(values))
+                additional_data["geometric_mean_intensity"].append(gmean(values))
+            
+            for prop in additional_props:
+                region_props[prop] = additional_data[prop]
+            
+            additional_df = pd.DataFrame(region_props)
+            mask_intensity_df = pd.concat([mask_intensity_df.reset_index(drop=True), additional_df.reset_index(drop=True)], axis=1)
+            
+            if homogeneity:
+                homogeneity_df = _calculate_homogeneity(label, channel, distances)
+                mask_intensity_df = pd.concat([mask_intensity_df.reset_index(drop=True), homogeneity_df], axis=1)
+
+            if periphery:
+                if ls[j] == 'nucleus' or ls[j] == 'pathogen':
+                    periphery_intensity_stats = _periphery_intensity(label, channel)
+                    mask_intensity_df = pd.concat([mask_intensity_df, pd.DataFrame(periphery_intensity_stats, columns=[f'periphery_{stat}' for stat in col_lables])],axis=1)
+
+            if outside:
+                if ls[j] == 'nucleus' or ls[j] == 'pathogen':
+                    outside_intensity_stats = _outside_intensity(label, channel)
+                    mask_intensity_df = pd.concat([mask_intensity_df, pd.DataFrame(outside_intensity_stats, columns=[f'outside_{stat}' for stat in col_lables])], axis=1)
+
+            # Adding GLCM features
+            glcm_features = _calculate_glcm_features(channel)
+            for k, v in glcm_features.items():
+                mask_intensity_df[f'{ls[j]}_channel_{i}_{k}'] = v
+
+            # Adding LBP features
+            lbp_features = _calculate_lbp_features(channel)
+            for k, v in lbp_features.items():
+                mask_intensity_df[f'{ls[j]}_channel_{i}_{k}'] = v
+            
+            # Adding Wavelet features
+            wavelet_features = _calculate_wavelet_features(channel)
+            for k, v in wavelet_features.items():
+                mask_intensity_df[f'{ls[j]}_channel_{i}_{k}'] = v
+
+            blur_col = [_estimate_blur(channel[label == region_label]) for region_label in mask_intensity_df['label']]
+            mask_intensity_df[f'{ls[j]}_channel_{i}_blur'] = blur_col
+
+            mask_intensity_df.columns = [f'{ls[j]}_channel_{i}_{col}' if col != 'label' else col for col in mask_intensity_df.columns]
+            df.append(mask_intensity_df)
+    
+    if radial_dist:
+        if np.max(nucleus_mask) != 0:
+            nucleus_radial_distributions = _calculate_radial_distribution(cell_mask, nucleus_mask, channel_arrays, num_bins=6)
+            nucleus_df = _create_dataframe(nucleus_radial_distributions, 'nucleus')
+            dfs[1].append(nucleus_df)
+            
+        if np.max(nucleus_mask) != 0:
+            pathogen_radial_distributions = _calculate_radial_distribution(cell_mask, pathogen_mask, channel_arrays, num_bins=6)
+            pathogen_df = _create_dataframe(pathogen_radial_distributions, 'pathogen')
+            dfs[2].append(pathogen_df)
+        
+    if calculate_correlation:
+        if channel_arrays.shape[-1] >= 2:
+            for i in range(channel_arrays.shape[-1]):
+                for j in range(i+1, channel_arrays.shape[-1]):
+                    chan_i = channel_arrays[:, :, i]
+                    chan_j = channel_arrays[:, :, j]
+                    for m, mask in enumerate(labels):
+                        coloc_df = _calculate_correlation_object_level(chan_i, chan_j, mask, settings)
+                        coloc_df.columns = [f'{ls[m]}_channel_{i}_channel_{j}_{col}' for col in coloc_df.columns]
+                        dfs[m].append(coloc_df)
+    
+    return pd.concat(cell_dfs, axis=1), pd.concat(nucleus_dfs, axis=1), pd.concat(pathogen_dfs, axis=1), pd.concat(cytoplasm_dfs, axis=1)
+

{spacr-0.0.20 → spacr-0.0.35}/spacr/annotate_app.py

@@ -13,7 +13,7 @@ from ttkthemes import ThemedTk
 
 from .logger import log_function_call
 
-from .gui_utils import ScrollableFrame, set_default_font, set_dark_style, create_dark_mode
+from .gui_utils import ScrollableFrame, set_default_font, set_dark_style, create_dark_mode, style_text_boxes, create_menu_bar
 
 class ImageApp:
     """
@@ -38,7 +38,7 @@ class ImageApp:
     - db_update_thread (threading.Thread): A thread for updating the database.
     """
 
-    def _init_(self, root, db_path, image_type=None, channels=None, grid_rows=None, grid_cols=None, image_size=(200, 200), annotation_column='annotate'):
+    def __init__(self, root, db_path, image_type=None, channels=None, grid_rows=None, grid_cols=None, image_size=(200, 200), annotation_column='annotate'):
         """
         Initializes an instance of the ImageApp class.
 
@@ -383,7 +383,7 @@ def annotate(db, image_type=None, channels=None, geom="1000x1100", img_size=(200
     root = tk.Tk()
     root.geometry(geom)
     app = ImageApp(root, db, image_type=image_type, channels=channels, image_size=img_size, grid_rows=rows, grid_cols=columns, annotation_column=annotation_column)
-    
+    #app = ImageApp()
     next_button = tk.Button(root, text="Next", command=app.next_page)
     next_button.grid(row=app.grid_rows, column=app.grid_cols - 1)
     back_button = tk.Button(root, text="Back", command=app.previous_page)
@@ -425,7 +425,8 @@ def initiate_annotation_app_root(width, height):
     root = ThemedTk(theme=theme)
     style = ttk.Style(root)
     set_dark_style(style)
-    set_default_font(root, font_name="Arial", size=10)
+    style_text_boxes(style)
+    set_default_font(root, font_name="Arial", size=8)
     root.geometry(f"{width}x{height}")
     root.title("Annotation App")
 
@@ -473,6 +474,7 @@ def initiate_annotation_app_root(width, height):
         new_root = tk.Tk()
         new_root.geometry(f"{width}x{height}")
         new_root.title("Mask Application")
+        
 
         # Start the annotation application in the new root window
         app_instance = annotate(db, image_type, channels, annotation_column, geom, img_size, rows, columns)
@@ -482,7 +484,7 @@ def initiate_annotation_app_root(width, height):
     create_dark_mode(root, style, console_output=None)
 
     run_button = ttk.Button(scrollable_frame.scrollable_frame, text="Run", command=run_app)
-    run_button.grid(row=row, column=0, columnspan=2, pady=10)
+    run_button.grid(row=row, column=0, columnspan=2, pady=10, padx=10)
 
     return root
 

spacr-0.0.35/spacr/chris.py

@@ -0,0 +1,50 @@
+import pandas as pd 
+import numpy as np
+from .core import _permutation_importance, _shap_analysis
+
+def join_measurments_and_annotation(src, tables = ['cell', 'nucleus', 'pathogen','cytoplasm']):
+    
+    from .io import _read_and_merge_data, _read_db
+    
+    db_loc = [src+'/measurements/measurements.db']
+    loc = src+'/measurements/measurements.db'
+    df, _ = _read_and_merge_data(db_loc, 
+                                 tables, 
+                                 verbose=True, 
+                                 include_multinucleated=True, 
+                                 include_multiinfected=True, 
+                                 include_noninfected=True)
+    
+    paths_df = _read_db(loc, tables=['png_list'])
+
+    merged_df = pd.merge(df, paths_df[0], on='prcfo', how='left')
+
+    return merged_df
+
+def plate_heatmap(src, model_type='xgboost', variable='predictions', grouping='mean', min_max='allq', cmap='viridis', channel_of_interest=3, min_count=25, n_estimators=100, col_to_compare='col', pos='c1', neg='c2', exclude=None, n_repeats=10, clean=True, nr_to_plot=20, verbose=False, n_jobs=-1):
+    from .io import _read_and_merge_data
+    from .plot import _plot_plates
+
+    db_loc = [src+'/measurements/measurements.db']
+    tables = ['cell', 'nucleus', 'pathogen','cytoplasm']
+    include_multinucleated, include_multiinfected, include_noninfected = True, 2.0, True
+    
+    df = join_measurments_and_annotation(src, tables=['cell', 'nucleus', 'pathogen', 'cytoplasm'])
+        
+    if not channel_of_interest is None:
+        df['recruitment'] = df[f'pathogen_channel_{channel_of_interest}_mean_intensity']/df[f'cytoplasm_channel_{channel_of_interest}_mean_intensity']
+        feature_string = f'channel_{channel_of_interest}'
+    else:
+        feature_string = None
+    
+    output = _permutation_importance(df, feature_string, col_to_compare, pos, neg, exclude, n_repeats, clean, nr_to_plot, n_estimators=n_estimators, random_state=42, model_type=model_type, n_jobs=n_jobs)
+    
+    _shap_analysis(output[3], output[4], output[5])
+
+    features = output[0].select_dtypes(include=[np.number]).columns.tolist()
+
+    if not variable in features:
+        raise ValueError(f"Variable {variable} not found in the dataframe. Please choose one of the following: {features}")
+    
+    plate_heatmap = _plot_plates(output[0], variable, grouping, min_max, cmap, min_count)
+    return [output, plate_heatmap]