text:
# %%
from pathlib import Path
from pyCompare import blandAltman
import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.transforms as transforms
import matplotlib.ticker as ticker

test_path = "/Users/akshay/Desktop/test-results - test-results.csv"
external_path = "/Users/akshay/Desktop/external_v3_model_assisted_vs_model.csv"
prospective_path = (
    "/Users/akshay/Desktop/prospective_v3-fixed_model_assisted_vs_model.csv"
)

def dice_plot(x, y, title):

    figureSize = (10, 7)
    dpi = 72
    fig, ax = plt.subplots(figsize=figureSize, dpi=dpi)
    draw = True
    meanColour = "#6495ED"
    loaColour = "coral"
    pointColour = "#6495ED"

    ax.scatter(x, y, alpha=0.5, c=pointColour)

    mean = np.mean(y)
    sd = np.std(y)
    sd95 = 1.96 * np.std(y)
    trans = transforms.blended_transform_factory(ax.transAxes, ax.transData)

    offset = 0.045
    ax.text(
        0.97,
        mean - 2 * offset,
        f"Mean {mean:.2f}",
        ha="right",
        va="bottom",
        transform=trans,
    )
    ax.text(
        0.97,
        mean - 3 * offset,
        f"SD {sd:.2f}",
        ha="right",
        va="bottom",
        transform=trans,
    )
    ax.text(
        0.97,
        mean - 4 * offset,
        f"±1.96 SD in coral",
        ha="right",
        va="bottom",
        transform=trans,
    )

    ax.axhspan(mean - sd95, mean + sd95, facecolor=loaColour, alpha=0.2)

    # Hide the right and top spines
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)

    ax.set_xlabel("Reference htTKV (mL)")
    ax.set_ylabel("Dice Similarity Coefficient")
    ax.set_title(title, pad=10)
    ax.set_ylim([0, 1])
    ax.axhline(np.mean(y), color=meanColour, linestyle="--")

    ax.set_yticks(np.arange(0, 1.1, 0.1))

# %%
external_ds = pd.read_csv(external_path)
prospective_ds = pd.read_csv(prospective_path)
test_ds = pd.read_csv(test_path)
# %%
blandAltman(
    prospective_ds.TKV_GT,
    prospective_ds.TKV_Pred,
    percentage=True,
    title="Bland-Altman Plot - Prospective dataset",
)
# %%
blandAltman(
    external_ds.TKV_GT,
    external_ds.TKV_Pred,
    percentage=True,
    title="Bland-Altman Plot - External dataset",
)
# %%
blandAltman(
    test_ds.TKV_GT,
    test_ds.TKV_Pred,
    percentage=True,
    title="Bland-Altman Plot - Hold-out-test dataset",
)

# %%
dice_plot(
    prospective_ds.TKV_GT,
    prospective_ds.patient_dice,
    title="Dice by htTKV - Prospective dataset",
)

dice_plot(
    external_ds.TKV_GT,
    external_ds.patient_dice,
    title="Dice by htTKV - External dataset",
)

dice_plot(
    test_ds.TKV_GT,
    test_ds.patient_dice,
    title="Dice by htTKV - Hold-out-test dataset",
)

# %%

max_stars_repo_name: kurtteichman/adpkd-segmentation-pytorch

max_issues_repo_name: kurtteichman/adpkd-segmentation-pytorch

max_forks_repo_name: kurtteichman/adpkd-segmentation-pytorch


text:
import os, sys, math
sys.path.insert(0, '../../..')

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import scipy.stats

class Attention(nn.Module):
    def __init__(self, encoder_size, decoder_size, device, type="additive"):
        """ Attention module.
                TODO description for each type

                TODO needed bias=False for KVQ transformations, as well for type?

                TODO need mask?

            Args:
                encoder_size (int): Size of the encoder's output (as input for the decoder).
                decoder_size (int): Size of the decoder's output.
                device (torch.device): Device (eg. torch.device("cpu"))
                type (string): One of several types of attention

            See: https://arxiv.org/pdf/1902.02181.pdf

            Notes:
                Self-Attention(Intra-attention) Relating different positions of the same input sequence. Theoretically the self-attention can adopt any score functions above, but just replace the target sequence with the same input sequence.
                Global/Soft	Attending to the entire input state space.
                Local/Hard	Attending to the part of input state space; i.e. a patch of the input image.
        """
        super(Attention, self).__init__()
        self.encoder_size = encoder_size
        self.decoder_size = decoder_size
        self.type = type

        # transforms encoder states into keys
        self.key_annotation_function = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
        # transforms encoder states into values
        self.value_annotation_function = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
        # transforms the hidden state into query
        self.query_annotation_function = nn.Linear(self.decoder_size, self.encoder_size, bias=False) # NOTE: transforming q to K size

        if type == "additive":
            # f(q, K) = wimp*tanh(W1K + W2q + b) , Bahdanau et al., 2015
            self.V = nn.Linear(self.encoder_size, 1, bias=False)
            self.W1 = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
            self.W2 = nn.Linear(self.encoder_size, self.encoder_size, bias=False) # encoder size because q is now K's size, otherwise dec_size to enc_size
            self.b = nn.Parameter(torch.zeros(self.encoder_size))

        elif type == "coverage": # https://arxiv.org/pdf/1601.04811.pdf
            # f(q, K) = wimp*tanh(W1K + W2q + b) , Bahdanau et al., 2015
            self.V = nn.Linear(self.encoder_size, 1, bias=False)
            self.W1 = nn.Linear(self.encoder_size, self.encoder_size, bias=False)
            self.W2 = nn.Linear(self.encoder_size, self.encoder_size, bias=False) # encoder size because q is now K's size, otherwise dec_size to enc_size
            self.b = nn.Parameter(torch.zeros(self.encoder_size))

            self.coverage_dim = 10
            self.coverage_input_size = self.coverage_dim + 1 + self.encoder_size + self.encoder_size
            self.cov_gru = nn.GRU(self.coverage_input_size, self.coverage_dim, batch_first=True)
            self.W3 = nn.Linear(self.coverage_dim, self.encoder_size)

        elif (type == "multiplicative" or type == "dot"):
            # f(q, K) = q^t K , Luong et al., 2015
            # direct dot product, nothing to declare here
            pass

        elif type == "scaled multiplicative" or type == "scaled dot":
            # f(q, K) = multiplicative / sqrt(dk) , Vaswani et al., 2017
            self.scale = math.sqrt(self.encoder_size)

        elif type == "general" or type == "bilinear":
            # f(q, K) = q^t WK , Luong et al., 2015
            self.W = nn.Linear(self.encoder_size, self.encoder_size, bias=False)

        elif type == "biased general":
            # f(q, K) = K|(W q + b) Sordoni et al., 2016
            pass
        elif type == "activated general":
            # f(q, K) = act(q|WK + b) Ma et al., 2017
            pass
        elif type == "concat":
            # f(q, K) = act(W[K;q] + b) , Luong et al., 2015
            pass
        elif type == "p":
            # https://arxiv.org/pdf/1702.04521.pdf pagina 3, de adaugat predict-ul in attention, KVP si Q
            pass
        else:
            raise Exception("Attention type not properly defined! (got type={})".format(self.type))
        self.device = device
        self.to(self.device)

    def _reshape_state_h(self, state_h):
        """
        Reshapes the hidden state to desired shape
        Input: [num_layers * 1, batch_size, decoder_hidden_size]
        Output: [batch_size, 1, decoder_hidden_state]

        Args:
            state_h (tensor): Hidden state of the decoder.
                [num_layers * 1, batch_size, decoder_hidden_size]

        Returns:
            The reshaped hidden state.
                [batch_size, 1, decoder_hidden_state]
        """
        num_layers, batch_size, hidden_size = state_h.size()
        # in case the decoder has more than 1 layer, take only the last one -> [1, batch_size, decoder_hidden_size]
        if num_layers > 1:
            state_h = state_h[num_layers-1:num_layers,:,:]

        # [1, batch_size, decoder_hidden_size] -> [batch_size, 1, decoder_hidden_size]
        return state_h.permute(1, 0, 2)

    def reset_coverage(self, batch_size, enc_seq_len, force_weight): # called on every new batch
        self.force_weight = force_weight

        if self.type == "coverage":
            self.C = torch.zeros(batch_size, enc_seq_len, self.coverage_dim, device=self.device)
            self.gru_input = torch.zeros(batch_size, 1, self.coverage_input_size, device=self.device)

    def _coverage_compute_next_C(self, attention_weights, enc_output, state_h):
        # attention_weights:        [batch_size, seq_len, 1]
        # enc_output :              [batch_size, seq_len, encoder_size]
        # state_h (after reshape):  [batch_size, 1, encoder_size]
        # self.C at prev timestep:  [batch_size, seq_len, coverage_dim]
        seq_len = enc_output.size[1]
        for i in range(seq_len):
            self.gru_input[:,:,0:self.coverage_dim] = self.C[:,i:i+1,:] # cat self.C at prev timestep for position i of source word
            self.gru_input[:,:,self.coverage_dim:self.coverage_dim+1] = attention_weights[:,i:i+1,:] # cat attention weight
            self.gru_input[:,:,self.coverage_dim+1:self.coverage_dim+1+self.encoder_size] = enc_output[:,i:i+1,:] # cat encoder output
            self.gru_input[:,:,self.coverage_dim+1+self.encoder_size:] = state_h[:,0:1,:] # cat state_h
            self.C[:,i:i+1,:] = self.cov_gru(self.gru_input)

    def _energy (self, K, Q):
        """
            Calculates the compatibility function f(query, keys)

            Args:
                K (tensor): Keys tensor of size [batch_size, seq_len, encoder_size]
                Q (tensor): Query tensor of size [batch_size, 1, decoder_size], but now dec_size is enc_size due to Q annotation

            Returns:
                energy tensor of size [batch_size, seq_len, 1]
        """
        if self.type == "additive":
            return self.V(torch.tanh(self.W1(K) + self.W2(Q) + self.b))

        elif self.type == "coverage":
            return self.V(torch.tanh(self.W1(K) + self.W2(Q) + self.W3(self.C) + self.b))

        elif self.type == "multiplicative" or self.type == "dot":
            # q^t K means batch matrix multiplying K with q transposed:
            # bmm( [batch_size, seq_len, enc_size] , [batch_size, enc_size, 1] ) -> [batch_size, seq_len, 1]
            return torch.bmm(K, Q.transpose(1,2))

        elif self.type == "scaled multiplicative" or self.type == "scaled dot":
            # same as multiplicative but scaled
            return torch.bmm(K, Q.transpose(1,2))/self.scale

        elif self.type == "general" or self.type == "bilinear":
            # f(q, K) = q^t WK , Luong et al., 2015
            return torch.bmm(self.W(K), Q.transpose(1,2))

    def forward(self, enc_output, state_h, decoder_step, dec_seq_len, mask=None):
        """
        This function calculates the context vector of the attention layer, given the hidden state and the encoder
        last lstm layer output.

        Args:
            state_h (tensor): The raw hidden state of the decoder's LSTM
                Shape: [num_layers * 1, batch_size, decoder_size].
            enc_output (tensor): The output of the last LSTM encoder layer.
                Shape: [batch_size, seq_len, encoder_size].
            mask (tensor): 1 and 0 as for encoder input
                Shape: [batch_size, seq_len].

        Returns:
            context (tensor): The context vector. Shape: [batch_size, encoder_size]
            attention_weights (tensor): Attention weights. Shape: [batch_size, seq_len, 1]
        """
        batch_size = enc_output.shape[0]
        seq_len = enc_output.shape[1]
        state_h = self._reshape_state_h(state_h) # [batch_size, 1, decoder_size]

        # get K, V, Q
        K = self.key_annotation_function(enc_output) # [batch_size, seq_len, encoder_size]
        V = self.value_annotation_function(enc_output) # [batch_size, seq_len, encoder_size]
        Q = self.query_annotation_function(state_h) # [batch_size, 1, encoder_size]

        # calculate energy
        energy = self._energy(K,Q) # [batch_size, seq_len, 1]

        # mask with -inf paddings
        if mask is not None:
            energy.masked_fill_(mask.unsqueeze(-1) == 0, -np.inf)

        # transform energy into probability distribution using softmax
        attention_weights = torch.softmax(energy, dim=1) # [batch_size, seq_len, 1]

        # calculate gaussian distribution as forced probability, assume that decoder len is roughly the same as encoder len
        """
        if self.force_weight > 0.:
            approx_position = decoder_step*seq_len/dec_seq_len
            x = np.linspace(0, seq_len, seq_len)
            y = scipy.stats.norm.pdf(x, approx_position, 3) # loc (mean) is decoder_step, scale (std dev) = 1.
            #print(decoder_step)
            #print(np.sum(y))
            y = y / np.sum(y) # rescale to make it a PDF
            gaussian_dist = torch.tensor(y, dtype = attention_weights.dtype, device = self.device) # make it a tensor, it's [seq_len]
            gaussian_dist = gaussian_dist.repeat(batch_size, 1) # same for all examples in batch, now it's [batch_size, seq_len]

            # calculate attention_weights as a PDF
            attention_weights = self.force_weight * gaussian_dist + (1-self.force_weight) * attention_weights.squeeze(2)
            attention_weights = attention_weights.unsqueeze(2) # recreate [batch_size, seq_len, 1]
        """

        # for coverage only, calculate the next C
        if type=="coverage":
            self._coverage_compute_next_C(attention_weights, enc_output, state_h)

        # calculate weighted values z (element wise multiplication of energy * values)
        # attention_weights is [batch_size, seq_len, 1], V is [batch_size, seq_len, encoder_size], z is same as V
        z = attention_weights*V # same as torch.mul(), element wise multiplication

        # finally, calculate context as the esum of z.
        # z is [batch_size, seq_len, encoder_size], context will be [batch_size, encoder_size]
        context = torch.sum(z, dim=1)

        return context, attention_weights # [batch_size, encoder_size], [batch_size, seq_len, 1]

