Add checks for removed foci

path_analysis/analyse.py

@@ -36,22 +36,25 @@ def get_paths_from_traces_file(traces_file):
         path_lengths.append(float(length))
     return all_paths, path_lengths
 
-def calculate_path_length(point_list, voxel_size=(1,1,1)):
-    # Simple calculation
-    l = 0
-    s = np.array(voxel_size)
-    for i in range(len(point_list)-1):
-        l += la.norm(s * (np.array(point_list[i+1]) - np.array(point_list[i])))
-    return l
 
 
 def calculate_path_length_partials(point_list, voxel_size=(1,1,1)):
+    """
+    Calculate the partial path length of a series of points.
+    
+    Args:
+    point_list (list of tuple): List of points, each represented as a tuple of coordinates (x, y, z).
+    voxel_size (tuple, optional): Size of the voxel in each dimension (x, y, z). Defaults to (1, 1, 1).
+    
+    Returns:
+    numpy.ndarray: Array of cumulative partial path lengths at each point.
+    """
     # Simple calculation
-    l = [0.0]
+    section_lengths = [0.0]
     s = np.array(voxel_size)
     for i in range(len(point_list)-1):
-        l.append(la.norm(s * (np.array(point_list[i+1]) - np.array(point_list[i]))))
-    return np.cumsum(l)
+        section_lengths.append(la.norm(s * (np.array(point_list[i+1]) - np.array(point_list[i]))))
+    return np.cumsum(section_lengths)
 
 
 def visualise_ordering(points_list, dim, wr=5, wc=5):
@@ -260,6 +263,7 @@ def measure_chrom2(path, hei10, config):
     measurements = measure_all_with_sphere(path, hei10, op='mean', R=sphere_xy_radius, z_scale_ratio=scale_ratio)
     measurements_max = measure_all_with_sphere(path, hei10, op='max', R=sphere_xy_radius, z_scale_ratio=scale_ratio)
 
+    
     return vis, measurements, measurements_max
 
 def extract_peaks(cell_id, all_paths, path_lengths, measured_traces, config):
@@ -286,7 +290,7 @@ def extract_peaks(cell_id, all_paths, path_lengths, measured_traces, config):
     n_paths = len(all_paths)
     
     data = []
-    foci_absolute_intensity, foci_position, foci_position_index, trace_median_intensities, trace_thresholds = analyse_traces(all_paths, path_lengths, measured_traces, config)
+    foci_absolute_intensity, foci_position, foci_position_index, dominated_foci_data, trace_median_intensities, trace_thresholds = analyse_traces(all_paths, path_lengths, measured_traces, config)
 
     foci_intensities = []
     for path_foci_abs_int, tmi in zip(foci_absolute_intensity, trace_median_intensities):
@@ -299,6 +303,8 @@ def extract_peaks(cell_id, all_paths, path_lengths, measured_traces, config):
 
         pl = calculate_path_length_partials(all_paths[i], (config['xy_res'], config['xy_res'], config['z_res']))
         
