nena.py

import h5py
import yaml
import numpy as np
import pandas as pd
import numba
import lmfit
import os
from scipy.special import iv
# import time


###############################################################################################
# HDF5 and YAML files loading and saving
###############################################################################################
    
def load_locs(path):
    """Load localizations from HDF5 file."""
    with h5py.File(path, "r") as locs_file:
        locs = locs_file["locs"][...]
    locs = np.rec.array(locs, dtype=locs.dtype)
    info = load_info(path)
    return locs, info

def load_info(path):
    """Load metadata from the corresponding YAML file."""
    path_base = path.rsplit(".", 1)[0]
    filename = path_base + ".yaml"
    try:
        with open(filename, "r") as info_file:
            info = list(yaml.load_all(info_file, Loader=yaml.UnsafeLoader))
    except FileNotFoundError as e:
        print(f"Could not find metadata file: {filename}")
        raise FileNotFoundError(e)
    return info

def save_locs(path, locs, info):
    """Save localizations to an HDF5 file."""
    with h5py.File(path, "w") as locs_file:
        locs_file.create_dataset("locs", data=locs)
    base = path.rsplit(".", 1)[0]
    save_info(base + ".yaml", info)

def save_info(path, info):
    """Save metadata information to a YAML file."""
    with open(path, "w") as file:
        yaml.dump_all(info, file, default_flow_style=False)

###############################################################################################
# 3. NeNA calculation
###############################################################################################        
        
def nena(input_data, info, callback=None):
    """ NeNA (nearest neighbor based analysis)
    
    Parameters:
        input_data: Either a path to HDF5 file or a structured numpy array of localizations
        info: metadata information
        callback: Optional callback function for progress updates
    """
    # Check if input_data is a string (file path) or structured array
    if isinstance(input_data, str):
        with h5py.File(input_data, "r") as f:
            locs = f["locs"][:]
    else:
        locs = input_data

    # Extract required fields
    frame = locs["frame"]
    x = locs["x"]
    y = locs["y"]
    lpx = locs["lpx"]
    lpy = locs["lpy"]
    
    sorted_indices = np.argsort(frame)
    frame = frame[sorted_indices]
    x = x[sorted_indices]
    y = y[sorted_indices]
    lpx = lpx[sorted_indices]
    lpy = lpy[sorted_indices]
    
    if 'group' in locs.dtype.names:
        group = locs['group']
    else:
        group = np.zeros(len(locs), dtype=np.int32)
    group = group[sorted_indices]

    bin_size = 0.001
    d_max = 1.0
    
    bin_centers, dnfl_ = _nfndh(frame, x, y, group, d_max, bin_size, callback)

    # Get the histogram of nearest neighbor distances
    # bin_centers, dnfl_ = next_frame_neighbor_distance_histogram(input_file, callback)

    def func(d, a, s, ac, dc, sc):
        """ Rayleigh distribution"""
        f = a * (d / s**2) * np.exp(-0.5 * d**2 / s**2)
        fc = (
            ac
            * (d / sc**2)
            * np.exp(-0.5 * (d**2 + dc**2) / sc**2)
            * iv(0, d * dc / sc)
        )
        return f + fc

    pdf_model = lmfit.Model(func)
    params = lmfit.Parameters()

    # Calculate area under the distribution curve
    area = np.trapz(dnfl_, bin_centers)

    # The median uncertainty in x/y (lpx, lpy) is used as an estimate of σ_SMLM
    median_lp = np.mean([np.median(lpx), np.median(lpy)])

    # Initialize parameters for fitting
    params.add("a", value=area / 2, min=0)
    params.add("s", value=median_lp, min=0)
    params.add("ac", value=area / 2, min=0)
    params.add("dc", value=2 * median_lp, min=0)
    params.add("sc", value=median_lp, min=0)

    # Fit the model to the data
    result = pdf_model.fit(dnfl_, params, d=bin_centers)

    # return result, result.best_values["s"]  # Return fitted result and NeNa value (s)
    return result.best_values["s"]  


