text:
#!/usr/bin/env python

import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split, GroupShuffleSplit

def test(a,b,cmp,cname=None):
  if cname is None: cname=cmp.__name__
  assert cmp(a,b),f"{cname}:\n{a}\n{b}"

def near(a,b): return np.allclose(a, b, rtol=1e-1, atol=1e-1)
def test_near(a,b): test(a,b,near)

def set_all_splits(df, val_pct, test_pct=0.0, seed=None):
  new_test_pct = np.around(test_pct / (val_pct + test_pct), 2)
  train_pct = 1 - (val_pct + test_pct)
  train_idxs, inter = train_test_split(np.arange(len(df)), test_size=(val_pct + test_pct), random_state=seed)
  val_idxs, test_idxs = train_test_split(inter, test_size=new_test_pct, random_state=seed)

  df['split'] = None
  df.iloc[train_idxs, df.columns.get_loc('split')] = 'train'
  df.iloc[val_idxs, df.columns.get_loc('split')] = 'val'
  df.iloc[test_idxs, df.columns.get_loc('split')] = 'test'

  test_near(round(len(df[df['split'] == 'train'])/len(df), 2), train_pct)
  test_near(round(len(df[df['split'] == 'val'])/len(df), 2), val_pct)
  test_near(round(len(df[df['split'] == 'test'])/len(df), 2), test_pct)

  return df

def set_splits_with_sample(df, val_pct, test_pct=0.0, sample_pct=0.0, seed=None):
  new_test_pct = np.around(test_pct / (val_pct + test_pct), 2)
  train_pct = 1 - (val_pct + test_pct)
  train_idxs, inter = train_test_split(np.arange(len(df)), test_size=(val_pct + test_pct), random_state=seed)
  val_idxs, test_idxs = train_test_split(inter, test_size=new_test_pct, random_state=seed)

  df['split'] = None
  df.iloc[train_idxs, df.columns.get_loc('split')] = 'train'
  df.iloc[val_idxs, df.columns.get_loc('split')] = 'val'
  df.iloc[test_idxs, df.columns.get_loc('split')] = 'test'

  if sample_pct > 0.0:
    df['is_sample'] = False
    _, sample_idxs = train_test_split(train_idxs, test_size=sample_pct)
    df.iloc[sample_idxs, df.columns.get_loc('is_sample')] = True
    test_near(round(len(df[df['is_sample']])/len(df), 2), round(sample_pct * train_pct, 2))

  test_near(round(len(df[df['split'] == 'train'])/len(df), 2), train_pct)
  test_near(round(len(df[df['split'] == 'val'])/len(df), 2), val_pct)
  test_near(round(len(df[df['split'] == 'test'])/len(df), 2), test_pct)

  return df

def set_two_splits(df, name, pct=0.15, seed=None):
  df['split'] = 'train'
  _, val_idxs = train_test_split(np.arange(len(df)), test_size=pct, random_state=seed)
  df.loc[val_idxs, 'split'] = name

  test_near(round(len(df[df['split'] == 'train'])/len(df), 2), 1-pct)
  test_near(round(len(df[df['split'] == name])/len(df), 2), pct)

  return df

def set_bool_split(df, pct=0.15, seed=None):
  df['is_valid'] = False
  _, val_idxs = train_test_split(np.arange(len(df)), test_size=pct, random_state=seed)
  df.loc[val_idxs, 'is_valid'] = True

  test_near(round(len(df[df['is_valid'] == False])/len(df), 2), 1-pct)
  test_near(round(len(df[df['is_valid'] == True])/len(df), 2), pct)

  return df

def set_group_splits(df, group_col, pct=0.15, name='test', seed=None):
  df['split'] = 'train'
  train_idxs, test_idxs = next(GroupShuffleSplit(test_size=pct, n_splits=2, random_state=seed).split(df, groups=df[group_col]))
  df.loc[test_idxs, 'split'] = name

  assert(set(df.loc[(df['split'] == 'train')][group_col].unique().tolist()).intersection(df[(df['split'] == name)][group_col].unique().tolist()) == set())
  test_near(len(df.loc[(df['split'] == 'train')])/len(df), 1-pct)
  test_near(len(df.loc[(df['split'] == name)])/len(df), pct)

  return df

def set_group_all_splits(df, group_col, val_pct, test_pct, seed=None):
  new_test_pct = np.around(test_pct / (val_pct + test_pct), 2)
  train_pct = 1 - (val_pct + test_pct)
  split_df = set_group_splits(df, group_col, (val_pct + test_pct), seed=seed)
  train_df = split_df.loc[(split_df['split'] == 'train')].reset_index()
  inter_df = split_df.loc[(split_df['split'] == 'test')].reset_index()
  split_df = set_group_splits(inter_df, group_col, new_test_pct, seed=seed)
  val_df = split_df.loc[(split_df['split'] == 'train')].reset_index()
  val_df['split'] = 'val'
  test_df = split_df.loc[(split_df['split'] == 'test')].reset_index()

  df = pd.concat([train_df, val_df, test_df], axis=0, sort=False)

  assert(set(df.loc[(df['split'] == 'train')][group_col].unique().tolist()).intersection(df[(df['split'] == 'val')][group_col].unique().tolist()) == set())
  assert(set(df.loc[(df['split'] == 'train')][group_col].unique().tolist()).intersection(df[(df['split'] == 'test')][group_col].unique().tolist()) == set())
  assert(set(df.loc[(df['split'] == 'val')][group_col].unique().tolist()).intersection(df[(df['split'] == 'test')][group_col].unique().tolist()) == set())

  test_near(round(len(df[df['split'] == 'train'])/len(df), 2), 1-(val_pct+test_pct))
  test_near(round(len(df[df['split'] == 'val'])/len(df), 2), val_pct)
  test_near(round(len(df[df['split'] == 'test'])/len(df), 2), test_pct)

  return df


text:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
import numpy as np
from timm.models.layers import DropPath, trunc_normal_

from mmcv.runner import load_checkpoint
from mmaction.utils import get_root_logger
from ..builder import BACKBONES

from functools import reduce, lru_cache
from operator import mul
from einops import rearrange

class Mlp(nn.Module):
    """ Multilayer perceptron."""

    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

def window_partition(x, window_size):
    """
    Args:
        x: (B, D, H, W, C)
        window_size (tuple[int]): window size

    Returns:
        windows: (B*num_windows, window_size*window_size, C)
    """
    B, D, H, W, C = x.shape
    x = x.view(B, D // window_size[0], window_size[0], H // window_size[1], window_size[1], W // window_size[2],
               window_size[2], C)
    windows = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous().view(-1, reduce(mul, window_size), C)
    return windows

def window_reverse(windows, window_size, B, D, H, W):
    """
    Args:
        windows: (B*num_windows, window_size, window_size, C)
        window_size (tuple[int]): Window size
        H (int): Height of image
        W (int): Width of image

    Returns:
        x: (B, D, H, W, C)
    """
    x = windows.view(B, D // window_size[0], H // window_size[1], W // window_size[2], window_size[0], window_size[1],
                     window_size[2], -1)
    x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(B, D, H, W, -1)
    return x

def get_window_size(x_size, window_size, shift_size=None):
    use_window_size = list(window_size)
    if shift_size is not None:
        use_shift_size = list(shift_size)
    for i in range(len(x_size)):
        if x_size[i] <= window_size[i]:
            use_window_size[i] = x_size[i]
            if shift_size is not None:
                use_shift_size[i] = 0

    if shift_size is None:
        return tuple(use_window_size)
    else:
        return tuple(use_window_size), tuple(use_shift_size)

class WindowAttention3D(nn.Module):
    """ Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.
    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The temporal length, height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(self, dim, window_size, num_heads, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):

        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wd, Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1) * (2 * window_size[2] - 1),
                        num_heads))  # 2*Wd-1 * 2*Wh-1 * 2*Ww-1, nH

        # get pair-wise relative position index for each token inside the window
        coords_d = torch.arange(self.window_size[0])
        coords_h = torch.arange(self.window_size[1])
        coords_w = torch.arange(self.window_size[2])
        coords = torch.stack(torch.meshgrid(coords_d, coords_h, coords_w))  # 3, Wd, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 3, Wd*Wh*Ww
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 3, Wd*Wh*Ww, Wd*Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wd*Wh*Ww, Wd*Wh*Ww, 3
        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 2] += self.window_size[2] - 1

        relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1)
        relative_coords[:, :, 1] *= (2 * self.window_size[2] - 1)
        relative_position_index = relative_coords.sum(-1)  # Wd*Wh*Ww, Wd*Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        """ Forward function.
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, N, N) or None
        """
        B_, N, C = x.shape
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # B_, nH, N, C

        q = q * self.scale
        attn = q @ k.transpose(-2, -1)

        relative_position_bias = self.relative_position_bias_table[
            self.relative_position_index[:N, :N].reshape(-1)].reshape(
            N, N, -1)  # Wd*Wh*Ww,Wd*Wh*Ww,nH
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wd*Wh*Ww, Wd*Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0)  # B_, nH, N, N

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class SwinTransformerBlock3D(nn.Module):
    """ Swin Transformer Block.

    Args:
        dim (int): Number of input channels.
        num_heads (int): Number of attention heads.
        window_size (tuple[int]): Window size.
        shift_size (tuple[int]): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, dim, num_heads, window_size=(2, 7, 7), shift_size=(0, 0, 0),
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_checkpoint=False):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        self.use_checkpoint = use_checkpoint

        assert 0 <= self.shift_size[0] < self.window_size[0], "shift_size must in 0-window_size"
        assert 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in 0-window_size"
        assert 0 <= self.shift_size[2] < self.window_size[2], "shift_size must in 0-window_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention3D(
            dim, window_size=self.window_size, num_heads=num_heads,
            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward_part1(self, x, mask_matrix):
        B, D, H, W, C = x.shape
        window_size, shift_size = get_window_size((D, H, W), self.window_size, self.shift_size)

        x = self.norm1(x)
        # pad feature maps to multiples of window size
        pad_l = pad_t = pad_d0 = 0
        pad_d1 = (window_size[0] - D % window_size[0]) % window_size[0]
        pad_b = (window_size[1] - H % window_size[1]) % window_size[1]
        pad_r = (window_size[2] - W % window_size[2]) % window_size[2]
        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b, pad_d0, pad_d1))
        _, Dp, Hp, Wp, _ = x.shape
        # cyclic shift
        if any(i > 0 for i in shift_size):
            shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), dims=(1, 2, 3))
            attn_mask = mask_matrix
        else:
            shifted_x = x
            attn_mask = None
        # partition windows
        x_windows = window_partition(shifted_x, window_size)  # B*nW, Wd*Wh*Ww, C
        # W-MSA/SW-MSA
        attn_windows = self.attn(x_windows, mask=attn_mask)  # B*nW, Wd*Wh*Ww, C
        # merge windows
        attn_windows = attn_windows.view(-1, *(window_size + (C,)))
        shifted_x = window_reverse(attn_windows, window_size, B, Dp, Hp, Wp)  # B D' H' W' C
        # reverse cyclic shift
        if any(i > 0 for i in shift_size):
            x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1], shift_size[2]), dims=(1, 2, 3))
        else:
            x = shifted_x

        if pad_d1 > 0 or pad_r > 0 or pad_b > 0:
            x = x[:, :D, :H, :W, :].contiguous()
        return x

    def forward_part2(self, x):
        return self.drop_path(self.mlp(self.norm2(x)))

    def forward(self, x, mask_matrix):
        """ Forward function.

        Args:
            x: Input feature, tensor size (B, D, H, W, C).
            mask_matrix: Attention mask for cyclic shift.
        """

        shortcut = x
        if self.use_checkpoint:
            x = checkpoint.checkpoint(self.forward_part1, x, mask_matrix)
        else:
            x = self.forward_part1(x, mask_matrix)
        x = shortcut + self.drop_path(x)

        if self.use_checkpoint:
            x = x + checkpoint.checkpoint(self.forward_part2, x)
        else:
            x = x + self.forward_part2(x)

        return x

