# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## LSH Prototype 2
#
# [SimHash](https://en.wikipedia.org/wiki/SimHash) implementation over Xarray/Dask.

# + tags=["parameters"]
h = None              # Number of individual signed projections to compute
g = None              # Number of composite hashes (groupings of signs)
sample_rate = None    # Variant sampling rate used in extraction
ds_name = None        # Dataset name 
n_workers = 8         # Number of dask workers
mem_fraction = .9     # Maximum fraction of system memory to use
projection = 'random'

# Example Settings
# h = 24
# g = 100        
# sample_rate = .05
# ds_name = 'hapmap'
# n_workers = 8
# mem_fraction = .9
# projection = 'orthogonal'
# -

import os
import pandas as pd
import numpy as np
from dask.distributed import Client
from gwas_analysis.dask import get_dask_client
import dask.array as da
import dask.dataframe as dd
import xarray as xr
# %run {os.environ['NB_DIR']}/nb.py
# %run $BENCHMARK_METHOD_DIR/common.py
assert h is not None
assert g is not None
assert sample_rate is not None
assert ds_name is not None
ds_config = DATASET_CONFIG[ds_name]
n_projections = h * g
ds_path = ld_prune_lsh.dataset_path(ds_name, sr=sample_rate)
ds_path

# ### Initialization

client = get_dask_client(n_workers=n_workers, max_mem_fraction=mem_fraction)
client

# +
# Load the coded call data
X = da.from_zarr(ds_path + '.zarr')

# Rechunk to ensure that more than one worker is used downstream
# (which is the case with a relatively small number of variants < 1M)
def blocks(n, n_workers):
    if n <= n_workers:
        return n
    return n // n_workers
X = X.rechunk(chunks=(blocks(X.shape[0], n_workers), -1))
X
# -

# Convert to xarray and center variant vectors (to make dot product equivalent to un-normalized cosine)
X = xr.DataArray(
    data=X, 
    dims=['variant', 'sample'],
    coords=dict(
        variant=da.arange(X.shape[0], dtype=np.int32),
        sample=da.arange(X.shape[1], dtype=np.int32)
    ),
    name='calls'
)
X -= X.mean(dim='sample')
X /= X.std(dim='sample')
X

# +
# Generate random vectors for projection
da.random.seed(seed=1)

def get_random_matrix(m, n):
    if projection == 'random':
        print('Using random projection')
        return da.random.normal(
            size=(m, n), 
            chunks='auto'
            # Split the random array into chunks here as well
            # chunks=(-1, blocks(n_projections, n_workers))
        )
    elif projection == 'orthogonal':
        print('Using orthogonal projection')
        n_batch = int(np.ceil(n / m))
        R = []
        for i in range(n_batch):
            rm = da.random.normal(
                size=(m, m), 
                chunks='auto'
            )
            R.append(da.linalg.qr(rm)[0])
        R = da.concatenate(R, axis=1)[:,:n]
        # rms = (rms - rms.mean(axis=1, keepdims=True)) / rms.std(axis=1, keepdims=True)
        return R
    else:
        raise ValueError(f'Projection type ${projection} not supported')

R = xr.DataArray(
    data=get_random_matrix(len(X['sample']), n_projections), 
    dims=('sample', 'projection'),
    coords=dict(
        sample=da.arange(len(X['sample']), dtype=np.int32),
        projection=da.arange(n_projections, dtype=np.int32)
    ),
    name='random'
)
R
# -

# ### Compute Hash Bits

# Apply projection and attach hash signature grouping
P = (
    (X @ R)
    .rename('projections')
    .pipe(lambda x: x.assign_coords(hash_group=('projection', x['projection']//h)))
)
P


# %%time
# Group by hash signature group (i.e. columns) and compute row-wise hashes
# within those columns, which will each be boolean vectors indicating sign
# TODO: This should operate on rows across hash group boundaries since the number
# of groups can be high and making many small chunks is inefficient
def hash_bits(x):
    return np.expand_dims(np.apply_along_axis(
        lambda r: hash(np.asarray(r).tobytes()), 
        axis=1, arr=x
    ), 1)
H = xr.DataArray(
    (P > 0).data.rechunk(chunks=(P.data.chunksize[0], h))
    .map_blocks(hash_bits, chunks=(P.data.chunksize[0], 1))
    .compute(),
    dims=('variant', 'hash_group'),
    coords=dict(
        variant=P['variant'],
        hash_group=np.arange(g)
    ),
    name='hash_value'
)
H

# Note that the above is unfortunately much faster than using the xarray API directly:

# +
# # %%time
# def hash_bits(x, axis=None):
#     # NOTE: It makes little difference if the np conversion is done
#     # per row or initially for this group (it must all be loaded into memory already)
#     return xr.DataArray(da.apply_along_axis(
#         lambda r: hash(np.asarray(r).tobytes()), 
#         axis=axis, arr=x
#     ))

# H = (
#     (P > 0)
#     .groupby('hash_group')
#     .reduce(hash_bits, dim='projection')
#     .rename('hash_value')
# )
# H

# For h = 24, g = 100:
# CPU times: user 24 s, sys: 1.64 s, total: 25.6 s
# Wall time: 2min 10s (compared to ~10s for Dask version w/ same parameters)
# -

# ### Compute Hash Bucket

L = H.stack(i=('variant', 'hash_group')).reset_index('i')
L


# %%time
def hash_bucket(x, axis=None):
    return xr.DataArray(da.apply_along_axis(
        lambda r: hash(np.asarray(r).tobytes()), 
        axis=axis, arr=x
    ))
L = L.assign_coords(hash_bucket=
    xr.concat([L, L['hash_group']], dim='component', coords='minimal').T
    .pipe(lambda x: x.reduce(hash_bucket, dim='component'))
    .rename('hash_bucket')
)
L

# ### Export

df = L.to_dataset().to_dask_dataframe().drop('i', axis='columns')
df.head(8)

# %%time
path = ld_prune_lsh.dataset_path(ds_name, sr=sample_rate, h=h, g=g, p=projection) + '.parquet'
df.to_parquet(path)
path

client.shutdown()

StarcoderdataJupyterScriptsDedupFiltered