+        print(i, len(all_paths[i]), len(pl))
+        
         path_data = { 'Cell_ID':cell_id,
                       'Trace': i+1,
                       'SNT_trace_length(um)': path_lengths[i],
@@ -315,7 +321,7 @@ def extract_peaks(cell_id, all_paths, path_lengths, measured_traces, config):
             path_data[f'Foci_{j+1}_relative_intensity'] = (v - trace_median_intensities[i])/mean_intensity
         data.append(path_data)
         trace_positions.append(pl)
-    return pd.DataFrame(data), foci_absolute_intensity, foci_position_index, trace_thresholds, trace_positions
+    return pd.DataFrame(data), foci_absolute_intensity, foci_position_index, dominated_foci_data, trace_thresholds, trace_positions
 
 
 def analyse_paths(cell_id,
@@ -351,25 +357,32 @@ def analyse_paths(cell_id,
         vis, m, _ = measure_chrom2(p,foci_stack.transpose(2,1,0), config)
         all_trace_vis.append(vis)
         all_m.append(m)
+        
 
-
-    extracted_peaks, foci_absolute_intensity, foci_pos_index, trace_thresholds, trace_positions = extract_peaks(cell_id, all_paths, path_lengths, all_m, config)
+    extracted_peaks, foci_absolute_intensity, foci_pos_index, dominated_foci_data, trace_thresholds, trace_positions = extract_peaks(cell_id, all_paths, path_lengths, all_m, config)
 
     
     n_cols = 2
     n_rows = (len(all_paths)+n_cols-1)//n_cols
-    fig, ax = plt.subplots(n_rows,n_cols)
+    fig, ax = plt.subplots(n_rows,n_cols, figsize=(5*n_cols, 3*n_rows))
     ax = ax.flatten()
 
     for i, m in enumerate(all_m):
         ax[i].set_title(f'Trace {i+1}')
         ax[i].plot(trace_positions[i], m)
-        print(foci_pos_index)
         if len(foci_pos_index[i]):
             ax[i].plot(trace_positions[i][foci_pos_index[i]], np.array(m)[foci_pos_index[i]], 'rx')
-            ax[i].set_xlabel('Distance from start (um)')
-            ax[i].set_ylabel('Intensity')
+
+        if len(dominated_foci_data[i]):
+
+            dominated_foci_pos_index = [u.idx for u in dominated_foci_data[i]]
+            ax[i].plot(trace_positions[i][dominated_foci_pos_index], np.array(m)[dominated_foci_pos_index], color=(0.5,0.5,0.5), marker='o', linestyle='None')
+
+            
+        if trace_thresholds[i] is not None:
             ax[i].axhline(trace_thresholds[i], c='r', ls=':')
+        ax[i].set_xlabel('Distance from start (um)')
+        ax[i].set_ylabel('Intensity')
     for i in range(len(all_m), n_cols*n_rows):
         ax[i].axis('off')
 

path_analysis/data_preprocess.py

@@ -9,45 +9,54 @@ from math import ceil
 
 
 
-def thin_points(point_list, dmin=10, voxel_size=(1,1,1)):
+def thin_peaks(peak_list, dmin=10, voxel_size=(1,1,1), return_larger_peaks=False):
     """
-    Remove points within a specified distance of each other, retaining the point with the highest intensity.
+    Remove peaks within a specified distance of each other, retaining the peak with the highest intensity.
 
     Args:
-    - point_list (list of tuples): Each tuple contains:
-        - x (list of float): 3D coordinates of the point.
-        - intensity (float): The intensity value of the point.
-        - idx (int): A unique identifier or index for the point.
-    - dmin (float, optional): Minimum distance between points. Points closer than this threshold will be thinned. Defaults to 10.
+    - peak_list (list of PeakData): Each element contains:
+        - pos (list of float): 3D coordinates of the peak.
+        - intensity (float): The intensity value of the peak.
+        - key (tuple): A unique identifier or index for the peak (#trace, #peak)
+    - dmin (float, optional): Minimum distance between peaks. peaks closer than this threshold will be thinned. Defaults to 10.
+    - return_larger_peaks (bool, optional): Indicate larger peak for each thinned peak
 
     Returns:
-    - list of int: A list containing indices of the removed points.
+    - list of tuples: A list containing keys of the removed peaks.
+    if return_larger_peaks
+    - list of tuples: A list containing the keys of the larger peak causing the peak to be removed
 
     Notes:
-    - The function uses the L2 norm (Euclidean distance) to compute the distance between points.
-    - When two points are within `dmin` distance, the point with the lower intensity is removed.
+    - The function uses the L2 norm (Euclidean distance) to compute the distance between peaks.
+    - When two peaks are within `dmin` distance, the peak with the lower intensity is removed.
     """
-    removed_points = []
-    for i in range(len(point_list)):
-        if point_list[i][2] in removed_points:
+    removed_peaks = []
+    removed_larger_peaks = []
+    for i in range(len(peak_list)):
+        if peak_list[i].key in removed_peaks:
             continue
-        for j in range(len(point_list)):
+        for j in range(len(peak_list)):
             if i==j:
                 continue
-            if point_list[j][2] in removed_points:
+            if peak_list[j].key in removed_peaks:
                 continue
-            d = (np.array(point_list[i][0]) - np.array(point_list[j][0]))*np.array(voxel_size)
+            d = (np.array(peak_list[i].pos) - np.array(peak_list[j].pos))*np.array(voxel_size)
             d = la.norm(d)
             if d<dmin:
-                hi = point_list[i][1]
-                hj = point_list[j][1]
+                hi = peak_list[i].intensity
+                hj = peak_list[j].intensity
                 if hi<hj:
-                    removed_points.append(point_list[i][2])
+                    removed_peaks.append(peak_list[i].key)
+                    removed_larger_peaks.append(peak_list[j].key)
                     break
                 else:
-                    removed_points.append(point_list[j][2])
-                
-    return removed_points
+                    removed_peaks.append(peak_list[j].key)
+                    removed_larger_peaks.append(peak_list[i].key)
+
+    if return_larger_peaks:
+        return removed_peaks, removed_larger_peaks
+    else:
+        return removed_peaks
 
 
 @dataclass
@@ -59,6 +68,17 @@ class CellData(object):
     """
     pathdata_list: list
 
+@dataclass
+class RemovedPeakData(object):
+    """Represents data related to a removed peak
+
+    Attributes:
+        idx (int): Index of peak along path
+        dominating_peak (tuple): (path_idx, position along path) for dominating peak
+    """
+    idx: int 
+    dominating_peak: tuple 
+
 @dataclass
 class PathData(object):
     """Represents data related to a specific path in the cell.
@@ -67,16 +87,23 @@ class PathData(object):
     the defining points, the fluorescence values, and the path length of a specific path.
 
     Attributes: peaks (list): List of peaks in the path (indicies of positions in points, o_hei10).
+        removed_peaks (list): List of peaks in the path which have been removed because of a nearby larger peak
         points (list): List of points defining the path.
         o_hei10 (list): List of (unnormalized) fluorescence intensity values along the path
         SC_length (float): Length of the path.
 
     """
     peaks: list
+    removed_peaks: list
     points: list
     o_hei10: list
     SC_length: float
 
+@dataclass
+class PeakData(object):
+    pos: tuple
+    intensity: float
+    key: tuple
 
 
 def find_peaks2(v, distance=5,  prominence=0.5):
@@ -136,7 +163,7 @@ def process_cell_traces(all_paths, path_lengths, measured_trace_fluorescence):
     
     cell_peaks = []
 
-    for points, path_length, o_hei10 in zip(all_paths, path_lengths, measured_trace_fluorescence):
+    for points, o_hei10 in zip(all_paths, measured_trace_fluorescence):
                      
         # For peak determination normalize each trace to have mean zero and s.d. 1
         hei10_normalized = (o_hei10 - np.mean(o_hei10))/np.std(o_hei10)
@@ -158,24 +185,34 @@ def process_cell_traces(all_paths, path_lengths, measured_trace_fluorescence):
     to_thin = []
     for k in range(len(cell_peaks)):
         for u in range(len(cell_peaks[k][0])):
-            to_thin.append((cell_peaks[k][1][u], cell_peaks[k][2][u], (k, u)))
+            to_thin.append(PeakData(pos=cell_peaks[k][1][u], intensity=cell_peaks[k][2][u], key=(k, u)))
     
     # Exclude any peak with a nearby brighter peak (on any SC)
-    removed_points = thin_points(to_thin)
+    removed_peaks, removed_larger_peaks = thin_peaks(to_thin, return_larger_peaks=True)
 
-    
     # Clean up and remove these peaks
     new_cell_peaks = []
-    for k in range(len(cell_peaks)):
-        cc = []
-        pp = cell_peaks[k][0]
-        for u in range(len(pp)):
-            if (k,u) not in removed_points:
-                cc.append(pp[u])
-        new_cell_peaks.append(cc)
+    removed_cell_peaks = []
+    removed_cell_peaks_larger = []
+    for path_idx in range(len(cell_peaks)):
+        path_retained_peaks = []
+        path_removed_peaks = []
+        path_peaks = cell_peaks[path_idx][0]
+
+        for peak_idx in range(len(path_peaks)):
+            if (path_idx, peak_idx) not in removed_peaks:
+                path_retained_peaks.append(path_peaks[peak_idx])
+            else:
+                # What's the larger point?
+                idx = removed_peaks.index((path_idx, peak_idx))
+                larger_path, larger_idx = removed_larger_peaks[idx]        
+                path_removed_peaks.append(RemovedPeakData(idx=path_peaks[peak_idx], dominating_peak=(larger_path, cell_peaks[larger_path][0][larger_idx])))
+                ###
+                 
+        new_cell_peaks.append(path_retained_peaks)
+        removed_cell_peaks.append(path_removed_peaks)
         
     cell_peaks = new_cell_peaks
-
     pd_list = []
     
     # Save peak positions, absolute HEI10 intensities, and length for each SC
@@ -184,8 +221,9 @@ def process_cell_traces(all_paths, path_lengths, measured_trace_fluorescence):
         points, o_hei10 = all_paths[k], measured_trace_fluorescence[k]
 
         peaks = cell_peaks[k]
+        removed_peaks = removed_cell_peaks[k]
         
-        pd = PathData(peaks=peaks, points=points, o_hei10=o_hei10, SC_length=path_lengths[k])
+        pd = PathData(peaks=peaks, removed_peaks=removed_peaks, points=points, o_hei10=o_hei10, SC_length=path_lengths[k])
         pd_list.append(pd)
 
     cd = CellData(pathdata_list=pd_list)
@@ -196,9 +234,9 @@ def process_cell_traces(all_paths, path_lengths, measured_trace_fluorescence):
 alpha_max = 0.4
 
 
-# Criterion used for identifying peak as a CO - normalized (with mean and s.d.)
+# Criterion used for identifying peak as a focus - normalized (with mean and s.d.)
 # hei10 levels being above 0.4 time maximum peak level
-def pc(pos, v, alpha=alpha_max):
+def focus_criterion(pos, v, alpha=alpha_max):
     """
     Identify and return positions where values in the array `v` exceed a certain threshold.
 
@@ -212,8 +250,11 @@ def pc(pos, v, alpha=alpha_max):
     Returns:
     - numpy.ndarray: Array of positions where corresponding values in `v` exceed the threshold.
     """
-    idx = (v>=alpha*np.max(v))
-    return np.array(pos[idx])
+    if len(v):
+        idx = (v>=alpha*np.max(v))
+        return np.array(pos[idx])
+    else:
+        return np.array([], dtype=np.int32)
 
 def analyse_celldata(cell_data, config):
     """
@@ -226,14 +267,18 @@ def analyse_celldata(cell_data, config):
                              'threshold_type' (str) = 'per-trace', 'per-foci'
 
     Returns:
-        tuple: A tuple containing three lists:
+        tuple: A tuple containing:
             - foci_rel_intensity (list): List of relative intensities for the detected foci.
             - foci_pos (list): List of absolute positions of the detected foci.
             - foci_pos_index (list): List of indices of the detected foci.
+            - dominated_foci_data (list): List of RemovedPeakData indicating positions of removed peaks and the index of the larger peak
+            - trace_median_intensities (list): Per-trace median intensity
+            - trace_thresholds (list): Per-trace absolute threshold for calling peaks as foci
     """
     foci_abs_intensity = []
     foci_pos = []
     foci_pos_index = []
+    dominated_foci_data = []
     trace_median_intensities = []
     trace_thresholds = []
     
@@ -254,15 +299,39 @@ def analyse_celldata(cell_data, config):
             h = np.array(path_data.o_hei10)
             h = h - np.mean(h)
             h = h/np.std(h)
-            # Extract peaks according to criterion
-            sig_peak_idx = pc(peaks, h[peaks], peak_threshold) 
-            trace_thresholds.append((1-peak_threshold)*np.mean(path_data.o_hei10) + peak_threshold*np.max(np.array(path_data.o_hei10)[peaks]))
+            # Extract foci according to criterion
+            foci_idx = focus_criterion(peaks, h[peaks], peak_threshold) 
+            print('peaks', peaks, h[peaks], foci_idx, np.mean(path_data.o_hei10))
+            
+            #
+            removed_peaks = path_data.removed_peaks
+            removed_peaks_idx = np.array([u.idx for u in removed_peaks], dtype=np.int32)
+
             
-            pos_abs = (sig_peak_idx/len(path_data.points))*path_data.SC_length
+            if len(peaks):
+                trace_thresholds.append((1-peak_threshold)*np.mean(path_data.o_hei10) + peak_threshold*np.max(np.array(path_data.o_hei10)[peaks]))
+            else:
+                trace_thresholds.append(None)
+
+            if len(removed_peaks):
+                if len(peaks):
+                    threshold = (1-peak_threshold)*np.mean(path_data.o_hei10) + peak_threshold*np.max(np.array(path_data.o_hei10)[peaks])
+                else:
+                    threshold = float('-inf')
+
+                
+                removed_peak_heights = np.array(path_data.o_hei10)[removed_peaks_idx]
+                dominated_foci_idx = np.where(removed_peak_heights>threshold)[0]
+                
+                dominated_foci_data.append([removed_peaks[i] for i in dominated_foci_idx])
+            else:
+                dominated_foci_data.append([])
+                
+            pos_abs = (foci_idx/len(path_data.points))*path_data.SC_length
             foci_pos.append(pos_abs)
-            foci_abs_intensity.append(np.array(path_data.o_hei10)[sig_peak_idx])
+            foci_abs_intensity.append(np.array(path_data.o_hei10)[foci_idx])
             
-            foci_pos_index.append(sig_peak_idx)
+            foci_pos_index.append(foci_idx)
             trace_median_intensities.append(np.median(path_data.o_hei10))
             
     elif threshold_type == 'per-cell':
@@ -286,22 +355,34 @@ def analyse_celldata(cell_data, config):
             h = np.array(path_data.o_hei10)
             h = h - np.mean(h)
 
-            sig_peak_idx = peaks[h[peaks]>peak_threshold*max_cell_intensity]
+            foci_idx = peaks[h[peaks]>peak_threshold*max_cell_intensity]
+
+            removed_peaks = path_data.removed_peaks
+            removed_peaks_idx = np.array([u.idx for u in removed_peaks], dtype=np.int32)
 
             trace_thresholds.append(np.mean(path_data.o_hei10) + peak_threshold*max_cell_intensity)
 
+            if len(removed_peaks):
+                threshold = np.mean(path_data.o_hei10) + peak_threshold*max_cell_intensity
+
+                removed_peak_heights = np.array(path_data.o_hei10)[removed_peaks_idx]
+                dominated_foci_idx = np.where(removed_peak_heights>threshold)[0]
+                
+                dominated_foci_data.append([removed_peaks[i] for i in dominated_foci_idx])
+            else:
+                dominated_foci_data.append([])
 
-            pos_abs = (sig_peak_idx/len(path_data.points))*path_data.SC_length
+            pos_abs = (foci_idx/len(path_data.points))*path_data.SC_length
             foci_pos.append(pos_abs)
-            foci_abs_intensity.append(np.array(path_data.o_hei10)[sig_peak_idx])
+            foci_abs_intensity.append(np.array(path_data.o_hei10)[foci_idx])
             
-            foci_pos_index.append(sig_peak_idx)
+            foci_pos_index.append(foci_idx)
             trace_median_intensities.append(np.median(path_data.o_hei10))            
             
     else:
         raise NotImplementedError
     
-    return foci_abs_intensity, foci_pos, foci_pos_index, trace_median_intensities, trace_thresholds
+    return foci_abs_intensity, foci_pos, foci_pos_index, dominated_foci_data, trace_median_intensities, trace_thresholds
 
 def analyse_traces(all_paths, path_lengths, measured_trace_fluorescence, config):
     

tests/test_analyse.py

@@ -1,9 +1,68 @@
 
+import pytest
 from path_analysis.analyse import *
 import numpy as np
 from math import pi
 import xml.etree.ElementTree as ET
+from PIL import ImageChops
 
+def test_draw_paths_no_error():
+    all_paths = [[[0, 0], [1, 1]], [[2, 2], [3, 3]]]
+    foci_stack = np.zeros((5, 5, 5))
+    foci_stack[0,0,0] = 1.0
+    foci_index = [[0], [1]]
+    r = 3
+
+    try:
+        im = draw_paths(all_paths, foci_stack, foci_index, r)
+    except Exception as e:
+        pytest.fail(f"draw_paths raised an exception: {e}")
+
+def test_draw_paths_image_size():
+    all_paths = [[[0, 0], [1, 1]], [[2, 2], [3, 3]]]
+    foci_stack = np.zeros((5, 5, 5))
+    foci_stack[0,0,0] = 1.0
+
+    foci_index = [[0], [1]]
+    r = 3
+
+    im = draw_paths(all_paths, foci_stack, foci_index, r)
+    assert im.size == (5, 5), f"Expected image size (5, 5), got {im.size}"
+
+def test_draw_paths_image_modified():
+    all_paths = [[[0, 0], [1, 1]], [[2, 2], [3, 3]]]
+    foci_stack = np.zeros((5, 5, 5))
+    foci_stack[0,0,0] = 1.0
+    foci_index = [[0], [1]]
+    r = 3
+
+    im = draw_paths(all_paths, foci_stack, foci_index, r)
+    blank_image = Image.new("RGB", (5, 5), "black")
+
+    # Check if the image is not entirely black (i.e., has been modified)
+    diff = ImageChops.difference(im, blank_image)
+    assert diff.getbbox() is not None, "The image has not been modified"
+
+
+
+def test_calculate_path_length_partials_default_voxel():
+    point_list = [(0, 0, 0), (1, 0, 0), (1, 1, 1)]
+    expected_result = np.array([0.0, 1.0, 1.0+np.sqrt(2)])
+    result = calculate_path_length_partials(point_list)
+    np.testing.assert_allclose(result, expected_result, atol=1e-5)
+
+def test_calculate_path_length_partials_custom_voxel():
+    point_list = [(0, 0, 0), (1, 0, 0), (1, 1, 0)]
+    voxel_size = (1, 2, 1)
+    expected_result = np.array([0.0, 1.0, 3.0])
+    result = calculate_path_length_partials(point_list, voxel_size=voxel_size)
+    np.testing.assert_allclose(result, expected_result, atol=1e-5)
+
+def test_calculate_path_length_partials_single_point():
+    point_list = [(0, 0, 0)]
+    expected_result = np.array([0.0])
+    result = calculate_path_length_partials(point_list)
+    np.testing.assert_allclose(result, expected_result, atol=1e-5)
 
 def test_get_paths_from_traces_file():
     # Mock the XML traces file content
@@ -221,3 +280,34 @@ def test_make_sphere_equal():
     assert abs(np.sum(sphere)-4/3*pi*R**3)<10, f"Expected approximate volume to be correct"
     assert (sphere[R,R,0] == 1), f"Expected centre point on top plane to be within sphere"
     assert (sphere[R+1,R,0] == 0), f"Expected point next to centre on top plane to be outside sphere"
+
+import pandas as pd
+
+
+# 1. Test basic functionality
+def test_extract_peaks_basic():
+    cell_id = 1
+    all_paths = [[[0, 0], [1, 1]]]
+    path_lengths = [1.41]  # length of the above path
+    measured_traces = [[100, 200]]  # fluorescence along the path
+    config = {'peak_threshold': 0.4, 'sphere_radius': 2, 'xy_res': 1, 'z_res': 1, 'use_corrected_positions': True}
+    
+    df, foci_abs_int, foci_pos_idx, _, _, _ = extract_peaks(cell_id, all_paths, path_lengths, measured_traces, config)
+    
+    # Now add your assertions to validate the result
+    assert len(df) == 1, "Expected one row in DataFrame"
+    assert df['Cell_ID'].iloc[0] == cell_id, "Unexpected cell_id"
+    # Add more assertions here based on expected values
+
+# 2. Test multiple paths
+def test_extract_peaks_multiple_paths():
+    cell_id = 1
+    all_paths = [[[0, 0], [1, 1]], [[1, 1], [2, 2]]]
+    path_lengths = [1.41, 1.41]
+    measured_traces = [[100, 200], [100, 150]]
+    config = {'peak_threshold': 0.4, 'sphere_radius': 2, 'xy_res': 1, 'z_res': 1, 'use_corrected_positions': True}
+
+    df, _, _, _, _, _ = extract_peaks(cell_id, all_paths, path_lengths, measured_traces, config)
+    
+    assert len(df) == 2, "Expected two rows in DataFrame"
+    # Add more assertions here

tests/test_preprocess.py

@@ -2,18 +2,19 @@ from path_analysis.data_preprocess import *
 import numpy as np
 import pytest
 
+
 def test_thin_points():
     # Define a sample point list
     points = [
-        ([0, 0, 0], 10, 0),
-        ([1, 1, 1], 8, 1),
-        ([10, 10, 10], 12, 2),
-        ([10.5, 10.5, 10.5], 5, 3),
-        ([20, 20, 20], 15, 4)
+        PeakData([0, 0, 0], 10, 0),
+        PeakData([1, 1, 1], 8, 1),
+        PeakData([10, 10, 10], 12, 2),
+        PeakData([10.5, 10.5, 10.5], 5, 3),
+        PeakData([20, 20, 20], 15, 4)
     ]
     
     # Call the thin_points function with dmin=5 (for example)
-    removed_indices = thin_points(points, dmin=5)
+    removed_indices = thin_peaks(points, dmin=5)
     
     # Check results
     # Point at index 1 ([1, 1, 1]) should be removed since it's within 5 units distance of point at index 0 and has lower intensity.
@@ -22,12 +23,12 @@ def test_thin_points():
 
     # Another simple test to check if function does nothing when points are far apart
     far_points = [
-        ([0, 0, 0], 10, 0),
-        ([100, 100, 100], 12, 1),
-        ([200, 200, 200], 15, 2)
+        PeakData([0, 0, 0], 10, 0),
+        PeakData([100, 100, 100], 12, 1),
+        PeakData([200, 200, 200], 15, 2)
     ]
     
-    removed_indices_far = thin_points(far_points, dmin=5)
+    removed_indices_far = thin_peaks(far_points, dmin=5)
     assert len(removed_indices_far) == 0  # Expect no points to be removed
 
 
@@ -72,29 +73,32 @@ def test_find_peaks2():
     assert peaks == [1]  # Only the peak at position 1 meets the prominence threshold
 
 
-def test_pc():
+def test_focus_criterion():
     pos = np.array([0, 1, 2, 3, 4, 6])
     values = np.array([0.1, 0.5, 0.2, 0.8, 0.3, 0.9])
 
     # Basic test
-    assert np.array_equal(pc(pos, values), np.array([1, 3, 6]))  # only values 0.8 and 0.9 exceed 0.4 times the max (which is 0.9)
+    assert np.array_equal(focus_criterion(pos, values), np.array([1, 3, 6]))  # only values 0.8 and 0.9 exceed 0.4 times the max (which is 0.9)
+
+    # Empty test
+    assert np.array_equal(focus_criterion(np.array([]), np.array([])), np.array([]))
 
     # Test with custom alpha
-    assert np.array_equal(pc(pos, values, alpha=0.5), np.array([1, 3, 6]))
+    assert np.array_equal(focus_criterion(pos, values, alpha=0.5), np.array([1, 3, 6]))
 
     # Test with a larger alpha
-    assert np.array_equal(pc(pos, values, alpha=1.0), [6])  # No values exceed the maximum value itself
+    assert np.array_equal(focus_criterion(pos, values, alpha=1.0), [6])  # No values exceed the maximum value itself
 
     # Test with all values below threshold
     values = np.array([0.1, 0.2, 0.3, 0.4])
 
-    assert np.array_equal(pc(pos[:4], values), [1,2,3])  # All values are below 0.4 times the max (which is 0.4)
+    assert np.array_equal(focus_criterion(pos[:4], values), [1,2,3])  # All values are below 0.4 times the max (which is 0.4)
  
 @pytest.fixture
 def mock_data():
     all_paths = [ [ (0,0,0), (0,2,0), (0,5,0), (0,10,0), (0,15,0), (0,20,0)], [ (1,20,0), (1,20,10), (1,20,20)  ] ] # Mock paths
     path_lengths = [ 2.2, 2.3 ]  # Mock path lengths
-    measured_trace_fluorescence = [ [100, 8, 3, 2, 3, 39], [38, 2, 20] ]  # Mock fluorescence data
+    measured_trace_fluorescence = [ [100, 8, 3, 2, 3, 49], [38, 2, 20] ]  # Mock fluorescence data
     return all_paths, path_lengths, measured_trace_fluorescence
 
 def test_process_cell_traces_return_type(mock_data):
@@ -117,23 +121,30 @@ def test_process_cell_traces_pathdata_path_lengths(mock_data):
 def test_process_cell_traces_peaks(mock_data):
     all_paths, path_lengths, measured_trace_fluorescence = mock_data
     result = process_cell_traces(all_paths, path_lengths, measured_trace_fluorescence)
+    print(result)
     peaks = [p.peaks for p in result.pathdata_list]
     assert peaks == [[0,5],[]]
     
 # Mock data
 @pytest.fixture
 def mock_celldata():
-    pathdata1 = PathData(peaks=[0, 5], points=[(0,0,0), (0,2,0), (0,5,0), (0,10,0), (0,15,0), (0,20,0)], o_hei10=[100, 8, 3, 2, 3, 39], SC_length=2.2)
-    pathdata2 = PathData(peaks=[0], points=[(1,20,0), (1,20,10), (1,20,20) ], o_hei10=[38, 2, 20], SC_length=2.3)
+    pathdata1 = PathData(peaks=[0, 5], points=[(0,0,0), (0,2,0), (0,5,0), (0,10,0), (0,15,0), (0,20,0)], removed_peaks=[], o_hei10=[100, 8, 3, 2, 3, 69], SC_length=2.2)
+    pathdata2 = PathData(peaks=[2], points=[(1,20,0), (1,20,10), (1,20,20) ], removed_peaks=[RemovedPeakData(0, (0,5))], o_hei10=[38, 2, 20], SC_length=2.3)
     return CellData(pathdata_list=[pathdata1, pathdata2])
 
-def test_analyse_celldata_output_length(mock_celldata):
-    rel_intensity, pos, pos_index, trace_median_intensity, trace_thresholds = analyse_celldata(mock_celldata, {'peak_threshold': 0.4, 'threshold_type':'per-trace'})
-    assert len(rel_intensity) == len(mock_celldata.pathdata_list), "Mismatch in relative intensities length"
-    assert len(pos) == len(mock_celldata.pathdata_list), "Mismatch in positions length"
-    assert len(pos_index) == len(mock_celldata.pathdata_list), "Mismatch in position indices length"
-
+def test_analyse_celldata(mock_celldata):
+    data_frame, foci_absolute_intensity, foci_position_index, dominated_foci_data, trace_median_intensity, trace_thresholds = analyse_celldata(mock_celldata, {'peak_threshold': 0.4, 'threshold_type':'per-trace'})
+    assert len(data_frame) == len(mock_celldata.pathdata_list), "Mismatch in dataframe length"
+    assert len(foci_absolute_intensity) == len(mock_celldata.pathdata_list), "Mismatch in relative intensities length"
+    assert len(foci_position_index) == len(mock_celldata.pathdata_list), "Mismatch in positions length"
 
+    assert list(map(list, foci_position_index)) == [[0, 5], [2]]
 
 
+def test_analyse_celldata_per_cell(mock_celldata):
+    data_frame, foci_absolute_intensity, foci_position_index, dominated_foci_data, trace_median_intensity, trace_thresholds = analyse_celldata(mock_celldata, {'peak_threshold': 0.4, 'threshold_type':'per-cell'})
+    assert len(data_frame) == len(mock_celldata.pathdata_list), "Mismatch in relative intensities length"
+    assert len(foci_absolute_intensity) == len(mock_celldata.pathdata_list), "Mismatch in positions length"
+    assert len(foci_position_index) == len(mock_celldata.pathdata_list), "Mismatch in position indices length"
+    assert list(map(list, foci_position_index)) == [[0, 5], []]