class PatchMerging(nn.Module):
    """ Patch Merging Layer

    Args:
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def forward(self, x):
        """ Forward function.

        Args:
            x: Input feature, tensor size (B, D, H, W, C).
        """
        B, D, H, W, C = x.shape

        # padding
        pad_input = (H % 2 == 1) or (W % 2 == 1)
        if pad_input:
            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))

        x0 = x[:, :, 0::2, 0::2, :]  # B D H/2 W/2 C
        x1 = x[:, :, 1::2, 0::2, :]  # B D H/2 W/2 C
        x2 = x[:, :, 0::2, 1::2, :]  # B D H/2 W/2 C
        x3 = x[:, :, 1::2, 1::2, :]  # B D H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)  # B D H/2 W/2 4*C

        x = self.norm(x)
        x = self.reduction(x)

        return x

# cache each stage results
@lru_cache()
def compute_mask(D, H, W, window_size, shift_size, device):
    img_mask = torch.zeros((1, D, H, W, 1), device=device)  # 1 Dp Hp Wp 1
    cnt = 0
    for d in slice(-window_size[0]), slice(-window_size[0], -shift_size[0]), slice(-shift_size[0], None):
        for h in slice(-window_size[1]), slice(-window_size[1], -shift_size[1]), slice(-shift_size[1], None):
            for w in slice(-window_size[2]), slice(-window_size[2], -shift_size[2]), slice(-shift_size[2], None):
                img_mask[:, d, h, w, :] = cnt
                cnt += 1
    mask_windows = window_partition(img_mask, window_size)  # nW, ws[0]*ws[1]*ws[2], 1
    mask_windows = mask_windows.squeeze(-1)  # nW, ws[0]*ws[1]*ws[2]
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
    return attn_mask

class BasicLayer(nn.Module):
    """ A basic Swin Transformer layer for one stage.

    Args:
        dim (int): Number of feature channels
        depth (int): Depths of this stage.
        num_heads (int): Number of attention head.
        window_size (tuple[int]): Local window size. Default: (1,7,7).
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
    """

    def __init__(self,
                 dim,
                 depth,
                 num_heads,
                 window_size=(1, 7, 7),
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=None,
                 drop=0.,
                 attn_drop=0.,
                 drop_path=0.,
                 norm_layer=nn.LayerNorm,
                 downsample=None,
                 use_checkpoint=False):
        super().__init__()
        self.window_size = window_size
        self.shift_size = tuple(i // 2 for i in window_size)
        self.depth = depth
        self.use_checkpoint = use_checkpoint

        # build blocks
        self.blocks = nn.ModuleList([
            SwinTransformerBlock3D(
                dim=dim,
                num_heads=num_heads,
                window_size=window_size,
                shift_size=(0, 0, 0) if (i % 2 == 0) else self.shift_size,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop,
                attn_drop=attn_drop,
                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
                norm_layer=norm_layer,
                use_checkpoint=use_checkpoint,
            )
            for i in range(depth)])

        self.downsample = downsample
        if self.downsample is not None:
            self.downsample = downsample(dim=dim, norm_layer=norm_layer)

    def forward(self, x):
        """ Forward function.

        Args:
            x: Input feature, tensor size (B, C, D, H, W).
        """
        # calculate attention mask for SW-MSA
        B, C, D, H, W = x.shape
        window_size, shift_size = get_window_size((D, H, W), self.window_size, self.shift_size)
        x = rearrange(x, 'b c d h w -> b d h w c')
        Dp = int(np.ceil(D / window_size[0])) * window_size[0]
        Hp = int(np.ceil(H / window_size[1])) * window_size[1]
        Wp = int(np.ceil(W / window_size[2])) * window_size[2]
        attn_mask = compute_mask(Dp, Hp, Wp, window_size, shift_size, x.device)
        for blk in self.blocks:
            x = blk(x, attn_mask)
        x = x.view(B, D, H, W, -1)

        if self.downsample is not None:
            x = self.downsample(x)
        x = rearrange(x, 'b d h w c -> b c d h w')
        return x

class PatchEmbed3D(nn.Module):
    """ Video to Patch Embedding.

    Args:
        patch_size (int): Patch token size. Default: (2,4,4).
        in_chans (int): Number of input video channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(self, patch_size=(2, 4, 4), in_chans=3, embed_dim=96, norm_layer=None):
        super().__init__()
        self.patch_size = patch_size

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        """Forward function."""
        # padding
        _, _, D, H, W = x.size()  # B C D H W
        if W % self.patch_size[2] != 0:
            x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
        if H % self.patch_size[1] != 0:
            x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
        if D % self.patch_size[0] != 0:
            x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))

        x = self.proj(x)  # B C D Wh Ww
        if self.norm is not None:
            D, Wh, Ww = x.size(2), x.size(3), x.size(4)
            # flatten:B, C, D*Wh*Ww
            # transpose:B, D*Wh*Ww, C
            x = x.flatten(2).transpose(1, 2)
            x = self.norm(x)
            x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)  # 拉回 B, C, D, Wh, Ww

        return x

