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

# +
import sys
import numpy as np
import matplotlib.pyplot as plt
import pickle
import os
import gzip
from torch.utils import data
import torch
import torch.optim as optim
from torch.autograd import Variable
from time import gmtime, strftime
import torch.nn as nn

import math

import h5py

import src.model as models
import src.utils as utils

plt.style.use('ggplot')
# -

# %load_ext autoreload
# %autoreload 2

#Constraint Pytoch to use the good GPU!
import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   
os.environ["CUDA_VISIBLE_DEVICES"]="0"

print torch.cuda.get_device_name(0)

# # Load Data

KeysSpec = ['ESC','CN'] #Keys Species
data_path = {}
data_path['ESC'] = '/home/invites/jmorlot/HDDTeam/data/Hi-C/Mouse/boyan/last_try2/ESC/'
data_path['CN'] = '/home/invites/jmorlot/HDDTeam/data/Hi-C/Mouse/boyan/last_try2/CN/'

# +
# import re
# InputFiles ={key: [data_path[key]+f for f in os.listdir(data_path[key]) if re.search('basic',f)] for key in data_path.keys()} 
# OutputFiles ={key: [data_path[key]+f.replace('basic','HiCPlusBoost') for f in os.listdir(data_path[key]) if re.search('basic',f)] for key in data_path.keys()} 
# -

KeysSub = ['0.1','1.0','10.0','100.0'] #Keys Subsampling
Data={}
for keySp in KeysSpec:
    print keySp
    Data[keySp]={}
    for keySu in KeysSub:
        print '\t' + keySu
#         path = data_path[keySp] + 'chr16_basic_' + keySu + '.hdf5'
        f = h5py.File(data_path[keySp] + 'chr16_basic_' + keySu + '.hdf5','r')
        Data[keySp][keySu] = np.array(f[f.keys()[0]])
        f.close()

DataBoostHiC = {}
for keySp in KeysSpec:
    print keySp
    DataBoostHiC[keySp]={}
    for keySu in KeysSub:
        print '\t' + keySu
#         path = data_path[keySp] + 'chr16_basic_' + keySu + '.hdf5'
        f = h5py.File(data_path[keySp] + 'chr16_boosted_' + keySu + '.hdf5','r')
        DataBoostHiC[keySp][keySu] = np.array(f[f.keys()[0]])
        f.close()

# # HiCPlus

# ## Internal variables

# +
use_gpu = 1

conv2d1_filters_numbers = 8
conv2d1_filters_size = 9
conv2d2_filters_numbers = 8
conv2d2_filters_size = 1
conv2d3_filters_numbers = 1
conv2d3_filters_size = 5


down_sample_ratio = 16
epochs = 10
HiC_max_value = 100

max_range = 3000 # in paper=201
batch_size = 32 # in paper = size of the dataset (silly...)


# -

# ## Rewritting functions to fit with our matrix format

# +
def divide(HiCsample):
    '''
        Subdivide the HiC matrix in an ensemble of subimages of size subImage_size
    '''
    subImage_size = 40
    step = 25
    result = []
    index = []
    
    total_loci = HiCsample.shape[0]
    for i in range(0, total_loci, step):
        for j in range(0, total_loci, ):
            if (abs(i-j) > max_range or i + subImage_size >= total_loci or j + subImage_size >= total_loci):
                continue
            subImage = HiCsample[i:i+subImage_size, j:j+subImage_size]

            result.append([subImage,])
            index.append((i, j))

    result = np.array(result)
#     print result.shape
    result = result.astype(np.double)
    index = np.array(index)
    return result, index

def HiCPlus(HiCsample):
    
    ## Subdivide the HiC matrix in subimages
    low_resolution_samples, index = divide(HiCsample)
    
    low_resolution_samples = np.minimum(HiC_max_value, low_resolution_samples)

#     batch_size = low_resolution_samples.shape[0]
    
    ## Reshape the high-quality Hi-C sample as the target value of the training.
    sample_size = low_resolution_samples.shape[-1]
    padding = conv2d1_filters_size + conv2d2_filters_size + conv2d3_filters_size - 3
    half_padding = padding / 2
    output_length = sample_size - padding

#     print low_resolution_samples.shape
    
    ## Data Loader
    lowres_set = data.TensorDataset(torch.from_numpy(low_resolution_samples), torch.from_numpy(np.zeros(low_resolution_samples.shape[0])))
    lowres_loader = torch.utils.data.DataLoader(lowres_set, batch_size=batch_size, shuffle=False)

    hires_loader = lowres_loader
    
    ## Get Model
    model = models.Net(40, 28)
    model.load_state_dict(torch.load('model/pytorch_model_12000'))
    if use_gpu:
        model = model.cuda()

    _loss = nn.MSELoss()
    
    ## Make predictions
    running_loss = 0.0
    running_loss_validate = 0.0 # WARNING : NOT USED
    reg_loss = 0.0 # WARNING : NOT USED

    y_prediction = []
    for i, _lowRes in enumerate(lowres_loader):

        _lowRes = Variable(_lowRes[0]).float()

        _lowRes = _lowRes.cuda()
        
        _hiRes = model(_lowRes).data.cpu()
        
        y_prediction.append(_hiRes)
        
        del _lowRes,_hiRes

    y_prediction = torch.cat(y_prediction)

    y_predict = y_prediction.numpy()


    print y_predict.shape
    
    ## Recombine samples
    length = int(y_predict.shape[2])
    y_predict = np.reshape(y_predict, (y_predict.shape[0], length, length))
