Python: Create submatrices from long/narrow table and save in HDF5

Question

Hello dear StackExchange Community,

I have written a small program which is doing the following. In short:

• it reads in a table in narrow format
• creates submatrices from it
• saves it in HDF5 format

The long table is actually representing a huge matrix and the data has the following shape (example not real):

POS1    POS2    VALUE
146000  146000  42
146000  147000  69
147000  147000  13
.       .       .
.       .       .
.       .       .

The data comes from experimental measurements and is not formatted user friendly. POS 1 and POS 2 represent coordinates of a small cluster on a big grid(matrix) and VALUE is a representing how strong these are associated. It's quite similar to a correlation matrix. The distance between positions is evenly sized and always has a distance of 1k (they represent bins). The goal is to compare every coordinate/bin vs all others. so if have n bins we get a matrix with n x n Values. Unfortunately not all comparisons are in the table. Some are missing.

For my project I need to subset this big matrix into smaller ones, while keeping track of the coordinates, so I can process them individually. I decided to go for 128x128 matrices.

Also I needed a solution how I can store the data efficiently to able to access it in a fast way for different ML applications. I decided to go for HDF5 using the h5py library. For every 128x128 submatrix I create a dataset in the HDF5 file with the name of POS1-POS2. This way I have the coordinates stored as the dataset name, which makes them easy to track, since all of them represent 1k bins.

While my code actually runs, it runs pretty slow. I could rewrite it in c++ using the eigen lib or try Rust, since I am currently learning it, but I thought I stick to Python and try to optimize it. Maybe someone here could help? Maybe numba or pypy could help here? Until now I did not make that much of a good experience with numba regarding more complicated code where one does not just crunch numbers (where it really shines).

Any suggestions and optimization are welcome! Many thanks in advance!

I separated the matrix class I wrote into another file:

from typing import List, Tuple

import numpy as np
import pandas as pd


class ImageMatrix:
    """ Class to store and fill 128x128 Image from long table """

    def __init__(self, x_axis_start: int = 0, y_axis_start: int = 0, mat_dims: Tuple[int, int] = (128, 128)):

        # axis' to know where submatrix belongs to in "whole" matrix
        self.x_axis: List[int] = [i for i in range(x_axis_start, x_axis_start + mat_dims[0] * 1_000, 1000)]
        self.y_axis: List[int] = [j for j in range(y_axis_start, y_axis_start + mat_dims[1] * 1_000, 1000)]


        self.mat_dims: Tuple[int, int] = mat_dims
        # create empty mat to fill up 
        self.image_matrix: np.ndarray = np.zeros(mat_dims).astype(int)


    def __str__(self) -> str:
        return(f"Matrix of size: {self.mat_dims}\n" 
            f"coordinates starting for x: {self.x_axis[0]} and y: {self.y_axis[0]}.")

    def print_matrix(self) -> None:
        """Print Matrix to the screen"""
        print(self.image_matrix)


    def add_value(self, value:int , x_coordinate:int, y_coordinate:int) -> None:
        """Adds a value to the matrix at given coordinates (x: column, y: row)"""

        self.image_matrix[self.y_axis.index(y_coordinate)][self.x_axis.index(x_coordinate)] = value

    def fill_from_pandas(self, pandas_df: pd.DataFrame) -> None:
        """
        Fills up the matrix from a pandas long table
        assuming the table has following structure:
        BIN1    BIN2    VALUE
        """
        for i in range(0,len(pandas_df)):
            self.add_value(pandas_df.iloc[i,:][2], pandas_df.iloc[i,:][0], pandas_df.iloc[i,:][1] ) 
    
    def fill_from_pandas_symmetric(self, pandas_df: pd.DataFrame) -> None:
        """
        Fills up the matrix and makes it symetric from a pandas long table
        assuming the table has following structure:
        BIN1    BIN2    VALUE
        """
        self.fill_from_pandas(pandas_df)

        tmp = np.tril(self.image_matrix) + np.triu(self.image_matrix.T,1)
        self.image_matrix = tmp 

Here's the main script:

import argparse
from time import time

import h5py
import numpy as np
import pandas as pd
from termcolor import colored

import imagematrix

def load_hic_from_csv(path_to_file: str) -> pd.DataFrame:
    """
    Input: path to T/CSV file
    Outout: pandas dataframe

    Expects t/csv to have the following format and no header:


    COL 1                  COL 2                  COL 3


    """
    return pd.read_csv(path_to_file, header=None, delim_whitespace=True, names=["POS1","POS2","VALUE"])

def create_hdf5_dataset(long_table: pd.DataFrame, file_name: str) -> None:
    """
    Input:
        - a pandas dataframe containing the hi-C counts
        - file name to write hdf5 file to
        - name of dataset group

    Output:
        None. Creates an HDF5 File in specified dir

    """

    # get first and last bin
    start = long_table['POS1'].min()
    end = long_table['POS2'].max()
    n_iters = int((end -start)/128_000)
    hf = h5py.File(file_name, 'a')

     # loop throug main matrix in 128k steps
    for counter, i in enumerate(range(start, end, 128_000)):
        # print steps every 10 steps
        if(counter % 10 == 0):
            print(f"[{counter}/{n_iters}]")

     # loop throug main matrix in 128k steps
        for j in range(start, end, 128_000):
                # select subregion from matrix 
                local_region = long_table[(long_table['POS1'] >= i) & (long_table['POS1'] <= i + 127_000 )&\
                    (long_table['POS2'] >= j) & (long_table['POS2'] <= j + 127_000 )]

                # instantiate the matrix
                My_matrix = imagematrix.ImageMatrix(x_axis_start=i, y_axis_start= j)

                # mirror matrix if it's symmetric -> long table does not contain that values
                if(i == j):
                    My_matrix.fill_from_pandas_symmetric(local_region)
                else:
                    My_matrix.fill_from_pandas(local_region)

                hf.create_dataset(f"{i}-{j}", data=My_matrix.image_matrix, compression="lzf")

    # close file
    hf.close()
    print(colored(f"Done processing: {file_name}!", 'green'))

