text:
# Copyright 2019 SAP SE
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#    http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch.nn as nn
import torch
import numpy as np
from torch.nn.parameter import Parameter
import math
import torch.nn.functional as F
from utils.utils import weights_init, weights_init_g

class Plastic_ConvTranspose2d(nn.ConvTranspose2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, dilation=1,
                 groups=1, bias=True,
                 num_tasks=1, out_size=None, batch_size=64, smax=800):
        super(Plastic_ConvTranspose2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding,
                                                      output_padding, groups, bias, dilation)
        if bias:
            self.ec = torch.nn.Embedding(num_tasks,
                                         out_channels * ((in_channels * kernel_size * kernel_size) + 1)).cuda()
        else:
            self.ec = torch.nn.Embedding(num_tasks, out_channels * ((in_channels * kernel_size * kernel_size))).cuda()

        self.smax = smax
        self.ec.weight.data.fill_(0)
        self.prev_weight_shape = self.weight.shape

    def forward(self, inputx, mask, d_in, d_out, output_size=None):
        self.prev_weight_shape = self.weight.shape
        output_padding = self._output_padding(inputx, output_size, self.stride, self.padding, self.kernel_size)
        bias = None
        if mask is not None:
            if not self.bias is None:
                bias = self.bias[:d_out] * mask[:, 0].contiguous().view(-1)
                mask_ = mask[:, 1:]
            else:
                mask_ = mask[:, :]
            out = F.conv_transpose2d(inputx, self.weight[:d_in, :d_out, :, :] * mask_.contiguous().view(
                self.weight.data.shape)[:d_in, :d_out, :, :],
                                     bias, self.stride, self.padding, output_padding, groups=self.groups,
                                     dilation=self.dilation)
        else:
            if not self.bias is None:
                bias = self.bias[:d_out]
            out = F.conv_transpose2d(inputx, self.weight[:d_in, :d_out, :, :], bias, self.stride, self.padding,
                                     output_padding, groups=self.groups, dilation=self.dilation)
        return out

    def expand(self, input_channels, out_channels):
        w_old = self.weight.data.clone()
        if self.bias is not None:
            b_old = self.bias.data.clone()
        self.out_channels += out_channels
        self.in_channels += input_channels
        self.weight = Parameter(
            torch.Tensor(self.in_channels, self.out_channels // self.groups, *self.kernel_size).cuda())
        if self.bias is not None:
            self.bias = Parameter(torch.Tensor(self.out_channels))
        self.weight.data.fill_(0)
        self.apply(weights_init)
        self.weight.data[:w_old.shape[0]:, :w_old.shape[1], :, :].copy_(w_old)
        if self.bias is not None:
            self.bias.data[:b_old.shape[0]].copy_(b_old)
        return self.weight.shape

    def expand_embeddings(self, n_new_classes, mask_pre=None):
        ec = self.ec.weight.view([-1] + list(self.prev_weight_shape)).data.clone()
        new_dim = self.out_channels * ((self.in_channels * self.kernel_size[0] * self.kernel_size[1]))
        if self.bias is not None:
            new_dim = self.out_channels * ((self.in_channels * self.kernel_size[0] * self.kernel_size[1]) + 1)
        self.ec = torch.nn.Embedding(ec.shape[0] + n_new_classes, new_dim).cuda()
        self.ec.weight.data.fill_(0)

        if ec.shape[0] > 0:
            self.ec.weight.data[:ec.shape[0], :].fill_(-90)  # for generating old samples do not use newly added parameters
            self.ec.weight.view([-1] + list(self.weight.shape)).data[:ec.shape[0], :ec.shape[1], :ec.shape[2], :,
            :].copy_(ec[:, :, :, :, :])  # but only the reserved once
        return self.prev_weight_shape

class BatchNorm2d(torch.nn.BatchNorm2d):
    def forward(self, input):
        self._check_input_dim(input)
        if self.training:
            momentum = self.momentum
        else:
            momentum = 0.
        return F.batch_norm(
            input, self.running_mean, self.running_var, self.weight, self.bias,
            self.training or not self.track_running_stats, momentum, self.eps)

class netG(nn.Module):
    def __init__(self, nz, ngf, nc, smax, n_classes=1):
        super(netG, self).__init__()

        self.gate = torch.nn.Sigmoid()

        self.nc = nc
        self.nz = nz
        self.ngf = ngf
        self.scalor = 1
        self.smax = smax
        self.ReLU = nn.ReLU(True)
        self.Tanh = nn.Tanh()

        self.conv1 = Plastic_ConvTranspose2d(nz, ngf * 4 * self.scalor, 4, 1, 0, bias=False, num_tasks=n_classes,
                                             smax=smax)
        # self.conv1 = Plastic_ConvTranspose2d(nz, ngf * 8*self.scalor, 3, 1, 0, bias=False, num_tasks=n_classes, smax=smax)

        self.cap_conv1 = [ngf * 4 * self.scalor]
        self.BatchNorms1 = torch.nn.ModuleList()
        self.BatchNorms1.append(torch.nn.BatchNorm2d(ngf * 4 * self.scalor).apply(weights_init))

        self.conv2 = Plastic_ConvTranspose2d(ngf * 4 * self.scalor, ngf * 2 * self.scalor, 4, 2, 1, bias=False,
                                             num_tasks=n_classes, smax=smax)
        self.cap_conv2 = [ngf * 2 * self.scalor]
        self.BatchNorms2 = torch.nn.ModuleList()
        self.BatchNorms2.append(torch.nn.BatchNorm2d(ngf * 2 * self.scalor).apply(weights_init))

        self.conv3 = Plastic_ConvTranspose2d(ngf * 2 * self.scalor, ngf * 1 * self.scalor, 4, 2, 1, bias=False,
                                             num_tasks=n_classes, smax=smax)
        self.cap_conv3 = [ngf * 1 * self.scalor]
        self.BatchNorms3 = torch.nn.ModuleList()
        self.BatchNorms3.append(torch.nn.BatchNorm2d(ngf * 1 * self.scalor).apply(weights_init))

        self.apply(weights_init_g)
        self.last = self.last = torch.nn.ModuleList()

    def extand(self, extention):
        ws_0 = self.conv1.expand(0, math.ceil(
            extention[0] / (self.nz * self.conv1.kernel_size[0] * self.conv1.kernel_size[1])))

        n_in_conv2 = self.conv1.weight.shape[1]
        n_out_conv2 = self.conv2.weight.shape[1]
        n_params_conv2 = n_in_conv2 * n_out_conv2 * self.conv2.kernel_size[0] * self.conv2.kernel_size[1]
        n_add_out_conv2 = max(math.ceil((extention[1] - (n_params_conv2 - np.prod(self.conv2.weight.size()).item())) / (
                    self.conv1.weight.data.shape[1] * self.conv2.kernel_size[0] * self.conv2.kernel_size[1])), 0)

        ws_1 = self.conv2.expand(
            math.ceil(extention[0] / (self.nz * self.conv1.kernel_size[0] * self.conv1.kernel_size[1])),
            n_add_out_conv2)
        n_in_conv3 = self.conv2.weight.shape[1]
        n_out_conv3 = self.conv3.weight.shape[1]
        n_params_conv3 = n_in_conv3 * n_out_conv3 * self.conv3.kernel_size[0] * self.conv3.kernel_size[1]
        n_add_out_conv3 = max(math.ceil((extention[2] - (n_params_conv3 - np.prod(self.conv3.weight.size()).item())) / (
                self.conv2.weight.data.shape[1] * self.conv3.kernel_size[0] * self.conv3.kernel_size[1])), 0)

        ws_2 = self.conv3.expand(n_add_out_conv2, n_add_out_conv3)
        self.cap_conv1.append(self.conv1.weight.shape[1])
        self.cap_conv2.append(self.conv2.weight.shape[1])
        self.cap_conv3.append(self.conv3.weight.shape[1])
        self.BatchNorms1.append(torch.nn.BatchNorm2d(self.conv1.weight.shape[1]).cuda())
        self.BatchNorms2.append(torch.nn.BatchNorm2d(self.conv2.weight.shape[1]).cuda())
        self.BatchNorms3.append(torch.nn.BatchNorm2d(self.conv3.weight.shape[1]).cuda())

        return [ws_0, ws_1, ws_2]

    def expand_embeddings(self, n_new_classes):
        prev_weight_shape_0 = self.conv1.expand_embeddings(n_new_classes)
        prev_weight_shape_1 = self.conv2.expand_embeddings(n_new_classes)
        prev_weight_shape_2 = self.conv3.expand_embeddings(n_new_classes)
        return [prev_weight_shape_0, prev_weight_shape_1, prev_weight_shape_2]

    def forward(self, input, t, lables=None, s=1, t_mix=None):
        task = torch.autograd.Variable(torch.LongTensor([t]).cuda())
        masks = self.mask(task, None, s=s)

        gc1, gc2, gc3 = masks
        # print(input.shape)
        x = self.conv1(input, gc1, self.nz, self.cap_conv1[t])
        x = self.BatchNorms1[t](x)
        x = self.ReLU(x)

        x = self.conv2(x, gc2, self.cap_conv1[t], self.cap_conv2[t])
        x = self.BatchNorms2[t](x)
        x = self.ReLU(x)

        x = self.conv3(x, gc3, self.cap_conv2[t], self.cap_conv3[t])
        x = self.BatchNorms3[t](x)
        x = self.ReLU(x)
        output = self.Tanh(self.last[t](x, None, self.cap_conv3[t], self.nc))
        return output, masks

    def mask(self, t, labels, s=1, test=False):
        gc1 = self.gate(s * self.conv1.ec(t))  # .view(self.conv1.out_channels, (
        # self.conv1.in_channels * self.conv1.kernel_size[0] * self.conv1.kernel_size[1]))# + 1))
        gc2 = self.gate(s * self.conv2.ec(t))  # .view(self.conv2.out_channels, (
        # self.conv2.in_channels * self.conv2.kernel_size[0] * self.conv2.kernel_size[1]))# + 1))
        gc3 = self.gate(s * self.conv3.ec(t))  # .view(self.conv3.out_channels, (
        # self.conv3.in_channels * self.conv3.kernel_size[0] * self.conv3.kernel_size[1]))# + 1))
        # gc4 = self.gate(s * self.conv4.ec(t))#.view(self.conv4.out_channels, (
        #         #self.conv4.in_channels * self.conv4.kernel_size[0] * self.conv4.kernel_size[1]))# + 1))
        return [gc1, gc2, gc3]  # ,gc4]

    def get_total_mask(self, t, labels, s):
        task = torch.autograd.Variable(torch.LongTensor(labels.data.cpu().numpy()).cuda())
        # t =  torch.autograd.Variable(torch.LongTensor([t]).cuda())
        masks = self.mask(task, None, s=s)
        m0 = torch.max(masks[0], 0)[0].view(1, -1)
        m1 = torch.max(masks[1], 0)[0].view(1, -1)
        m2 = torch.max(masks[2], 0)[0].view(1, -1)
        return [m0, m1, m2]  # + m3

    def get_view_for(self, n, masks):
        gc1, gc2, gc3 = masks
        if n == 'conv1.weight':
            return gc1.data[:, :].contiguous().view(self.conv1.weight.shape)
        elif n == 'conv1.bias':
            return gc1.data[:, 0].contiguous().view(-1)

        elif n == 'conv2.weight':
            return gc2.data[:, :].contiguous().view(self.conv2.weight.shape)
        elif n == 'conv2.bias':
            return gc2.data[:, 0].contiguous().view(-1)

        elif n == 'conv3.weight':
            return gc3.data[:, :].contiguous().view(self.conv3.weight.shape)
        elif n == 'conv3.bias':
            return gc3.data[:, 0].contiguous().view(-1)
        return None

class netD(nn.Module):
    def __init__(self, ndf, nc, out_class=1):
        super(netD, self).__init__()
        self.nc = nc
        self.momentum = 0.1
        self.LeakyReLU = nn.LeakyReLU(0.2, inplace=True)
        self.conv1 = nn.Conv2d(nc, ndf, 5, 2, 2, bias=False)
        self.conv2 = nn.Conv2d(ndf, ndf * 2, 5, 2, 2, bias=False)
        self.BatchNorm2 = BatchNorm2d(ndf * 2, momentum=self.momentum, track_running_stats=True)
        # self.conv3 = nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False)
        self.conv3 = nn.Conv2d(ndf * 2, ndf * 4, 5, 2, 2, bias=False)

        self.BatchNorm3 = BatchNorm2d(ndf * 4, momentum=self.momentum, track_running_stats=True)
        self.output_size = ndf * 4 * 4 * 4
        self.disc_linear = nn.Linear(self.output_size, 1)  # .append(nn.Linear(ndf, 1))
        self.aux_linear = nn.Linear(self.output_size, out_class)

        self.softmax = nn.Softmax(dim=1)
        self.sigmoid = nn.Sigmoid()
        self.ndf = ndf
        self.apply(weights_init)

    def forward(self, input):
        batch_size = input.size()[0]
        x = self.conv1(input)
        x = self.LeakyReLU(x)
        x = self.conv2(x)
        x = self.BatchNorm2(x)
        x = self.LeakyReLU(x)
        x = self.conv3(x)
        x = self.BatchNorm3(x)
        x = self.LeakyReLU(x)
        x = x.view(batch_size, -1)
        c = self.aux_linear(x)
        s = self.disc_linear(x)
        return s.view(-1), c

max_stars_repo_name: DavideA/machine-learning-dgm

max_issues_repo_name: DavideA/machine-learning-dgm

max_forks_repo_name: DavideA/machine-learning-dgm


text:
import numpy as np
import math

# TODO: somehow make nnrt import also take care of importing open3d, if that hasn't yet been imported
import open3d as o3d  # necessary to import nnrt!

from nnrt import backproject_depth_ushort as backproject_depth_ushort_c
from nnrt import backproject_depth_float as backproject_depth_float_c
from nnrt import warp_flow as warp_flow_c
from nnrt import warp_3d as warp_3d_c

from image_processing.numba_cuda.preprocessing import cuda_compute_normal
from image_processing.numpy_cpu.preprocessing import cpu_compute_normal

def in_bounds(point, h, w):
    return point[0] >= 0.0 and point[0] < w and point[1] >= 0.0 and point[1] < h

def warp_flow_py(image, flow, mask):
    # We assume:
    #               image shape (3, 480, 640)
    #               flow shape  (2, 480, 640)
    #               mask shape  (2, 480, 640)

    image_copy = np.copy(image)
    flow_copy = np.copy(flow)
    mask_copy = np.copy(mask)

    # print(type(image_copy), image_copy.dtype)

    image_copy = np.moveaxis(image_copy, 0, -1)  # (h, w, 3)
    flow_copy = np.moveaxis(flow_copy, 0, -1)
    mask_copy = np.moveaxis(mask_copy, 0, -1)

    h, w = flow_copy.shape[0], flow_copy.shape[1]

    assert image_copy.shape[0] == flow_copy.shape[0] and image_copy.shape[1] == flow_copy.shape[1]
    assert mask_copy.shape[0] == flow_copy.shape[0] and mask_copy.shape[1] == flow_copy.shape[1]

    image_warped = np.zeros(image_copy.shape, dtype=np.float32)
    weights_warped = np.zeros((h, w, 1), dtype=np.float32)

    for v in range(h):
        for u in range(w):

            if mask_copy is not None and (not mask_copy[v, u, 0] or not mask_copy[v, u, 1]):
                continue

            u_warped = u + flow_copy[v, u, 0]
            v_warped = v + flow_copy[v, u, 1]

            u0 = math.floor(u_warped)
            u1 = u0 + 1
            v0 = math.floor(v_warped)
            v1 = v0 + 1

            if not in_bounds((u0, v0), h, w) or not in_bounds((u0, v1), h, w) \
                    or not in_bounds((u1, v0), h, w) or not in_bounds((u1, v1), h, w):
                continue

            du = u_warped - u0
            dv = v_warped - v0

            w00 = (1 - du) * (1 - dv)
            w01 = (1 - du) * dv
            w10 = du * (1 - dv)
            w11 = du * dv

            c = image_copy[v, u]
            image_warped[v0, u0] += w00 * c
            image_warped[v0, u1] += w10 * c
            image_warped[v1, u0] += w01 * c
            image_warped[v1, u1] += w11 * c
            weights_warped[v0, u0] += w00
            weights_warped[v0, u1] += w10
            weights_warped[v1, u0] += w01
            weights_warped[v1, u1] += w11

    zero_weights = weights_warped == 0
    weights_warped[zero_weights] = 1.0
    image_warped /= weights_warped
    image_warped[np.repeat(zero_weights, 3, axis=2)] = 1.0

    return np.moveaxis(image_warped, -1, 0)  # (3, h, w)

def warp_flow(image, flow, mask):
    # We assume:
    #               image shape (3, h, w)
    #               flow shape  (2, h, w)
    #               mask shape  (2, h, w)

    assert image.shape[1] == flow.shape[1] and image.shape[2] == flow.shape[2]
    assert mask.shape[1] == flow.shape[1] and mask.shape[2] == flow.shape[2]

    image_warped = warp_flow_c(image, flow, mask)

    return image_warped

def warp_deform_py(image, pixel_anchors, pixel_weights, node_positions, node_rotations, node_translations, fx, fy, cx, cy):
    # We assume:
    #               image               shape (6, h, w)
    #               pixel_anchors       shape (h, w, 4)
    #               pixel_weights       shape (h, w, 4)
    #               canonical_node_positions      shape (num_nodes, 3)
    #               node_rotations      shape (num_nodes, 3, 3)
    #               node_translations   shape (num_nodes, 3)

    image_copy = np.copy(image)

    h = image.shape[1]
    w = image.shape[2]
    num_pixels = h * w

    fx, fy, cx, cy = modify_intrinsics_due_to_cropping(fx, fy, cx, cy, h, w, original_h=480, original_w=640)

    invalid_pixel_anchors = pixel_anchors < 0
    filtered_pixel_anchors = np.copy(pixel_anchors)
    filtered_pixel_anchors[invalid_pixel_anchors] = 0

    # Warp the image pixels using graph poses.
    image_points = image[3:, :, :]
    image_points = np.moveaxis(image_points, 0, -1).reshape(num_pixels, 3, 1)
    deformed_points = np.zeros((num_pixels, 3, 1), dtype=image.dtype)

    num_nodes = node_translations.shape[0]
    node_translations = node_translations.reshape(num_nodes, 3, 1)

    for k in range(4):
        node_idxs_k = pixel_anchors.reshape(num_pixels, 4)[:, k]
        nodes_k = node_positions[node_idxs_k].reshape(num_pixels, 3, 1)

        # Compute deformed point contribution.
        rotated_points_k = np.matmul(node_rotations[node_idxs_k], image_points - nodes_k)  # (num_pixels, 3, 1)
        deformed_points_k = rotated_points_k + nodes_k + node_translations[node_idxs_k]
        interpolation_weights = pixel_weights.reshape(num_pixels, 4)[:, k].reshape(num_pixels, 1, 1)
        deformed_points += np.repeat(interpolation_weights, 3, axis=1) * deformed_points_k  # (num_pixels, 3, 1)

    deformed_points = deformed_points.reshape(h, w, 3)
    deformed_points = np.moveaxis(deformed_points, -1, 0)

    # Compute the warped image.
    image_warped = np.zeros((h, w, 3), dtype=np.float32)
    weights_warped = np.zeros((h, w, 1), dtype=np.float32)

    for v in range(h):
        for u in range(w):
            p = image[3:, v, u].reshape(3, 1)
            if p[2] <= 0.0: continue

            p_def = deformed_points[:, v, u]
            if p_def[2] <= 0.0: continue

            if invalid_pixel_anchors[v, u, 0] and invalid_pixel_anchors[v, u, 1] and invalid_pixel_anchors[v, u, 2] and invalid_pixel_anchors[
                v, u, 3]:
                continue

            u_warped = fx * p_def[0] / p_def[2] + cx
            v_warped = fy * p_def[1] / p_def[2] + cy

            u0 = math.floor(u_warped)
            u1 = u0 + 1
            v0 = math.floor(v_warped)
            v1 = v0 + 1

            if not in_bounds((u0, v0), h, w) or not in_bounds((u0, v1), h, w) \
                    or not in_bounds((u1, v0), h, w) or not in_bounds((u1, v1), h, w):
                continue

            du = u_warped - u0
            dv = v_warped - v0

            w00 = (1 - du) * (1 - dv)
            w01 = (1 - du) * dv
            w10 = du * (1 - dv)
            w11 = du * dv

            c = image_copy[:3, v, u]

            image_warped[v0, u0] += w00 * c
            image_warped[v0, u1] += w10 * c
            image_warped[v1, u0] += w01 * c
            image_warped[v1, u1] += w11 * c
            weights_warped[v0, u0] += w00
            weights_warped[v0, u1] += w10
            weights_warped[v1, u0] += w01
            weights_warped[v1, u1] += w11

    zero_weights = weights_warped == 0
    weights_warped[zero_weights] = 1.0
    image_warped /= weights_warped
    image_warped[np.repeat(zero_weights, 3, axis=2)] = 1.0

    return np.moveaxis(image_warped, -1, 0)

def warp_deform(image, pixel_anchors, pixel_weights, node_positions, node_rotations, node_translations, fx, fy, cx, cy):
    # We assume:
    #               image               shape (6, h, w)
    #               pixel_anchors       shape (h, w, 4)
    #               pixel_weights       shape (h, w, 4)
    #               canonical_node_positions      shape (num_nodes, 3)
    #               node_rotations      shape (num_nodes, 3, 3)
    #               node_translations   shape (num_nodes, 3)

    image_copy = np.copy(image)

    h = image_copy.shape[1]
    w = image_copy.shape[2]
    num_pixels = h * w

    fx, fy, cx, cy = modify_intrinsics_due_to_cropping(fx, fy, cx, cy, h, w, original_h=480, original_w=640)

    # invalid_pixel_anchors = pixel_anchors < 0
    # filtered_pixel_anchors = np.copy(pixel_anchors)
    # filtered_pixel_anchors[invalid_pixel_anchors] = 0

    # Warp the image pixels using graph poses.
    image_points = image_copy[3:, :, :]
    image_points = np.moveaxis(image_points, 0, -1).reshape(num_pixels, 3, 1)
    deformed_points = np.zeros((num_pixels, 3, 1), dtype=image_copy.dtype)

    num_nodes = node_translations.shape[0]
    node_translations = node_translations.reshape(num_nodes, 3, 1)

    for k in range(4):
        node_idxs_k = pixel_anchors.reshape(num_pixels, 4)[:, k]
        nodes_k = node_positions[node_idxs_k].reshape(num_pixels, 3, 1)

        # Compute deformed point contribution.
        rotated_points_k = np.matmul(node_rotations[node_idxs_k], image_points - nodes_k)  # (num_pixels, 3, 1)
        deformed_points_k = rotated_points_k + nodes_k + node_translations[node_idxs_k]
        interpolation_weights = pixel_weights.reshape(num_pixels, 4)[:, k].reshape(num_pixels, 1, 1)
        deformed_points += np.repeat(interpolation_weights, 3, axis=1) * deformed_points_k  # (num_pixels, 3, 1)

    deformed_points = deformed_points.reshape(h, w, 3)
    deformed_points = np.moveaxis(deformed_points, -1, 0)

    # point_validity = np.logical_not(np.all(invalid_pixel_anchors, axis=2))
    point_validity = ~np.any(pixel_anchors == -1, axis=2)

    # Compute the warped image.
    image_warped = warp_3d_c(image_copy, deformed_points, point_validity, fx, fy, cx, cy)

    return image_warped

def warp_deform_3d(image, pixel_anchors, pixel_weights, node_positions, node_rotations, node_translations):
    # We assume:
    #               image               shape (6, h, w)
    #               pixel_anchors       shape (h, w, 4)
    #               pixel_weights       shape (h, w, 4)
    #               canonical_node_positions      shape (num_nodes, 3)
    #               node_rotations      shape (num_nodes, 3, 3)
    #               node_translations   shape (num_nodes, 3)

    assert image.shape[0] == 6 and len(image.shape) == 3, image.shape
    assert pixel_anchors.shape[-1] == 4 and len(pixel_anchors.shape) == 3
    assert pixel_weights.shape[-1] == 4 and len(pixel_weights.shape) == 3
    assert node_positions.shape[-1] == 3
    assert node_rotations.shape[1] == 3 and node_rotations.shape[2] == 3
    assert node_translations.shape[-1] == 3

    image_copy = np.copy(image)

    h = image_copy.shape[1]
    w = image_copy.shape[2]
    num_pixels = h * w

    # Warp the image pixels using graph poses.
    image_points = image_copy[3:, :, :]
    image_points = np.moveaxis(image_points, 0, -1).reshape(num_pixels, 3, 1)
    deformed_points = np.zeros((num_pixels, 3, 1), dtype=image_copy.dtype)

    num_nodes = node_translations.shape[0]
    node_translations = node_translations.reshape(num_nodes, 3, 1)

    for k in range(4):
        node_idxs_k = pixel_anchors.reshape(num_pixels, 4)[:, k]
        nodes_k = node_positions[node_idxs_k].reshape(num_pixels, 3, 1)

        # Compute deformed point contribution.
        rotated_points_k = np.matmul(node_rotations[node_idxs_k], image_points - nodes_k)  # (num_pixels, 3, 1)
        deformed_points_k = rotated_points_k + nodes_k + node_translations[node_idxs_k]
        interpolation_weights = pixel_weights.reshape(num_pixels, 4)[:, k].reshape(num_pixels, 1, 1)
        deformed_points += np.repeat(interpolation_weights, 3, axis=1) * deformed_points_k  # (num_pixels, 3, 1)

    deformed_points = deformed_points.reshape(h, w, 3)
    deformed_points = np.moveaxis(deformed_points, -1, 0)

    point_validity = np.all(pixel_anchors != -1, axis=2)
    point_validity = np.repeat(point_validity.reshape(1, h, w), 3, axis=0)

    deformed_points[~point_validity] = 0.0

    return deformed_points

def modify_intrinsics_due_to_cropping(fx, fy, cx, cy, h, w, original_h=480, original_w=640):
    # Modify intrinsics
    delta_height = (h - original_h) / 2
    cy_modified = cy + delta_height

    delta_width = (w / original_w) / 2
    cx_modified = cx + delta_width

    return fx, fy, cx_modified, cy_modified

def backproject_depth_py(depth_image, fx, fy, cx, cy, normalizer=1000.0):
    assert len(depth_image.shape) == 2
    width = depth_image.shape[1]
    height = depth_image.shape[0]

    point_image = np.zeros((3, height, width))
    for y in range(height):
        for x in range(width):
            depth = depth_image[y, x] / normalizer
            if depth > 0:
                pos_x = depth * (x - cx) / fx
                pos_y = depth * (y - cy) / fy
                pos_z = depth

                point_image[0, y, x] = pos_x
                point_image[1, y, x] = pos_y
                point_image[2, y, x] = pos_z

    return point_image

def backproject_depth(depth_image, fx, fy, cx, cy, depth_scale=1000.0):
    assert len(depth_image.shape) == 2
    width = depth_image.shape[1]
    height = depth_image.shape[0]

    if depth_image.dtype == np.float32:
        point_image = np.zeros((height, width, 3), dtype=np.float32)
        backproject_depth_float_c(depth_image, point_image, fx, fy, cx, cy)
    else:
        point_image = np.zeros((height, width, 3), dtype=np.float32)
        backproject_depth_ushort_c(depth_image, point_image, fx, fy, cx, cy, depth_scale)
    return point_image

def compute_boundary_mask(point_image: np.ndarray, max_distance: float) -> np.ndarray:
    """
    Pont image expected shape: (3, h, w)
    Resulting mask filters out points where either:
        distance between pixel-wise left & right neighbors exceeds max_distance
        - or -
        distance between pixel-wise above & below neighbors exceeds max_distance.

    In the following diagram, a, b, c, d, and e are pixel-neighboring points.

           [d]
            ||
    [a] == [b] == [c]
            ||
           [e]
    If dist(x,y) yields distance between point x & y and coord(x) gives
    the pixel-coordinate of point x:

    mask[coord[b]] = dist(a,c) > max_distance | dist(e,d) > max_distance

    :param point_image: input point image
    :param max_distance: maximum distance threshold
    :return: binary mask based on the distances, as described above
    """
    shift_right = np.zeros_like(point_image)
    shift_left = np.zeros_like(point_image)
    shift_down = np.zeros_like(point_image)
    shift_up = np.zeros_like(point_image)

    # points shifted to the right (x-1)
    shift_right[:, 1:, :] = point_image[:, :-1, :]
    shift_left[:, :-1, :] = point_image[:, 1:, :]
    shift_down[1:, :, :] = point_image[:-1, :, :]
    shift_up[:-1, :, :] = point_image[1:, :, :]

    horizontal_dist = np.linalg.norm(shift_left - shift_right, axis=2)
    vertical_dist = np.linalg.norm(shift_up - shift_down, axis=2)

    assert np.isfinite(horizontal_dist).all()
    assert np.isfinite(vertical_dist).all()

    boundary_horizontal_mask = horizontal_dist > max_distance
    boundary_vertical_mask = vertical_dist > max_distance

    boundary_mask = boundary_horizontal_mask | boundary_vertical_mask

    # the slice with [None, ...] here reshapes the mask to shape ( 1, h, w )
    return boundary_mask[None, ...]

def compute_normals(device: o3d.core.Device, vertex_map: np.ndarray) -> np.ndarray:
    if device.get_type() == o3d.core.Device.DeviceType.CPU:
        return cpu_compute_normal(vertex_map)
    elif device.get_type() == o3d.core.Device.DeviceType.CUDA:
        return cuda_compute_normal(vertex_map)
    else:
        raise ValueError(f"Unsupported device type: {device.get_type()}")

ext: py

max_stars_repo_name: Algomorph/NeuralTracking

max_issues_repo_name: Algomorph/NeuralTracking

max_forks_repo_name: Algomorph/NeuralTracking

reason: import numpy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.1.5
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # s_checklist_scenariobased_step10 [<img src="https://www.arpm.co/lab/icons/icon_permalink.png" width=30 height=30 style="display: inline;">](https://www.arpm.co/lab/redirect.php?code=s_checklist_scenariobased_step10&codeLang=Python)
# For details, see [here](https://www.arpm.co/lab/redirect.php?permalink=ex-vue-10).

# +
import numpy as np
import pandas as pd
from scipy.stats import norm
import matplotlib.pyplot as plt

from arpym.statistics import meancov_sp
from arpym.tools import add_logo
# -

# ## [Input parameters](https://www.arpm.co/lab/redirect.php?permalink=s_checklist_scenariobased_step10-parameters)

# +
# indicates which projection to continue from
# True: use copula-marginal projections
# False: use historical projections
copula_marginal = True

q_now = 0  # initial volume time
q_end = 1  # final volume time
k_ = 300  # number of elements in the q grid
l_ = 500  # number of elements in the beta grid
alpha = 1  # parameter of the power slippage component
gamma = 3.14e-5 # permanent impact parameter
eta = 1.42e-6  # temporary impact parameter
c = 0.95  # confidence level for quantile satisfaction measure

n_plot = 2  # index of instrument to plot
# -

# ## [Step 0](https://www.arpm.co/lab/redirect.php?permalink=s_checklist_scenariobased_step10-implementation-step00): Load data

# +
path = '../../../databases/temporary-databases/'

# Risk drivers identification
db_riskdrivers_tools = pd.read_csv(path+'db_riskdrivers_tools.csv')
n_stocks = int(db_riskdrivers_tools['n_stocks'][0])
n_bonds = int(db_riskdrivers_tools.n_bonds[0])
n_ = n_stocks+n_bonds+3

# Aggregation
db_holdings = pd.read_csv(path+'db_holdings.csv')
h = np.squeeze(db_holdings.values)  # initial holdings

if copula_marginal:
    # Projection
    db_scenprob = pd.read_csv(path+'db_scenario_probs.csv')
    p = db_scenprob.p.values

    # Pricing
    # import daily P&Ls computed in step 5 with m_=1
    db_pi_oneday = pd.read_csv(path+'db_oneday_pl.csv')
    pi_oneday = db_pi_oneday.values

    # Construction
    db_final_portfolio = pd.read_csv(path+'db_final_portfolio.csv')
    # the final portfolio is the one obtained in the construction step,
    # that optimizes the cVaR satisfaction measure
    h_qsi = np.squeeze(db_final_portfolio.values)

else:
    # Projection
    db_scenprob = pd.read_csv(path+'db_scenario_probs_bootstrap.csv')
    p = db_scenprob.p.values

    # Pricing
    # import daily P&Ls computed in step 5 with m_=1
    db_pi_oneday = pd.read_csv(path+'db_oneday_pl_historical.csv')
    pi_oneday = db_pi_oneday.values

    # Construction
    db_final_portfolio = pd.read_csv(path+'db_final_portfolio_historical.csv')
    # the final portfolio is the one obtained in the construction step,
    # that optimizes the cVaR satisfaction measure
    h_qsi = np.squeeze(db_final_portfolio.values)

# start portfolio
h_qnow = h
# final portfolio
h_qend = h_qsi
# -

# ## [Step 1](https://www.arpm.co/lab/redirect.php?permalink=s_checklist_scenariobased_step10-implementation-step01): Find trajectory

# +
sigma2 = np.zeros(n_)
variance_pihat = np.zeros((n_, l_))
mean_pihat = np.zeros((n_, l_))
xi = np.zeros(l_)
traj = np.zeros((n_, l_, k_))

# size of parent order
delta_h_parent = (h_qend - h_qnow).astype('int')
# beta grid
beta = np.linspace(alpha/(1+alpha), 1, l_+1, endpoint=True)
beta = beta[1:]
# q grid
q_grid = np.linspace(q_now, q_end, k_)

for n in range(n_):
    if delta_h_parent[n] == 0:
        # no change in holdings
        traj[n, :, :] = np.tile(h_qend[n], (l_, k_))
    else:
        _, sigma2[n] = meancov_sp(pi_oneday[:, n], p)
        for l in range(l_):
            # expected P&L
            xi[l] = beta[l]**(alpha+1)/(beta[l]+beta[l]*alpha-alpha)
            mean_pihat[n, l] = gamma/2*(h_qend[n]**2 - h_qnow[n]**2) - \
                eta*xi[l]*np.abs(delta_h_parent[n])**(1+alpha) * \
                (q_end-q_now)**(-alpha)
            # P&L variance
            variance_pihat[n, l] = sigma2[n] * (q_end-q_now) * \
                (h_qnow[n]**2 + 2*h_qnow[n]*delta_h_parent[n]/(beta[l]+1) +
                (delta_h_parent[n]**2)/(2*beta[l]+1))
            # trajectory
            traj[n, l, :] = h_qnow[n] + \
                ((q_grid-q_now)/(q_end-q_now))**beta[l]*delta_h_parent[n]
# -

# ## [Step 2](https://www.arpm.co/lab/redirect.php?permalink=s_checklist_scenariobased_step10-implementation-step02): Maximization

q_satis = np.zeros((n_, l_))
beta_star = np.zeros(n_)
l_star = np.zeros(n_)
for n in range(n_):
    if delta_h_parent[n] == 0:
        # no change in holdings
        beta_star[n] = beta[-1]
    else:
        # quantile satisfaction measure
        for l in range(l_):
            q_satis[n, l] = mean_pihat[n, l] + \
                            np.sqrt(variance_pihat[n, l])*norm.ppf(1-c)
        # beta corresponding to the optimal liquidation trajectory
        l_star[n] = \
            np.where(q_satis[n, :] == np.max(q_satis[n, :]))[0]
        beta_star[n] = beta[np.int(l_star[n])]

# ## Plots

# plot execution trajectories
plt.style.use('arpm')
fig = plt.figure(figsize=(1280.0/72.0, 720.0/72.0), dpi = 72.0)
for i in range(0, l_, 50):
    plt.plot(q_grid, traj[n_plot-1, i, :]*1e-6, color='grey')
plt.plot(q_grid, traj[n_plot-1, np.int(l_star[n_plot-1]), :]*1e-6,
         color='red')
plt.title('Optimal trading trajectory - ' + db_pi_oneday.columns[n_plot-1],
         fontsize=20, fontweight='bold')
plt.xlabel('Volume time', fontsize=17)
plt.ylabel('Holdings (million units)', fontsize=17)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlim(0,1)
add_logo(fig, location=1, set_fig_size=False)

max_stars_repo_name: dpopadic/arpmRes

max_issues_repo_name: dpopadic/arpmRes

max_forks_repo_name: dpopadic/arpmRes

reason: import numpy,from scipy


text:
# ----------------------------------------------------------------------------
# -                        Open3D: www.open3d.org                            -
# ----------------------------------------------------------------------------
# The MIT License (MIT)
#
# Copyright (c) 2018-2021 www.open3d.org
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
# ----------------------------------------------------------------------------

import os
import open3d as o3d
import numpy as np
import pytest

def test_set_mesh_data_deserialization():
    """Tests the deserialization of messages created with the set_mesh_data_function.
    """

    def set_mesh_data_to_geometry(*args, **kwargs):
        bc = o3d.io.rpc.BufferConnection()
        o3d.io.rpc.set_mesh_data(*args, **kwargs, connection=bc)
        return o3d.io.rpc.data_buffer_to_meta_geometry(bc.get_buffer())

    rng = np.random

    verts = rng.rand(100, 3).astype(np.float32)
    tris = rng.randint(0, 100, size=[71, 3]).astype(np.int32)
    lines = rng.randint(0, 100, size=[82, 2]).astype(np.int32)

    dtypes = [np.uint8, np.int16, np.int32, np.float32, np.float64]
    vert_attrs = {
        'a':
            rng.uniform(0, 256, size=verts.shape).astype(rng.choice(dtypes)),
        'b':
            rng.uniform(0, 256,
                        size=verts.shape[:1] + (7,)).astype(rng.choice(dtypes)),
        'c':
            rng.uniform(0, 256,
                        size=verts.shape + (2,)).astype(rng.choice(dtypes)),
    }

    tri_attrs = {
        'a':
            rng.uniform(0, 256,
                        size=tris.shape[:1] + (1,)).astype(rng.choice(dtypes)),
        'b':
            rng.uniform(0, 256,
                        size=tris.shape[:1] + (3,)).astype(rng.choice(dtypes)),
        'c':
            rng.uniform(0, 256,
                        size=tris.shape[:1] + (5,)).astype(rng.choice(dtypes)),
    }

    line_attrs = {
        'a':
            rng.uniform(0, 256,
                        size=lines.shape[:1] + (1,)).astype(rng.choice(dtypes)),
        'b':
            rng.uniform(0, 256, size=lines.shape[:1] + (1, 2, 3)).astype(
                rng.choice(dtypes)),
        'c':
            rng.uniform(0, 256,
                        size=lines.shape[:1] + (2,)).astype(rng.choice(dtypes)),
    }

    path, time, geom = set_mesh_data_to_geometry(vertices=verts,
                                                 vertex_attributes=vert_attrs,
                                                 path="pcd",
                                                 time=123)
    assert path == "pcd"
    assert time == 123
    assert isinstance(geom, o3d.t.geometry.PointCloud)
    np.testing.assert_equal(geom.point['positions'].numpy(), verts)
    for k, v in vert_attrs.items():
        np.testing.assert_equal(geom.point[k].numpy(), v)

    path, time, geom = set_mesh_data_to_geometry(vertices=verts,
                                                 faces=tris,
                                                 vertex_attributes=vert_attrs,
                                                 face_attributes=tri_attrs,
                                                 path="trimesh",
                                                 time=123)
    assert path == "trimesh"
    assert time == 123
    assert isinstance(geom, o3d.t.geometry.TriangleMesh)
    np.testing.assert_equal(geom.vertex['positions'].numpy(), verts)
    np.testing.assert_equal(geom.triangle['indices'].numpy(), tris)
    for k, v in vert_attrs.items():
        np.testing.assert_equal(geom.vertex[k].numpy(), v)
    for k, v in tri_attrs.items():
        np.testing.assert_equal(geom.triangle[k].numpy(), v)

    path, time, geom = set_mesh_data_to_geometry(vertices=verts,
                                                 lines=lines,
                                                 vertex_attributes=vert_attrs,
                                                 line_attributes=line_attrs,
                                                 path="lines",
                                                 time=123)
    assert path == "lines"
    assert time == 123
    assert isinstance(geom, o3d.t.geometry.LineSet)
    np.testing.assert_equal(geom.point['positions'].numpy(), verts)
    np.testing.assert_equal(geom.line['indices'].numpy(), lines)
    for k, v in vert_attrs.items():
        np.testing.assert_equal(geom.point[k].numpy(), v)
    for k, v in line_attrs.items():
        np.testing.assert_equal(geom.line[k].numpy(), v)

    #
    # Test partial data
    #
    path, time, geom = set_mesh_data_to_geometry(vertex_attributes=vert_attrs,
                                                 path="pcd",
                                                 time=123,
                                                 o3d_type="PointCloud")
    assert path == "pcd"
    assert time == 123
    assert isinstance(geom, o3d.t.geometry.PointCloud)
    for k, v in vert_attrs.items():
        np.testing.assert_equal(geom.point[k].numpy(), v)

    path, time, geom = set_mesh_data_to_geometry(vertex_attributes=vert_attrs,
                                                 face_attributes=tri_attrs,
                                                 path="trimesh",
                                                 time=123,
                                                 o3d_type="TriangleMesh")
    assert path == "trimesh"
    assert time == 123
    assert isinstance(geom, o3d.t.geometry.TriangleMesh)
    for k, v in vert_attrs.items():
        np.testing.assert_equal(geom.vertex[k].numpy(), v)
    for k, v in tri_attrs.items():
        np.testing.assert_equal(geom.triangle[k].numpy(), v)

    path, time, geom = set_mesh_data_to_geometry(vertex_attributes=vert_attrs,
                                                 line_attributes=line_attrs,
                                                 path="lines",
                                                 time=123,
                                                 o3d_type="LineSet")
    assert path == "lines"
    assert time == 123
    assert isinstance(geom, o3d.t.geometry.LineSet)
    for k, v in vert_attrs.items():
        np.testing.assert_equal(geom.point[k].numpy(), v)
    for k, v in line_attrs.items():
        np.testing.assert_equal(geom.line[k].numpy(), v)

    # Without o3d_type and no primary key data the returned object is None
    path, time, geom = set_mesh_data_to_geometry(vertex_attributes=vert_attrs,
                                                 path="unknown",
                                                 time=123,
                                                 o3d_type="")
    assert path == "unknown"
    assert time == 123
    assert geom is None


text:
#  Copyright (c) 2022, Chair of Software Technology
#  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 University Mannheim 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 OWNER 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.
import unittest

import numpy as np
from PIL import Image

from insynth.perturbators.image import ImageFlipPerturbator
from insynth.perturbators.pipeline import PipelinePerturbator

class TestPipeline(unittest.TestCase):
    def _generate_random_image(self):
        imarray = np.random.rand(100, 100, 3) * 255
        im = Image.fromarray(imarray.astype('uint8')).convert('RGBA')
        return im

    def test_PipelinePerturbator(self):
        input_image = self._generate_random_image()

        flip_perturbator = ImageFlipPerturbator(p=1.0, transformation_type='flip')
        mirror_perturbator = ImageFlipPerturbator(p=1.0, transformation_type='mirror')
        pipeline = PipelinePerturbator([flip_perturbator, mirror_perturbator])

        output_image = pipeline.apply(input_image)

        np.testing.assert_array_equal(output_image, np.flip(input_image, axis=(0, 1)))

if __name__ == '__main__':
    unittest.main()

max_stars_repo_name: mlxyz/insynth

max_issues_repo_name: mlxyz/insynth

max_forks_repo_name: mlxyz/insynth


text:
from __future__ import division, print_function

import logging
import numpy as np

FORMAT = '[%(asctime)s] %(name)-15s %(message)s'
DATEFMT = "%H:%M:%S"
logging.basicConfig(format=FORMAT, datefmt=DATEFMT, level=logging.INFO)

import os
import theano
import theano.tensor as T
import fuel
import ipdb
import time
import cPickle as pickle

from argparse import ArgumentParser
from theano import tensor

from fuel.streams import DataStream
from fuel.schemes import SequentialScheme
from fuel.transformers import Flatten

from blocks.algorithms import GradientDescent, CompositeRule, StepClipping, RMSProp, Adam
from blocks.bricks import Tanh, Identity
from blocks.bricks.cost import BinaryCrossEntropy
from blocks.bricks.recurrent import SimpleRecurrent, LSTM
from blocks.initialization import Constant, IsotropicGaussian, Orthogonal
from blocks.filter import VariableFilter
from blocks.graph import ComputationGraph
from blocks.roles import PARAMETER
from blocks.monitoring import aggregation
from blocks.extensions import FinishAfter, Timing, Printing, ProgressBar
from blocks.extensions.saveload import Checkpoint
from blocks.extensions.monitoring import DataStreamMonitoring, TrainingDataMonitoring
from blocks.main_loop import MainLoop
from blocks.model import Model

try:
    from blocks.extras import Plot
except ImportError:
    pass

import draw.datasets as datasets
from draw.draw import *
from draw.samplecheckpoint import SampleCheckpoint
from draw.partsonlycheckpoint import PartsOnlyCheckpoint

import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from sklearn.manifold import TSNE, MDS
from sklearn import cross_validation

from PIL import Image

import unittest

from draw.attention import *

import SVRT_analysis_helper_functions
from draw.datasets import package_sketch_images

def main(dataset,model_dir,model_file):

	#dataset = 'sketch_uint8_shuffle'
	#data_dir = '/Users/drewlinsley/Desktop/res_results_problem_4'
	data_dir = os.path.join('/Users/drewlinsley/Desktop/',dataset)
    #data_dir = '/home/ubuntu/svrt_data/'+dataset

	#Get shuffled data
	test_X, train_X, test_y, train_y = package_sketch_images.import_sketch(data_dir)
	image_size = (int(np.sqrt(test_X.shape[1])),int(np.sqrt(test_X.shape[1])))
	channels = 1
	target_categories = np.unique(train_y).shape[0]

	#Load old model
	#model_file = 'new_class_test-16_it/new_class_test_model'
	model_pointer = os.path.join(model_dir,model_file)
	with open(model_pointer,"rb") as f:
	    model = pickle.load(f)
	draw = model.get_top_bricks()[0]

	#Get attention model extractions
	images = T.matrix('images')

	probs, h_enc, c_enc, center_y, center_x, delta = draw.reconstruct(images)
	do_sample = theano.function(inputs=[images],outputs=[probs, h_enc, c_enc, center_y, center_x, delta],allow_input_downcast=True)
	tprobs, th_enc, tc_enc, tcenter_y, tcenter_x, tdelta = do_sample(test_X)
	correct_ims = tprobs.argmax(axis=1) == test_y.flatten()
	acc = np.sum(correct_ims) / test_X.shape[0]
	incorrect_ims = ~correct_ims

	#Save everything
	outfile = os.path.join(model_dir,'classification_data')
	np.savez(outfile,tprobs=tprobs, correct_ims=correct_ims, test_y=test_y, acc=acc)

	#Plot everything
	max_ims = 100
	from matplotlib.collections import LineCollection
	from SVRT_analysis_helper_functions import scale_range
	from SVRT_analysis_helper_functions import plot_ims
	# get the colors from the color map
	test_image_res = test_X#[0:batch_size,:].reshape(batch_size,image_size[0],image_size[1])
	#Make a folder
	subdir = os.path.join(model_dir,"classification_trace_errors")
	if not os.path.exists(subdir):
	    os.makedirs(subdir)
	plot_ims(test_image_res,tcenter_x,tcenter_y,tdelta,incorrect_ims,image_size,max_ims,subdir)
	#Make a folder
	subdir = os.path.join(model_dir,"classification_trace_correct")
	if not os.path.exists(subdir):
	    os.makedirs(subdir)
	plot_ims(test_image_res,tcenter_x,tcenter_y,tdelta,correct_ims,image_size,max_ims,subdir)
	#Which vars to save? tprobs, correct_ims, test_y, acc
	print("finished")

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--dataset", type=str, dest="dataset",
                default=False, help="directory pointer to images")
    parser.add_argument("--model_dir", type=str, dest="model_dir",
                default=False, help="directory pointer to model")
    parser.add_argument("--model_file", type=str, dest="model_file",
                default=False, help="name of model file")
    args = parser.parse_args()
    main(**vars(args))

reason: import numpy,from scipy,import theano,from theano


text:
import functools
import os
from typing import Any, Callable, Dict, Literal, Optional, Sequence, Tuple

import numpy as np
import tensorflow as tf
import torch
import GPUtil

from rlo import rewrites, search_tree, utils
from rlo.best_results import TypeBestCost
from rlo.config_utils import kwargs_from_config
from rlo.expr_sets import ExpressionSet, get_expression_set
from rlo.expression import Expression

from rlo.model.layers import global_pooling_from_str
from rlo.model.model import TorchModelWrapper, Model, ModelState
from rlo.model.state_value_estimator import StateValueModel
from rlo.pipelines.graph_pipeline import EdgeType
from rlo.expression_util import NamedExprWithEnv
from rlo.search_ops import AbstractSearcher
from rlo.tf_model import ModelWrapper, DualModelWrapper, Weights
from rlo.torch_graph_data import DataConverter
from rlo.training import TrainingConfig, train_model

ConfigType = Dict[str, Any]

def get_train_and_eval_exprs(
    config: ConfigType,
) -> Tuple[ExpressionSet, ExpressionSet]:
    """Return a tuple of
    (1) the training expressions
    (2) the evaluation expressions, i.e. test + train
    """
    train_exprs = get_expression_set(
        config["train_exprs"], take_defs=config.get("train_on_defs")
    )
    test_exprs = get_expression_set(
        config["test_exprs"], take_defs=config.get("test_on_defs")
    )

    # Take a union so that we evaluate on both train and test expressions
    # using the same searcher. Success metrics are calculated and plotted
    # separatedly.
    return train_exprs, train_exprs.union(test_exprs)

def alpha_schedule_from_config(config: ConfigType):
    from rlo.rollouts import AlphaSchedule

    kwargs = dict(
        init_alpha=1.0,
        max_alpha=10.0,
        alpha_scaling_factor=1.1,
        alpha_scaling_factor_fail=1.0,
    )
    kwargs.update(
        {
            k: config[k]
            for k in (
                "init_alpha",
                "max_alpha",
                "alpha_scaling_factor",
                "alpha_scaling_factor_fail",
            )
            if k in config
        }
    )
    return AlphaSchedule(
        init_value=kwargs["init_alpha"],
        max_value=kwargs["max_alpha"],
        scale_factor_success=kwargs["alpha_scaling_factor"],
        scale_factor_fail=kwargs["alpha_scaling_factor_fail"],
    )

PhaseType = Literal["train", "eval"]

def rollout_searcher_from_config(
    config: ConfigType, phase: PhaseType
) -> AbstractSearcher:
    required = (
        "num_episode_clusters",
        "num_positive_examples",
        "alpha_test",
    )
    from rlo.rollouts import ValueBasedRolloutSearcher

    optional = ("batch_size",)
    renames = (
        ("search_batch_size", "batch_size"),
        (f"max_num_episodes_{phase}", "max_num_episodes"),
        (f"simulation_depth_{phase}", "simulation_depth"),
    )
    kwargs = kwargs_from_config(config, required, optional, renames)
    if "maxing" in config:
        kwargs["maxing"] = maxing_algs[config["maxing"]]
    return ValueBasedRolloutSearcher(
        rules=rewrites.get_rules(config["rules"]),
        alpha_schedule=alpha_schedule_from_config(config),
        **kwargs,
    )

def searcher_kwargs_from_config(config: ConfigType, phase: PhaseType) -> Dict[str, Any]:
    """ Arguments general to most searchers, with the exception of Rollout. """
    required = ()
    optional = (
        "num_episode_clusters",
        "cost_per_step",
    )
    renames = (
        ("search_batch_size", "batch_size"),
        (f"max_gnn_{phase}", "max_gnn"),
        (f"simulation_depth_{phase}", "simulation_depth"),
    )
    kwargs = kwargs_from_config(config, required, optional, renames)
    if "maxing" in config:
        kwargs["maxing"] = maxing_algs[config["maxing"]]
    kwargs["rules"] = rewrites.get_rules(config["rules"])
    return kwargs

def astar_searcher_from_config(
    config: ConfigType, phase: PhaseType
) -> AbstractSearcher:
    from rlo.astar_search import AStarSearcher

    return AStarSearcher(**searcher_kwargs_from_config(config, phase))

def beam_searcher_from_config(config: ConfigType, phase: PhaseType) -> AbstractSearcher:
    from rlo.beam_search import PseudoBeamSearcher

    return PseudoBeamSearcher(**searcher_kwargs_from_config(config, phase))

def hybrid_searcher_from_config(
    config: ConfigType, phase: PhaseType
) -> AbstractSearcher:
    from rlo.hybrid_search import HybridSearcher, MergeHandling

    required = ()
    optional = ()
    renames = (
        ("hybrid_prob_rollout", "prob_rollout"),
        ("hybrid_alpha", "alpha"),
    )
    kwargs = {
        **searcher_kwargs_from_config(config, phase),
        **kwargs_from_config(config, required, optional, renames),
    }
    if "hybrid_merge_handling" in config:
        kwargs["merge_handling"] = MergeHandling[config["hybrid_merge_handling"]]
    return HybridSearcher(**kwargs)

def searcher_from_config(
    config: ConfigType, phase: PhaseType = "train"
) -> AbstractSearcher:
    search = config[f"{phase}_search"]
    return search_algs[search](config, phase)

def tf_distillator_from_config(config):
    from rlo.distillator import Distillator
    from rlo.distillator_v2 import DistillatorV2

    kwargs = kwargs_from_config(
        config, ("max_epochs", "min_epochs", "patience_epochs"), ("split", "verbose"),
    )
    sparse = config.get("sparse_gnn", False)
    if sparse:
        kwargs["max_nodes_per_batch"] = config["max_nodes_per_batch"]
    else:
        kwargs["batch_size"] = config.get("batch_size", 16)
    if config.get("v2", False):
        kwargs["tb_dir"] = config.get("tb_dir")
        return DistillatorV2(**kwargs)
    else:
        if "verbose" in kwargs:
            del kwargs["verbose"]
        return Distillator(**kwargs)

def distillator_from_config(config):
    if config["tensorflow"]:
        return tf_distillator_from_config(config)

    # Pytorch. Sparse GNN only, so always batch by num_nodes.
    kwargs = kwargs_from_config(
        config,
        ("max_epochs", "min_epochs", "max_nodes_per_batch", "patience_epochs"),
        ("split",),
    )
    training_config = TrainingConfig(**kwargs)
    return functools.partial(train_model, config=training_config)

def dataset_refiner_factory_from_config(config: ConfigType):
    from rlo import dataset_refiner

    return dataset_refiner.construct_dataset_refiner_factory(
        [
            dataset_refiner.get_dataset_refiner_creator_by_name(class_name)
            for class_name in config["dataset_refiners"]
        ]
    )

def model_save_dir_from_config(config: ConfigType) -> str:
    # Support legacy configs with separate model_save_path
    return config.get(
        "model_save_path", os.path.join(config["result_save_path"], "saved_models")
    )

def rng_factory_from_config(
    config: ConfigType,
) -> Callable[[int], np.random.Generator]:
    def generation_rng(seed: int, generation: int) -> np.random.Generator:
        root_rng = utils.rng(seed)
        for _ in range(generation):
            # Skip generations. If this becomes expensive, we could consider
            # something like utils.rng(nth_prime(seed == repetition) * (generation+1))
            utils.seed(root_rng)
        return utils.rng(root_rng)

    return functools.partial(generation_rng, seed_from_config(config))

def expression_shuffler_from_config(
    config: ConfigType, train_exprs: Sequence[NamedExprWithEnv]
):
    from rlo.expr_shuffler import ExpressionShuffler

    return ExpressionShuffler(
        rng_factory_from_config(config), train_exprs, config["exprs_per_generation"]
    )

def best_cost_from_config(config):
    """ Returns a function that returns the best oracle or declared cost
        for a given expression taking into account time_budget if applicable.
        The function will first try to look up the oracle results in best_results.py.
        If that fails, it will try to look up the best costs declared in
        train / test expression files. If that fails too, None is returned.
    """
    rules_name = config["rules"]
    _, eval_set = get_train_and_eval_exprs(config)

    @functools.lru_cache(maxsize=None)
    def best_cost_func(expr: str, time_budget: int) -> TypeBestCost:
        res = eval_set.best_cost_for_expression(expr, rules_name, time_budget)
        if res is None:
            raise ValueError(
                f"Could not find the best cost for {expr} at time budget of {time_budget}"
            )
        return res

    return best_cost_func

def azureml_reporter_from_config(config, train_set: ExpressionSet):
    from rlo.reporting.azureml import AzuremlReporter
    from rlo.reporting.utils import SearchOptimalityMetricFunction

    return AzuremlReporter(
        config["repetition"],
        config["num_repetitions"],
        SearchOptimalityMetricFunction(best_cost_from_config(config)),
        [name for name, _ in train_set.named_exprenvs()],
    )

def simul_search_curriculum_from_config(
    config: ConfigType, train_set: ExpressionSet, eval_set: ExpressionSet,
):
    from rlo.simul_search import SimultaneousSearchCurriculum

    return SimultaneousSearchCurriculum(
        train_exprs_per_generation=expression_shuffler_from_config(
            config, train_set.named_exprenvs()
        ),
        train_searcher=searcher_from_config(config, "train"),
        eval_expressions=eval_set.named_exprenvs(),
        eval_searcher=searcher_from_config(config, "eval"),
        distillator=distillator_from_config(config),
        dataset_refiner_factory=dataset_refiner_factory_from_config(config),
        rng_for_gen_fn=rng_factory_from_config(config),
        model_save_dir=model_save_dir_from_config(config),
        reporter=azureml_reporter_from_config(config, train_set),
        model_state_factory=ModelStateFactory.from_config(config),
        **{
            k: config[k]
            for k in (
                "num_generations",
                "total_train_time",
                "save_all_models",
                "repetition",
            )
            if k in config
        },
    )

class NoGPUError(ValueError):
    """Raise this when user demands GPU but there is no GPU."""

def torch_device_from_config(config: ConfigType) -> torch.device:
    """Get the CPU or GPU device to run on.  Optionally, also set the fraction of GPU memory to be used."""
    force_gpu = config.get("force_gpu", False)
    force_cpu = config.get("force_cpu", False)

    if force_cpu and force_gpu:
        raise ValueError("You cannot force both CPU and GPU")

    if (not force_cpu) and torch.cuda.is_available():
        # Default to GPU when available
        memory_frac = config.get("gpu_memory_fraction")
        if memory_frac is not None:
            torch.cuda.set_per_process_memory_fraction(memory_frac)
        return torch.device("cuda:0")
    else:
        if force_gpu:
            raise NoGPUError()
        return torch.device("cpu")

def tf_device_name_from_config(config: Dict) -> Optional[str]:
    """Get GPU device name, or None if GPU is not requested. Optionally, also set the fraction of GPU memory to be used."""
    force_gpu = config.get("force_gpu", False)
    if config.get("force_cpu", False):
        raise NotImplementedError
    if not force_gpu:
        return None
    physical_devices = tf.config.list_physical_devices("GPU")
    if len(physical_devices) == 0:
        raise NoGPUError("No GPUs visible - cannot use force_gpu")
    physical_device = physical_devices[0]
    memory_frac = config.get("gpu_memory_fraction")
    if memory_frac is not None:
        total_memory = GPUtil.getGPUs()[0].memoryTotal
        memory_limit = memory_frac * total_memory
        tf.config.set_logical_device_configuration(
            physical_device, [tf.config.LogicalDeviceConfiguration(memory_limit)]
        )
    return tf.config.list_logical_devices("GPU")[0].name

def num_time_heads_from_config(config: ConfigType):
    # At the moment, this is the only configuration we support,
    # although we expect these to diverge in the future.
    return config["simulation_depth_train"] + 1

def data_converter_from_config(config: ConfigType) -> DataConverter:
    kwargs = kwargs_from_config(
        config, required_keys=("use_subtree_match_edges",), optional_keys=(),
    )
    kwargs["num_time_heads"] = num_time_heads_from_config(config)
    kwargs["cost_norm"] = config["cost_normalization"]
    kwargs["device"] = torch_device_from_config(config)
    return DataConverter(**kwargs)

def torch_model_from_config(config: ConfigType, seed: Optional[int] = None) -> Model:
    # TODO edit flags so we don't need all these renames.
    model_kwargs = kwargs_from_config(
        config,
        required_keys=("num_gnn_blocks", "loss", "lr", "aggregation_over_edge_types"),
        optional_keys=("num_propagations",),
        renames=[
            ("hidden_dim", "gnn_hidden_dim"),
            ("output_hidden_dim", "regressor_hidden_dim"),
            ("decoder_readout", "global_pooling"),
        ],
    )
    global_pooling = model_kwargs.pop("global_pooling")
    model_kwargs["global_pooling"] = global_pooling_from_str(global_pooling)
    model_kwargs["gnn_dropout"] = 1 - config["graph_state_keep_prob"]
    model_kwargs["regressor_dropout"] = 1 - config["output_keep_prob"]
    model_kwargs["num_time_heads"] = num_time_heads_from_config(config)

    model_kwargs["num_node_types"] = Expression.num_node_types
    model_kwargs["num_edge_types"] = EdgeType.num_edge_types(
        use_subtree_match_edges=config["use_subtree_match_edges"]
    )
    seed = seed_from_config(config) if seed is None else seed
    with utils.random_state_context(seed):
        model = StateValueModel(**model_kwargs)
    return model

def tf_single_regressor_from_config(seed, config: ConfigType) -> ModelWrapper:
    from rlo.cost_normalizers import cost_normalizer
    from rlo.regressor import DenseGNNStateValueRegressor, SparseGNNStateValueRegressor

    sparse = config.get("sparse_gnn", False)
    cost_norm_str: str = config["cost_normalization"]
    cost_norm = cost_normalizer(cost_norm_str)

    # Model parameters go into model_hypers
    model_hypers = kwargs_from_config(
        config,
        ("output_hidden_dim", "num_propagations", "hidden_dim",),
        (
            "lr",
            "loss",
            "cumsum",
            "one_hot_embedding",
            "nonlinear_messages",
            "decoder_readout",
            "aggregation_over_edge_types",
            "message_from_sender_receiver",
            "grad_clip_value",
        ),
    )
    # All other parameters including pipeline parameters go into kwargs
    kwargs = kwargs_from_config(config, ("use_subtree_match_edges",), ("eager",),)
    kwargs["device"] = tf_device_name_from_config(config)
    kwargs["cost_norm"] = cost_norm
    kwargs["seed"] = seed_from_config(config) if seed is None else seed

    kwargs["num_time_heads"] = num_time_heads_from_config(config)
    if "graph_state_keep_prob" in config:
        model_hypers["recurrent_dropout"] = 1 - config["graph_state_keep_prob"]
    if "output_keep_prob" in config:
        model_hypers["output_dropout"] = 1 - config["output_keep_prob"]
    model_hypers["num_gnn_blocks"] = config["num_gnn_blocks"]
    model_hypers["stacked_gnn_double_hidden"] = config["stacked_gnn_double_hidden"]
    kwargs["model_hypers"] = model_hypers

    if sparse:
        return SparseGNNStateValueRegressor(**kwargs)
    else:
        return DenseGNNStateValueRegressor(**kwargs)

def single_regressor_from_config(config: ConfigType, seed=None):

    if config["tensorflow"]:
        return tf_single_regressor_from_config(seed=seed, config=config)

    data_converter = data_converter_from_config(config)
    model = torch_model_from_config(config)
    device = torch_device_from_config(config)
    return TorchModelWrapper(model=model, data_converter=data_converter, device=device)

def regressor_from_config(config: ConfigType):

    if config.get("two_value_func") is not None:
        if not config["tensorflow"]:
            raise NotImplementedError(
                "DualModelWrapper is not implemented for torch models."
            )
        # Store seed_2 in config just so we have a record.
        first_seed = seed_from_config(config)
        second_seed = DualModelWrapper.seed_for_second_model(first_seed)
        config["seed_2"] = second_seed
        model_1 = single_regressor_from_config(config, first_seed)
        model_2 = single_regressor_from_config(config, second_seed)

        if config["two_value_func"] == "train":
            return DualModelWrapper(
                model_1, model_2, config["two_value_func_var_frac_train"], 0
            )
        else:
            raise ValueError("For two value functions only train is supported atm.")
    else:
        return single_regressor_from_config(config)

search_algs = {
    "rollout": rollout_searcher_from_config,
    "astar": astar_searcher_from_config,
    "beam": beam_searcher_from_config,
    "hybrid": hybrid_searcher_from_config,
}

maxing_algs = {
    "episodes": search_tree.max_over_episodes,
    "accumulator": search_tree.max_with_accumulator,
}

def seed_from_config(config: ConfigType):
    return int(
        config[
            "seed_all_reps" if config.get("seed_all_reps") is not None else "repetition"
        ]
    )

def framework_from_config(config: ConfigType):
    # Field tensorflow doesn't exist before PyTorch port
    tensorflow = config.get("tensorflow", True)
    return "tensorflow" if tensorflow else "torch"

class ModelStateFactory:
    def __init__(self, model_state_class):
        self._model_state_class = model_state_class

    @staticmethod
    def from_config(config: ConfigType):
        framework = framework_from_config(config)
        return ModelStateFactory(Weights if framework == "tensorflow" else ModelState)

    def from_path(self, path):
        return self._model_state_class.from_path(path)

max_stars_repo_name: tomjaguarpaw/knossos-ksc

max_issues_repo_name: tomjaguarpaw/knossos-ksc

max_forks_repo_name: tomjaguarpaw/knossos-ksc


text:
import networkx as nx
from networkx.algorithms import tree

from .random import random
from collections import deque, namedtuple

def predecessors(h, root):
    return {a: b for a, b in nx.bfs_predecessors(h, root)}

def random_spanning_tree(graph):
    for edge in graph.edges:
        graph.edges[edge]["weight"] = random.random()

    spanning_tree = tree.maximum_spanning_tree(
        graph, algorithm="kruskal", weight="weight"
    )
    return spanning_tree

class PopulatedGraph:
    def __init__(self, graph, populations, ideal_pop, epsilon):
        self.graph = graph
        self.subsets = {node: {node} for node in graph}
        self.population = populations.copy()
        self.ideal_pop = ideal_pop
        self.epsilon = epsilon
        self._degrees = {node: graph.degree(node) for node in graph}

    def __iter__(self):
        return iter(self.graph)

    def degree(self, node):
        return self._degrees[node]

    def contract_node(self, node, parent):
        self.population[parent] += self.population[node]
        self.subsets[parent] |= self.subsets[node]
        self._degrees[parent] -= 1

    def has_ideal_population(self, node):
        return (
            abs(self.population[node] - self.ideal_pop) < self.epsilon * self.ideal_pop
        )

def contract_leaves_until_balanced_or_none(h, choice=random.choice):
    # this used to be greater than 2 but failed on small grids:(
    root = choice([x for x in h if h.degree(x) > 1])
    # BFS predecessors for iteratively contracting leaves
    pred = predecessors(h.graph, root)

    leaves = deque(x for x in h if h.degree(x) == 1)
    while len(leaves) > 0:
        leaf = leaves.popleft()
        if h.has_ideal_population(leaf):
            return h.subsets[leaf]
        # Contract the leaf:
        parent = pred[leaf]
        h.contract_node(leaf, parent)
        if h.degree(parent) == 1 and parent != root:
            leaves.append(parent)
    return None

Cut = namedtuple("Cut", "edge subset")

def find_balanced_edge_cuts(h, choice=random.choice):
    # this used to be greater than 2 but failed on small grids:(
    root = choice([x for x in h if h.degree(x) > 1])
    # BFS predecessors for iteratively contracting leaves
    pred = predecessors(h.graph, root)

    cuts = []
    leaves = deque(x for x in h if h.degree(x) == 1)
    while len(leaves) > 0:
        leaf = leaves.popleft()
        if h.has_ideal_population(leaf):
            cuts.append(Cut(edge=(leaf, pred[leaf]), subset=h.subsets[leaf].copy()))
        # Contract the leaf:
        parent = pred[leaf]
        h.contract_node(leaf, parent)
        if h.degree(parent) == 1 and parent != root:
            leaves.append(parent)
    return cuts

def bipartition_tree(
    graph,
    pop_col,
    pop_target,
    epsilon,
    node_repeats=1,
    spanning_tree=None,
    choice=random.choice,
):
    """This function finds a balanced 2 partition of a graph by drawing a
    spanning tree and finding an edge to cut that leaves at most an epsilon
    imbalance between the populations of the parts. If a root fails, new roots
    are tried until node_repeats in which case a new tree is drawn.

    Builds up a connected subgraph with a connected complement whose population
    is ``epsilon * pop_target`` away from ``pop_target``.

    Returns a subset of nodes of ``graph`` (whose induced subgraph is connected).
    The other part of the partition is the complement of this subset.

    :param graph: The graph to partition
    :param pop_col: The node attribute holding the population of each node
    :param pop_target: The target population for the returned subset of nodes
    :param epsilon: The allowable deviation from  ``pop_target`` (as a percentage of
        ``pop_target``) for the subgraph's population
    :param node_repeats: A parameter for the algorithm: how many different choices
        of root to use before drawing a new spanning tree.
    :param spanning_tree: The spanning tree for the algorithm to use (used when the
        algorithm chooses a new root and for testing)
    :param choice: :func:`random.choice`. Can be substituted for testing.
    """
    populations = {node: graph.nodes[node][pop_col] for node in graph}

    balanced_subtree = None
    if spanning_tree is None:
        spanning_tree = random_spanning_tree(graph)
    restarts = 0
    while balanced_subtree is None:
        if restarts == node_repeats:
            spanning_tree = random_spanning_tree(graph)
            restarts = 0
        h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
        balanced_subtree = contract_leaves_until_balanced_or_none(h, choice=choice)
        restarts += 1

    return balanced_subtree

def bipartition_tree_random(
    graph,
    pop_col,
    pop_target,
    epsilon,
    node_repeats=1,
    spanning_tree=None,
    choice=random.choice,
):
    """This is like :func:`bipartition_tree` except it chooses a random balanced
    cut, rather than the first cut it finds.

    This function finds a balanced 2 partition of a graph by drawing a
    spanning tree and finding an edge to cut that leaves at most an epsilon
    imbalance between the populations of the parts. If a root fails, new roots
    are tried until node_repeats in which case a new tree is drawn.

    Builds up a connected subgraph with a connected complement whose population
    is ``epsilon * pop_target`` away from ``pop_target``.

    Returns a subset of nodes of ``graph`` (whose induced subgraph is connected).
    The other part of the partition is the complement of this subset.

    :param graph: The graph to partition
    :param pop_col: The node attribute holding the population of each node
    :param pop_target: The target population for the returned subset of nodes
    :param epsilon: The allowable deviation from  ``pop_target`` (as a percentage of
        ``pop_target``) for the subgraph's population
    :param node_repeats: A parameter for the algorithm: how many different choices
        of root to use before drawing a new spanning tree.
    :param spanning_tree: The spanning tree for the algorithm to use (used when the
        algorithm chooses a new root and for testing)
    :param choice: :func:`random.choice`. Can be substituted for testing.
    """
    populations = {node: graph.nodes[node][pop_col] for node in graph}

    possible_cuts = []
    if spanning_tree is None:
        spanning_tree = random_spanning_tree(graph)

    while len(possible_cuts) == 0:
        spanning_tree = random_spanning_tree(graph)
        h = PopulatedGraph(spanning_tree, populations, pop_target, epsilon)
        possible_cuts = find_balanced_edge_cuts(h, choice=choice)

    return choice(possible_cuts).subset

def recursive_tree_part(
    graph, parts, pop_target, pop_col, epsilon, node_repeats=1, method=bipartition_tree
):
    """Uses :func:`~gerrychain.tree.bipartition_tree` recursively to partition a tree into
    ``len(parts)`` parts of population ``pop_target`` (within ``epsilon``). Can be used to
    generate initial seed plans or to implement ReCom-like "merge walk" proposals.

    :param graph: The graph
    :param parts: Iterable of part labels (like ``[0,1,2]`` or ``range(4)``
    :param pop_target: Target population for each part of the partition
    :param pop_col: Node attribute key holding population data
    :param epsilon: How far (as a percentage of ``pop_target``) from ``pop_target`` the parts
        of the partition can be
    :param node_repeats: Parameter for :func:`~gerrychain.tree_methods.bipartition_tree` to use.
    :return: New assignments for the nodes of ``graph``.
    :rtype: dict
    """
    flips = {}
    remaining_nodes = set(graph.nodes)
    # We keep a running tally of deviation from ``epsilon`` at each partition
    # and use it to tighten the population constraints on a per-partition
    # basis such that every partition, including the last partition, has a
    # population within +/-``epsilon`` of the target population.
    # For instance, if district n's population exceeds the target by 2%
    # with a +/-2% epsilon, then district n+1's population should be between
    # 98% of the target population and the target population.
    debt = 0

    for part in parts[:-1]:
        min_pop = max(pop_target * (1 - epsilon), pop_target * (1 - epsilon) - debt)
        max_pop = min(pop_target * (1 + epsilon), pop_target * (1 + epsilon) - debt)
        nodes = method(
            graph.subgraph(remaining_nodes),
            pop_col=pop_col,
            pop_target=(min_pop + max_pop) / 2,
            epsilon=(max_pop - min_pop) / (2 * pop_target),
            node_repeats=node_repeats,
        )

        part_pop = 0
        for node in nodes:
            flips[node] = part
            part_pop += graph.nodes[node][pop_col]
        debt += part_pop - pop_target
        remaining_nodes -= nodes

    # All of the remaining nodes go in the last part
    for node in remaining_nodes:
        flips[node] = parts[-1]

    return flips

max_stars_repo_name: OpenPrecincts/verification

max_issues_repo_name: OpenPrecincts/verification

max_forks_repo_name: OpenPrecincts/verification

reason: import networkx,from networkx


text:
#!/usr/bin/env python2.7
"""Get R and G bands, average and display as greyscale"""
from __future__ import (absolute_import, division,  # Python3 compatibility
                        print_function)
import time
import picamera
from picamera.array import PiRGBArray
import cv2
import numpy as np

with picamera.PiCamera() as camera:
    width = 800
    height = 480
    camera.resolution = (width,height)
    camera.framerate = 30
    cv2.namedWindow("window", cv2.WND_PROP_FULLSCREEN)
    cv2.setWindowProperty("window",cv2.WND_PROP_FULLSCREEN, cv2.cv.CV_WINDOW_FULLSCREEN)

    with picamera.array.PiRGBArray(camera) as output:
        for foo in  camera.capture_continuous(output, 'bgr', use_video_port=True):
                image = output.array
                output.truncate(0)
                # Do something
                RpG = image[:,:,2].astype(np.int16) + image[:,:,1].astype(np.int16)
                avgd = RpG / 2
                cv2.imshow("window", avgd.astype(np.uint8))
                # Sleep until next frame
                #cv2.waitKey(int(1000/30))
                cv2.waitKey(1)
    cv2.destroyAllWindows()


text:
from itertools import product
import numpy as np
from numpy.testing import assert_allclose
from pytest import raises
from scipy.spatial.transform import Rotation, RotationSpline
from scipy.spatial.transform._rotation_spline import (
    _angular_rate_to_rotvec_dot_matrix,
    _rotvec_dot_to_angular_rate_matrix,
    _matrix_vector_product_of_stacks,
    _angular_acceleration_nonlinear_term,
    _create_block_3_diagonal_matrix)

def test_angular_rate_to_rotvec_conversions():
    np.random.seed(0)
    rv = np.random.randn(4, 3)
    A = _angular_rate_to_rotvec_dot_matrix(rv)
    A_inv = _rotvec_dot_to_angular_rate_matrix(rv)

    # When the rotation vector is aligned with the angular rate, then
    # the rotation vector rate and angular rate are the same.
    assert_allclose(_matrix_vector_product_of_stacks(A, rv), rv)
    assert_allclose(_matrix_vector_product_of_stacks(A_inv, rv), rv)

    # A and A_inv must be reciprocal to each other.
    I_stack = np.empty((4, 3, 3))
    I_stack[:] = np.eye(3)
    assert_allclose(np.matmul(A, A_inv), I_stack, atol=1e-15)

def test_angular_rate_nonlinear_term():
    # The only simple test is to check that the term is zero when
    # the rotation vector
    np.random.seed(0)
    rv = np.random.rand(4, 3)
    assert_allclose(_angular_acceleration_nonlinear_term(rv, rv), 0,
                    atol=1e-19)

def test_create_block_3_diagonal_matrix():
    np.random.seed(0)
    A = np.empty((4, 3, 3))
    A[:] = np.arange(1, 5)[:, None, None]

    B = np.empty((4, 3, 3))
    B[:] = -np.arange(1, 5)[:, None, None]
    d = 10 * np.arange(10, 15)

    banded = _create_block_3_diagonal_matrix(A, B, d)

    # Convert the banded matrix to the full matrix.
    k, l = list(zip(*product(np.arange(banded.shape[0]),
                             np.arange(banded.shape[1]))))
    k = np.asarray(k)
    l = np.asarray(l)

    i = k - 5 + l
    j = l
    values = banded.ravel()
    mask = (i >= 0) & (i < 15)
    i = i[mask]
    j = j[mask]
    values = values[mask]
    full = np.zeros((15, 15))
    full[i, j] = values

    zero = np.zeros((3, 3))
    eye = np.eye(3)

    # Create the reference full matrix in the most straightforward manner.
    ref = np.block([
        [d[0] * eye, B[0], zero, zero, zero],
        [A[0], d[1] * eye, B[1], zero, zero],
        [zero, A[1], d[2] * eye, B[2], zero],
        [zero, zero, A[2], d[3] * eye, B[3]],
        [zero, zero, zero, A[3], d[4] * eye],
    ])

    assert_allclose(full, ref, atol=1e-19)

def test_spline_2_rotations():
    times = [0, 10]
    rotations = Rotation.from_euler('xyz', [[0, 0, 0], [10, -20, 30]],
                                    degrees=True)
    spline = RotationSpline(times, rotations)

    rv = (rotations[0].inv() * rotations[1]).as_rotvec()
    rate = rv / (times[1] - times[0])
    times_check = np.array([-1, 5, 12])
    dt = times_check - times[0]
    rv_ref = rate * dt[:, None]

    assert_allclose(spline(times_check).as_rotvec(), rv_ref)
    assert_allclose(spline(times_check, 1), np.resize(rate, (3, 3)))
    assert_allclose(spline(times_check, 2), 0, atol=1e-16)

def test_constant_attitude():
    times = np.arange(10)
    rotations = Rotation.from_rotvec(np.ones((10, 3)))
    spline = RotationSpline(times, rotations)

    times_check = np.linspace(-1, 11)
    assert_allclose(spline(times_check).as_rotvec(), 1, rtol=1e-15)
    assert_allclose(spline(times_check, 1), 0, atol=1e-17)
    assert_allclose(spline(times_check, 2), 0, atol=1e-17)

    assert_allclose(spline(5.5).as_rotvec(), 1, rtol=1e-15)
    assert_allclose(spline(5.5, 1), 0, atol=1e-17)
    assert_allclose(spline(5.5, 2), 0, atol=1e-17)

def test_spline_properties():
    times = np.array([0, 5, 15, 27])
    angles = [[-5, 10, 27], [3, 5, 38], [-12, 10, 25], [-15, 20, 11]]

    rotations = Rotation.from_euler('xyz', angles, degrees=True)
    spline = RotationSpline(times, rotations)

    assert_allclose(spline(times).as_euler('xyz', degrees=True), angles)
    assert_allclose(spline(0).as_euler('xyz', degrees=True), angles[0])

    h = 1e-8
    rv0 = spline(times).as_rotvec()
    rvm = spline(times - h).as_rotvec()
    rvp = spline(times + h).as_rotvec()
    assert_allclose(rv0, 0.5 * (rvp + rvm), rtol=1e-15)

    r0 = spline(times, 1)
    rm = spline(times - h, 1)
    rp = spline(times + h, 1)
    assert_allclose(r0, 0.5 * (rm + rp), rtol=1e-14)

    a0 = spline(times, 2)
    am = spline(times - h, 2)
    ap = spline(times + h, 2)
    assert_allclose(a0, am, rtol=1e-7)
    assert_allclose(a0, ap, rtol=1e-7)

def test_error_handling():
    raises(ValueError, RotationSpline, [1.0], Rotation.random())

    r = Rotation.random(10)
    t = np.arange(10).reshape(5, 2)
    raises(ValueError, RotationSpline, t, r)

    t = np.arange(9)
    raises(ValueError, RotationSpline, t, r)

    t = np.arange(10)
    t[5] = 0
    raises(ValueError, RotationSpline, t, r)

    t = np.arange(10)

    s = RotationSpline(t, r)
    raises(ValueError, s, 10, -1)

    raises(ValueError, s, np.arange(10).reshape(5, 2))

max_stars_repo_name: Ennosigaeon/scipy

max_issues_repo_name: Ennosigaeon/scipy

max_forks_repo_name: Ennosigaeon/scipy

reason: import numpy,from numpy,from scipy