def _nfndh(frame, x, y, group, d_max, bin_size, callback=None):
    """ Compute histogram of nearest neighbor distances """
    N = len(frame)
    bins = np.arange(0, d_max, bin_size)
    dnfl = np.zeros(len(bins))
    
    one_percent = int(N / 100)
    starts = one_percent * np.arange(100)
    
    for k, start in enumerate(starts):
        for i in range(start, start + one_percent):
            _fill_dnfl(N, frame, x, y, group, i, d_max, dnfl, bin_size)
        if callback is not None:
            callback(k + 1)
    
    bin_centers = bins + bin_size / 2
    return bin_centers, dnfl

# @numba.jit(nopython=True)
@numba.jit(nopython=True)
def _fill_dnfl(N, frame, x, y, group, i, d_max, dnfl, bin_size):
    """ Helper function to fill the nearest neighbor distance histogram """
    frame_i = frame[i]
    x_i = x[i]
    y_i = y[i]
    group_i = group[i]
    
    min_frame = frame_i + 1
    for min_index in range(i + 1, N):
        if frame[min_index] >= min_frame:
            break
    
    max_frame = frame_i + 1
    for max_index in range(min_index, N):
        if frame[max_index] > max_frame:
            break
    
    d_max_2 = d_max**2
    for j in range(min_index, max_index):
        if group[j] == group_i:
            dx2 = (x_i - x[j]) ** 2
            if dx2 <= d_max_2:
                dy2 = (y_i - y[j]) ** 2
                if dy2 <= d_max_2:
                    d = np.sqrt(dx2 + dy2)
                    if d <= d_max:
                        bin = int(d / bin_size)
                        dnfl[bin] += 1



def process_nena(input_folder, input_extension, pixelsize):
    """Process all HDF5 files in the input folder for NeNA calculation."""
    
    # Dictionary to store NeNA values with clean base names
    nena_values = {}
    results = []
    
    hdf5_files = []
    for root, _, files in os.walk(input_folder):
        for f in files:
            if f.endswith(f'{input_extension}.hdf5'):
                hdf5_files.append((root, f))
    
    print(f"Found {len(hdf5_files)} files to process")
    
    for root, filename in hdf5_files:
        # Get base name by removing the input extension and .hdf5
        base_name = filename.replace(f"{input_extension}.hdf5", "")
        input_file = os.path.join(root, filename)
        
        print(f"\nProcessing file: {filename}")
        # print(f"Using base name: {base_name}")
        
        try:
            # Load localizations and info
            with h5py.File(input_file, "r") as f:
                locs = f["locs"][...]
            info = load_info(input_file)
            
            # Calculate NeNA value
            nena_value_px = nena(input_file, info)
            nena_value_nm = nena_value_px * pixelsize
            
            nena_values[base_name] = nena_value_px
            
            # Add to results for CSV
            results.append({
                'Filename': base_name,
                # 'Input_filename': filename,
                'NeNA_px': nena_value_px,
                'NeNA_nm': nena_value_nm
            })
            
            print(f"NeNA values for {base_name}:")
            print(f"  Pixels: {nena_value_px:.4f}")
            print(f"  Nanometers: {nena_value_nm:.4f}")
            
            
            
        except Exception as e:
            print(f"Error processing {filename}: {str(e)}")
            continue
    
        
    # Save NeNA values to YAML file with additional metadata
    nena_data = {
        'metadata': {
            'Filename': base_name,
            'pixelsize': pixelsize
        },
        'values': nena_values
    }
    
    output_dir = os.path.dirname(input_folder)
    nena_yaml_file = os.path.join(output_dir, "nena_values.yaml")
    nena_csv_file = os.path.join(output_dir, "nena_values.csv")
    
    with open(nena_yaml_file, 'w') as f:
        yaml.dump(nena_data, f)
    print(f"\nNeNA values saved to: {nena_yaml_file}")
    
    # Save results to CSV
    df = pd.DataFrame(results)
    df.to_csv(nena_csv_file, index=False)
    print(f"NeNA values saved to: {nena_csv_file}")
    
    # print("\nAll NeNA values calculated and saved")
    return nena_data
        

if __name__ == "__main__":

    with open('config.yaml', 'r') as f:
        config = yaml.safe_load(f)
    
    process_nena(
        input_folder=config['paths']['input_folder'],
        input_extension=config['nena']['input_extension'],
        pixelsize=config['pixelsize']['pixelsize']
    )