class PatchUnEmbed3D(nn.Module):
    """ Video to Patch UnEmbedding.

    Args:
        patch_size (int): Patch token size. Default: (2,4,4).
        in_chans (int): Number of input video channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(self, patch_size=(2, 4, 4), in_chans=3, embed_dim=96, norm_layer=None):
        super().__init__()
        self.embed_dim = embed_dim

    def forward(self, x, x_size):
        """Forward function."""
        # padding
        B, D, HW, C = x.size()  # B C D H W
        x = x.transpose(1, 2).view(B, D, self.embed_dim, D, x_size[0],x_size[1])  # B,D,C,ph, pw

        return x
class RSTB(nn.Module):
    """Residual Swin Transformer Block (RSTB).

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
        img_size: Input image size.
        patch_size: Patch size.
        resi_connection: The convolutional block before residual connection.
    """

    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
                 img_size=224, patch_size=4, resi_connection='1conv'):
        super(RSTB, self).__init__()

        self.dim = dim
        self.input_resolution = input_resolution
        # print('>>>>>>input_resolution=',input_resolution)

        self.residual_group = BasicLayer(dim=dim,
                                         input_resolution=input_resolution,
                                         depth=depth,
                                         num_heads=num_heads,
                                         window_size=window_size,
                                         mlp_ratio=mlp_ratio,
                                         qkv_bias=qkv_bias, qk_scale=qk_scale,
                                         drop=drop, attn_drop=attn_drop,
                                         drop_path=drop_path,
                                         norm_layer=norm_layer,
                                         downsample=downsample,
                                         use_checkpoint=use_checkpoint)

        if resi_connection == '1conv':
            self.conv = nn.Conv3d(dim, dim, 3, 1, 1)
        elif resi_connection == '3conv':
            # to save parameters and memory
            self.conv = nn.Sequential(nn.Conv3d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),
                                      nn.Conv3d(dim // 4, dim // 4, 1, 1, 0),
                                      nn.LeakyReLU(negative_slope=0.2, inplace=True),
                                      nn.Conv3d(dim // 4, dim, 3, 1, 1))

        self.patch_embed = PatchEmbed3D(
            patch_size=patch_size, in_chans=dim, embed_dim=dim, norm_layer=None)

        self.patch_unembed = PatchUnEmbed3D(
            patch_size=patch_size, in_chans=dim, embed_dim=dim, norm_layer=None)

    def forward(self, x, x_size):
        # print('>>>>>>RSTB input shape=', x.shape) # [4,65536,64]
        # print('>>>>>>RSTB output shape=', (self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x).shape) # [4,65536,64]
        return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x

    def flops(self):
        flops = 0
        flops += self.residual_group.flops()
        H, W = self.input_resolution
        flops += H * W * self.dim * self.dim * 9
        flops += self.patch_embed.flops()
        flops += self.patch_unembed.flops()

        return flops

@BACKBONES.register_module()
class SwinIR3D(nn.Module):
    """ SwinIR Transformer backbone.

    Args:
        patch_size (int | tuple(int)): Patch size. Default: (4,4,4).
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        depths (tuple[int]): Depths of each Swin Transformer stage.
        num_heads (tuple[int]): Number of attention head of each stage.
        window_size (int): Window size. Default: 7.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: Truee
        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
        drop_rate (float): Dropout rate.
        attn_drop_rate (float): Attention dropout rate. Default: 0.
        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
        norm_layer: Normalization layer. Default: nn.LayerNorm.
        patch_norm (bool): If True, add normalization after patch embedding. Default: False.
        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
            -1 means not freezing any parameters.
    """

    def __init__(self,
                 pretrained=None,
                 pretrained2d=True,
                 patch_size=(4, 4, 4),
                 in_chans=3,
                 embed_dim=96,
                 depths=[2, 2, 6, 2],
                 num_heads=[3, 6, 12, 24],
                 window_size=(2, 7, 7),
                 mlp_ratio=4.,
                 qkv_bias=True,
                 qk_scale=None,
                 drop_rate=0.,
                 attn_drop_rate=0.,
                 drop_path_rate=0.2,
                 norm_layer=nn.LayerNorm,
                 patch_norm=False,
                 frozen_stages=-1,
                 use_checkpoint=False,
                 img_range=1.,
                 resi_connection='1conv'):
        super().__init__()
        num_out_ch = in_chans
        num_feat = 64
        self.pretrained = pretrained
        self.pretrained2d = pretrained2d
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.patch_norm = patch_norm
        self.frozen_stages = frozen_stages
        self.window_size = window_size
        self.patch_size = patch_size
        self.img_range = img_range
        if in_chans == 3:
            rgb_mean = (0.4488, 0.4371, 0.4040)
            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
        else:
            self.mean = torch.zeros(1, 1, 1, 1)
        #####################################################################################################
        ################################### 1, shallow feature extraction ###################################
        self.conv_first = nn.Conv3d(in_chans, embed_dim, 3, 1, 1)

        #####################################################################################################
        ################################### 2, deep feature extraction ######################################
        self.ape = ape
        self.num_features = embed_dim
        self.mlp_ratio = mlp_ratio
        # split image into non-overlapping patches
        self.patch_embed = PatchEmbed3D(
            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None)

        self.pos_drop = nn.Dropout(p=drop_rate)
        self.patch_unembed = PatchUnEmbed3D(
            patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None)
        # stochastic depth
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule

        # build layers
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = RSTB(
                dim=int(embed_dim * 2 ** i_layer),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
                norm_layer=norm_layer,
                downsample=None,
                use_checkpoint=use_checkpoint,
                resi_connection=resi_connection)
            self.layers.append(layer)

        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))

        # add a norm layer for each output
        self.norm = norm_layer(self.num_features)
        # build the last conv layer in deep feature extraction
        if resi_connection == '1conv':
            self.conv_after_body = nn.Conv3d(embed_dim, embed_dim, 3, 1, 1)
        elif resi_connection == '3conv':
            # to save parameters and memory
            self.conv_after_body = nn.Sequential(nn.Conv3d(embed_dim, embed_dim // 4, 3, 1, 1),
                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),
                                                 nn.Conv3d(embed_dim // 4, embed_dim // 4, 1, 1, 0),
                                                 nn.LeakyReLU(negative_slope=0.2, inplace=True),
                                                 nn.Conv3d(embed_dim // 4, embed_dim, 3, 1, 1))

        #####################################################################################################
        ################################ 3, high quality image reconstruction ################################
        # for classical SR
        self.conv_before_upsample = nn.Sequential(nn.Conv3d(embed_dim, num_feat, 3, 1, 1),
                                                  nn.LeakyReLU(inplace=True))
        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)

        self._freeze_stages()

    def _freeze_stages(self):
        if self.frozen_stages >= 0:
            self.patch_embed.eval()
            for param in self.patch_embed.parameters():
                param.requires_grad = False

        if self.frozen_stages >= 1:
            self.pos_drop.eval()
            for i in range(0, self.frozen_stages):
                m = self.layers[i]
                m.eval()
                for param in m.parameters():
                    param.requires_grad = False

    def inflate_weights(self, logger):
        """Inflate the swin2d parameters to swin3d.

        The differences between swin3d and swin2d mainly lie in an extra
        axis. To utilize the pretrained parameters in 2d model,
        the weight of swin2d models should be inflated to fit in the shapes of
        the 3d counterpart.

        Args:
            logger (logging.Logger): The logger used to print
                debugging infomation.
        """
        checkpoint = torch.load(self.pretrained, map_location='cpu')
        state_dict = checkpoint['model']

        # delete relative_position_index since we always re-init it
        relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k]
        for k in relative_position_index_keys:
            del state_dict[k]

        # delete attn_mask since we always re-init it
        attn_mask_keys = [k for k in state_dict.keys() if "attn_mask" in k]
        for k in attn_mask_keys:
            del state_dict[k]

        state_dict['patch_embed.proj.weight'] = state_dict['patch_embed.proj.weight'].unsqueeze(2).repeat(1,1,self.patch_size[0],1,1) / self.patch_size[0]

        # bicubic interpolate relative_position_bias_table if not match
        relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k]
        for k in relative_position_bias_table_keys:
            relative_position_bias_table_pretrained = state_dict[k]
            relative_position_bias_table_current = self.state_dict()[k]
            L1, nH1 = relative_position_bias_table_pretrained.size()
            L2, nH2 = relative_position_bias_table_current.size()
            L2 = (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1)
            wd = self.window_size[0]
            if nH1 != nH2:
                logger.warning(f"Error in loading {k}, passing")
            else:
                if L1 != L2:
                    S1 = int(L1 ** 0.5)
                    relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
                        relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1),
                        size=(2 * self.window_size[1] - 1, 2 * self.window_size[2] - 1),
                        mode='bicubic')
                    relative_position_bias_table_pretrained = relative_position_bias_table_pretrained_resized.view(nH2,
                                                                                L2).permute(
                        1, 0)
            state_dict[k] = relative_position_bias_table_pretrained.repeat(2 * wd - 1, 1)

        msg = self.load_state_dict(state_dict, strict=False)
        logger.info(msg)
        logger.info(f"=> loaded successfully '{self.pretrained}'")
        del checkpoint
        torch.cuda.empty_cache()

    def init_weights(self, pretrained=None):
        """Initialize the weights in backbone.

        Args:
            pretrained (str, optional): Path to pre-trained weights.
                Defaults to None.
        """

        def _init_weights(m):
            if isinstance(m, nn.Linear):
                trunc_normal_(m.weight, std=.02)
                if isinstance(m, nn.Linear) and m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)

        if pretrained:
            self.pretrained = pretrained
        if isinstance(self.pretrained, str):
            self.apply(_init_weights)
            logger = get_root_logger()
            logger.info(f'load model from: {self.pretrained}')

            if self.pretrained2d:
                # Inflate 2D model into 3D model.
                self.inflate_weights(logger)
            else:
                # Directly load 3D model.
                load_checkpoint(self, self.pretrained, strict=False, logger=logger)
        elif self.pretrained is None:
            self.apply(_init_weights)
        else:
            raise TypeError('pretrained must be a str or None')

    def check_image_size(self, x):
        _, _, h, w = x.size()
        mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
        mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
        return x

    def forward_features(self, x):
        x_size = (x.shape[2], x.shape[3])
        x = self.patch_embed(x)
        #if self.ape:
            #x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        for layer in self.layers:
            x = layer(x, x_size)

        #x = rearrange(x, 'n c d h w -> n d h w c')
        x = self.norm(x)
        #x = rearrange(x, 'n d h w c -> n c d h w')
        x = self.patch_unembed(x, x_size)

        return x

    def forward(self, x):
        """Forward function."""
        print(f'input shape is : {x.size()}')
        x = self.check_image_size(x)
        print(f'input shape after padding is : {x.size()}')
        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        x = self.conv_first(x)
        print(f'feature shape after first  conv is : {x.size()}')
        x = self.conv_after_body(self.forward_features(x)) + x
        print(f'feature shape after rstb : {x.size()}')
        x = self.conv_before_upsample(x)
        print(f'feature shape after last second conv : {x.size()}')
        x = self.conv_last(x)
        print(f'output shape : {x.size()}')

        x = x / self.img_range + self.mean

        return x

    def train(self, mode=True):
        """Convert the model into training mode while keep layers freezed."""
        super(SwinIR3D, self).train(mode)
        self._freeze_stages()

ext: py

max_stars_repo_name: heliossun/Video-Swin-Transformer

max_issues_repo_name: heliossun/Video-Swin-Transformer

max_forks_repo_name: heliossun/Video-Swin-Transformer

reason: import numpy


text:
import numpy as np
import math
import sys

'''
Retorna la entropia de un conjunto de datos para un determinado atributo objetivo.
'''
def entropy(attributes, data, targetAttr):

    i = attributes.index(targetAttr)
    targetValues = data[:,i]

    #Frecuencia de cada valor del atributo objetivo dentro del conjunto de datos.
    valFreq = {}
    entropy = 0.0

    for targetValue in targetValues:
        if (valFreq.has_key(targetValue)):
            valFreq[targetValue] += 1.0
        else:
            valFreq[targetValue]  = 1.0

    for freq in valFreq.values():
        proportion = freq/len(data)
        entropy += (-proportion) * math.log(proportion, 2)

    return entropy

'''
Retorna el information gain para el atributo attr y el atributo objetivo targetAttr.
'''
def informationGain(attributes, data, attr, targetAttr):

    i = attributes.index(attr)
    attrValues = data[:,i]

    valFreq = {}
    for attrValue in attrValues:
        if (valFreq.has_key(attrValue)):
            valFreq[attrValue] += 1.0
        else:
            valFreq[attrValue]    = 1.0

    attrEntropy = 0.0
    dataLength = len(data)
    for val in valFreq.keys():
        valProb           = valFreq[val] / dataLength
        attrData     = np.array([entry for entry in data if entry[i] == val])
        attrEntropy += valProb * entropy(attributes, attrData, targetAttr)

    return (entropy(attributes, data, targetAttr) - attrEntropy)

'''
Obtiene el atributo que aporta mas information gain.
'''
def getBestAttr(data, attributes, target):
    best = attributes[0]
    maxGain = 0;
    nonTargetAttributes = attributes[:]
    nonTargetAttributes.remove(target)
    for attr in nonTargetAttributes:
        newGain = informationGain(attributes, data, attr, target)
        if newGain>maxGain:
            maxGain = newGain
            best = attr
    return best

'''
Retorna el valor mas comun del atributo objetivo.
'''
def mostCommonValue(targetValues):
    valFreq = {}
    maxFreq = 0
    commonValue = 0

    for value in targetValues:
        if (valFreq.has_key(value)):
            valFreq[value] += 1
        else:
            valFreq[value] = 1
        if valFreq[value] > maxFreq:
            maxFreq = valFreq[value]
            commonValue = value

    return commonValue

'''
Construye y retorna el arbol de decision.
'''
def genDecisionTree(data, attributes, target,maxHeight,currentHeight):
    data = np.array(data[:])
    #Obtengo un arreglo con todos los valores del atributo target
    #en cada instancia.
    targetValues = data[:,attributes.index(target)]
    mcv = mostCommonValue(targetValues)

    #Si todas las instancias tienen el mismo valor para el atributo,
    #etiquetar con ese valor.
    if len(set(targetValues)) == 1:
        return targetValues[0]

    #Si no quedan atributos, etiquetar con el valor mas comun.
    if (len(attributes) -1) == 0:
        return mcv

    #Si llego al largo maximo establecido para las ramas, etiquetar
    #con el valor mas comun.
    if currentHeight == maxHeight:
        return mcv

    #Si no alcance el largo maximo, continua creciendo la rama.
    currentHeight = currentHeight + 1

    # Se elige el mejor atributo segun su information gain.
    bestAttr = getBestAttr(data, attributes, target)
    bestAttrIndex = attributes.index(bestAttr)

    bestAttrTargetValues = data[:,bestAttrIndex]
    bestAttrMostCommonValue = mostCommonValue(bestAttrTargetValues)

    #Genero una rama.
    #Cada nodo tendra su valor mas comun (de acuerdo a ejemplos de
    #entrenamiento) asociado
    tree = {(bestAttr, bestAttrMostCommonValue):{}}

    # Genero una rama para cada valor del atributo elegido.
    for val in set(data[:,bestAttrIndex]):

        #Subconjunto de datos para cada valor del atributo.
        #Obtengo los ejemplos que comparten el mismo valor del atributo.
        valSubset = np.array([x for x in data if x[bestAttrIndex] == val])
        #A cada ejemplo le quito la coordenada correspondiente al atributo
        #elegido, porque no interesa para las proximas iteraciones.
        valSubset = np.delete(valSubset, bestAttrIndex, 1)

        #Remuevo el atributo de la lista de atributos.
        newAttr = attributes[:]
        newAttr.remove(bestAttr)

        #Llamada recursiva.
        subtree = genDecisionTree(valSubset, newAttr, target, maxHeight, currentHeight)

        #Agrego el subarbol al nodo actual.
        tree[(bestAttr,bestAttrMostCommonValue)][val] = subtree

    return tree

'''
Clasifica una instancia de acuerdo a un arbol de decision.
'''
def evalInstance(instance, attributes, decisionTree):
    if (isinstance(decisionTree,dict)):
        #Obtengo nodo.
        node = decisionTree.keys()[0]
        #Obtengo indice del atributo en el arreglo de atributos posibles.
        index = attributes.index(node[0])
        #Si el valor del atributo de interes no es contemplado por el arbol,
        #se lo intercambia por el valor mas comun para ese atributo.

        attributeValue = instance[index]
        #Si no es una etiqueta soportada.
        if (not attributeValue in decisionTree[node].keys()):
            #Valor mas comun asociado en construccion.
            attributeValue = node[1]
        #De otra manera, simplemente sigo la rama que corresponda con
        #el valor original del atributo para la instancia.
        nextBranchOrValue = decisionTree[node][attributeValue]
        #Si estoy frente a una nueva rama, me sigo moviendo en ella
        if (isinstance(nextBranchOrValue,dict)):
            return evalInstance(instance, attributes, nextBranchOrValue)
        #De otra manera, devuelvo la clasificacion obtenida.
        else:
            return nextBranchOrValue
    else:
        return sys.maxint

'''
Obtiene un subconjunto a partir de un conjunto de datos.
'''
def getFold(i,data,fold_size):
    return data[(i-1)*fold_size:i*fold_size,:]

'''
Quita un subconjunto de un conjunto de datos.
'''
def getSampleWithoutFold(i,data,fold_size,k):
    return np.concatenate((data[0:(i-1)*fold_size,:],data[i*fold_size:k*fold_size,:]))

'''
Funcion que evalua el error cometido al intentar clasificar las instancias
del conjunto de validacion pasado por parametro, no tomado en cuenta para realizar
el entrenamiento.
'''
def error(tree,attrs,targetAttr,validation_set):
    tmpAttrs = attrs[:-1]
    e_aux = 0
    for instance in validation_set:
        #Valor asignado por hipotesis obtenida mediante entrenamiento.
        h_x = float(evalInstance(instance, tmpAttrs, tree))
        #Valor asignado por funcion objetivo.
        f_x = float(instance[attrs.index(targetAttr)])
        if (h_x != f_x):
            e_x = 1.0
        else:
            e_x = 0.0
        e_aux += e_x

    return (1.0/len(validation_set)*e_aux)

max_stars_repo_name: marcciosilva/maa

max_issues_repo_name: marcciosilva/maa

max_forks_repo_name: marcciosilva/maa


text:
"""
    Provides functions to compute flooded surface.
