3DVA analysis

extract_data.py

import os

import numpy as np
import pandas as pd
import mrcfile
from pathlib import Path


def components_from_cryosparc(cs_file):
    data = np.load(cs_file)
    coeff_headings = [
        name
        for name in data.dtype.names
        if ('components_mode_' in name and 'value' in name)
     ]
    coeffs = {
        f'component_{i}': data[heading] for i, heading in enumerate(coeff_headings)
    }
    return pd.DataFrame.from_dict(coeffs)


def map_from_particle_subset(component_df, subset_idx):
    consensus_file = 'J27/cryosparc_P21_J27_map.mrc'
    consensus = mrcfile.open(consensus_file).data
    volume_files = list(Path('J27').glob('*_component*.mrc'))
    volume_files.reverse()
    n_coeffs = len(volume_files)
    print(volume_files)
    variability_volumes = np.stack(
        [mrcfile.open(volume_file).data for volume_file in volume_files],
        axis=0
    )
    subset_coeffs = component_df.to_numpy()[subset_idx, :].mean(axis=0)
    subset_coeffs = subset_coeffs.reshape((n_coeffs, 1, 1, 1))
    map = np.sum(consensus + subset_coeffs * variability_volumes, axis=0)
    return map


def volumes_along_component(component_df, n_volumes, component_idx, qcut=True):
    """generate a set of n volumes along a given component

    qcut=True means bins will have equal numbers of particles
    qcut=False means distance between each bin is identical
    """
    coefficients = component_df[f'component_{component_idx}']
    if qcut:
        subsets = pd.qcut(coefficients, n_volumes)
    else:
        subsets = pd.cut(coefficients, n_volumes)
    volumes = [
        map_from_particle_subset(component_df, subsets == subset)
        for subset
        in subsets.array.categories
        ]
    return np.stack(volumes)



if __name__ == '__main__':
    import napari
    df = components_from_cryosparc('J27/cryosparc_P21_J27_particles.cs')
    viewer = napari.Viewer(ndisplay=3)
    viewer.add_image(volumes_along_component(df, 20, 0))
    napari.run()

plots.ipynb