if __name__ == "__main__":
    import numpy as np

    # debug stuff:
    q = torch.tensor([ [ [2.,2.,2.] ] ])
    K = torch.tensor([ [ [1.,1.,1.] , [5.,5.,5.] ] ])
    #result would be ([ [ [2.,2.,2.] , [2.5,2.5,2.5] ] ])

    print(K.size())
    print(q)
    print(q.size())
    #qt = q.expand(1,3,3)#q.transpose(1,2)
    qt = q.transpose(1,2)
    print(qt)
    print(qt.size())

    print()
    r = torch.bmm(K,qt)
    print(r)
    print(r.size())
    print()
    #print(e1.size())
    #print(v1.size())
    #qq = e1*v1
    #print(qq)

    # prep inputs
    batch_size = 2
    seq_len = 10
    enc_size = 4
    dec_layers = 5
    dec_size = 3

    encoder_outputs = torch.tensor(np.random.rand(batch_size, seq_len, enc_size), dtype=torch.float)
    decoder_hidden_state = torch.tensor(np.random.rand(dec_layers*1, batch_size, dec_size), dtype=torch.float) # 1 for unidirectional

    # prep layer
    device = torch.device("cpu")
    #type = "additive"
    type = "general"
    att = Attention(enc_size, dec_size, device, type)

    # run
    context, attention_weights = att(encoder_outputs, decoder_hidden_state)
    print("Output is:")
    print(context)
    print("Attention weights size:" + str(attention_weights.size()))

    # debug stuff:
    #e1 = torch.tensor([[[2],[0.5]]])
    #v1 = torch.tensor([ [ [1.,1.,1.] , [5.,5.,5.] ] ])
    #result would be ([ [ [2.,2.,2.] , [2.5,2.5,2.5] ] ])
    #print(e1.size())
    #print(v1.size())
    #qq = e1*v1
    #print(qq)

ext: py

reason: import numpy,import scipy


text:
import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0, 2.0, 20)

plt.plot(x, x, 'b:')
plt.plot(x, x**1.5, 'rs-')
plt.plot(x, x**2, color='#30F050', marker='^', markerfacecolor='blue', markeredgecolor='blue')

plt.tight_layout()

plt.savefig("data/dst/matplotlib_example_single.png")

reason: import numpy


text:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

from .fusions.factory import factory as factory_fusion
from .mlp import MLP

class VRDNet(nn.Module):

    def __init__(self, opt):
        super(VRDNet, self).__init__()
        self.opt = opt
        self.classeme_embedding = nn.Embedding(
            self.opt['nb_classeme'],
            self.opt['classeme_dim'])
        self.fusion_c = factory_fusion(self.opt['classeme'])
        self.fusion_s = factory_fusion(self.opt['spatial'])
        self.fusion_f = factory_fusion(self.opt['feature'])
        self.predictor = MLP(**self.opt['predictor'])

    def forward(self, batch):
        bsize = batch['subject_boxes'].size(0)
        x_c = [self.classeme_embedding(batch['subject_cls_id']),
               self.classeme_embedding(batch['object_cls_id'])]
        x_s = [batch['subject_boxes'], batch['object_boxes']]
        x_f = [batch['subject_features'], batch['object_features']]

        x_c = self.fusion_c(x_c)
        x_s = self.fusion_s(x_s)
        x_f = self.fusion_f(x_f)

        x = torch.cat([x_c, x_s, x_f], -1)

        if 'aggreg_dropout' in self.opt:
            x = F.dropout(x, self.opt['aggreg_dropout'], training=self.training)
        y = self.predictor(x)

        out = {
            'rel_scores': y
        }
        return out


text:
import os

from .. import ConvergenceMap

from .. import dataExtern

import numpy as np
from astropy.units import deg,rad

import matplotlib.pyplot as plt

test_map = ConvergenceMap.load(os.path.join(dataExtern(),"conv.fit"))

#Set bin edges
l_edges = np.arange(200.0,50000.0,200.0)
thresholds_mf = np.arange(-2.0,2.0,0.2)
thresholds_pk = np.arange(-1.0,5.0,0.2)

def test_visualize():

	assert test_map.data.dtype == np.float

	test_map.setAngularUnits(deg)
	test_map.visualize()
	test_map.savefig("map.png")
	test_map.setAngularUnits(deg)

def test_save():
	test_map.save("conv_save.fits")

def test_power():

	#Compute
	l,Pl = test_map.powerSpectrum(l_edges)
	assert type(l)==np.ndarray
	assert type(Pl)==np.ndarray

	#Visualize
	fig,ax = plt.subplots()
	ax.plot(l,l*(l+1)*Pl/(2.0*np.pi))
	ax.set_xscale("log")
	ax.set_yscale("log")
	ax.set_xlabel(r"$l$")
	ax.set_ylabel(r"$l(l+1)P_l/2\pi$")

	plt.savefig("power_spectrum.png")
	plt.clf()

def test_cross():

	#Load
	conv1 = ConvergenceMap.load(os.path.join(dataExtern(),"conv1.fit"))
	conv2 = ConvergenceMap.load(os.path.join(dataExtern(),"conv2.fit"))

	#Cross
	l,Pl = conv1.cross(conv2,l_edges=l_edges)

	#Visualize
	fig,ax = plt.subplots()
	ax.plot(l,np.abs(l*(l+1)*Pl/(2.0*np.pi)))
	ax.set_xscale("log")
	ax.set_yscale("log")
	ax.set_xlabel(r"$l$")
	ax.set_ylabel(r"$l(l+1)P_l/2\pi$")

	plt.savefig("cross_spectrum.png")
	plt.clf()

def test_bispectrum():

	#Measure bispectrum
	l,be = test_map.bispectrum(l_edges,configuration="equilateral",scale=lambda x:(x**4)/(2.*np.pi)**2)
	l,bf = test_map.bispectrum(l_edges,configuration="folded",ratio=0.5,scale=lambda x:(x**4)/(2.*np.pi)**2)

	#Plot
	fig,ax = plt.subplots()
	ax.plot(l,np.abs(be),label=r"$\ell_1=\ell_2=\ell_3=\ell$")
	ax.plot(l,np.abs(bf),label=r"$\ell_1=\ell,\ell_2=\ell_3=\ell/2$")
	ax.set_xscale("log")
	ax.set_yscale("log")
	ax.set_xlabel(r"$\ell$")
	ax.set_ylabel(r"$(\ell_1\ell_2\ell_3)^{4/3}B(\ell_1,\ell_2,\ell_3)/(2\pi)^2$")
	ax.legend()

	plt.savefig("kappa_bispectrum.png")
	plt.clf()

def test_pdf():

	#Compute
	v,p = test_map.pdf(thresholds_mf,norm=True)

	#Visualize
	fig,ax = plt.subplots()
	ax.plot(v,p)
	ax.set_xlabel(r"$\nu=\kappa/\sigma$")
	ax.set_ylabel(r"$P(\nu)$")

	plt.savefig("pdf.png")
	plt.clf()

def test_minkowski():

	#Compute
	nu,V0,V1,V2 = test_map.minkowskiFunctionals(thresholds_mf,norm=True)

	#Assert computation went OK
	assert hasattr(test_map,"gradient_x")
	assert hasattr(test_map,"gradient_y")
	assert hasattr(test_map,"hessian_xx")
	assert hasattr(test_map,"hessian_yy")
	assert hasattr(test_map,"hessian_xy")

	#Visualize
	fig,ax = plt.subplots(1,3,figsize=(24,8))
	ax[0].plot(nu,V0)
	ax[1].plot(nu,V1)
	ax[2].plot(nu,V2)

	ax[0].set_xlabel(r"$\nu=\kappa/\sigma$")
	ax[0].set_ylabel(r"$V_0(\nu)$")

	ax[1].set_xlabel(r"$\nu=\kappa/\sigma$")
	ax[1].set_ylabel(r"$V_1(\nu)$")

	ax[2].set_xlabel(r"$\nu=\kappa/\sigma$")
	ax[2].set_ylabel(r"$V_2(\nu)$")

	fig.tight_layout()

	plt.savefig("minkowski.png")
	plt.clf()