"""

import os
import numpy as np
import re
import argparse
import sys
import flopy.utils.binaryfile as fpu
import math
import csv
import pandas as pd
import math
import flopy
import pickle
import statistics
#from utils import utils as u

__author__ = "June Sallou"
__maintainer__ = "June Sallou"
__credits__ = ["June Sallou"]
__license__ = "MIT"
__version__ = "0.0.4"
__date__ = "07/02/2019"
__email__ = "june.benvegnu-sallou@univ-rennes1.fr"

mainAppRepo = os.path.dirname(os.path.abspath(__file__)) + '/'

dc = 0.3  # depth of threshold for the undeground vulnerable zone
alpha = float(1/3)

def get_soil_surface_values_for_a_simulation(repo_simu, model_name):
    """
        Retrieve the matrix of the topographical altitude.
    """
    mf = flopy.modflow.Modflow.load(repo_simu + "/" + model_name + '.nam')
    dis = flopy.modflow.ModflowDis.load(
        repo_simu + "/" + model_name + '.dis', mf)
    topo = dis.top._array

    return topo

def get_model_name(site_number, chronicle, approx, rate, ref, perm):
    model_name = "model_time_0_geo_0_thick_1_K_86.4_Sy_0.1_Step1_site" + str(site_number) + "_Chronicle" + str(chronicle)
    if perm:
        model_name += "_SteadyState"
    elif not ref:
        model_name += "_Approx" + str(approx)
        if approx == 0:
            model_name += "_Period" + str(rate)
        elif approx==1:
            model_name += "_RechThreshold" + str(rate)
    return model_name

def get_site_name_from_site_number(site_number):
    sites = pd.read_csv(mainAppRepo + 'data/study_sites.txt',
                        sep=',', header=0, index_col=0) #\\s+
    site_name = sites.index._data[site_number] + '/'
    return site_name

def get_path_to_simulation_directory(site_number, chronicle, approx, rate, ref, perm):
    model_name = get_model_name(site_number, chronicle, approx, rate, ref, perm)
    site_name = get_site_name_from_site_number(site_number)
    return os.path.join(mainAppRepo, "outputs/", site_name, model_name)

def getNonDryCellHdsValue(hds, nrow, ncol, nlayer):
    layer = 0
    h = hds[layer][nrow][ncol]
    while (math.isclose(abs(h)/1e+30, 1, rel_tol=1e-3)) and layer < nlayer:
        if layer == nlayer-1:
            print("cell completely dry")
        else:
            h = hds[layer+1][nrow][ncol]
            layer += 1
    return h

def getWeightToSurface(zs, h, dc, alpha):
    """
        zs : value of soil surface at the point x
        h : head value for watertable at the x point
        dc : critical depth
        alpha : ratio of critical depth for calculating the width of transition zone
    """
    ddc = alpha * dc  # Alpha must be not null

    borneInf = zs - (dc + (ddc / 2))
    borneSup = zs - (dc - (ddc / 2))

    if h <= borneInf:
        Ws = 0
    elif h >= borneSup:
        Ws = 1
    else:
        Ws = math.sin((math.pi * (h - borneSup)) / (2*ddc))

    return Ws

def getSaturatedZoneArea(site_number, chronicle, folder, ref, dc, alpha, timestep=1):

    model_name = get_model_name(site_number, chronicle, None, None, ref, perm=True)

    if folder is None:
        repo_simu = get_path_to_simulation_directory(site_number, chronicle, None, None, ref=ref, perm=True)
    else :
        site_name = get_site_name_from_site_number(site_number)
        print(site_name)
        repo_simu = folder + site_name + model_name

    topoSimu = get_soil_surface_values_for_a_simulation(repo_simu, model_name)

    simuHds = fpu.HeadFile(repo_simu + '/' + model_name + '.hds')
    # simuTimes = simuHds.get_times()
    # simuKstpkper = simuHds.get_kstpkper()
    # print(simuTimes, simuKstpkper)
    hds_data = simuHds.get_data(kstpkper=(0,0))
    # print(hds_data)

    nbrowtot = hds_data.shape[1]
    nbcoltot = hds_data.shape[2]

    size_site = nbrowtot * nbcoltot
    Ws_sum = 0
    saturated_area = 0

    for nrow in range(nbrowtot):
        for ncol in range(nbcoltot):

            Zs = topoSimu[nrow][ncol]
            h = getNonDryCellHdsValue(hds_data, nrow, ncol, hds_data.shape[0])
            d =  Zs - h
            if d<= dc:
                saturated_area += 1
            Ws_sum += getWeightToSurface(Zs, h, dc, alpha)

    with open(repo_simu + "/" + model_name + '_extracted_features.csv', 'w') as f:
        writer = csv.writer(f, delimiter=';')
        writer.writerow(['Satured Zone Area', 'Vulnerability Sum', 'Vulnerability Rate', 'Global Area of site (in cells)'])
        writer.writerow([saturated_area, Ws_sum, Ws_sum/size_site, size_site])
    print("Saturated Area : ", saturated_area)
    print("Vulnerability Sum : ", Ws_sum)
    print("Vulnerability Rate : ", Ws_sum/size_site)
    print("Global Area of site :", size_site)

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-site", "--sitenumber", type=int, required=True)
    parser.add_argument("-chr", "--chronicle", type=int, required=True)
    parser.add_argument("-approx", "--approximation", type=int, required=False)
    parser.add_argument("-rate", "--rate", type=float, required=False)
    parser.add_argument("-dc", "--criticaldepth", type=float, required=False)
    parser.add_argument("-alpha", "--alpha", type=float, required=False)
    parser.add_argument("-ref", "--ref", action='store_true')
    parser.add_argument("-f", "--folder", type=str, required=False)

    parser.add_argument("-ts", "--timestep", type=int, required=False)

    args = parser.parse_args()

    site_number = args.sitenumber
    timestep = args.timestep
    dc = args.criticaldepth
    alpha = args.alpha
    chronicle= args.chronicle
    approx=args.approximation
    rate = args.rate
    ref = args.ref
    folder = args.folder

    if timestep is None:
        timestep = 1

    if alpha is None:
        alpha = float(1/3)

    getSaturatedZoneArea(site_number, chronicle, folder, ref=False, dc=0.3, alpha=alpha, timestep=timestep)


text:
# Copyright (c) 2018, NECOTIS
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
#  - Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#  - Neither the name of the copyright holder nor the names of its contributors
#    may be used to endorse or promote products derived from this software
#    without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
# OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# Authors: Simon Brodeur, Jean Rouat (advisor)
# Date: April 22th, 2019
# Organization: Groupe de recherche en Neurosciences Computationnelles et Traitement Intelligent des Signaux (NECOTIS),
# Université de Sherbrooke, Canada

import random
import logging
import numpy as np
import matplotlib.pyplot as plt

from brian2.units import ms, Hz
from brian2.synapses.synapses import Synapses
from brian2.core.clocks import defaultclock
from brian2.monitors.spikemonitor import SpikeMonitor
from brian2.core.network import Network
from brian2.units.allunits import second
from brian2.monitors.statemonitor import StateMonitor
from brian2.input.spikegeneratorgroup import SpikeGeneratorGroup
from brian2.groups.neurongroup import NeuronGroup
from brian2.core.functions import implementation
from brian2.units.fundamentalunits import check_units
from brian2.core.base import BrianObject
from brian2.devices.cpp_standalone.codeobject import CPPStandaloneCodeObject
from brian2.core.preferences import prefs

from critical.microcircuit import Microcircuit

logger = logging.getLogger(__name__)

class AvalancheEndDetector(BrianObject):

    def __init__(self, G, maxQuietTime=15 * ms, network=None, *args, **kwargs):
        super(AvalancheEndDetector, self).__init__(*args, **kwargs)

        prefs.codegen.cpp.headers += ['"run.h"']  # This is necessary to use brian_end()

        namespace = None
        if network is not None:
            # Stop simulation when avalanche end is detected
            @implementation(CPPStandaloneCodeObject, '''
            double stop_on_avalanche_end(double lastspiketime_input, double maxQuietTime) {
                if (lastspiketime_input > maxQuietTime) {
                    brian_end();
                    std::exit(0);
                }
                return 0.0;
            }
            ''')
            @implementation('numpy', discard_units=True)
            @check_units(lastspiketime_input=second, maxQuietTime=second, result=1)
            def stop_on_avalanche_end(lastspiketime_input, maxQuietTime):
                if lastspiketime_input > maxQuietTime:
                    network.stop()
                return 0.0

            namespace = {'stop_on_avalanche_end': stop_on_avalanche_end, 'maxQuietTime': maxQuietTime}

        # Spike monitor used to compute the avalanche length
        M = SpikeMonitor(G, record=True)

        # Dummy neuron and synapses used to monitor the end of the avalanche from the spiking activity
        P = NeuronGroup(1, '''
        lastspiketime_input : second
        ''', namespace=namespace)
        P.lastspiketime_input = 0.0 * second
        Sae = Synapses(G, P, on_pre='''
        lastspiketime_input = 0.0 * second
        ''')
        Sae.connect('True')
        P.run_regularly('lastspiketime_input += dt', when='synapses')
        if network is not None:
            P.run_regularly('dummy = stop_on_avalanche_end(lastspiketime_input, maxQuietTime)', when='after_synapses')

        self.P = P
        self.M = M
        self.contained_objects.extend([P, M, Sae])

    def getAvalancheSize(self):
        return max(int(self.M.num_spikes - 1), 0)

    def getAvalancheLength(self):
        if len(self.M.t) > 1:
            minTime, maxTime = np.min(self.M.t), np.max(self.M.t)
            avalancheLength = (maxTime - minTime)
        else:
            avalancheLength = 0.0 * ms
        return avalancheLength

def estimatePowerLawScaling(net, microcircuit, nbSamples=1000, maxAvalancheTime=150 * ms):

    # Disable plasticity and spontaneous activity
    microcircuit.S.plastic = False
    microcircuit.G.noise.active = False

    # Spike generator used for input stimulation
    # NOTE: use a high weight to force spiking of the postsynaptic neuron
    nbInputs = len(microcircuit.G)
    G = SpikeGeneratorGroup(nbInputs, indices=[0], times=[0 * ms])
    Si = Synapses(G, microcircuit.G, model='w : 1', on_pre='''v_post += w * int(not_refractory_post)
                                                              c_in_tot_post += w * int(not_refractory_post)''')
    Si.connect(j='i')
    Si.w = 10.0

    # Detect the end of avalanche and stop the simulation if no spike occured in the last 5 ms
    D = AvalancheEndDetector(microcircuit.G, maxQuietTime=5 * ms, network=net)
    defaultclock.dt = 0.1 * ms
    net.add([G, Si, D])
    net.store('before_testing')

    # Generate multiple avalanches
    # NOTE: only select excitatory neurons
    targets = np.arange(len(microcircuit.G), dtype=np.int)
    validTargets = targets[microcircuit.G.ntype > 0]
    avalancheSizes = []
    for n in range(nbSamples):
        if (n + 1) % 100 == 0:
            logger.debug('Generating avalanche no.%d (out of %d)' % (n + 1, nbSamples))

        # Reinitialization
        net.restore('before_testing')

        # Chose a target neuron in the population
        target = np.random.choice(validTargets)
        G.set_spikes(indices=[target], times=[defaultclock.t + 1 * ms])
        net.run(maxAvalancheTime)

        # Get the number of elicited spikes
        logger.debug('Avalanche no.%d: size of %d, length of %4.3f ms' % (n + 1, D.getAvalancheSize(), D.getAvalancheLength() / ms))
        avalancheSizes.append(D.getAvalancheSize())

    avalancheSizes = np.array(avalancheSizes)

    # Compute the histogram of avalanche sizes
    sizes = np.arange(0, np.max(avalancheSizes) + 1)
    bins = np.concatenate((sizes, [np.max(avalancheSizes) + 1, ])) - 0.5
    hist, _ = np.histogram(avalancheSizes, bins)
    pdf = hist.astype(np.float) / np.sum(hist)
    assert len(pdf) == len(sizes)

    # Show histogram of avalanche sizes
    fig = plt.figure(facecolor='white')
    ax = fig.add_subplot(111)
    ax.scatter(sizes[pdf > 0], pdf[pdf > 0], color='k', marker='.')
    ax.set_xlabel('Avalanche size')
    ax.set_ylabel('Probability')
    ax.set_xscale('log')
    ax.set_yscale('log')

    # Fitting the power-law line in log-log domain
    mask = np.logical_and(sizes > 0, pdf > 0)
    logX = np.log(sizes[mask])
    logY = np.log(pdf[mask])
    p = np.polyfit(logX, logY, 1)
    a, b = p[0], p[1]
    logger.info('Estimated exponent using curve fitting: %f' % (a))

    # Plot the fitted powerlaw curve
    npl = np.arange(0, np.max(sizes) + 1)
    pl = np.exp(b) * (npl ** (a))
    ax.plot(npl, pl, color='k', linestyle='--')

    return fig

def main():

    # Choose the duration of the training
    duration = 60 * second
    targetCbf = 1.0

    logger.info('Simulating for target branching factor of %f' % (targetCbf))

    # Create the microcircuit
    # NOTE: p_max is chosen so to have an out-degree of N=16
    m = Microcircuit(connectivity='small-world', macrocolumnShape=[2, 2, 2], minicolumnShape=[4, 4, 4],
                     p_max=0.056, srate=1 * Hz, excitatoryProb=0.8, delay='1*ms + 2*ms * rand()')

    # Configure CRITICAL learning rule
    m.G.c_out_ref = targetCbf          # target critical branching factor
    m.S.alpha = 0.05                   # learning rate

    logger.info('Number of neurons in the population: %d' % (len(m.G)))

    # Configure the monitors and simulation
    # NOTE: setting a high time resolution increase the stability of the learning rule
    M = SpikeMonitor(m.G, record=True)
    Mg = StateMonitor(m.G, variables=['cbf'], record=True, dt=10 * ms)
    defaultclock.dt = 0.1 * ms
    net = Network(m.getBrianObjects())
    net.store('initialized')

    # Add inputs and monitors
    net.add([M, Mg])

    # Run the simulation with input stimuli and plasticity enabled
    m.S.plastic = True
    net.run(duration, report='text')

    # Compute population average firing rate
    avgOutputFiringRate = len(M.i) / (len(m.G) * duration)
    logger.info('Average output firing rate: %4.2f Hz' % (avgOutputFiringRate))

    # NOTE: compute statistics on excitatory neurons only
    meanCbf = np.mean(Mg.cbf.T[:, m.G.ntype > 0], axis=-1)
    fig = plt.figure(facecolor='white', figsize=(6, 5))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlabel('Time [sec]')
    ax.set_ylabel('Average output contributions')
    ax.plot(Mg.t, meanCbf, color='k')
    ax.set_ylim((0.0, 2.0))
    fig.tight_layout()
    fig.savefig('convergence.eps')

    # Visualization of the simulation
    # NOTE: show only the last 10 sec of the simulation
    fig = plt.figure(facecolor='white', figsize=(6, 5))
    plt.subplot(111)
    plt.title('Spiking activity (output)')
    plt.plot(M.t / ms, M.i, '.', color='b')
    plt.ylabel('Neurons')
    plt.xlabel('Time [ms]')
    plt.xlim([0.0, duration / ms])
    fig.tight_layout()

    # Disable spontaneous activity and let the population reach a resting state
    m.S.plastic = False
    m.G.noise.active = False
    net.remove([M, Mg])
    net.run(10 * second)
    net.store('after_learning')

    fig = estimatePowerLawScaling(net, m, nbSamples=1000, maxAvalancheTime=250 * ms)
    fig.savefig('avalanches_distribution.eps')

    plt.show()

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)

    # Fix the seed of all random number generator
    seed = 0
    random.seed(seed)
    np.random.seed(seed)

    main()
    logger.info('All done.')

max_stars_repo_name: NECOTIS/CRITICAL

max_issues_repo_name: NECOTIS/CRITICAL

max_forks_repo_name: NECOTIS/CRITICAL


text:
import numpy as np

from util.mcts import MCTS

def generate(state, nn, callback, args):
    generator = Generator()
    generator.generate(state, nn, callback, args)

class Generator:
    def __init__(self):
        pass

    def generate(self, state, nn, cb, args):
        self.mcts = MCTS(nn)

        iterator = range(args['selfplay'])
        if args['progress']:
            from tqdm import tqdm
            iterator = tqdm(iterator, ncols = 50)

        # self play
        for pi in iterator:
            result = self._selfplay(state, args)
            if cb != None:
                cb(result)

    def _selfplay(self, state, args):
        self.mcts.resetStats()
        turn = 0
        hard_random_turn = args['hard_random'] if 'hard_random' in args else 0
        soft_random_turn = (args['soft_random'] if 'soft_random' in args else 30) + hard_random_turn
        history = []
        while state.getWinner() == None:
            if turn < hard_random_turn:
                # random action
                action_list = state.getAction()
                index = np.random.choice(len(action_list))
                (action, key) = action_list[index]
            else:
                action_list = self.mcts.getActionInfo(state, args['simulation'])
                if turn < soft_random_turn:
                    # random action by visited count
                    visited = [1.0 * a.visited for a in action_list]
                    sum_visited = sum(visited)
                    assert(sum_visited > 0)
                    p = [v / sum_visited for v in visited]
                    index = np.random.choice(len(action_list), p = p)
                else:
                    # select most visited count
                    index = np.argmax([a.visited for a in action_list])
                choice = action_list[index]
                key = choice.key
                action = choice.action

            # collect history
            history.append(state.actionToString(action))

            # next state
            state = state.getNextState(action)
            turn += 1

        winner = state.getWinner()
        win_str = 'x' if winner > 0 else 'o' if winner < 0 else '='

        return win_str + ' ' + ' '.join(history)

max_issues_repo_name: ksekimoto/Mobile-Deep-Learning-Projects

max_forks_repo_name: ksekimoto/Mobile-Deep-Learning-Projects


text:
from __future__ import annotations

import logging
import dataclasses
import typing as t

import numpy as np
from OpenGL import GL

from spyke.exceptions import GraphicsException
from spyke.graphics import gl
from spyke.enums import (
    _SizedInternalFormat,
    _CompressedInternalFormat,
    _TextureFormat,
    MinFilter,
    MagFilter,
    WrapMode,
    SizedInternalFormat,
    SwizzleTarget,
    SwizzleMask,
    PixelType,
    TextureFormat,
    TextureParameter,
    TextureTarget,)

_logger = logging.getLogger(__name__)

@dataclasses.dataclass
class TextureSpec:
    width: int = 0
    height: int = 0
    mipmaps: int = 3
    min_filter: MinFilter = MinFilter.LinearMipmapLinear
    mag_filter: MagFilter = MagFilter.Linear
    wrap_mode: WrapMode = WrapMode.Repeat
    texture_swizzle: t.Optional[SwizzleTarget] = None
    internal_format: _SizedInternalFormat = SizedInternalFormat.Rgba8
    _swizzle_mask: t.Optional[np.ndarray] = None

    @property
    def swizzle_mask(self) -> np.ndarray:
        return self._swizzle_mask or np.zeros((4,), dtype=np.int32)

    @swizzle_mask.setter
    def swizzle_mask(self, value: t.Sequence[SwizzleMask]) -> None:
        self._swizzle_mask = np.array(value, dtype=np.int32)

class Texture(gl.GLObject):
    _white_texture: t.Optional[Texture] = None
    _invalid_texture: t.Optional[Texture] = None

    @staticmethod
    def set_pixel_alignment(alignment: int) -> None:
        assert alignment in [
            1, 2, 4, 8], f'Invalid pixel alignment: {alignment}'
        GL.glPixelStorei(GL.GL_UNPACK_ALIGNMENT, alignment)

    @classmethod
    def empty(cls):
        '''
        Returns empty white texture of size 1x1,
        creating one if necessary.
        '''

        if cls._white_texture:
            return cls._white_texture

        data = np.array([255, 255, 255, 255], dtype=np.ubyte)

        spec = TextureSpec()
        spec.width = 1
        spec.height = 1
        spec.internal_format = SizedInternalFormat.Rgba8
        spec.min_filter = MinFilter.Nearest
        spec.mag_filter = MagFilter.Nearest
        spec.mipmaps = 1

        tex = cls(spec)
        tex.upload(None, 0, TextureFormat.Rgba, PixelType.UnsignedByte, data)
        tex.check_immutable()

        cls._white_texture = tex
        return cls._white_texture

    @classmethod
    def invalid(cls):
        '''
        Returns texture used to indicate that resource loading
        failed, creating texture if necessary.
        '''

        if cls._invalid_texture:
            return cls._invalid_texture

        data = np.array([
            255, 0, 255, 255,
            0, 0, 0, 255,
            0, 0, 0, 255,
            255, 0, 255, 255], dtype=np.ubyte)

        spec = TextureSpec()
        spec.width = 2
        spec.height = 2
        spec.internal_format = SizedInternalFormat.Rgba8
        spec.min_filter = MinFilter.Nearest
        spec.mag_filter = MagFilter.Nearest
        spec.mipmaps = 1

        tex = cls(spec)
        tex.upload(None, 0, TextureFormat.Rgba, PixelType.UnsignedByte, data)
        tex.check_immutable()

        cls._invalid_texture = tex
        return cls._invalid_texture

    def __init__(self, specification: TextureSpec):
        super().__init__()

        self.width: int = specification.width
        self.height: int = specification.height
        self.mipmaps: int = specification.mipmaps
        self.is_compressed: bool = isinstance(
            specification.internal_format, _CompressedInternalFormat)

        # TODO: Check if texture width and height are less than GL_MAX_TEXTURE_SIZE

        self._id = gl.create_texture(TextureTarget.Texture2d)

        GL.glTextureStorage2D(self.id, specification.mipmaps,
                              specification.internal_format, self.width, self.height)

        self.set_parameter(TextureParameter.WrapS, specification.wrap_mode)
        self.set_parameter(TextureParameter.WrapR, specification.wrap_mode)
        self.set_parameter(TextureParameter.WrapT, specification.wrap_mode)
        self.set_parameter(TextureParameter.MinFilter,
                           specification.min_filter)
        self.set_parameter(TextureParameter.MagFilter,
                           specification.mag_filter)
        self.set_parameter(TextureParameter.BaseLevel, 0)
        self.set_parameter(TextureParameter.MaxLevel, self.mipmaps - 1)

        if specification.texture_swizzle:
            assert specification.swizzle_mask is not None, 'Texture swizzle target was set but swizzle mask was not specified'

            GL.glTextureParameteriv(
                self.id, specification.texture_swizzle, specification.swizzle_mask)

        _logger.debug('Texture object with id %d initialized succesfully.', self.id)

    def upload(self,
               size: t.Optional[t.Tuple[int, int]],
               level: int,
               _format: _TextureFormat,
               pixel_type: PixelType,
               data: np.ndarray) -> None:
        assert not self.is_compressed, 'Cannot use Texture.upload on compressed texture. Use Texture.upload_compressed instead.'

        # TODO: Add check for weird texture conversion from formats that differ much
        # i.e. GL.GL_RG8 <= GL.GL_RGBA
        if size is None:
            size = (self.width, self.height)

        GL.glTextureSubImage2D(self.id, level, 0, 0, *size, _format, pixel_type, data)

        _logger.debug('Texture data of size (%d, %d) uploaded to texture object %d as level %d.', size[0], size[1], self.id, level)

    def upload_compressed(self,
                          size: t.Optional[t.Tuple[int, int]],
                          level: int,
                          _format: _TextureFormat,
                          image_size: int,
                          data: np.ndarray) -> None:
        assert self.is_compressed, 'Cannot use Texture.upload_compressed on non-compressed texture. Use Texture.upload instead.'

        if size is None:
            size = (self.width, self.height)

        GL.glCompressedTextureSubImage2D(self.id, level, 0, 0, size[0], size[1], _format, image_size, data)

        _logger.debug('Compressed texture data of size (%d, %d) uploaded to texture object %d as level %d. Texture format: %s', size[0], size[1], self.id, level, _format.name)

    def set_parameter(self, parameter: TextureParameter, value: t.Any) -> None:
        GL.glTextureParameteri(self.id, parameter, value)

    def generate_mipmap(self) -> None:
        GL.glGenerateTextureMipmap(self.id)

    def bind_to_unit(self, slot) -> None:
        GL.glBindTextureUnit(slot, self.id)

    def _delete(self) -> None:
        GL.glDeleteTextures(1, [self.id])

    def check_immutable(self) -> None:
        success = GL.GLint()
        GL.glGetTextureParameteriv(
            self.id, GL.GL_TEXTURE_IMMUTABLE_FORMAT, success)

        if success.value != 1:
            raise GraphicsException('Cannot create immutable texture.')


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

                stateless.py

Stateless training and active learning wrapper functions
for integrating with an app.

"""
import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow.keras.backend as K
import scipy.stats as st