#     print y_predict.shape

#     chrs_length = [249250621,243199373,198022430,191154276,180915260,171115067,159138663,146364022,141213431,135534747,135006516,133851895,115169878,107349540,102531392,90354753,81195210,78077248,59128983,63025520,48129895,51304566]

#     chrN = 21

    length = HiCsample.shape[0]

    prediction_1 = np.zeros((length, length))
    
    for i in range(index.shape[0]):
        x = int(index[i][0])
        y = int(index[i][1])
        prediction_1[x+6:x+34, y+6:y+34] = y_predict[i]
        prediction_1[x+6:x+34, y+6:y+34] = y_predict[i]
    
    return prediction_1
# -

# ## Apply HiCPlus to our datasets

DataBoosted={}
for keySp in KeysSpec:
    DataBoosted[keySp]={}
    print keySp
    for keySu in KeysSub:
        print '\t' + keySu
        DataBoosted[keySp][keySu] = HiCPlus(Data[keySp][keySu])


# # Display Images

# +
# keySp = 'ESC'
# keySu = '1'
figpath = 'Figures/'
if not os.path.exists(figpath):
    os.makedirs(figpath)
for keySp in KeysSpec:
    print keySp
    for keySu in KeysSub:
        print '\t' + keySu
        for Wm in [100,300,3000]:
            print '\t' + '\t' + str(Wm)
            W = [0,Wm]
            f,ax = plt.subplots(1,4,figsize=(20,5))
            ax[0].imshow(np.log10(Data[keySp][keySu])[W[0]:W[1],W[0]:W[1]],cmap='jet')
            ax[0].set_title('Data Subsampled')
            ax[1].imshow(np.log10(DataBoosted[keySp][keySu])[W[0]:W[1],W[0]:W[1]],cmap='jet')
            ax[1].set_title('Data Enhanced with HiCPlus')
            ax[2].imshow(np.log10(DataBoostHiC[keySp][keySu])[W[0]:W[1],W[0]:W[1]],cmap='jet')
            ax[2].set_title('Data Enhanced with BoostHiC')
            ax[3].imshow(np.log10(Data[keySp]['100'])[W[0]:W[1],W[0]:W[1]],cmap='jet')
            ax[3].set_title('Original Data')
            plt.savefig(figpath + 'CompareBoost_W'+ str(W[1]) + '_Spec' + keySp + '_Res' + keySu + '.png',dpi=300,format='png')
            
plt.show()
# -

# # Saving

for keySp in KeysSpec:
    path = data_path[keySp] + '/HiCPlus/'
    if not os.path.exists(path):
        os.makedirs(path)
        
    print keySp
    for keySu in KeysSub:
        print '\t' + keySu
        f = h5py.File(path + 'chr16_basic_' + keySu + '_HiCPlus.hdf5','w')
        f['data'] = DataBoosted[keySp][keySu]
        f.close()
        

# # Debugging

# +
HiCsample = Data[keySp][keySu]
## Subdivide the HiC matrix in subimages
low_resolution_samples, index = divide(HiCsample)

low_resolution_samples = np.minimum(HiC_max_value, low_resolution_samples)

batch_size = 32

## Reshape the high-quality Hi-C sample as the target value of the training.
sample_size = low_resolution_samples.shape[-1]
padding = conv2d1_filters_size + conv2d2_filters_size + conv2d3_filters_size - 3
half_padding = padding / 2
output_length = sample_size - padding

print low_resolution_samples.shape

## Data Loader
lowres_set = data.TensorDataset(torch.from_numpy(low_resolution_samples), torch.from_numpy(np.zeros(low_resolution_samples.shape[0])))
lowres_loader = torch.utils.data.DataLoader(lowres_set, batch_size=batch_size, shuffle=False)

hires_loader = lowres_loader

## Get Model
model = models.Net(40, 28)
model.load_state_dict(torch.load('model/pytorch_model_12000'))
if use_gpu:
    model = model.cuda()

_loss = nn.MSELoss()

## Make predictions
running_loss = 0.0
running_loss_validate = 0.0 # WARNING : NOT USED
reg_loss = 0.0 # WARNING : NOT USED
# -

ll = []
for i,l in enumerate(lowres_loader):
    l = Variable(l[0]).float()
    l = l.cuda()
    l = l.cpu()
    ll.append(l)
    
    if i==3:
        break

l = torch.cat(ll)

l[0]



StarcoderdataJupyterScriptsDedupFiltered