def test_peaks():

	#Compute
	nu,pk = test_map.peakCount(thresholds_pk,norm=True)

	#Check if computation went OK
	assert type(nu)==np.ndarray
	assert type(pk)==np.ndarray

	#Visualize
	fig,ax = plt.subplots()
	ax.plot(nu,pk)
	ax.set_xlabel(r"$\nu=\kappa/\sigma$")
	ax.set_ylabel(r"$dN/d\nu$")

	plt.savefig("peaks.png")

def test_peak_locations():

	#Thresholds for high peaks
	high_thresholds = np.arange(0.3,0.6,0.01)

	#Find the peak locations
	values,locations = test_map.locatePeaks(high_thresholds)

	#Visualize the result
	fig,ax = plt.subplots(1,2,figsize=(16,8))
	test_map.visualize(fig=fig,ax=ax[0],colorbar=True)
	test_map.visualize(fig=fig,ax=ax[1])
	ax[1].scatter(locations[:,0].value,locations[:,1].value,color="black")
	ax[1].set_xlim(0.0,test_map.side_angle.value)
	ax[1].set_ylim(0.0,test_map.side_angle.value)

	#And save it
	fig.tight_layout()
	fig.savefig("peak_locations.png")

def test_getValues():

	b = np.linspace(0.0,test_map.side_angle.value,test_map.data.shape[0])
	xx,yy = np.meshgrid(b,b) * deg

	new_values = test_map.getValues(xx,yy)
	assert (new_values==test_map.data)[:-1,:-1].all()

def test_gradient_partial():

	b = np.linspace(0.0,test_map.side_angle.value,test_map.data.shape[0])
	xx,yy = np.meshgrid(b,b) * deg

	gx,gy = test_map.gradient()
	gxp,gyp = test_map.gradient(xx,yy)

	assert (gx==gxp)[:-1,:-1].all()
	assert (gy==gyp)[:-1,:-1].all()

def test_hessian_partial():

	b = np.linspace(0.0,test_map.side_angle.value,test_map.data.shape[0])
	xx,yy = np.meshgrid(b,b) * deg

	hxx,hyy,hxy = test_map.hessian()
	hxxp,hyyp,hxyp = test_map.hessian(xx,yy)

	assert (hxx==hxxp)[:-1,:-1].all()
	assert (hyy==hyyp)[:-1,:-1].all()
	assert (hxy==hxyp)[:-1,:-1].all()

def test_cut():

	b = np.array([0.0,test_map.side_angle.value/2,0.0,test_map.side_angle.value/2]) * deg
	cut_map = test_map.cutRegion(b)
	cut_map.visualize()
	cut_map.savefig("map_cut.png")

def test_translate():

	b = np.array([0.5,test_map.side_angle.value+0.5,0.0,test_map.side_angle.value]) * deg
	translated_map = test_map.cutRegion(b)
	translated_map.visualize()
	translated_map.savefig("map_translated.png")

max_stars_repo_name: asabyr/LensTools

max_issues_repo_name: asabyr/LensTools

max_forks_repo_name: asabyr/LensTools

reason: import numpy,from astropy


text:
import numpy as np
import pandas as pd
import os

from operator import itemgetter

from abc import ABCMeta, abstractmethod

from flow_equation_parser import FlowEquationParser

class IntermediateVectorManager:
    def __init__(self, couplings):
        self.couplings = couplings

        self.num_intermediate_vectors = 0
        self.intermediate_vectors = np.array([])
        self.global_iterator_map = dict()
        for coupling_index, key in enumerate(self.couplings):
            self.global_iterator_map[key] = "variables[" + str(coupling_index) + "]"
        self.num_couplings = len(self.couplings)

    def get_intermediate_vector(self):
        # Generate new intermediate vector
        if len(self.intermediate_vectors) == len(np.array(list(self.global_iterator_map.values()))) - self.num_couplings:
            self.intermediate_vectors = np.append(
                self.intermediate_vectors, "inter_med_vec" + str(self.num_intermediate_vectors))
            self.num_intermediate_vectors += 1
            return self.intermediate_vectors[-1]
        # Return existing intermediate vector
        else:
            used_intermediate_vectors = list(itemgetter(*self.get_actual_intermediate_vector_keys())(
                self.global_iterator_map))
            available_intermediate_vectors = [item for item in self.intermediate_vectors
                                              if item not in used_intermediate_vectors]
            return available_intermediate_vectors[-1]

    def get_actual_intermediate_vector_keys(self):
        keys = list(self.global_iterator_map.keys())
        return [item for item in keys if item not in self.couplings]

    def free_intermediate_vector(self, name):
        self.global_iterator_map.pop(name)