from patchwork._sample import stratified_sample, find_labeled_indices
from patchwork._sample import PROTECTED_COLUMN_NAMES
from patchwork._badge import KPlusPlusSampler, _build_output_gradient_function
from patchwork._losses import masked_binary_crossentropy, masked_binary_focal_loss
from patchwork._util import shannon_entropy

import sklearn.preprocessing, sklearn.decomposition, sklearn.cluster
from sklearn.metrics import roc_auc_score

def _mlflow(server_uri, experiment_name, params=None, metrics=None,
            model=None, run_id=None, model_name="classification_model",
            log_model=False):
    """
    Wrapper for MLflow integration

    :server_uri: string; URI for tracking server
    :experiment_name: string; name of experiment to log to
    :params: dict; will be passed to mlflow.log_params()
    :metrics: dict of dicts; will be passed to mlflow.log_metrics()
    :model: tf.keras.Model; will be passed to mlflow.keras.log_model()
    :run_id: ID of existing run to connect to
    :model_name: relative path for logging model as an artifact
    :log_model: Boolean for whether to log model as an artifact
    """
    # if Nones were passed, don't do anything
    if server_uri is None or experiment_name is None:
        return None
    import mlflow
    mlflow.set_tracking_uri(server_uri)
    mlflow.set_experiment(experiment_name)
    run = mlflow.start_run(run_id)

    if params is not None:
        mlflow.log_params(params)
    # Metrics are stored in a dictionary of dictionaries- one set
    # of metrics per category. Flatten the dictionary and save
    # to the tracking server.
    if metrics is not None:
        metric_dict = {}
        for c in metrics:
            for m in ["accuracy", "base_rate", "prob_above_base_rate", "auc"]:
                if m in metrics[c]:
                    # accuracy is reported with an interval, e.g. 0.80 (0.75-0.85)
                    # so split on the space and just take the first value, so
                    # MLflow server can treat it like a float
                    metric_dict[c+"_"+m] = metrics[c][m].split(" ")[0]
        if len(metric_dict) > 0:
            mlflow.log_metrics(metric_dict)
    if log_model and (model is not None):
        mlflow.keras.log_model(model, model_name)

    mlflow.end_run()

