Source code for dipy.stats.sketching

import os
import tempfile

import numpy as np




def count_sketch(matrixa_name, matrixa_dtype, matrixa_shape, sketch_rows, tmp_dir):
    """Count Sketching algorithm to reduce the size of the matrix.

    Parameters
    ----------
    matrixa_name : str
        The name of the memmap file containing the matrix A.
    matrixa_dtype : dtype
        The dtype of the matrix A.
    matrixa_shape : tuple
        The shape of the matrix A.
    sketch_rows : int
        The number of rows in the sketch matrix.
    tmp_dir : str
        The directory to save the temporary files.

    Returns
    -------
    matrixc_file.name : str
        The name of the memmap file containing the sketch matrix.
    matrixc.dtype : dtype
        The dtype of the sketch matrix.
    matrixc.shape : tuple
        The shape of the sketch matrix.

    """
    matrixa = np.squeeze(
        np.memmap(matrixa_name, dtype=matrixa_dtype, mode="r+", shape=matrixa_shape)
    ).reshape(np.prod(matrixa_shape[:-1]), matrixa_shape[-1])

    with tempfile.NamedTemporaryFile(
        delete=False, dir=tmp_dir, suffix="matrix_t"
    ) as matrixt_file:
        matrixt = np.memmap(
            matrixt_file.name, dtype=matrixa_dtype, mode="w+", shape=matrixa.shape
        )
        hashed_indices = np.random.choice(sketch_rows, matrixa.shape[0], replace=True)
        rand_signs = np.random.choice(2, matrixa.shape[0], replace=True) * 2 - 1
        for i in range(0, matrixa.shape[0], matrixa.shape[0] // 20):
            end_index = min(i + matrixa.shape[0] // 20, matrixa.shape[0])
            matrixt[i:end_index, :] = (
                matrixa[i:end_index, :] * rand_signs[i:end_index, np.newaxis]
            )

    with tempfile.NamedTemporaryFile(
        delete=False, dir=tmp_dir, suffix="matrix_C"
    ) as matrixc_file:
        matrixc = np.memmap(
            matrixc_file.name,
            dtype=matrixa_dtype,
            mode="w+",
            shape=(sketch_rows, matrixa.shape[1]),
        )
        np.add.at(matrixc, hashed_indices, matrixt)
        matrixc.flush()
        matrixt.flush()
    del matrixt
    os.unlink(matrixt_file.name)
    return matrixc_file.name, matrixc.dtype, matrixc.shape