class ThrustMetaProgrammer:
    __metaclass__ = ABCMeta

    comp_functor_counter = 0

    def __init__(self, **kwargs):
        super(ThrustMetaProgrammer, self).__init__()

        self.theory_name = kwargs.pop("theory_name")
        self.dim = kwargs.pop("dim")

        self.base_struct_name = kwargs.pop("base_struct_name")

        self.class_name = ''.join([item.capitalize() for item in self.theory_name.split(sep="_")])

        self.intermediate_vector_manager = None
        self.couplings = None

    @abstractmethod
    def write_header(self):
        pass

    @abstractmethod
    def write_footer(self):
        pass

    def write_flow_equation(self, dim_index, flow_equation_parser):
        constant_expressions = flow_equation_parser.operation_tree_dataframe.query("type == 'constant expression'").value
        unique_constant_expressions = pd.unique(constant_expressions)
        self.couplings = np.array(flow_equation_parser.couplings)
        self.intermediate_vector_manager = IntermediateVectorManager(couplings=self.couplings)

        #  generate const expressions
        const_expr_map = dict()
        for idx, const_expr in enumerate(unique_constant_expressions):
            const_expr_map[const_expr] = 'const_expr' + str(idx)

        # generate transformations
        tree_root_index = flow_equation_parser.operation_tree_dataframe.query("depth == 0").iloc[0].name
        transformations, comp_functors = self.write_transformations(
            tree_index=tree_root_index,
            flow_equation_parser=flow_equation_parser,
            const_expr_map=const_expr_map)

        with open(self.filename, "a") as f:
            # write computation functors
            for comp_functor in reversed(comp_functors):
                f.write(comp_functor)

            # Struct name + constructor
            f.write('struct ' + self.class_name + self.base_struct_name + str(dim_index) + ' : public ' + self.base_struct_name + '\n'
                    '{\n'
                    '\t' + self.class_name + self.base_struct_name + str(dim_index) + '(const cudaT k_) : k(k_)')

            for idx, const_expr in enumerate(unique_constant_expressions):
                f.write(',\n\t\tconst_expr' + str(idx) + const_expr)
            f.write('\n\t{}\n\n')

            # void operator()
            f.write('\tvoid operator() (DimensionIteratorC &derivatives, const DevDatC &variables) override\n'
                    '\t{\n')

            # write initialization of intermediate vectors
            for intermediate_vector in self.intermediate_vector_manager.intermediate_vectors:
                # f.write('\t\tauto * ' + intermediate_vector + ' = new dev_vec(derivatives.size());\n')
                f.write('\t\tdev_vec ' + intermediate_vector + '(derivatives.size());\n')

            # write transformations
            for transformation in reversed(transformations):
                f.write(transformation)

            f.write('\t}\n\n')

            # private variables
            f.write('private:\n'
                    '\tconst cudaT k;\n')
            for idx, _ in enumerate(unique_constant_expressions):
                f.write('\n\tconst cudaT const_expr' + str(idx) + ';')
            f.write('\n};\n\n')        # replace values of constant expressions by corresponding variable name

    def write_transformations(self, tree_index, flow_equation_parser, const_expr_map):
        current_entry = flow_equation_parser.operation_tree_dataframe.loc[tree_index]

        # Refer to the variables it will be iterated over in the transformation
        iterators = current_entry.eff_terms
        # Keys refer the keys of the local_iterator_map,
        if current_entry.type == "coupling":
            # Identity
            coupling = current_entry.value[:-len(FlowEquationParser.coupling_appendix)]
            local_iterator_map = {coupling: "val"}
            keys = [coupling]
            iterators = [coupling]
        elif current_entry.type in ["numeric constant", "constant expression"]:
            local_iterator_map = None
            keys = None
            iterators = []
        else:
            # The local_iterator_map contains a dictionary for the assigment of keys for the repsective iterators
            # to variables in the comp_functors or in the lambda expressions
            local_iterator_map, keys = self.prepare_transformation(iterators)
        # Determine the output_iterators of the resulting transformation
        if current_entry.value is None or current_entry.type in ["coupling", "numeric constant", "constant expression"]:
            output_iterator = 'derivatives.begin()'
            output_end_iterator = 'derivatives.end()'
        else:
            output_iterator = self.intermediate_vector_manager.global_iterator_map[current_entry.value] + '.begin()'
            output_end_iterator = self.intermediate_vector_manager.global_iterator_map[current_entry.value] + '.end()'

        # Determine the return expression of the resulting lambda expression
        return_expression, const_expressions, open_tree_ids = self.write_return_expression(
            tree_index=tree_index,
            flow_equation_parser=flow_equation_parser,
            local_iterator_map=local_iterator_map,
            const_expr_map=const_expr_map,
            initial=True
        )

        if current_entry.type not in ["coupling", "numeric constant", "constant expression"]:
            # strip ()-braces
            return_expression = return_expression[1:-1] + ";"

        # ToDo: Check if this if condition has any implications on other outputs!!!
        if return_expression == "val":
            return_expression = return_expression + ";"

        # Determine attributes of the resulting lambda expression
        if current_entry.type in ["numeric constant", "constant expression"]:
            attributes = ''
        elif len(iterators) <= 2:
            attributes = 'const cudaT &' + ', const cudaT &'.join(local_iterator_map.values())
        else:
            attributes = 'Tuple t'

        # Determine captures of the resulting lambda expression
        if len(const_expressions) > 0:
            # const_expressions = [const_expr + "=this." + const_expr for const_expr in const_expressions]
            # captures = ', &'.join(const_expressions)
            # captures = '&' + captures
            # captures = 'this'
            captures = const_expressions
        else:
            captures = ''

        # Write the actual transformation
        transformation, comp_functor = self.get_and_write_transformation(len(iterators))(
            output_iterator=output_iterator,
            output_end_iterator=output_end_iterator,
            captures=captures,
            attributes=attributes,
            return_expression=return_expression,
            keys=keys,
        )

        if current_entry.value in self.intermediate_vector_manager.get_actual_intermediate_vector_keys():
            self.intermediate_vector_manager.free_intermediate_vector(current_entry.value)
        # remove evaluated keys from global iterator map

        transformations = [transformation]
        comp_functors = [comp_functor] if comp_functor is not None else []
        for open_tree_id in open_tree_ids:
            transformation, comp_functor = self.write_transformations(
                tree_index=open_tree_id,
                flow_equation_parser=flow_equation_parser,
                const_expr_map=const_expr_map
            )
            transformations += transformation
            comp_functors += comp_functor
        return transformations, comp_functors

        # f.write('\t\tthrust::transform(')
        # variables[' + str(dim_index) + '].begin(), '
        #        'variables[' + str(dim_index) + '].end(), '
        # f.write('derivatives.begin(), )[](const cudaT &vertex) { return vertex; });\n')

    def prepare_transformation(self, iterators):
        if len(iterators) <= 2:
            local_iterator_map = {key: "val" + str(idx + 1) for idx, key in enumerate(iterators)}
        else:
            local_iterator_map = {key: "thrust::get<" + str(idx) + ">(t)" for idx, key in enumerate(iterators)}
        for key in iterators:
            if key not in self.couplings:
                self.intermediate_vector_manager.global_iterator_map[key] = \
                    self.intermediate_vector_manager.get_intermediate_vector()
        # assert np.equal(self.intermediate_vector_manager.global_iterator_map.keys(), local_iterator_map.keys()),\
        #     "Keys do not coincide"
        return local_iterator_map, list(local_iterator_map.keys())

    @staticmethod
    def write_comp_functor(theory_name, captures, attributes, return_expression, return_index=None):
        if len(captures) == 0:
            string = str('\nstruct comp_func_' + theory_name + str(ThrustMetaProgrammer.comp_functor_counter) + '\n'
                         '{\n'
                         '\tcomp_func_' + theory_name + str(ThrustMetaProgrammer.comp_functor_counter) + '()\n'
                         '\t{}\n\n')
        else:
            string = str('\nstruct comp_func_' + theory_name + str(ThrustMetaProgrammer.comp_functor_counter) + '\n'
                         '{\n'
                         '\tconst cudaT ' + ';\n\tconst cudaT '.join(captures) + ';\n\n'
                         '\tcomp_func_' + theory_name + str(ThrustMetaProgrammer.comp_functor_counter) + '(const cudaT ' + '_, const cudaT '.join(captures) + '_)\n'
                         '\t\t: ' + ', '.join([const_expr + '(' + const_expr + '_)' for const_expr in captures]) + ' {}\n\n')
        if return_index is not None:
            string += '\ttemplate <typename Tuple>\n'
            string += str('\t__host__ __device__\n'
                          '\tvoid operator()(' + attributes + ')\n'
                          '\t{\n')
        else:
            string += str('\t__host__ __device__\n'
                          '\tcudaT operator()(' + attributes + ')\n'
                          '\t{\n')
        if return_index is not None:
            string += str('\t\tthrust::get<' + str(return_index) + '>(t) = ' + return_expression + '\n')
        else:
            string += str('\t\treturn ' + return_expression + '\n')
        string += str('\t}\n'
                      '};\n\n')
        comp_functor_name = 'comp_func_' + theory_name + str(ThrustMetaProgrammer.comp_functor_counter)
        ThrustMetaProgrammer.comp_functor_counter += 1
        return comp_functor_name, string

    def write_scalar_transformation(self, output_iterator, output_end_iterator, captures, attributes,
                                   return_expression, keys):
        comp_functor_string = None
        string = str('\t\tthrust::fill(thrust::device, ' + output_iterator + ", " + output_end_iterator + ", " + return_expression + ');\n')
        return string, comp_functor_string

    def write_unary_transformation(self, output_iterator, output_end_iterator, captures, attributes,
                                   return_expression, keys):
        # identity operator
        # f.write('\t\tthrust::transform(variables[' + str(dim_index) + '].begin(), '
        #         'variables[' + str(dim_index) + '].end(), '
        #         'derivatives.begin(), [](const cudaT &vertex) { return vertex; });\n')
        if len(captures) == 0:
            comp_functor_string = None
            string = str('\t\tthrust::transform(' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.begin(), ' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.end(), ' +
                         output_iterator +
                         ', [] __host__ __device__ (' + attributes + ') { return ' + return_expression + ' });\n')
        else:
            comp_functor_name, comp_functor_string = ThrustMetaProgrammer.write_comp_functor(self.theory_name, captures, attributes, return_expression)
            string = str('\t\tthrust::transform(' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.begin(), ' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.end(), ' +
                         output_iterator +
                         ', ' + comp_functor_name + '(' + ', '.join(captures) + '));\n')
        return string, comp_functor_string

    def write_binary_transformation(self, output_iterator, output_end_iterator, captures, attributes,
                                    return_expression, keys):
        if len(captures) == 0:
            comp_functor_string = None
            string = str('\t\tthrust::transform(' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.begin(), ' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.end(), ' +
                         self.intermediate_vector_manager.global_iterator_map[keys[1]] + '.begin(), ' +
                         output_iterator +
                         ', [] __host__ __device__ (' + attributes + ') { return ' + return_expression + ' });\n')
        else:
            comp_functor_name, comp_functor_string = ThrustMetaProgrammer.write_comp_functor(self.theory_name, captures, attributes, return_expression)
            string = str('\t\tthrust::transform(' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.begin(), ' +
                         self.intermediate_vector_manager.global_iterator_map[keys[0]] + '.end(), ' +
                         self.intermediate_vector_manager.global_iterator_map[keys[1]] + '.begin(), ' +
                         output_iterator +
                         ', ' + comp_functor_name + '(' + ', '.join(captures) + '));\n')
        return string, comp_functor_string

    def write_arbitrary_transformation(self, output_iterator, output_end_iterator, captures, attributes,
                                       return_expression, keys):
        comp_functor_name, comp_functor_string = ThrustMetaProgrammer.write_comp_functor(self.theory_name, captures, attributes, return_expression, len(keys))
        string = str('\t\tthrust::for_each(thrust::make_zip_iterator(thrust::make_tuple(')
        string += ('.begin(), '.join(self.intermediate_vector_manager.global_iterator_map[key] for key in keys) +
                   '.begin(), ' + output_iterator + ')),')
        string += 'thrust::make_zip_iterator(thrust::make_tuple('
        string += ('.end(), '.join(self.intermediate_vector_manager.global_iterator_map[key] for key in keys) +
                   '.end(), ' + output_end_iterator + '))'
                                                       ', ' + comp_functor_name + '(' + ', '.join(captures) + '));\n')
        return string, comp_functor_string

    def get_and_write_transformation(self, number_of_terms):
        if number_of_terms == 0:
            return self.write_scalar_transformation
        if number_of_terms == 1:
            return self.write_unary_transformation
        elif number_of_terms == 2:
            return self.write_binary_transformation
        else:
            return self.write_arbitrary_transformation

    def write_return_expression(self, tree_index, flow_equation_parser, local_iterator_map, const_expr_map, initial=False):
        const_expressions = []
        open_tree_ids = []

        current_entry = flow_equation_parser.operation_tree_dataframe.loc[tree_index]
        operation, child_ids, value, type = current_entry[["operation", "child_ids", "value", "type"]].values

        # Found leaf
        if value is not None:
            if type == "coupling":
                return local_iterator_map[value[:-len(FlowEquationParser.coupling_appendix)]], const_expressions, open_tree_ids
            elif type == "numeric constant":
                return value, const_expressions, open_tree_ids
            elif type == "constant expression":
                const_expressions.append(const_expr_map[value])
                return const_expr_map[value], const_expressions, open_tree_ids
            elif type == "intermediate expression" and value in local_iterator_map.keys():
                open_tree_ids.append(tree_index)
                return local_iterator_map[value], const_expressions, open_tree_ids

        args = []
        for child_id in child_ids:
            content, new_const_expressions, new_open_tree_ids = self.write_return_expression(
                tree_index=child_id,
                flow_equation_parser=flow_equation_parser,
                local_iterator_map=local_iterator_map,
                const_expr_map=const_expr_map,
                initial=False
            )
            args.append(content)
            const_expressions += new_const_expressions
            open_tree_ids += new_open_tree_ids
        return FlowEquationParser.plain_operations[operation](args), list(np.unique(const_expressions)), open_tree_ids

class FlowEquationMetaProgrammer(ThrustMetaProgrammer):

    def __init__(self, **kwargs):
        super(FlowEquationMetaProgrammer, self).__init__(**kwargs)

        self.custom = kwargs.pop("custom")
        if self.custom:
            self.filename = kwargs.get("project_path") + "/flow_equations/" +\
                            self.theory_name + "/" + self.theory_name + "_flow_equation.hpp"
            self.flow_equation_src_filename = kwargs.get("project_path") + "/flow_equations/flow_equation.cu"
            self.flow_equations_directory = kwargs.get("project_path") + "/flow_equations/"
        else:
            self.filename = "../flow_equations/" + self.theory_name + "_flow_equation.hpp"
            self.flow_equation_src_filename = None
            self.flow_equations_directory = kwargs.get("project_path") + "/flow_equations/"

        self.ode_solver_path = kwargs.pop("ode_solver_path")

        self.write_header()

    def write_header(self):
        ifndef_name = ''.join([item.upper() for item in self.theory_name.split(sep="_")])

        with open(self.filename, "w") as f:
            f.write('#ifndef PROJECT_' + ifndef_name + 'FLOWEQUATION_HPP\n'
                    '#define PROJECT_' + ifndef_name + 'FLOWEQUATION_HPP\n\n'
                    '#include <math.h>\n'
                    '#include <tuple>\n\n'
                    '#include "' + "../" + self.ode_solver_path + "/" +\
                    'thrust_code/include/flow_equation_interface/flow_equation.hpp"\n\n')

    def write_footer(self):
        ifndef_name = ''.join([item.upper() for item in self.theory_name.split(sep="_")])

        with open(self.filename, "a") as f:
            f.write('class ' + self.class_name + 'FlowEquations : public FlowEquationsWrapper\n'
                    '{\n'
                    'public:\n'
                    '\t' + self.class_name + 'FlowEquations(const cudaT k_) : k(k_)\n'
                    '\t{\n'
                    '\t\tflow_equations = std::vector< ' + self.base_struct_name + '* > {\n')
            for dim_index in range(self.dim-1):
                f.write('\t\t\tnew ' + self.class_name + self.base_struct_name + str(dim_index) + '(k),\n')
            f.write('\t\t\tnew ' + self.class_name + self.base_struct_name + str(self.dim-1) + '(k)\n\t\t};\n')
            f.write('\t}\n\n'
                    '\tvoid operator() (DimensionIteratorC &derivatives, const DevDatC &variables, '
                    'const int dim_index) override\n'
                    '\t{\n'
                    '\t\t(*flow_equations[dim_index])(derivatives, variables);\n'
                    '\t}\n\n'
                    '\tuint8_t get_dim() override\n'
                    '\t{\n'
                    '\t\treturn dim;\n'
                    '\t}\n\n'
                    '\tbool pre_installed_theory()\n'
                    '\t{\n')
            if self.custom:
                f.write('\t\treturn false;\n')
            else:
                f.write('\t\treturn true;\n')
            f.write('\t}\n\n'
                    '\tstatic std::string name()\n'
                    '\t{\n'
                    '\t\treturn "' + self.theory_name + '";\n'
                    '\t}\n\n'
                    '\tconst static uint8_t dim = ' + str(self.dim) + ';\n'
                    '\n'
                    'private:\n'
                    '\tconst cudaT k;\n'
                    '\tstd::vector < ' + self.base_struct_name + '* > flow_equations;\n'
                    '};\n\n'
                    '# endif //PROJECT_' + ifndef_name + 'FLOWEQUATION_HPP\n')

    def write_src_file(self):
        import os
        theories = [f.name for f in os.scandir(self.flow_equations_directory) if f.is_dir()]
        class_names = [''.join([item.capitalize() for item in theory.split(sep="_")]) for theory in theories]
        with open(self.flow_equation_src_filename, "w") as f:
            f.write('#include "' + self.ode_solver_path + "/" +\
                    'thrust_code/include/flow_equation_interface/flow_equation.hpp"\n\n')
            for theory in theories:
                f.write('#include "' + theory + "/" + theory + '_flow_equation.hpp"\n')
            f.write('\nFlowEquationsWrapper *FlowEquationsWrapper::make_flow_equation(const std::string theory)\n'
                    '{\n'
                    '\tif (theory == "' + theories[0] + '")\n'
                    '\t\treturn new ' + class_names[0] + 'FlowEquations(0);\n')
            for theory, class_name in zip(theories[1:], class_names[1:]):
                f.write('\telse if (theory == "' + theory + '")\n'
                        '\t\treturn new ' + class_name + 'FlowEquations(0);\n')
            f.write('\telse\n'
                    '\t{\n'
                    '\t\tstd:: cout << "ERROR: Flow equation not known" << std::endl;\n'
                    '\t\tstd::exit(EXIT_FAILURE);\n'
                    '\t}\n'
                    '}')