def _build_model(input_dim, num_classes, num_hidden_layers=0,
                 hidden_dimension=128,
                 normalize_inputs=False, dropout=0):
    """
    Macro to generate a Keras classification model
    """
    inpt = tf.keras.layers.Input((input_dim))
    net = inpt

    # if we're normalizing inputs:
    if normalize_inputs:
        norm = tf.keras.layers.Lambda(lambda x:K.l2_normalize(x,axis=1))
        net = norm(net)

    # for each hidden layer
    for _ in range(num_hidden_layers):
        if dropout > 0:
            net = tf.keras.layers.Dropout(dropout)(net)
        net = tf.keras.layers.Dense(hidden_dimension, activation="relu")(net)

    # final layer
    if dropout > 0:
        net = tf.keras.layers.Dropout(dropout)(net)
    net = tf.keras.layers.Dense(num_classes, activation="relu")(net)

    return tf.keras.Model(inpt, net)

def _build_training_dataset(features, df, num_classes, num_samples,
                            batch_size):
    indices, labels = stratified_sample(df, num_samples,
                                            return_indices=True)
    ds = tf.data.Dataset.from_tensor_slices((indices, labels))

    def _load_feature(x,y):
        return tf.gather(features, x, axis=0), y
    ds = ds.map(_load_feature)
    ds = ds.batch(batch_size)
    ds = ds.prefetch(1)
    return ds