def main() -> None:

    argsParser = argparse.ArgumentParser(description = "Write diagonal sub matrices to hdf5 file in 128 x 128 ")
    argsParser.add_argument('--input_file', '-i', type = str, required = True, help="Long table containing POS1, POS2 and VALUES")
    argsParser.add_argument('--output_file', '-o', type = str, required= True, help="A string with the path to the output file")

    # parse arguments
    args =argsParser.parse_args()

    DF = load_hic_from_csv(args.input_file)
    create_hdf5_dataset(DF, args.output_file)

if __name__ == "__main__":
    main()

Here's a longer example input file. The ones I am currently working with are actually tab separated and not csvs anymore.

Furthermore maybe i should mention that these files can get quite big - a few GB.

7320000        17351000        3.0
17321000        17351000        1.0
17322000        17351000        2.0
17323000        17351000        2.0
17324000        17351000        4.0
17325000        17351000        4.0
17326000        17351000        1.0
17327000        17351000        1.0
17328000        17351000        2.0
17329000        17351000        1.0
17332000        17351000        1.0
17333000        17351000        5.0
17340000        17351000        4.0
17341000        17351000        6.0
17342000        17351000        6.0
17343000        17351000        5.0
17344000        17351000        3.0
17345000        17351000        4.0
17346000        17351000        4.0
17347000        17351000        9.0
17348000        17351000        1.0
17349000        17351000        4.0
17350000        17351000        77.0
17351000        17351000        48.0
16848000        17352000        1.0
16864000        17352000        1.0
16865000        17352000        1.0
16874000        17352000        1.0
16877000        17352000        1.0
16884000        17352000        4.0
16895000        17352000        1.0
16900000        17352000        1.0
16910000        17352000        1.0
16913000        17352000        1.0
16928000        17352000        1.0
16936000        17352000        1.0
16939000        17352000        1.0
16943000        17352000        2.0
16946000        17352000        1.0
16954000        17352000        1.0
17003000        17352000        1.0
17008000        17352000        1.0
17013000        17352000        1.0
17018000        17352000        1.0
17024000        17352000        1.0
17030000        17352000        1.0
17031000        17352000        1.0
17034000        17352000        1.0
17035000        17352000        2.0
17037000        17352000        1.0
17042000        17352000        1.0
17055000        17352000        1.0
17062000        17352000        1.0
17064000        17352000        1.0
17065000        17352000        1.0
17066000        17352000        1.0
17068000        17352000        1.0
17072000        17352000        1.0
17076000        17352000        2.0
17078000        17352000        1.0
17079000        17352000        1.0
17085000        17352000        1.0
17086000        17352000        1.0
17089000        17352000        1.0
17091000        17352000        2.0
17092000        17352000        1.0
17094000        17352000        3.0

Answer 1

POS 1 and POS 2 represent coordinates of a small cluster on a big grid (matrix)

Unfortunately not all comparisons are in the table. Some are missing.

This is a job for sparse matrices. The selection of which sparse matrix format to use is non-trivial but BSR lends itself well to direct construction from your data.

I'd cut out Pandas and use Numpy file loading directly; loadtxt is trivially easy for processing your format. I assume that your third column is all int; if not, you can use a struct-dtype.

Divide out the factor of 1000 from your coordinates.

Consider ejecting HDF5 and using standard Numpy npz. Think: you had already been loading the entire dataframe of coordinates into memory all at once, which means that it's reasonable to assume you should be able to do the same for a sparse matrix which scales in the same way. You can save and load this sparse matrix in one shot, and when you load it and do (mystery) processing, you can slice it for segmentation then. The slicing operation will be much easier in the sparse representation and you can lose most of your code.

Suggested

import argparse

import numpy as np
import scipy.sparse
from scipy.sparse import spmatrix


def load_hic(path: str) -> np.ndarray:
    return np.loadtxt(fname=path, dtype=int)


def create_dataset(long_table: np.ndarray) -> spmatrix:
    data = long_table[:, 2]
    ij = long_table[:, :2].T // 1_000
    return scipy.sparse.bsr_matrix((data, ij))


def parse_args() -> argparse.Namespace:
    args_parser = argparse.ArgumentParser(
        description="Convert a coordinate-list file to a sparse NPZ file")
    args_parser.add_argument('--input-file', '-i', type=str, required=True,
                             help="Long table containing POS1, POS2 and VALUES")
    args_parser.add_argument('--output-file', '-o', type=str, required=True,
                             help="A string with the path to the output file")
    return args_parser.parse_args()


def main() -> None:
    args = parse_args()

    array = load_hic(args.input_file)
    sparse = create_dataset(array)
    scipy.sparse.save_npz(args.output_file, sparse)


if __name__ == "__main__":
    main()