class JacobianEquationMetaProgrammer(ThrustMetaProgrammer):

    def __init__(self, **kwargs):
        super(JacobianEquationMetaProgrammer, self).__init__(**kwargs)

        self.custom = kwargs.pop("custom")
        if self.custom:
            self.filename = kwargs.get("project_path") + "/flow_equations/" +\
                            self.theory_name + "/" + self.theory_name + "_jacobian.hpp"
            self.jacobian_src_filename = kwargs.get("project_path") + "/flow_equations/jacobian.cu"
            self.jacobian_equations_directory = kwargs.get("project_path") + "/flow_equations/"
        else:
            self.filename = "../flow_equations/source/" + self.theory_name + "_jacobian.hpp"
            self.jacobian_src_filename = None
            self.jacobian_equations_directory = kwargs.get("project_path") + "/flow_equations/"

        self.ode_solver_path = kwargs.pop("ode_solver_path")

        self.write_header()

    def write_header(self):
        ifndef_name = ''.join([item.upper() for item in self.theory_name.split(sep="_")])

        with open(self.filename, "w") as f:
            f.write('#ifndef PROJECT_' + ifndef_name + 'JACOBIAN_HPP\n'
                    '#define PROJECT_' + ifndef_name + 'JACOBIAN_HPP\n\n'
                    '#include <math.h>\n'
                    '#include <tuple>\n\n'
                    '#include "' + "../" + self.ode_solver_path + "/" +\
                    'thrust_code/include/flow_equation_interface/jacobian_equation.hpp"\n\n')

    def write_footer(self):
        ifndef_name = ''.join([item.upper() for item in self.theory_name.split(sep="_")])

        with open(self.filename, "a") as f:
            f.write('class ' + self.class_name + 'JacobianEquations : public JacobianWrapper\n'
                    '{\n'
                    'public:\n'
                    '\t' + self.class_name + 'JacobianEquations(const cudaT k_) : k(k_)\n'
                    '\t{\n'
                    '\t\tjacobian_equations = std::vector< ' + self.base_struct_name + '* > {\n')
            for dim_index in range(pow(self.dim, 2)-1):
                f.write('\t\t\tnew ' + self.class_name + self.base_struct_name + str(dim_index) + '(k),\n')
            f.write('\t\t\tnew ' + self.class_name + self.base_struct_name + str(pow(self.dim, 2)-1) + '(k)\n\t\t};\n')
            f.write('\t}\n\n'
                    '\tvoid operator() (DimensionIteratorC &derivatives, const DevDatC &variables, '
                    'const int row_idx, const int col_idx) override\n'
                    '\t{\n'
                    '\t\t(*jacobian_equations[row_idx * dim + col_idx])(derivatives, variables);\n'
                    '\t}\n\n'
                    '\tvoid operator() (DimensionIteratorC &derivatives, const DevDatC &variables, '
                    'const int matrix_idx) override\n'
                    '\t{\n'
                    '\t\t(*jacobian_equations[matrix_idx])(derivatives, variables);\n'
                    '\t}\n\n'
                    '\tuint8_t get_dim() override\n'
                    '\t{\n'
                    '\t\treturn dim;\n'
                    '\t}\n\n'
                    '\tstatic std::string name()\n'
                    '\t{\n'
                    '\t\treturn "' + self.theory_name + '";\n'
                    '\t}\n\n'
                    '\tconst static uint8_t dim = ' + str(self.dim) + ';\n'
                    '\n'
                    'private:\n'
                    '\tconst cudaT k;\n'
                    '\tstd::vector < ' + self.base_struct_name + '* > jacobian_equations;\n'
                    '};\n\n'
                    '# endif //PROJECT_' + ifndef_name + 'JACOBIAN_HPP\n')

    def write_src_file(self):
        import os
        theories = [f.name for f in os.scandir(self.jacobian_equations_directory) if f.is_dir()]
        class_names = [''.join([item.capitalize() for item in theory.split(sep="_")]) for theory in theories]
        with open(self.jacobian_src_filename, "w") as f:
            f.write('#include "' + self.ode_solver_path + "/" +\
                    'thrust_code/include/flow_equation_interface/jacobian_equation.hpp"\n\n')
            for theory in theories:
                f.write('#include "' + theory + "/" + theory + '_jacobian.hpp"\n')
            f.write('\nJacobianWrapper *JacobianWrapper::make_jacobian(const std::string theory)\n'
                    '{\n'
                    '\tif (theory == "' + theories[0] + '")\n'
                    '\t\treturn new ' + class_names[0] + 'JacobianEquations(0);\n')
            for theory, class_name in zip(theories[1:], class_names[1:]):
                f.write('\telse if (theory == "' + theory + '")\n'
                        '\t\treturn new ' + class_name + 'JacobianEquations(0);\n')
            f.write('\telse\n'
                    '\t{\n'
                    '\t\tstd:: cout << "ERROR: Jacobian equation not known" << std::endl;\n'
                    '\t\tstd::exit(EXIT_FAILURE);\n'
                    '\t}\n'
                    '}')


text:
#!/home/a/anaconda3/envs/tinysleepnet/bin/python

#!/bin/env python

import argparse
import glob
import importlib
import os
import numpy as np
import shutil
import sklearn.metrics as skmetrics
import tensorflow as tf

''' todo:
    1) look around for a sleep dataset with electrodes similar to muse
    2) add additional data for each epoch:
        spectrogram?
        respiratory rate?
        spo2
        (noise, light, spo2, pulse...)
    3) generate a daily report notebook
    4) questions:
        -why is AF7 and AF8 poor at predicting rems
        -why am I not registering much deep sleep?
        -how well can TP9 and TP10 substitute for Cz and Pz
    '''

from data import load_data, get_subject_files
from model import TinySleepNet
from minibatching import (iterate_minibatches,
                          iterate_batch_seq_minibatches,
                          iterate_batch_multiple_seq_minibatches)
from utils import (get_balance_class_oversample,
                   print_n_samples_each_class,
                   save_seq_ids,
                   load_seq_ids)
from logger import get_logger

import scriptine

config = {
    # Train
    "n_epochs": 200,
    "learning_rate": 1e-4,
    "adam_beta_1": 0.9,
    "adam_beta_2": 0.999,
    "adam_epsilon": 1e-8,
    "clip_grad_value": 5.0,
    "evaluate_span": 50,
    "checkpoint_span": 50,

    # Early-stopping
    "no_improve_epochs": 50,

    # Model
    "model": "model-mod-8",
    "n_rnn_layers": 1,
    "n_rnn_units": 128,
    "sampling_rate": 100.0,
    "input_size": 3000,
    "n_classes": 5,
    "l2_weight_decay": 1e-3,

    # Dataset
    "dataset": "sleepedf",
    "data_dir": "./data/sleepedf/sleep-cassette/eeg_fpz_cz",
    "n_folds": 20,
    "n_subjects": 20,

    # Data Augmentation
    "augment_seq": True,
    "augment_signal_full": True,
    "weighted_cross_ent": True,
}

config.update({
    "seq_length": 20,
    "batch_size": 15,
})

config.update({
    "batch_size": 1,
    "seq_length": 1,
})
config["class_weights"] = np.ones(config["n_classes"], dtype=np.float32)

import pandas as pd

def getAccData(h5fname):
    x = pd.read_hdf(h5fname, 'Muse/ACC').reset_index(drop=True)
    x = x[x.ts.cummax().diff()>0]
    x['orientation'] = 180/np.pi*np.arctan(x.Y/x.X)
    x['motion'] = x.Y.rolling(1000).std().fillna(0)
    return x

def getPpgData(h5fname):
    #read the data
    x = pd.read_hdf(h5fname, 'Muse/PPG').reset_index(drop=True)
    x = x[x.ts.cummax().diff()>0]
    x.index = x.ts

    #label with heartbeat cycle (based on PPG1)
    x['PPG1m'] = x.PPG1.rolling(5).mean()
    isPeak = x['PPG1m'] == x['PPG1m'].rolling(int(.5*64), center=True).max()
    isPeak -= isPeak.rolling(20).max().shift(1).fillna(0)    #remove consecutive isPeak trues
    isPeak = isPeak.clip(0,1)
    x['cycle'] = isPeak.cumsum()

    #compute PPG1 stats on each cycle
    b = x['PPG1m'].groupby(x['cycle']).agg(['count','min','first','last','idxmin', ('idxfirst',lambda y: y.index.values[0]), ('idxlast',lambda y: y.index.values[-1])])
    b['max'] = b['first'] + (b['last']-b['first']) * (b['idxmin']-b['idxfirst']) / (b['idxlast']-b['idxfirst'])
    ppg = pd.DataFrame()
    ppg['ts'] = b['idxfirst']
    ppg['bpm'] = 1 / (b['count']/64/60)
    ppg['ampl1'] = (b['max']-b['min']) / (b['max']+b['min'])

    #compute PPG2 stats on each cycle
    x['PPG2m'] = x.PPG2.rolling(5).mean()
    b = x['PPG2m'].groupby(x['cycle']).agg(['count','min','first','last','idxmin', ('idxfirst',lambda y: y.index.values[0]), ('idxlast',lambda y: y.index.values[-1])])
    b['max'] = b['first'] + (b['last']-b['first']) * (b['idxmin']-b['idxfirst']) / (b['idxlast']-b['idxfirst'])
    ppg['ampl2'] = (b['max']-b['min']) / (b['max']+b['min'])
    ppg['spo2'] = 0
    return ppg