def _labels_to_dataframe(labels, classes=None):
    """
    Input labels as a list of dictionaries; return same information as
    a pandas dataframe.

    :classes: if a list of strings is passed here, will prune the columns
        to just these categories
    """
    # convert to dataframe
    df = pd.DataFrame(labels)

    # if exclude and validation not present, add them
    for c in ["exclude", "validation"]:
        if c not in df.columns:
            df[c] = False

    if "filepath" not in df.columns:
        df["filepath"] = ""

    # prune if necessary
    if classes is not None:
        for c in df.columns:
            if c not in classes+PROTECTED_COLUMN_NAMES:
                df = df.drop(c, axis=1)
    # make sure columns are in the same order
    if classes is not None:
        df = df.reindex(columns=["filepath", "exclude",
                                 "validation"]+classes)
    return df

def _estimate_accuracy(tp, fp, tn, fn, alpha=5, beta=5):
    """
    Compute a set of estimates around the accuracy, using a
    beta-binomial model.

    :tp: count of true positives
    :fp: count of false positives
    :tn: count of true negatives
    :fn: count of false negatives
    :alpha, beta: parameters of Beta prior

    Returns dictionary containing STRING-FORMATTED RESULTS
    :accuracy: point estimate of accuracy
    :interval_low: low end of 90% credible interval of accuracy
    :interval_high: high end of 90% credible interval of accuracy
    :base_rate: frequentist point estimate of the base rate
    :prob_higher_than_base_rate: estimate of the probability that
        accuracy is higher than the base rate
    """
    # total number
    N = tp + tn + fp + fn
    # accuracy point estimate
    #acc = (tp+tn)/N
    acc = (alpha+tp+tn)/(alpha + beta + tp + tn + fp + fn)
    # how many right/wrong
    num_right = tp+tn
    num_wrong = fp+fn
    # estimate base rate
    frac_pos = (tp+fn)/N
    base_rate = max(frac_pos, 1-frac_pos)

    interval_low = st.beta.ppf(0.05, alpha+num_right, beta+num_wrong)
    interval_high = st.beta.ppf(0.95, alpha+num_right, beta+num_wrong)
    prob_above_base_rate = 1-st.beta.cdf(base_rate, alpha+num_right,
                                         beta+num_wrong)

    return {"accuracy":"{:0.3f} ({:0.3f}-{:0.3f})".format(acc, interval_low,
                                                              interval_high),
                "base_rate":"{:0.3f}".format(base_rate),
                "prob_above_base_rate":"{:0.3f}".format(prob_above_base_rate)}

def _eval(features, df, classes, model, threshold=0.5, alpha=5,beta=5):
    """
    Macro for evaluating a classification model. does the following:

        1) find the subset of tags marked "validation"
        2) generate model predictions on the validation features
        3) round predictions to 1 or 0
        4) for each category:
            a) check to see whether the labels include at least one positive
                and one negative example
            b) if so, call _estimate_accuracy() to get performance measures
        5) return everything as a dict
    """
    outdict = {}
    # boolean array
    val_subset = (df["validation"] == True).values
    # if there aren't any validation points, return empty dictionaries
    if val_subset.sum() == 0:
        return {c:{} for c in classes}
    # (num_val, num_classes) array
    val_labels = df[classes].values[val_subset,:]
    # get model sigmoid predictions
    preds = model.predict(features[val_subset,:])
    # round to 1 or 0
    predictions = (preds >= threshold).astype(int)

    # for each class
    for i,c in enumerate(classes):
        # only bother computing if we have at least one positive
        # and one negative val example in this category
        if (1 in val_labels[:,i])&(0 in val_labels[:,i]):
            # true positives
            tp = np.sum((val_labels[:,i]==1)&(predictions[:,i] == 1))
            # true negatives
            tn = np.sum((val_labels[:,i]==0)&(predictions[:,i] == 0))
            # false positives
            fp = np.sum((val_labels[:,i]==0)&(predictions[:,i] == 1))
            # false negatives
            fn = np.sum((val_labels[:,i]==1)&(predictions[:,i] == 0))
            outdict[c] = _estimate_accuracy(tp,fp,tn,fn, alpha, beta)
            # some people are really into AUC. we need to pull out the subset
            # of validation points that are nonempty for this class (in case
            # there are partially-labeled val points)
            nonempty = pd.notnull(df[c]).values[val_subset]
            outdict[c]["auc"] = "{:0.3f}".format(roc_auc_score(val_labels[nonempty,i],
                                                               preds[nonempty,i]))

            # also include raw values of model outputs for positive/negative
            # ground truth for each class, broken out by train and test
            raw_preds = {}
            raw_preds["validation_positive"] = [float(x) for x in
                                                preds[val_labels[:,i] == 1][:,i]]
            raw_preds["validation_negative"] = [float(x) for x in
                                                preds[val_labels[:,i] == 0][:,i]]

            # positive and negative training subsets
            trainpos = ((df["validation"] != True)&(df[c] == 1)).values
            trainneg = ((df["validation"] != True)&(df[c] == 0)).values
            raw_preds["training_positive"] = [float(x) for x in
                                              model.predict(features[trainpos,:])[:,i]]
            raw_preds["training_negative"] = [float(x) for x in
                                              model.predict(features[trainneg,:])[:,i]]

            outdict[c]["predictions"] = raw_preds

        else:
            outdict[c] = {}

    return outdict

def train(features, labels, classes, training_steps=1000,
          batch_size=32, learning_rate=1e-3, focal_loss=False,
          server_uri=None, experiment_name=None, log_model=False,
          model_name="classification_model", **model_kwargs):
    """
    Hey now, you're an all-star. Get your train on.

    :features: (num_samples, input_dim) array of feature vectors
    :labels:
    :classes: list of strings; names of categories to include in model
    :training_steps: how many steps to train for
    :batch_size: number of examples per batch
    :learning_rate: learning rate for Adam optimizer
    :focal_loss: Boolean; whether to use focal loss (with gamma=2) instead
        of binary crossentropy loss
    :model_kwargs: keyword arguments for model construction
    :server_uri: string; URI for mlflow tracking server
    :experiment_name: string; name of mlflow experiment to log to
    :log_model: whether to log model as an mlflow artifact after training
    :model_name: if logging, relative path to store artifact at

    Returns
    :model: trained Keras model
    :training_loss: array of length (training_steps) giving the
        loss function value at each training step
    :validation_metrics: dictionary of validation metrics
    """
    # convert labels to a dataframe
    df = _labels_to_dataframe(labels, classes)

    num_classes = len(classes)
    input_dim = features.shape[1]

    # create a model and optimizer
    model = _build_model(input_dim, num_classes, **model_kwargs)
    opt = tf.keras.optimizers.Adam(learning_rate)

    # build a dataset for training
    num_samples = training_steps*batch_size
    ds = _build_training_dataset(features, df, num_classes, num_samples,
                                 batch_size)

    # train the model, recording loss at each step
    training_loss = []
    @tf.function
    def training_step(x,y):
        with tf.GradientTape() as tape:
            pred = model(x, training=True)
            if focal_loss:
                loss = masked_binary_focal_loss(y, pred,2)
            else:
                loss = masked_binary_crossentropy(y, pred)
        variables = model.trainable_variables
        grads = tape.gradient(loss, variables)
        opt.apply_gradients(zip(grads, variables))
        return loss

    for x, y in ds:
        training_loss.append(training_step(x,y).numpy())

    # finally run some performance metrics
    acc = _eval(features, df, classes, model)

    # log results to mlflow
    num_val_examples = np.sum(df.validation)
    num_train_examples = len(labels) - num_val_examples
    params = {"batch_size":batch_size, "learning_rate":learning_rate,
              "focal_loss":focal_loss, "training_steps":training_steps,
              "num_val_examples":num_val_examples,
              "num_train_examples":num_train_examples}
    for k in model_kwargs:
        params[k] = model_kwargs[k]

    _mlflow(server_uri, experiment_name, params=params,
                metrics=acc, model=model, log_model=log_model,
                model_name=model_name)

    return model, np.array(training_loss), acc