def predict_command(h5fname):
    # Add dummy class weights
    trues = []
    preds = []
    #tf.random.set_random_seed(123)

    model = TinySleepNet(
        config=config,
        output_dir='out_sleepedf/train/0',
        use_rnn=True,
        testing=True,
        use_best=True,
    )

    x = pd.read_hdf(h5fname, 'Muse/EEG').reset_index(drop=True)

    #resample to 100Hz
    x = x[x.ts==x.ts.cummax()]
    ts = np.arange(x.ts.values[0], x.ts.values[-1], 10**7)
    ts -= ts%1000
    x = x.set_index('ts').reindex(ts, method='nearest')

    #
    n_epochs = len(x)//3000
    x = x.iloc[:n_epochs*3000]
    p = pd.DataFrame()
    p['ts'] = x.index[::3000]
    for c in 'AF7 AF8 TP9 TP10'.split():
        night_x = x[c].values.reshape( (n_epochs,3000,1,1))
        print('night_x', night_x.shape, night_x.mean(), night_x.std())
        night_y = np.zeros(n_epochs)

        # Create minibatches for testing
        test_minibatch_fn = iterate_batch_multiple_seq_minibatches(
            [night_x],
            [night_y],
            batch_size=config["batch_size"],
            seq_length=config["seq_length"],
            shuffle_idx=None,
            augment_seq=False,
        )

        # Evaluate
        test_outs = model.evaluate(test_minibatch_fn)
        #from IPython import embed; embed()
        print('test_outs', c, np.array(test_outs['test/preds']).mean())
        p[c] = test_outs['test/preds']

        # Filter bad data
        isBad = night_x[:,:,0,0].std(axis=1) > 500
        print(c, isBad.sum(), p[c].mean(), p[c].std())
        p[c][isBad] = -1

    #acc data
    y = getAccData(h5fname)
    y = y.set_index('ts').reindex(p.ts, method='ffill')
    for c in ['orientation','motion']:
        p[c] = y[c].values

    #ppg data
    ppg = getPpgData(h5fname)

    #save results
    store = pd.HDFStore(h5fname, complib='zlib',complevel=5)
    store.put('tinysleepnet', p)
    store.put('spo2', ppg)

    #from IPython import embed; embed()
    tf.reset_default_graph()

if __name__ == "__main__":
    scriptine.run()


text:
# -*- coding: utf-8 -*-

# This code is part of Kaleidoscope.
#
# (C) Copyright IBM 2020.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""Interactive Qsphere"""

import numpy as np
import scipy.linalg as la
import scipy.special as spsp
import matplotlib as mpl
import colorcet as cc
from plotly.subplots import make_subplots
import plotly.graph_objects as go
from kaleidoscope.errors import KaleidoscopeError
from .plotly_wrapper import PlotlyWidget, PlotlyFigure
from .bloch.primitives import BSPHERE

def qsphere(state, state_labels=True,
            state_labels_kind='bits',
            as_widget=False):
    """Plots a statevector of qubits using the qsphere
    representation.

    Parameters:
        state (ndarray): Statevector as 1D NumPy array.
        state_labels (bool): Show state labels.
        state_labels_kind (str): 'bits' (default) or 'ints'.
        as_widget (bool): Return a widget instance.

    Returns:
        PlotlyFigure or PlotlyWidget: Figure instance.

    Raises:
        KaleidoscopeError: Invalid statevector input.

    Example:

        .. jupyter-execute::

            from qiskit import QuantumCircuit
            from qiskit.quantum_info import Statevector
            import kaleidoscope.qiskit
            from kaleidoscope.interactive import qsphere

            qc = QuantumCircuit(3)
            qc.h(range(3))
            qc.ch(0,1)
            qc.s(2)
            qc.cz(2,1)
            state = qc.statevector()

            qsphere(state)
    """
    if state.__class__.__name__ in ['Statevector'] \
            and 'qiskit' in state.__class__.__module__:
        state = state.data

    if state.__class__.__name__ in ['DensityMatrix'] \
            and 'qiskit' in state.__class__.__module__:

        if not abs(1-state.data.dot(state.data).trace().real) < 1e-6:
            raise KaleidoscopeError('Input density matrix is not a pure state.')
        # pylint: disable=unexpected-keyword-arg
        _, evecs = la.eig(state.data)
        state = evecs[0].ravel()

    if len(state.shape) == 2:
        if not abs(1-state.dot(state).trace().real) < 1e-6:
            raise KaleidoscopeError('Input density matrix is not a pure state.')
        # pylint: disable=unexpected-keyword-arg
        _, evecs = la.eig(state.data)
        state = evecs[0].ravel()

    if len(state.shape) != 1:
        raise KaleidoscopeError('Input state is not 1D array.')

    if np.log2(state.shape[0]) % 1:
        raise KaleidoscopeError('Input is not a valid statevector of qubits.')

    eps = 1e-8
    norm = mpl.colors.Normalize(vmin=0, vmax=2*np.pi)
    cmap = cc.cm.CET_C1s
    num_qubits = int(np.log2(state.shape[0]))

    xvals = []
    yvals = []
    zvals = []
    colors = []
    bases = []
    probs = []
    marker_sizes = []

    for idx in range(2**num_qubits):
        prob = (state[idx]*state[idx].conj()).real
        if prob > eps:
            elem = bin(idx)[2:].zfill(num_qubits)
            weight = elem.count("1")
            zvalue = -2 * weight / num_qubits + 1
            number_of_divisions = spsp.comb(num_qubits, weight)
            weight_order = _bit_string_index(elem)
            angle = (float(weight) / num_qubits) * (np.pi * 2) + \
                    (weight_order * 2 * (np.pi / number_of_divisions))

            if (weight > num_qubits / 2) or (((weight == num_qubits / 2) and
                                              (weight_order >= number_of_divisions / 2))):
                angle = np.pi - angle - (2 * np.pi / number_of_divisions)

            xvalue = np.sqrt(1 - zvalue ** 2) * np.cos(angle)
            yvalue = np.sqrt(1 - zvalue ** 2) * np.sin(angle)

            bases.append(elem)
            probs.append(prob)
            xvals.append(xvalue)
            yvals.append(yvalue)
            zvals.append(zvalue)

            phase = np.arctan2(state[idx].imag, state[idx].real)
            phase = phase if phase >= 0 else phase+2*np.pi
            colors.append(mpl.colors.rgb2hex(cmap(norm(phase))))
            marker_sizes.append(np.sqrt(prob) * 40)

    if state_labels_kind == 'ints':
        bases = [int(kk, 2) for kk in bases]

    # Output figure instance
    fig = make_subplots(rows=5, cols=5,
                        specs=[[{"type": "scene", "rowspan": 5, "colspan": 5},
                                None, None, None, None],
                               [None, None, None, None, None],
                               [None, None, None, None, None],
                               [None, None, None, None, None],
                               [None, None, None, None,
                                {"rowspan": 1, "colspan": 1, "type": "domain"}
                                ]
                               ]
                        )

    figsize = (350, 350)

    # List for vector annotations, if any
    fig_annotations = []

    fig.add_trace(BSPHERE(), row=1, col=1)

    # latitudes
    for kk in _qsphere_latitudes(zvals):
        fig.add_trace(kk, row=1, col=1)

    fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0],
                               mode='markers',
                               opacity=0.6,
                               marker=dict(size=4,
                                           color='#555555'),
                               ),
                  row=1, col=1)

    for kk, _ in enumerate(xvals):
        fig.add_trace(go.Scatter3d(x=[0, xvals[kk]], y=[0, yvals[kk]], z=[0, zvals[kk]],
                                   mode="lines",
                                   hoverinfo=None,
                                   opacity=0.5,
                                   line=dict(color=colors[kk], width=3)
                                   ),
                      row=1, col=1
                      )

        if state_labels:
            xanc = 'center'
            if xvals[kk] != 0:
                if xvals[kk] < 0:
                    xanc = 'right'
                else:
                    pass

            yanc = 'middle'
            if zvals[kk] != 0:
                if zvals[kk] < 0:
                    yanc = 'top'
                else:
                    yanc = 'bottom'

            fig_annotations.append(dict(showarrow=False,
                                        x=xvals[kk]*1.1,
                                        y=yvals[kk]*1.1,
                                        z=zvals[kk]*1.1,
                                        text="<b>|{}\u3009</b>".format(bases[kk]),
                                        align='left',
                                        opacity=0.7,
                                        xanchor=xanc,
                                        yanchor=yanc,
                                        xshift=10,
                                        bgcolor="#ffffff",
                                        font=dict(size=10,
                                                  color="#000000",
                                                  ),
                                        )
                                   )

    fig.add_trace(go.Scatter3d(x=xvals, y=yvals, z=zvals,
                               mode='markers',
                               opacity=1,
                               marker=dict(size=marker_sizes,
                                           color=colors),
                               ),
                  row=1, col=1)

    slices = 128
    labels = ['']*slices
    values = [1]*slices

    phase_colors = [mpl.colors.rgb2hex(cmap(norm(2*np.pi*kk/slices))) for kk in range(slices)]

    fig.add_trace(go.Pie(labels=labels, values=values, hole=.6,
                         showlegend=False,
                         textinfo='none',
                         hoverinfo='none',
                         textposition="outside",
                         rotation=90,
                         textfont_size=12,
                         marker=dict(colors=phase_colors)
                         ),
                  row=5, col=5)

    pie_x = fig.data[-1]['domain']['x']
    pie_y = fig.data[-1]['domain']['y']

    fig['layout'].update(annotations=[
        dict(
            xref='paper',
            yref='paper',
            x=(pie_x[1]-pie_x[0])/2+pie_x[0],
            y=(pie_y[1]-pie_y[0])/2+pie_y[0],
            text='Phase',
            xanchor="center",
            yanchor="middle",
            showarrow=False,
            font=dict(size=9),
        ),
        dict(
            xref='paper',
            yref='paper',
            x=pie_x[0]-0.03,
            y=(pie_y[1]-pie_y[0])/2+pie_y[0],
            text='\U0001D70B',
            xanchor="left",
            yanchor="middle",
            showarrow=False,
            font=dict(size=14),
        ),
        dict(
            xref='paper',
            yref='paper',
            x=pie_x[1]+0.03,
            y=(pie_y[1]-pie_y[0])/2+pie_y[0],
            text='0',
            xanchor="right",
            yanchor="middle",
            showarrow=False,
            font=dict(size=12),
        ),
        dict(
            xref='paper',
            yref='paper',
            x=(pie_x[1]-pie_x[0])/2+pie_x[0],
            y=pie_y[1]+0.05,
            text='\U0001D70B/2',
            xanchor="center",
            yanchor="top",
            showarrow=False,
            font=dict(size=12),
        ),
        dict(
            xref='paper',
            yref='paper',
            x=(pie_x[1]-pie_x[0])/2+pie_x[0],
            y=pie_y[0]-0.05,
            text='3\U0001D70B/2',
            xanchor="center",
            yanchor="bottom",
            showarrow=False,
            font=dict(size=12),
        )
    ])

    fig.update_layout(width=figsize[0],
                      height=figsize[1],
                      autosize=False,
                      hoverdistance=50,
                      showlegend=False,
                      scene_aspectmode='cube',
                      margin=dict(r=15, b=15, l=15, t=15),
                      scene=dict(annotations=fig_annotations,
                                 xaxis=dict(showbackground=False,
                                            range=[-1.2, 1.2],
                                            showspikes=False,
                                            visible=False),
                                 yaxis=dict(showbackground=False,
                                            range=[-1.2, 1.2],
                                            showspikes=False,
                                            visible=False),
                                 zaxis=dict(showbackground=False,
                                            range=[-1.2, 1.2],
                                            showspikes=False,
                                            visible=False)),
                      scene_camera=dict(eye=dict(x=0,
                                                 y=-1.4,
                                                 z=0.3)
                                        )
                      )

    if as_widget:
        return PlotlyWidget(fig)

    return PlotlyFigure(fig, modebar=True)