def get_indices_of_tiles_in_predicted_class(features, model, category_index,
                                            threshold=0.5):
    """
    :features: (num_samples, input_dim) array of feature vectors
    :model: trained Keras classifier model
    :category_index: integer index of category to query
    :threshold: float between 0 and 1; minimum probability assessed by
        classifier
    """
    predictions = model.predict(features)
    category_predictions = predictions[:,category_index]
    return np.arange(features.shape[0])[category_predictions >= threshold]

def sample_random(labels, max_to_return=None):
    """
    Generate a random sample of indices

    :labels: list of dictionaries containing labels
    :max_to_return: if not None; max number of indices to return
    """
    N = len(labels)
    if max_to_return is None:
        max_to_return = N
    # create a list of unlabeled indices
    df = _labels_to_dataframe(labels)
    labeled = list(find_labeled_indices(df))
    indices_to_sample_from = [n for n in range(N) if n not in labeled]
    # update num_to_return in case not many unlabeled are left
    max_to_return = min(max_to_return, len(indices_to_sample_from))

    return np.random.choice(indices_to_sample_from,size=max_to_return,
                            replace=False)

def sample_uncertainty(labels, features, model, max_to_return=None):
    """
    Return indices sorted by decreasing entropy.

    :features: (num_samples, input_dim) array of feature vectors
    :model: trained Keras classifier model
    :max_to_return: if not None; max number of indices to return
    """
    N = features.shape[0]
    # get model predictions
    predictions = model.predict(features)
    # compute entropies for each prediction
    entropy = shannon_entropy(predictions)

    # create a list of unlabeled indices
    df = _labels_to_dataframe(labels)
    labeled = list(find_labeled_indices(df))

    # order by decreasing entropy
    ordering = entropy.argsort()[::-1]
    ordered_indices = np.arange(N)[ordering]
    # prune out labeled indices
    ordered_indices = [i for i in ordered_indices if i not in labeled]
    # and clip list if necessary
    if max_to_return is not None:
        ordered_indices = ordered_indices[:max_to_return]
    return ordered_indices

def sample_diversity(labels, features, max_to_return=20):
    """
    Attempt to create a diverse, model-independent sample of
    images to tag. PCA-reduce feature vectors an then use k-means
    for sampling.

    :features: (num_samples, input_dim) array of feature vectors
    :max_to_return: if not None; max number of indices to return
    """
    N = features.shape[0]

    # figure out which indices are yet labeled
    df = _labels_to_dataframe(labels)
    labeled = find_labeled_indices(df)

    # rescale, PCA reduce, and cluster feature vectors
    scaled = sklearn.preprocessing.StandardScaler().fit_transform(features)
    reduced = sklearn.decomposition.PCA(16).fit_transform(scaled)
    kmeans = sklearn.cluster.KMeans(n_clusters=max_to_return).fit(reduced)

    # now use this for sample- pull up to one not-yet-labeled
    # data point from each cluster.
    all_indices = np.arange(N)
    indices = []
    for i in range(max_to_return):
        subset = [j for j in all_indices[kmeans.labels_ == i]
                  if j not in labeled]
        if len(subset) > 0:
            indices.append(np.random.choice(subset))
    return np.array(indices)

def sample_badge(labels, features, model, max_to_return=None):
    """
    Use a trained model to compute output-gradient vectors for the
    BADGE algorithm for active learning. Return sorted indices.

    :features: (num_samples, input_dim) array of feature vectors
    :model: trained Keras classifier model
    :max_to_return: if not None; max number of indices to return
    """
    N = features.shape[0]
    if max_to_return is None:
        max_to_return = N

    # compute badge embeddings- define a tf.function for it
    compute_output_gradients = _build_output_gradient_function(model)
    # then run that function across all the images.
    output_gradients = compute_output_gradients(features).numpy()

    # figure out which indices are yet labeled
    df = _labels_to_dataframe(labels)
    labeled = list(find_labeled_indices(df))
    # update max_to_return in case there aren't very many
    # unlabeled indices left
    max_to_return = min(max_to_return, N-len(labeled))

    # initialize a K++ sampler
    badge_sampler = KPlusPlusSampler(output_gradients, indices=labeled)
    return badge_sampler(max_to_return)

max_stars_repo_name: mediocretech/patchwork

max_issues_repo_name: mediocretech/patchwork

max_forks_repo_name: mediocretech/patchwork

reason: import numpy,import scipy


text:
import numpy
import time
import sys
import subprocess
import os
import random

from is13.data import load
from is13.rnn.jordan import model
from is13.metrics.accuracy import conlleval
from is13.utils.tools import shuffle, minibatch, contextwin

if __name__ == '__main__':

    s = {'fold':3, # 5 folds 0,1,2,3,4
         'lr':0.0627142536696559,
         'verbose':1,
         'decay':False, # decay on the learning rate if improvement stops
         'win':7, # number of words in the context window
         'bs':9, # number of backprop through time steps
         'nhidden':100, # number of hidden units
         'seed':345,
         'emb_dimension':100, # dimension of word embedding
         'nepochs':50}

    folder = os.path.basename(__file__).split('.')[0]
    if not os.path.exists(folder): os.mkdir(folder)

    # load the dataset
    train_set, valid_set, test_set, dic = load.atisfold(s['fold'])
    idx2label = dict((k,v) for v,k in dic['labels2idx'].iteritems())
    idx2word  = dict((k,v) for v,k in dic['words2idx'].iteritems())

    train_lex, train_ne, train_y = train_set
    valid_lex, valid_ne, valid_y = valid_set
    test_lex,  test_ne,  test_y  = test_set

    vocsize = len(set(reduce(\
                       lambda x, y: list(x)+list(y),\
                       train_lex+valid_lex+test_lex)))

    nclasses = len(set(reduce(\
                       lambda x, y: list(x)+list(y),\
                       train_y+test_y+valid_y)))

    nsentences = len(train_lex)

    # instanciate the model
    numpy.random.seed(s['seed'])
    random.seed(s['seed'])
    rnn = model(    nh = s['nhidden'],
                    nc = nclasses,
                    ne = vocsize,
                    de = s['emb_dimension'],
                    cs = s['win'] )

    # train with early stopping on validation set
    best_f1 = -numpy.inf
    s['clr'] = s['lr']
    for e in xrange(s['nepochs']):
        # shuffle
        shuffle([train_lex, train_ne, train_y], s['seed'])
        s['ce'] = e
        tic = time.time()
        for i in xrange(nsentences):
            cwords = contextwin(train_lex[i], s['win'])
            words  = map(lambda x: numpy.asarray(x).astype('int32'),\
                         minibatch(cwords, s['bs']))
            labels = train_y[i]

            for word_batch , label_last_word in zip(words, labels):
                rnn.train(word_batch, label_last_word, s['clr'])
                rnn.normalize()

            if s['verbose']:
                print '[learning] epoch %i >> %2.2f%%'%(e,(i+1)*100./nsentences),'completed in %.2f (sec) <<\r'%(time.time()-tic),
                sys.stdout.flush()

        # evaluation // back into the real world : idx -> words
        predictions_test = [ map(lambda x: idx2label[x], \
                             rnn.classify(numpy.asarray(contextwin(x, s['win'])).astype('int32')))\
                             for x in test_lex ]
        groundtruth_test = [ map(lambda x: idx2label[x], y) for y in test_y ]
        words_test = [ map(lambda x: idx2word[x], w) for w in test_lex]

        predictions_valid = [ map(lambda x: idx2label[x], \
                             rnn.classify(numpy.asarray(contextwin(x, s['win'])).astype('int32')))\
                             for x in valid_lex ]
        groundtruth_valid = [ map(lambda x: idx2label[x], y) for y in valid_y ]
        words_valid = [ map(lambda x: idx2word[x], w) for w in valid_lex]

        # evaluation // compute the accuracy using conlleval.pl
        res_test  = conlleval(predictions_test, groundtruth_test, words_test, folder + '/current.test.txt')
        res_valid = conlleval(predictions_valid, groundtruth_valid, words_valid, folder + '/current.valid.txt')

        if res_valid['f1'] > best_f1:
            rnn.save(folder)
            best_f1 = res_valid['f1']
            if s['verbose']:
                print 'NEW BEST: epoch', e, 'valid F1', res_valid['f1'], 'best test F1', res_test['f1'], ' '*20
            s['vf1'], s['vp'], s['vr'] = res_valid['f1'], res_valid['p'], res_valid['r']
            s['tf1'], s['tp'], s['tr'] = res_test['f1'],  res_test['p'],  res_test['r']
            s['be'] = e
            subprocess.call(['mv', folder + '/current.test.txt', folder + '/best.test.txt'])
            subprocess.call(['mv', folder + '/current.valid.txt', folder + '/best.valid.txt'])
        else:
            print ''

        # learning rate decay if no improvement in 10 epochs
        if s['decay'] and abs(s['be']-s['ce']) >= 10: s['clr'] *= 0.5
        if s['clr'] < 1e-5: break

    print 'BEST RESULT: epoch', e, 'valid F1', s['vf1'], 'best test F1', s['tf1'], 'with the model', folder