def _lex_index(n, k, lst):
    """Return  the lex index of a combination..

    Args:
        n (int): the total number of options .
        k (int): The number of elements.
        lst (list): list

    Returns:
        int: returns int index for lex order

    Raises:
        KaleidoscopeError: if length of list is not equal to k
    """
    if len(lst) != k:
        raise KaleidoscopeError("list should have length k")
    comb = list(map(lambda x: n - 1 - x, lst))
    dualm = sum([spsp.comb(comb[k - 1 - i], i + 1) for i in range(k)])
    return int(dualm)

def _bit_string_index(s):
    """Return the index of a string of 0s and 1s.

    Parameters:
        s (str): Bitstring.

    Returns:
        int: Index.

    Raises:
        KaleidoscopeError: If string is not binary.
    """
    n = len(s)
    k = s.count("1")
    if s.count("0") != n - k:
        raise KaleidoscopeError("s must be a string of 0 and 1")
    ones = [pos for pos, char in enumerate(s) if char == "1"]
    return _lex_index(n, k, ones)

def _qsphere_latitudes(zvals):
    """Latitude lines for sphere.

    Parameters:
        zvals (int): Input zvals

    Returns:
        list: List of Plotly traces.
    """
    lats = []
    u = np.linspace(0, 2*np.pi, 100)

    if 0 not in zvals:
        zvals = [0] + zvals
    for zv in zvals:
        th = np.arctan2(np.sqrt(1 - zv ** 2), zv)
        xvals = np.sin(th)*np.cos(u)
        yvals = np.sin(th)*np.sin(u)
        zvals = zv*np.ones_like(u)
        lats.append(go.Scatter3d(
            x=xvals, y=yvals, z=zvals,
            mode="lines",
            hoverinfo='skip',
            line=dict(
                color='#1e1e1e' if th == np.pi/2 else '#373737',
                width=2 if th == np.pi/2 else 1
            )
        ))
    return lats

max_stars_repo_name: QuSTaR/kaleidoscope

max_issues_repo_name: nonhermitian/kaleidoscope

max_forks_repo_name: QuSTaR/kaleidoscope


text:
import os
import mmap
from glob import glob
import itertools
import numpy as np
import argparse as ap
import concurrent.futures as cf
import time
from queue import Queue as Queue
import torch.cuda.nvtx as nvtx
import copy

from mpi4py import MPI
from mpi4py.util import dtlib

import torch
import io_helpers as ioh

def numpy_integrity_check(src_dir, dst_dir, files):
    checklist = []
    issue_found = False
    for fname in files:
        src_arr = np.load(os.path.join(src_dir, fname))
        dst_arr = np.load(os.path.join(dst_dir, fname))
        compare = np.allclose(src_arr, dst_arr,
                              rtol=1e-07,
                              atol=1e-08)
        src_nan = np.any(np.isnan(src_arr))
        dst_nan = np.any(np.isnan(dst_arr))

        issue_found = issue_found or not compare or src_nan or dst_nan

        checklist.append({"file": fname, "equal": compare, "src_nan": src_nan, "dst_nan": dst_nan})

    return checklist, issue_found

def get_shard(files, num_shards, shard_id, cycle_dist=0):
    num_files = len(files)

    # shard files into bulk and remainder:
    num_files_per_shard = num_files // num_shards
    files_bulk = files[:num_files_per_shard * num_shards]
    files_remainder = files[num_files_per_shard * num_shards:]

    # get my shard
    shard_start = shard_id * num_files_per_shard
    shard_end = shard_start + num_files_per_shard
    files_shard = files_bulk[shard_start:shard_end].copy()

    # deal with remainder: round robin with some offset for even distribution
    cycle_offset = 0
    for idf, fname in enumerate(files_remainder):
        if ((idf + cycle_offset) % num_shards == shard_id):
            files_shard.append(fname)
        cycle_offset += cycle_dist

    return files_shard

def allgather_safe(comm, fdata):
    #total size
    comm_size = comm.Get_size()
    num_bytes = len(fdata)
    total_bytes = num_bytes * comm_size

    #chunk by ~1GB:
    gigabyte = 1024*1024*1024

    # determine number of chunks
    num_chunks = (total_bytes + gigabyte - 1) // gigabyte

    # determine local chunksize
    chunksize = (num_bytes + num_chunks - 1) // num_chunks

    # datatype stuff
    datatype = MPI.BYTE
    np_dtype = dtlib.to_numpy_dtype(datatype)

    # gather stuff
    if True:
        # prepare buffers:
        sendbuff = np.frombuffer(memoryview(fdata), dtype=np_dtype, count=num_bytes)
        recvbuff = np.empty((comm_size * chunksize), dtype=np_dtype)
        resultbuffs = np.split(np.empty(num_bytes * comm_size, dtype=np_dtype), comm_size)

        # do subsequent gathers
        for i in range(0, num_chunks):
            # create buffer views
            start = i * chunksize
            end = min(start + chunksize, num_bytes)
            eff_bytes = end - start
            sendbuffv = sendbuff[start:end]
            recvbuffv = recvbuff[0:eff_bytes*comm_size]

            # perform allgather on views
            comm.Allgather([sendbuffv, datatype], [recvbuffv, datatype])

            # split result buffer for easier processing
            recvbuff_split = np.split(recvbuffv, comm_size)
            for j in range(comm_size):
                resultbuffs[j][start:end] = recvbuff_split[j][...]
        results = [x.tobytes() for x in resultbuffs]
    else:
        recvbuff = np.empty((total_bytes), dtype=np_dtype)
        for i in range(0, num_chunks):
            # prepare local chunks
            local_start = i * chunksize
            local_end = min(local_start + chunksize, num_bytes)
            eff_bytes = local_end - local_start

            # set up send buffer and recv buffer specv
            sendbuff = np.frombuffer(memoryview(fdata[local_start:local_end]), dtype=np_dtype, count=eff_bytes)
            counts = [eff_bytes for _ in range(comm_size)]
            recv_displacements = [local_start + j * num_bytes for j in range(comm_size)]

            # perform the gather
            comm.Allgatherv([sendbuff, datatype], [recvbuff, (counts, recv_displacements), datatype])

        # create the output vector
        recvbuff_split = np.split(recvbuff, comm_size)
        results = [x.tobytes() for x in recvbuff_split]

    return results

def load_file(filename):
    with open(filename, "rb") as f:
        token = f.read()

    return filename, token, len(token)

def load_file_direct(filename, filesize=None, blocksize=None):
    io_blocksize = 4096 if blocksize is None else blocksize
    token = ioh.load_file_direct(filename, blocksize=io_blocksize, filesize=0 if filesize is None else filesize)

    return filename, token, len(token)

def save_file(ofname, fdata):
    with open(ofname, "wb") as f:
        f.write(fdata)
    return len(fdata)

def save_file_direct(ofname, fdata, blocksize=512, sync=True):
    wbytes = ioh.save_file_direct(ofname, fdata, blocksize, sync)
    return wbytes

class FileStager(object):

    def _load_batch_parallel(self, source_directory, files, num_shards, shard_id, filesize=None, blocksize=None):

        # profiling tag
        nvtx.range_push("load_batch_parallel")

        # split the data across ranks
        if num_shards > 1:
            files_load = get_shard(files, num_shards, shard_id, cycle_dist=0)
        else:
            files_load = files

        # submit loads:
        queue = Queue()
        for filename in files_load:
            fname = os.path.join(source_directory, filename)
            if self.use_direct_read:
                queue.put(self.executor.submit(load_file_direct, fname, filesize, blocksize))
            else:
                queue.put(self.executor.submit(load_file, fname))
        nvtx.range_pop()

        return queue

    def _finalize_save_local(self, save_queue):
        nvtx.range_push("finalize_save_local")
        wbytes = 0
        while not save_queue.empty():
            handle = save_queue.get()
            if handle.done():
                fbytes = handle.result()
                wbytes += fbytes
            else:
                save_queue.put(handle)
        nvtx.range_pop()

        return wbytes

    # this uses the data load queue to communicate the loaded files across ranks and saves them:
    def _distribute_save_batch_parallel(self, comm, load_queue, target_directory):

        # profile region
        nvtx.range_push("distribute_save_batch_parallel")

        # we need these
        rbytes = 0
        save_queue = Queue()

        # iterate till queue is empty
        while not load_queue.empty():

            # get handle
            handle = load_queue.get()

            # if load is not done, requeue
            if not handle.done():
                load_queue.put(handle)
                continue

            # load is done, extract data
            fname, fdata, fbytes = handle.result()

            # communicate if necessary and store
            if comm is not None:
                fname_all = comm.allgather(fname)

                # we need to do that to check integrity
                if self.extended_verify:
                    from hashlib import blake2b
                    checksum = blake2b(fdata).hexdigest()
                    checksums_io = comm.allgather(checksum)

                # issue allgather
                fdata_all = allgather_safe(comm, fdata)

                # compare checksums
                if self.extended_verify:
                    checksums_ag = [blake2b(x).hexdigest() for x in fdata_all]
                    if checksums_io != checksums_ag:
                        print(checksums_io, checksums_ag)
                    assert(checksums_io == checksums_ag)

                for fname, fdata in zip(fname_all, fdata_all):
                    ofn = os.path.join(target_directory, os.path.basename(fname))
                    if self.use_direct_write:
                        save_queue.put(self.executor.submit(save_file_direct, ofn, copy.copy(fdata)))
                    else:
                        save_queue.put(self.executor.submit(save_file, ofn, copy.copy(fdata)))
                    rbytes += len(fdata)
            else:
                ofn = os.path.join(target_directory, os.path.basename(fname))
                if self.use_direct_write:
                    save_queue.put(self.executor.submit(save_file_direct, ofn, copy.copy(fdata)))
                else:
                    save_queue.put(self.executor.submit(save_file, ofn, copy.copy(fdata)))
                rbytes += len(fdata)

        nvtx.range_pop()

        return save_queue, rbytes

    def _stage_batch_data(self, shardinfo,
                          source_directory, target_directory,
                          files, filesize, blocksize):

        # profile region
        nvtx.range_push("stage_batch_data")

        # extract some shard info
        start = shardinfo["start"]
        end = shardinfo["end"]
        files_shard = files[start:end]

        # queue up the loads
        load_queue = self._load_batch_parallel(source_directory,
                                               files_shard,
                                               shardinfo["num_shards"],
                                               shardinfo["shard_id"],
                                               filesize = filesize,
                                               blocksize = blocksize)

        # batch loop: go one step further to process the last batch
        comm = None if (shardinfo["num_shards"] == 1) else self.stage_comm
        save_queue, rbytes = self._distribute_save_batch_parallel(comm, load_queue, target_directory)

        # wait for all the stores to complete
        wbytes = self._finalize_save_local(save_queue)

        # close region
        nvtx.range_pop()

        return rbytes, wbytes

    def __init__(self,
                 global_comm,
                 num_instances,
                 instance_id,
                 instance_comm,
                 local_size,
                 local_rank,
                 batch_size = -1,
                 num_workers = 1,
                 stage_mode="global",
                 verify=False,
                 full_dataset_per_node=True,
                 use_direct_io=False,
                 seed=333):

        # global chunking parameters
        self.global_comm = global_comm
        self.num_instances = num_instances
        self.instance_id = instance_id
        self.instance_comm = instance_comm
        self.local_size = local_size
        self.local_rank = local_rank

        # stage optimization info
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.stage_mode = stage_mode
        self.full_dataset_per_node = full_dataset_per_node
        self.use_direct_read = use_direct_io
        self.use_direct_write = use_direct_io
        self.seed = seed

        # debug
        self.verify = verify
        # warning, this is slow!
        self.extended_verify = False

        # extract comm info
        self.gsize = self.global_comm.Get_size()
        self.grank = self.global_comm.Get_rank()
        self.isize = self.instance_comm.Get_size()
        self.irank = self.instance_comm.Get_rank()
        self.lsize = self.local_size
        self.lrank = self.local_rank

        # create new helper comms
        self.stage_comm = self.global_comm.Split(color=self.irank, key=self.instance_id)
        self.ssize = self.stage_comm.Get_size()
        self.srank = self.stage_comm.Get_rank()
        # split the instance by nodes and create a comm with all matching local ranks by node
        self.num_nodes_per_instance = self.isize // self.lsize
        self.instance_node_id = self.irank // self.lsize
        self.instance_node_comm = self.instance_comm.Split(color=self.lrank, key=self.instance_node_id)
        # get a local communicator too
        self.local_comm = self.instance_comm.Split(color=(self.irank // self.lsize), key=self.lrank)

        # create stage executor
        self.executor = cf.ThreadPoolExecutor(max_workers = self.num_workers)

    # helper to prepare staging for the instance
    def _prepare_instance_stage(self, files, target_directory):
        if (self.stage_mode == "node") and self.full_dataset_per_node:
            files_shard = get_shard(files, self.lsize, self.lrank)
        else:
            # here we need to make sure the data is evenly distributed
            # across nodes, otherwise one node might have longer epochs
            files_shard = get_shard(files, self.isize, self.irank, cycle_dist=self.lsize)

        if self.full_dataset_per_node:
            files_print = files
        else:
            files_print = self.local_comm.allgather(files_shard)
            files_print = list(itertools.chain(*files_print))

        # create tags
        tag = os.path.basename(files[0]).split("-")[0]
        fname = os.path.join(target_directory, f"files_{tag}.lst")
        if self.lrank == 0:
            with open(fname, "w") as f:
                f.write("\n".join(files_print))

        return fname

    # prepare staging the instance
    def prepare(self, data_dir_prefix, stage_dir_prefix, stage_filter_list):

        # append instance ID to target dir
        target_directory = os.path.join(stage_dir_prefix, f"instance{self.instance_id}")

        if self.grank == 0:
            print("Copying stats.h5", flush=True)
            with open(os.path.join(data_dir_prefix, "stats.h5"), "rb") as f:
                statsfile = f.read()
        else:
            statsfile = None

        # broadcast the statsfile
        statsfile = self.global_comm.bcast(statsfile, 0)

        # save it
        if self.lrank == 0:
            os.makedirs(target_directory, exist_ok = True)
            with open(os.path.join(target_directory, "stats.h5"), "wb") as f:
                f.write(statsfile)

        # iterate over staging filters
        self.file_stats = {}
        for stage_filter in stage_filter_list:

            nvtx.range_push(f"stage {stage_filter}")

            if (self.grank == 0):
                print(f"Preparing file lists for {stage_filter}", flush=True)

            # get directories
            stage_source_directory = os.path.join(data_dir_prefix, os.path.dirname(stage_filter))
            stage_target_directory = os.path.join(target_directory, os.path.dirname(stage_filter))

            # create target directory if not exist:
            if self.local_rank == 0:
                os.makedirs(stage_target_directory, exist_ok = True)

            # get file info to everybody
            if self.grank == 0:
                allfiles = sorted([os.path.basename(x) for x in glob(os.path.join(stage_source_directory, os.path.basename(stage_filter)))])
            else:
                allfiles = None

            # shuffle files if requested
            if (self.grank == 0) and (not self.full_dataset_per_node) and (self.seed is not None):
                rng = np.random.default_rng(self.seed)
                rng.shuffle(allfiles)

            # communicate list of files
            allfiles = self.global_comm.bcast(allfiles, 0)

            # now stage the data so that each rank in each instance has the relevant data
            stage_start = time.perf_counter()
            list_file = self._prepare_instance_stage(allfiles, stage_target_directory)
            stage_stop = time.perf_counter()

            # updating file stats buffer
            self.file_stats[stage_filter] = {"num_files": len(allfiles),
                                             "source_directory": stage_source_directory,
                                             "target_directory": stage_target_directory,
                                             "list_file": list_file}

        return

    def _stage_instance_data(self, stage_filter):

        # comm parameters
        num_files = self.file_stats[stage_filter]["num_files"]
        source_directory = self.file_stats[stage_filter]["source_directory"]
        target_directory = self.file_stats[stage_filter]["target_directory"]

        # this is the tag of the file list
        fname = self.file_stats[stage_filter]["list_file"]

        # load metadata
        with open(fname, "r") as f:
            files_shard = f.read().splitlines()

        # shard locally
        files_shard = get_shard(files_shard, self.lsize, self.lrank)

        # get file stats:
        stats = os.stat(os.path.join(source_directory, files_shard[0]))
        filesize = stats.st_size
        blocksize = stats.st_blksize
        self.file_stats[stage_filter]["filesize"] = filesize
        self.file_stats[stage_filter]["blocksize"] = blocksize

        # now, let's take care of the data: update the number of files because of remainder
        batch_size = self.batch_size
        if self.stage_mode == "global":
            batch_size *= self.ssize
        num_files_per_shard = len(files_shard)
        num_batches_bulk = num_files_per_shard // batch_size
        num_files_remainder = num_files_per_shard - num_batches_bulk * batch_size

        # create list of batch sizes, shard sizes, etc:
        stage_info = [{"start": i*batch_size,
                       "end": (i+1)*batch_size,
                       "num_shards": self.ssize if self.stage_mode == "global" else 1,
                       "shard_id": self.srank if self.stage_mode == "global" else 0}
                      for i in range(0, num_batches_bulk)]

        # deal with the remainder:
        remainder_start = num_batches_bulk * batch_size
        if self.stage_mode == "global":
            # see if we can squeeze in one more batch with reduced size
            eff_batchsize = (num_files_remainder // self.ssize) * self.ssize

            if eff_batchsize > 0:
                stage_info.append({"start": remainder_start,
                                   "end": remainder_start + eff_batchsize,
                                   "num_shards": self.ssize,
                                   "shard_id": self.srank})

            remainder_start += eff_batchsize

        # remainder:
        if (num_files_per_shard - remainder_start > 0):
            stage_info.append({"start": remainder_start,
                               "end": num_files_per_shard,
                               "num_shards": 1,
                               "shard_id": 0})

        # do the staging
        total_bytes_read = 0
        total_bytes_write = 0
        for shardinfo in stage_info:
            rbytes, wbytes = self._stage_batch_data(shardinfo,
                                                    source_directory,
                                                    target_directory,
                                                    files_shard,
                                                    filesize, blocksize)
            total_bytes_read += rbytes
            total_bytes_write += wbytes

        # global barrier
        self.instance_comm.Barrier()

        return total_bytes_read, total_bytes_write

    def execute_stage(self):

        # some useful variables
        # unit conversion
        unit_convert_gb = 1./float(1024*1024*1024)

        # iterate over all the prepared stage filters
        for stage_filter in self.file_stats.keys():

            nvtx.range_push(f"stage {stage_filter}")

            stage_start = time.perf_counter()
            total_read, total_write = self._stage_instance_data(stage_filter)
            stage_stop = time.perf_counter()

            # allreduce:
            total_read = self.global_comm.allreduce(total_read)
            total_write = self.global_comm.allreduce(total_write)

            # convert units
            total_read *= unit_convert_gb
            total_write *= unit_convert_gb

            # stage duration:
            stage_duration = stage_stop - stage_start

            # print
            if self.grank == 0:
                print(f"""Staging {stage_filter} done.
                          Total number of files: {self.file_stats[stage_filter]["num_files"]}.
                          Elapsed time {stage_duration:.2f}s.
                          Read {total_read:.2f} GB (bandwidth: {total_read/stage_duration:.2f} GB/s).
                          Write {total_write:.2f} GB (bandwidth: {total_write/stage_duration:.2f} GB/s).
                       """, flush=True)

            # verify staging results if requested
            if self.verify:
                nvtx.range_push(f"stage_verify")
                if self.lrank == 0:
                    stage_target_directory = self.file_stats[stage_filter]["target_directory"]
                    staged_files = glob(os.path.join(stage_target_directory, os.path.basename(stage_filter)))
                else:
                    staged_files = []

                if not self.full_dataset_per_node:
                    # if every node hosts a shard, we need to sum the results, if not we need to make sure everybody has the same
                    staged_files_full = self.instance_comm.allgather(staged_files)
                    staged_files_full = set(itertools.chain(*staged_files_full))
                else:
                    staged_files_full = set(staged_files)
                staged_num_files = len(staged_files_full)

                global_files_full = None
                if self.irank == 0:
                    stage_source_directory = self.file_stats[stage_filter]["source_directory"]
                    global_files_full = glob(os.path.join(stage_source_directory, os.path.basename(stage_filter)))
                global_files_full = self.instance_comm.bcast(global_files_full, 0)
                global_num_files = len(global_files_full)

                # strip off the directory
                checkfiles1 = sorted([os.path.basename(x) for x in staged_files_full])
                checkfiles2 = sorted([os.path.basename(x) for x in global_files_full])

                assert(staged_num_files == global_num_files)
                assert(checkfiles1 == checkfiles2)

                if self.extended_verify:
                    if self.lrank == 0:
                        checks, issue = numpy_integrity_check(self.file_stats[stage_filter]["source_directory"],
                                                              self.file_stats[stage_filter]["target_directory"],
                                                              [os.path.basename(x) for x in staged_files])
                        if issue:
                            print(f"Instance {self.instance_id}, local rank {self.irank}: Verification Error. Results:", checks, flush=True)
                        else:
                            print(f"Instance {self.instance_id}, local rank {self.irank}: Verification OK", flush=True)
                    self.instance_comm.Barrier()

                if self.irank == 0:
                    print(f'Staged data for {stage_filter}: {staged_num_files}, expected: {global_num_files}', flush=True)
                nvtx.range_pop()

            # close range
            nvtx.range_pop()

        return