text:
"""Definition of HLO Instructions"""

from collections import defaultdict
from enum import Enum, auto

import numpy as np

from common import compute_bytes, append_flatten_elements, transpose_flatten, reshape_flatten

class ShardingSpecType(Enum):
    REPLICATED = auto()
    MAXIMAL = auto()
    OTHER = auto()
    TUPLE = auto()
    PARTIAL_REDUCTION = auto()

INF_COST = 1e10  # infinity cost

class ShardingSpec:
    def __init__(self, type_, tile_assignment_dimensions, tile_assignment_devices,
                 replicate_on_last_tile_dim, partial_reduce_replication):
        self.type = type_
        self.tile_assignment_dimensions = tuple(tile_assignment_dimensions)
        self.tile_assignment_devices = tuple(tile_assignment_devices)
        self.replicate_on_last_tile_dim = replicate_on_last_tile_dim
        self.partial_reduce_replication = partial_reduce_replication

    def num_tile_devices(self):
        if self.type == ShardingSpecType.REPLICATED:
            return 1

        assert self.type == ShardingSpecType.OTHER
        ret = np.prod(self.tile_assignment_dimensions)
        if self.replicate_on_last_tile_dim:
            ret /= self.tile_assignment_dimensions[-1]
        return ret

    def transpose(self, dimensions):
        if self.type == ShardingSpecType.REPLICATED:
            return self

        assert self.type == ShardingSpecType.OTHER

        spec_trans_dims = list(dimensions)
        if self.replicate_on_last_tile_dim:
            spec_trans_dims.append(len(dimensions))

        tile_assignment_dimensions = [self.tile_assignment_dimensions[i]
            for i in spec_trans_dims]
        tile_assignment_devices = transpose_flatten(self.tile_assignment_devices,
            self.tile_assignment_dimensions, spec_trans_dims)

        ret = ShardingSpec(self.type,
                           tile_assignment_dimensions,
                           tile_assignment_devices,
                           self.replicate_on_last_tile_dim,
                           self.partial_reduce_replication)
        return ret

    def broadcast(self, new_shape, dimensions):
        if self.type == ShardingSpecType.REPLICATED:
            return self

        assert self.type == ShardingSpecType.OTHER

        tile_assignment_dimensions = []
        for i in range(len(new_shape)):
            if i in dimensions:
                tile_assignment_dimensions.append(
                    self.tile_assignment_dimensions[dimensions.index(i)])
            else:
                tile_assignment_dimensions.append(1)

        if self.replicate_on_last_tile_dim:
            tile_assignment_dimensions.append(self.tile_assignment_dimensions[-1])

        output_spec = ShardingSpec(self.type,
                                   tile_assignment_dimensions,
                                   self.tile_assignment_devices,
                                   self.replicate_on_last_tile_dim,
                                   self.partial_reduce_replication)
        return output_spec

    def reshape(self, old_shape, new_shape):
        if self.type == ShardingSpecType.REPLICATED:
            return self

        assert self.type == ShardingSpecType.OTHER

        # Construct a map that maps an old dimension to its corresponding new dimension
        dim_mapping = {}
        new_pt = -1
        old_pt = -1
        old_prod = 1
        new_prod = 1
        while True:
            move_new = False
            move_old = False

            if new_prod == old_prod:
                dim_mapping[old_pt + 1] = new_pt + 1
                move_new = move_old = True
            elif new_prod < old_prod:
                move_new = True
            else:
                move_old = True

            if move_new:
                new_pt += 1
                if new_pt < len(new_shape):
                    new_prod *= new_shape[new_pt]
                else:
                    break
            if move_old:
                old_pt += 1
                if old_pt < len(old_shape):
                    old_prod *= old_shape[old_pt]
                else:
                    break

        tile_assignment_dimensions = []
        cur_prod = 1
        state = 1  # 0: start  1: middle
        i = 0

        failed = False
        while i < len(old_shape) and not failed:
            if state == 0:
                assert i in dim_mapping
                while len(tile_assignment_dimensions) < dim_mapping[i]:
                    tile_assignment_dimensions.append(1)
                tile_assignment_dimensions.append(
                    self.tile_assignment_dimensions[i])
                state = 1
                i += 1
            elif state == 1:
                if i in dim_mapping:
                    state = 0
                else:
                    if self.tile_assignment_dimensions[i] == 1:
                        i += 1
                    else:
                        failed = True

        if failed:
            return None

        while len(tile_assignment_dimensions) < len(new_shape):
            tile_assignment_dimensions.append(1)

        if self.replicate_on_last_tile_dim:
            tile_assignment_dimensions.append(self.tile_assignment_dimensions[-1])
        output_spec = ShardingSpec(self.type,
                                   tile_assignment_dimensions,
                                   self.tile_assignment_devices,
                                   self.replicate_on_last_tile_dim,
                                   self.partial_reduce_replication)
        return output_spec

    @staticmethod
    def tile_internal(shape, tensor_dims, mesh_dims, cluster_env, partial_reduce_replication):
        assert len(tensor_dims) == len(mesh_dims)

        tile_assignment_dimensions = [1] * len(shape)

        # Split on certain mesh dimensions
        split_prod = 1
        for tensor_dim, mesh_dim in zip(tensor_dims, mesh_dims):
            tile_assignment_dimensions[tensor_dim] = cluster_env.device_mesh.shape[mesh_dim]
            split_prod *= cluster_env.device_mesh.shape[mesh_dim]

        if split_prod == 1:
            return ShardingSpec.replicated(cluster_env)

        # Replicate on reminding mesh dimensions
        if split_prod < cluster_env.num_devices:
            tile_assignment_dimensions.append(cluster_env.num_devices // split_prod)
            replicate_on_last_tile_dim = True
        else:
            replicate_on_last_tile_dim = False

        # Map device ids from device_mesh to tile_assignment_devices
        tile_assignment_devices = []
        tmp_indices = [None] * len(cluster_env.device_mesh.shape)
        def generate_tile_assignment_devices(tensor_dim, mesh_indices):
            if tensor_dim == len(shape) - 1:
                append_flatten_elements(tile_assignment_devices, cluster_env.device_mesh,
                                        mesh_indices, -1, tmp_indices)
            else:
                next_tensor_dim = tensor_dim + 1
                next_mesh_dim = -1

                if next_tensor_dim in tensor_dims:
                    next_mesh_dim = mesh_dims[tensor_dims.index(next_tensor_dim)]

                for i in range(tile_assignment_dimensions[next_tensor_dim]):
                    if next_mesh_dim != -1:
                        mesh_indices[next_mesh_dim] = i
                    generate_tile_assignment_devices(next_tensor_dim, mesh_indices)

        generate_tile_assignment_devices(-1, [-1] * len(cluster_env.device_mesh.shape))

        return ShardingSpec(ShardingSpecType.OTHER,
                            tile_assignment_dimensions, tile_assignment_devices,
                            replicate_on_last_tile_dim,
                            False)

    @staticmethod
    def tile(shape, tensor_dims, mesh_dims, cluster_env):
        return ShardingSpec.tile_internal(shape, tensor_dims, mesh_dims, cluster_env, False)

    @staticmethod
    def tile_partial_reduce(shape, tensor_dims, mesh_dims, cluster_env):
        return ShardingSpec.tile_internal(shape, tensor_dims, mesh_dims, cluster_env, True)

    @staticmethod
    def replicated(cluster_env):
        tile_assignment_devices = range(cluster_env.num_devices)
        return ShardingSpec(ShardingSpecType.REPLICATED, (), tile_assignment_devices,
                            False, False)

    @staticmethod
    def split(shape, dim, cluster_env):
        tile_assignment_dimensions = [1] * len(shape)
        tile_assignment_dimensions[dim] = cluster_env.num_devices
        tile_assignment_devices = range(cluster_env.num_devices)
        return ShardingSpec(ShardingSpecType.OTHER,
                            tile_assignment_dimensions, tile_assignment_devices,
                            False, False)

    @staticmethod
    def tuple():
        return ShardingSpec(ShardingSpecType.TUPLE, (), (), False, False)

    def __str__(self):
        return f"{self.tile_assignment_dimensions}"\
               f"{list(self.tile_assignment_devices)}"

    def __eq__(self, other):
        return (self.type == other.type and
                self.tile_assignment_dimensions == other.tile_assignment_dimensions and
                self.tile_assignment_devices == other.tile_assignment_devices and
                self.replicate_on_last_tile_dim == other.replicate_on_last_tile_dim and
                self.partial_reduce_replication == other.partial_reduce_replication)

def resharding_cost_vector(cluster_env, source_ins, required_spec):
    cost_vector = []
    for strategy in source_ins.strategies:
        cost_vector.append(cluster_env.resharding_cost(source_ins.shape,
            strategy.output_spec, required_spec))
    return cost_vector

def follow_ins_cost_vector(source_ins, index):
    ret = [INF_COST] * len(source_ins.strategies)
    ret[index] = 0
    return ret

class InstructionStrategy:
    def __init__(self, name, output_spec):
        self.name = name
        self.output_spec = output_spec

class OpCode(Enum):
    PARAMETER = auto()
    CONSTANT = auto()
    BROADCAST = auto()
    RESHAPE = auto()
    TRANSPOSE = auto()
    IDENTITY = auto()
    EXP = auto()
    FORCE_REPLICATED = auto()
    ADD = auto()
    SUBTRACT = auto()
    MULTIPLY = auto()
    DIV = auto()
    COMPARE = auto()
    SELECT = auto()
    REDUCE = auto()
    DOT = auto()
    TUPLE = auto()

op_code_ct = defaultdict(int)

class HloInstruction:
    def __init__(self, op_code, shape, operands=[]):
        # Attributes
        self.op_code = op_code
        self.shape = shape
        self.operands = operands
        self.name = f"{str(op_code)[7:].lower()}.{op_code_ct[op_code]}"
        op_code_ct[op_code] += 1

        # Cost
        self.strategies = []
        self.compute_costs = []
        self.communication_costs = []
        self.memory_costs = []
        self.resharding_costs = []
        self.follow_ins = None
        self.depth = None

        # The index in HloComputation
        self.index = HloComputation.cur_env.append(self)
        self.batch_dim = None

    def build_strategy_and_cost(self, cluster_env, solver_option):
        raise NotImplementedError(f"{self.op_code}")

    def propagate_batch_dim(self, operand):
        raise NotImplementedError(f"{self.op_code}")

class HloParameter(HloInstruction):
    def __init__(self, shape, fix_strategy=None):
        super().__init__(OpCode.PARAMETER, shape, [])
        self.fix_strategy = fix_strategy

    def build_strategy_and_cost(self, cluster_env, solver_option):
        for i in range(len(self.shape)):
            for j in range(len(cluster_env.device_mesh.shape)):
                if (cluster_env.device_mesh.shape[j] == 1 or
                    self.shape[i] < cluster_env.device_mesh.shape[j]):
                    continue

                name = f"S{i} @ {j}"
                output_spec = ShardingSpec.tile(self.shape, [i], [j], cluster_env)
                self.strategies.append(InstructionStrategy(name, output_spec))
                self.compute_costs.append(0)
                self.communication_costs.append(0)
                self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())

        self.strategies.append(InstructionStrategy("R", ShardingSpec.replicated(cluster_env)))
        self.compute_costs.append(2)
        self.communication_costs.append(0)
        self.memory_costs.append(compute_bytes(self.shape))

        if self.fix_strategy:
            new_strategies = []
            new_compute_costs = []
            new_communication_costs = []
            new_memory_costs = []

            # filter strategies
            for i in range(len(self.strategies)):
                if self.strategies[i].name == self.fix_strategy:
                    new_strategies.append(self.strategies[i])
                    new_compute_costs.append(self.compute_costs[i])
                    new_communication_costs.append(self.communication_costs[i])
                    new_memory_costs.append(self.memory_costs[i])

            self.strategies = new_strategies
            self.compute_costs = new_compute_costs
            self.communication_costs = new_communication_costs
            self.memory_costs = new_memory_costs

    def __str__(self):
        return f"{self.name} {self.shape} = parameter()"

class HloConstant(HloInstruction):
    def __init__(self, value):
        super().__init__(OpCode.CONSTANT, (), [])
        self.value = value

    def build_strategy_and_cost(self, cluster_env, solver_option):
        self.strategies.append(InstructionStrategy("R", ShardingSpec.replicated(cluster_env)))
        self.compute_costs.append(0)
        self.communication_costs.append(0)
        self.memory_costs.append(compute_bytes(self.shape))

    def __str__(self):
        return f"{self.name} {self.shape} = constant({self.value})"

class HloBroadcast(HloInstruction):
    def __init__(self, operand, shape, dimensions=()):
        for i in dimensions:
            assert shape[i] == operand.shape[dimensions.index(i)]
        super().__init__(OpCode.BROADCAST, shape, [operand])
        self.dimensions = dimensions

    def build_strategy_and_cost(self, cluster_env, solver_option):
        follow = self.operands[0]
        self.follow_ins = follow

        for sid in range(len(follow.strategies)):
            output_spec = follow.strategies[sid].output_spec.broadcast(
                    self.shape, self.dimensions)
            name = f"{output_spec.tile_assignment_dimensions}"
            self.strategies.append(InstructionStrategy(name, output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(0)
            self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())
            self.resharding_costs.append([follow_ins_cost_vector(follow, sid)])

    def __str__(self):
        return f"{self.name} {self.shape} = broadcast({self.operands[0].name})"

class HloReshape(HloInstruction):
    def __init__(self, operand, new_shape):
        # todo: mark this as inplace
        assert np.prod(operand.shape) == np.prod(new_shape)
        super().__init__(OpCode.RESHAPE, new_shape, [operand])
        self.new_shape = new_shape

    def build_strategy_and_cost(self, cluster_env, solver_option):
        follow = self.operands[0]
        self.follow_ins = follow
        old_shape = self.operands[0].shape
        new_shape = self.new_shape

        for sid in range(len(follow.strategies)):
            output_spec = follow.strategies[sid].output_spec.reshape(
                    follow.shape, self.shape)
            if output_spec is None:
                continue

            name = f"{output_spec.tile_assignment_dimensions}"
            self.strategies.append(InstructionStrategy(name, output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(0)
            self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())
            self.resharding_costs.append([follow_ins_cost_vector(follow, sid)])

    def __str__(self):
        return f"{self.name} {self.shape} = reshape({self.operands[0].name})"

class HloTranspose(HloInstruction):
    def __init__(self, operand, dimensions):
        assert len(dimensions) == len(operand.shape)
        new_shape = tuple(operand.shape[i] for i in dimensions)
        super().__init__(OpCode.TRANSPOSE, new_shape, [operand])
        self.dimensions = dimensions

    def build_strategy_and_cost(self, cluster_env, solver_option):
        follow = self.operands[0]
        self.follow_ins = follow

        for sid in range(len(follow.strategies)):
            output_spec = follow.strategies[sid].output_spec.transpose(self.dimensions)
            name = f"{output_spec.tile_assignment_dimensions}"
            self.strategies.append(InstructionStrategy(name, output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(0)
            self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())
            self.resharding_costs.append([follow_ins_cost_vector(follow, sid)])

    def __str__(self):
        return f"{self.name} {self.shape} = transpose({self.operands[0].name}) " +\
               f"dimensions={self.dimensions}"

class HloElementwise(HloInstruction):
    def __init__(self, op_code, operands):
        for i in range(0, len(operands)):
            assert operands[0].shape == operands[i].shape
        super().__init__(op_code, operands[0].shape, operands)

    def build_strategy_and_cost(self, cluster_env, solver_option):
        depths = [operand.depth for operand in self.operands]
        follow_idx = np.argmax(depths)

        follow = self.operands[follow_idx]
        self.follow_ins = follow

        for sid in range(len(follow.strategies)):
            output_spec = follow.strategies[sid].output_spec

            name = f"{output_spec.tile_assignment_dimensions}"
            self.strategies.append(InstructionStrategy(name, output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(0)
            self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())

            resharding_costs = []
            for k in range(len(self.operands)):
                if k == follow_idx:
                    resharding_costs.append(
                        follow_ins_cost_vector(follow, sid))
                else:
                    resharding_costs.append(
                    resharding_cost_vector(cluster_env, self.operands[k], output_spec))
            self.resharding_costs.append(resharding_costs)

    def propagate_batch_dim(self, ins):
        self.batch_dim = ins.batch_dim
        return True

    def __str__(self):
        fun_name = str(self.op_code)[7:].lower()
        args = ", ".join(f"{self.operands[i].name}" for i in range(len(self.operands)))
        return f"{self.name} {self.shape} = {fun_name}({args})"

class HloIdentity(HloElementwise):
    def __init__(self, operand):
        super().__init__(OpCode.IDENTITY, [operand])

class HloExp(HloElementwise):
    def __init__(self, operand):
        super().__init__(OpCode.EXP, [operand])

class HloForceReplicated(HloElementwise):
    def __init__(self, operand):
        super().__init__(OpCode.FORCE_REPLICATED, [operand])

    def build_strategy_and_cost(self, cluster_env, solver_option):
        self.strategies.append(InstructionStrategy("R",
            ShardingSpec.replicated(cluster_env)))
        self.compute_costs.append(0)
        self.communication_costs.append(0)
        self.memory_costs.append(0)
        self.resharding_costs.append([
            resharding_cost_vector(cluster_env, self.operands[0],
                ShardingSpec.replicated(cluster_env))
        ])

class HloAdd(HloElementwise):
    def __init__(self, lhs, rhs):
        super().__init__(OpCode.ADD, [lhs, rhs])

class HloSubtract(HloElementwise):
    def __init__(self, lhs, rhs):
        super().__init__(OpCode.SUBTRACT, [lhs, rhs])

class HloMutiply(HloElementwise):
    def __init__(self, lhs, rhs):
        super().__init__(OpCode.MULTIPLY, [lhs, rhs])

class HloDiv(HloElementwise):
    def __init__(self, lhs, rhs):
        super().__init__(OpCode.DIV, [lhs, rhs])

class HloCompare(HloElementwise):
    def __init__(self, lhs, rhs):
        super().__init__(OpCode.COMPARE, [lhs, rhs])

class HloSelect(HloElementwise):
    def __init__(self, pred, true_value, false_value):
        super().__init__(OpCode.SELECT, [pred, true_value, false_value])

class HloReduce(HloInstruction):
    def __init__(self, operand, dimensions):
        new_shape = tuple(operand.shape[i] for i in range(len(operand.shape)) if i not in dimensions)
        super().__init__(OpCode.REDUCE, new_shape, [operand])
        self.dimensions = dimensions

    def build_strategy_and_cost(self, cluster_env, solver_option):
        operand = self.operands[0]
        self.follow_ins = operand

        # Map old dims to new dim
        old_dim_to_new_dim = []
        pt = 0
        for old_dim in range(len(operand.shape)):
            if old_dim in self.dimensions:
                old_dim_to_new_dim.append(-1)
            else:
                old_dim_to_new_dim.append(pt)
                pt += 1
        assert pt == len(self.shape)

        # Create follow strategies
        for sid in range(len(operand.strategies)):
            tensor_dim_to_mesh = cluster_env.get_tensor_dim_to_mesh_dim(
                operand.shape, operand.strategies[sid].output_spec)

            tile_tensor_dims = []
            tile_mesh_dims = []
            all_reduce_dims = []

            for tensor_dim in range(len(operand.shape)):
                mesh_dim = tensor_dim_to_mesh[tensor_dim]
                if tensor_dim in self.dimensions:
                    if mesh_dim == -1:  # reduce on a replicated dim
                        continue
                    else:               # reduce on a split dim
                        all_reduce_dims.append(mesh_dim)
                else:
                    if mesh_dim == -1: # follow replicated dim
                        pass
                    else:              # follow split dim
                        tile_tensor_dims.append(old_dim_to_new_dim[tensor_dim])
                        tile_mesh_dims.append(mesh_dim)

            output_spec = ShardingSpec.tile(self.shape, tile_tensor_dims, tile_mesh_dims, cluster_env)

            mem_cost = compute_bytes(self.shape) / output_spec.num_tile_devices()
            comm_cost = 0
            for mesh_dim in all_reduce_dims:
                comm_cost += cluster_env.all_reduce_cost(mem_cost, mesh_dim)

            reduce_dims_str = "".join([str(x) for x in all_reduce_dims])
            if reduce_dims_str:
                name = f"follow (allreduce @ {reduce_dims_str})"
            else:
                name = f"{output_spec.tile_assignment_dimensions}"

            self.strategies.append(InstructionStrategy(name, output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(comm_cost)
            self.memory_costs.append(mem_cost)
            self.resharding_costs.append([follow_ins_cost_vector(operand, sid)])

    def __str__(self):
        return f"{self.name} {self.shape} = reduce({self.operands[0].name}) " +\
               f"dimensions={self.dimensions}"

class HloDot(HloInstruction):
    def __init__(self, lhs, rhs,
                 lhs_batch_dims=(), lhs_contracting_dims=(1,),
                 rhs_batch_dims=(), rhs_contracting_dims=(0,)):
        # shape inference
        lhs_space_shape = \
            tuple(lhs.shape[i] for i in range(len(lhs.shape))
                  if i not in lhs_contracting_dims and i not in lhs_batch_dims)
        rhs_space_shape = \
            tuple(rhs.shape[i] for i in range(len(rhs.shape))
                  if i not in rhs_contracting_dims and i not in rhs_batch_dims)
        lhs_batch_shape = tuple(lhs.shape[i] for i in lhs_batch_dims)

        shape = lhs_batch_shape + lhs_space_shape + rhs_space_shape

        for i, j in zip(lhs_contracting_dims, rhs_contracting_dims):
            assert lhs.shape[i] == rhs.shape[j]
        for i, j in zip(lhs_batch_dims, rhs_batch_dims):
            assert lhs.shape[i] == rhs.shape[j]

        super().__init__(OpCode.DOT, shape, [lhs, rhs])
        self.lhs = lhs
        self.lhs_batch_dims = lhs_batch_dims
        self.lhs_contracting_dims = lhs_contracting_dims
        self.lhs_space_dims = tuple(set(range(len(lhs.shape))) - set(self.lhs_batch_dims) - set(self.lhs_contracting_dims))
        assert len(self.lhs_contracting_dims) == 1
        assert len(self.lhs_space_dims) == 1
        self.rhs = rhs
        self.rhs_batch_dims = rhs_batch_dims
        self.rhs_contracting_dims = rhs_contracting_dims
        self.rhs_space_dims = tuple(set(range(len(rhs.shape))) - set(self.rhs_batch_dims) - set(self.rhs_contracting_dims))
        assert len(self.rhs_contracting_dims) == 1
        assert len(self.rhs_space_dims) == 1

    def build_strategy_and_cost(self, cluster_env, solver_option):
        lhs = self.lhs
        lhs_batch_dims = self.lhs_batch_dims
        lhs_space_dim = self.lhs_space_dims[0]
        lhs_con_dim = self.lhs_contracting_dims[0]

        rhs = self.rhs
        rhs_batch_dims = self.rhs_batch_dims
        rhs_space_dim = self.rhs_space_dims[0]
        rhs_con_dim = self.rhs_contracting_dims[0]

        space_base_dim = len(self.lhs_batch_dims)

        assert len(cluster_env.device_mesh.shape) == 2

        # Split lhs space dim + rhs space dim
        # @ {0, 1}
        output_spec =\
            ShardingSpec.tile(self.shape, [space_base_dim, space_base_dim + 1], [0, 1], cluster_env)
        self.strategies.append(InstructionStrategy("SS = SR x RS @ {0,1}", output_spec))
        self.compute_costs.append(0)
        self.communication_costs.append(0)
        self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())
        self.resharding_costs.append([
            resharding_cost_vector(cluster_env, lhs,
                ShardingSpec.tile(lhs.shape, [lhs_space_dim], [0], cluster_env)),
            resharding_cost_vector(cluster_env, rhs,
                ShardingSpec.tile(rhs.shape, [rhs_space_dim], [1], cluster_env))
        ])

        # @ {1, 0}
        output_spec =\
            ShardingSpec.tile(self.shape, [space_base_dim, space_base_dim + 1], [1, 0], cluster_env)
        self.strategies.append(InstructionStrategy("SS = SR x RS @ {1,0}", output_spec))
        self.compute_costs.append(0)
        self.communication_costs.append(0)
        self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())
        self.resharding_costs.append([
            resharding_cost_vector(cluster_env, lhs,
                ShardingSpec.tile(lhs.shape, [lhs_space_dim], [1], cluster_env)),
            resharding_cost_vector(cluster_env, rhs,
                ShardingSpec.tile(rhs.shape, [rhs_space_dim], [0], cluster_env))
        ])

        # Split lhs space dim + contracting dim
        # @ {0, 1}
        if cluster_env.device_mesh.shape[1] > 1:
            output_spec = ShardingSpec.tile(self.shape, [space_base_dim], [0], cluster_env)
            memory_cost = compute_bytes(self.shape) / output_spec.num_tile_devices()
            self.strategies.append(
                InstructionStrategy("SR = SS x SR @ {0,1} (allreduce @ 1)", output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(cluster_env.all_reduce_cost(memory_cost, 1))
            self.memory_costs.append(memory_cost)
            self.resharding_costs.append([
                resharding_cost_vector(cluster_env, lhs,
                    ShardingSpec.tile(lhs.shape, [lhs_space_dim, lhs_con_dim], [0, 1], cluster_env)),
                resharding_cost_vector(cluster_env, rhs,
                    ShardingSpec.tile(rhs.shape, [rhs_con_dim], [1], cluster_env))
            ])

        # @ {1, 0}
        if cluster_env.device_mesh.shape[0] > 1:
            output_spec = ShardingSpec.tile(self.shape, [space_base_dim], [1], cluster_env)
            memory_cost = compute_bytes(self.shape) / output_spec.num_tile_devices()
            self.strategies.append(
                InstructionStrategy("SR = SS x SR @ {1,0} (allreduce @ 0)", output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(cluster_env.all_reduce_cost(memory_cost, 0))
            self.memory_costs.append(memory_cost)
            self.resharding_costs.append([
                resharding_cost_vector(cluster_env, lhs,
                    ShardingSpec.tile(lhs.shape, [lhs_space_dim, lhs_con_dim], [1, 0], cluster_env)),
                resharding_cost_vector(cluster_env, rhs,
                    ShardingSpec.tile(rhs.shape, [rhs_con_dim], [0], cluster_env))
            ])

        # Split rhs space dim + contracting dim
        # @ {0, 1}
        if cluster_env.device_mesh.shape[0] > 1 and cluster_env.device_mesh.shape[1] > 1:
            output_spec = ShardingSpec.tile(self.shape, [space_base_dim+1], [1], cluster_env)
            memory_cost = compute_bytes(self.shape) / output_spec.num_tile_devices()
            self.strategies.append(
                InstructionStrategy("RS = RS x SS @ {0,1} (allreduce @ 0)", output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(cluster_env.all_reduce_cost(memory_cost, 0))
            self.memory_costs.append(memory_cost)
            self.resharding_costs.append([
                resharding_cost_vector(cluster_env, lhs,
                    ShardingSpec.tile(lhs.shape, [lhs_con_dim], [0], cluster_env)),
                resharding_cost_vector(cluster_env, rhs,
                    ShardingSpec.tile(rhs.shape, [rhs_con_dim, rhs_space_dim], [0, 1], cluster_env))
            ])

        # @ {1, 0}
        if cluster_env.device_mesh.shape[0] > 1 and cluster_env.device_mesh.shape[1] > 1:
            output_spec = ShardingSpec.tile(self.shape, [space_base_dim+1], [0], cluster_env)
            memory_cost = compute_bytes(self.shape) / output_spec.num_tile_devices()
            self.strategies.append(
                InstructionStrategy("RS = RS x SS @ {1,0} (allreduce @ 1)", output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(cluster_env.all_reduce_cost(memory_cost, 1))
            self.memory_costs.append(memory_cost)
            self.resharding_costs.append([
                resharding_cost_vector(cluster_env, lhs,
                    ShardingSpec.tile(lhs.shape, [lhs_con_dim], [1], cluster_env)),
                resharding_cost_vector(cluster_env, rhs,
                    ShardingSpec.tile(rhs.shape, [rhs_con_dim, rhs_space_dim], [1, 0], cluster_env))
            ])

        # Split one batch dim
        for i in range(len(self.lhs_batch_dims)):
            for j in range(len(cluster_env.device_mesh.shape)):
                if (cluster_env.device_mesh.shape[j] == 1 or
                    self.shape[i] < cluster_env.device_mesh.shape[j]):
                    continue

                output_spec = ShardingSpec.tile(self.shape, [i], [j], cluster_env)
                self.strategies.append(InstructionStrategy(f"Sb_{i} = Sb x Sb @ {j}", output_spec))
                self.compute_costs.append(0)
                self.communication_costs.append(0)
                self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())
                self.resharding_costs.append([
                    resharding_cost_vector(cluster_env, lhs,
                        ShardingSpec.tile(lhs.shape, [lhs_batch_dims[i]], [j], cluster_env)),
                    resharding_cost_vector(cluster_env, rhs,
                        ShardingSpec.tile(rhs.shape, [rhs_batch_dims[i]], [j], cluster_env))
                ])

        # Split two batch dims
        if len(self.lhs_batch_dims) == 2 and cluster_env.device_mesh.shape[0] > 1\
                and cluster_env.device_mesh.shape[1] > 1:

            self.strategies = []
            self.compute_costs = []
            self.communication_costs = []
            self.memory_costs = []
            self.resharding_costs = []

            # Split two batch dims
            output_spec = ShardingSpec.tile(self.shape, [0, 1], [0, 1], cluster_env)
            self.strategies.append(InstructionStrategy("Sb = Sb x Sb @ {0,1}", output_spec))
            self.compute_costs.append(0)
            self.communication_costs.append(0)
            self.memory_costs.append(compute_bytes(self.shape) / output_spec.num_tile_devices())
            self.resharding_costs.append([
                resharding_cost_vector(cluster_env, lhs,
                    ShardingSpec.tile(lhs.shape, [lhs_batch_dims[0], lhs_batch_dims[1]], [0, 1], cluster_env)),
                resharding_cost_vector(cluster_env, rhs,
                    ShardingSpec.tile(rhs.shape, [rhs_batch_dims[0], rhs_batch_dims[1]], [0, 1], cluster_env))
            ])

        # If force batch dim to a mesh dim, filter out invalid strategies
        if solver_option.force_batch_dim_to_mesh_dim is not None and self.batch_dim is not None:
            filter_indices = []
            for i in range(len(self.strategies)):
                tensor_dim_to_mesh_dim = cluster_env.get_tensor_dim_to_mesh_dim(
                    self.shape, self.strategies[i].output_spec)
                if tensor_dim_to_mesh_dim[self.batch_dim] == solver_option.force_batch_dim_to_mesh_dim:
                    filter_indices.append(i)

            self.strategies = [self.strategies[i] for i in filter_indices]
            self.compute_costs = [self.compute_costs[i] for i in filter_indices]
            self.communication_costs = [self.communication_costs[i] for i in filter_indices]
            self.memory_costs = [self.memory_costs[i] for i in filter_indices]
            self.resharding_costs = [self.resharding_costs[i] for i in filter_indices]

    def propagate_batch_dim(self, operand):
        index = self.operands.index(operand)

        if index == 0:
            for i in range(len(self.lhs_batch_dims)):
                if operand.batch_dim == self.lhs_batch_dims[i]:
                    self.batch_dim = i
                    return True
            if operand.batch_dim == self.lhs_space_dims[0]:
                self.batch_dim = len(self.lhs_batch_dims)
                return True
            if operand.batch_dim in self.lhs_contracting_dims:
                return False
        else:
            for i in range(len(self.rhs_batch_dims)):
                if operand.batch_dim == self.rhs_batch_dims[i]:
                    self.batch_dim = i
                    return True
            if operand.batch_dim == self.rhs_space_dims[0]:
                self.batch_dim = len(self.rhs_batch_dims)
                return True
            if operand.batch_dim in self.rhs_contracting_dims:
                return False

    def __str__(self):
        return f"{self.name} {self.shape} = dot({self.lhs.name}, {self.rhs.name}) "\
               f" lhs_con_dim={self.lhs_contracting_dims},"\
               f" rhs_con_dim={self.rhs_contracting_dims}"

class HloTuple(HloInstruction):
    def __init__(self, operands):
        super().__init__(OpCode.TUPLE, (), operands)

    def build_strategy_and_cost(self, cluster_env, solver_option):
        self.strategies.append(InstructionStrategy("tuple", ShardingSpec.tuple()))
        self.memory_costs.append(0)
        self.compute_costs.append(0)
        self.communication_costs.append(0)
        self.resharding_costs.append([np.zeros(len(operand.strategies))
            for operand in self.operands])

    def __str__(self):
        names = tuple(x.name for x in self.operands)
        return f"{self.name} {self.shape} = tuple{names}"

class HloComputation:
    cur_env = None

    def __init__(self):
        self.ct = 0
        self.instructions = []
        self.alias_list = []
        self.alias_cost_vector = []

        self.parameters = []

        self.strategy_built = False

    def append(self, instruction):
        ct = len(self.instructions)
        self.instructions.append(instruction)

        if instruction.op_code == OpCode.PARAMETER:
            self.parameters.append(instruction)

        return ct

    def liveness_analysis(self):
        liveness_dict = dict()

        live_set = set()

        for t in range(len(self.instructions)-1, -1, -1):
            inst = self.instructions[t]

            live_set.add(inst)
            for operand in inst.operands:
                live_set.add(operand)

            liveness_dict[t] = set(live_set)

            live_set.remove(inst)

        return liveness_dict

    def set_alias(self, alias_list):
        self.alias_list = alias_list

    def concurrency_analysis(self):
        frontier_list = []
        edge_dict = defaultdict(list)

        # Build degree dict
        #out_degree = defaultdict(lambda : 0)
        #for ins in self.instructions:
        #    for operand in ins.operands:
        #        out_degree[operand] += 1

        degree = defaultdict(lambda : 0)
        for ins in self.instructions:
            for operand in ins.operands:
                degree[ins] += 1
                edge_dict[operand].append(ins)

        # Init frontier
        collected = 0
        current_frontier = []
        for ins in self.instructions:
            if degree[ins] == 0:
                current_frontier.append(ins)
                collected += 1
        frontier_list.append(current_frontier)

        # Push forward frontier
        while collected < len(self.instructions):
            current_frontier = frontier_list[-1]
            next_frontier = []
            for ins in current_frontier:
                for node in edge_dict[ins]:
                    degree[node] -= 1
                    if degree[node] == 0:
                        next_frontier.append(node)
                        collected += 1
            frontier_list.append(next_frontier)

        for i, frontier in enumerate(frontier_list):
            print(i)
            for ins in frontier:
                print(ins)

    def forward_backward_analysis(self):
        used_by = defaultdict(list)
        for ins in self.instructions:
            for operand in ins.operands:
                used_by[operand].append(ins.index)

        sep_id = 0
        for param in self.parameters:
            if len(used_by[param]) > 2:
                backward_id = used_by[param][0]
                sep_id = max(sep_id, backward_id + 1)

        return sep_id

    def batch_dim_analysis(self):
        # Build used by dict
        used_by = defaultdict(list)
        for ins in self.instructions:
            for operand in ins.operands:
                used_by[operand].append(ins)

        # Find source.
        # Rule: The first dim of parameters that are only used once
        #possible_inputs = []
        #for param in self.parameters:
        #    if len(used_by[param]) == 1:
        #        possible_inputs.append(param)
        #source = possible_inputs[0]
        source = self.instructions[0]
        source.batch_dim = 0

        # Dim propagation
        queue = [source]
        visited = set([source])

        while len(queue) > 0:
            ins = queue.pop(0)

            # Propagate to operand

            # Propagate to used_by
            for consumer in used_by[ins]:
                #print(f"Propagate from {ins} to {consumer}")
                success = consumer.propagate_batch_dim(ins)
                if not success:
                    continue
                if consumer.index not in visited:
                    visited.add(consumer)
                    queue.append(consumer)

    def depth_analysis(self):
        edge_dict = defaultdict(list)

        degree = defaultdict(lambda : 0)
        for ins in self.instructions:
            for operand in ins.operands:
                degree[ins] += 1
                edge_dict[operand].append(ins)

        # Init frontier
        collected = 0
        current_frontier = []
        for ins in self.instructions:
            if degree[ins] == 0:
                ins.depth = 0
                current_frontier.append(ins)
                collected += 1

        # Push forward frontier
        depth = 0
        while collected < len(self.instructions):
            next_frontier = []
            for ins in current_frontier:
                for node in edge_dict[ins]:
                    degree[node] -= 1
                    if degree[node] == 0:
                        next_frontier.append(node)
                        collected += 1

            depth += 1
            current_frontier = next_frontier
            for ins in current_frontier:
                ins.depth = depth

    def build_strategy_and_cost(self, cluster_env, solver_option):
        if self.strategy_built:
            for ins in self.instructions:
                ins.strategies = []
                ins.compute_costs = []
                ins.communication_costs = []
                ins.memory_costs = []
                ins.resharding_costs = []
                ins.follow_ins = None

            self.alias_cost_vector = []

        # Analyze depth for all instructions
        self.depth_analysis()

        # Analyze batch dim
        if solver_option.force_batch_dim_to_mesh_dim is not None:
            batch_dim = self.batch_dim_analysis()
            print("===== Batch Dim Analysis =====")
            for i in range(len(self.instructions)):
                print(f"Time {i:2d}: {self.instructions[i]}  Batch: {self.instructions[i].batch_dim}")

        # Build strategies and costs for each instruction
        for ins in self.instructions:
            ins.build_strategy_and_cost(cluster_env, solver_option)

        # Build alias costs
        for (ins_a, ins_b) in self.alias_list:
            assert ins_a.shape == ins_b.shape
            cost_vector = []
            for stra_a in ins_a.strategies:
                for stra_b in ins_b.strategies:
                    if stra_a.output_spec == stra_b.output_spec:
                        cost_vector.append(0)
                    else:
                        cost_vector.append(1)
            self.alias_cost_vector.append(cost_vector)

        self.strategy_built = True

    def __enter__(self):
        assert HloComputation.cur_env is None
        HloComputation.cur_env = self

    def __exit__(self, *args, **kwargs):
        HloComputation.cur_env = None

    def __str__(self):
        strs = []
        for i, ins in enumerate(self.instructions):
            strs.append(f"{i:2d}: " + str(ins))
        return "\n".join(strs)


text:
from .context import assert_equal, _Add, _Mul, _Pow
import pytest
from sympy import binomial, Symbol, Rational

x = Symbol('x', real=True, positive=True)
y = Symbol('y', real=True, positive=True)
theta = Symbol('theta', real=True, positive=True)
gamma = Symbol('gamma', real=True, positive=True)

def test_binomial_numeric():
    assert_equal("\\binom{16}{2}", binomial(16, 2))

def test_binomial_symbols():
    assert_equal("\\binom{x}{y}", binomial(x, y))

def test_binomial_greek_symbols():
    assert_equal("\\binom{\\theta}{\\gamma}", binomial(theta, gamma))

def test_binomial_expr():
    assert_equal("\\binom{16+2}{\\frac{4}{2}}", binomial(_Add(16, 2), Rational(4, 2)))

def test_choose_numeric():
    assert_equal("\\choose{16}{2}", binomial(16, 2))

def test_choose_symbols():
    assert_equal("\\choose{x}{y}", binomial(x, y))

def test_choose_greek_symbols():
    assert_equal("\\choose{\\theta}{\\gamma}", binomial(theta, gamma))

ext: py

reason: from sympy


text:
import cv2
import numpy as np
import glob
import json
import urllib.request
import tempfile
import requests
import requests
import shutil
import os

'''
def imread_web(url):
    # 画像をリクエストする
    res = requests.get(url)
    img = None
    # Tempfileを作成して即読み込む
    with tempfile.NamedTemporaryFile(dir='./') as fp:
        fp.write(res.content)
        fp.file.seek(0)
        img = cv2.imread(fp.name)
    return img
'''

def download_img(url, file_name):
    r = requests.get(url, stream=True)
    if r.status_code == 200:
        with open(file_name, 'wb') as f:
            r.raw.decode_content = True
            shutil.copyfileobj(r.raw, f)

# files = glob.glob("/Users/nakamura/Dropbox/U-PARL_Moeller/input/*.jpg")
# files = glob.glob("input/*.jpg")

manifest = "https://iiif.dl.itc.u-tokyo.ac.jp/repo/iiif/56653a59-0d55-4d1a-a7e3-2242e02859a1/manifest"
num = "2"

res = urllib.request.urlopen(manifest)
# json_loads() でPythonオブジェクトに変換
data = json.loads(res.read())

canvases = data["sequences"][0]["canvases"]

th = 30

rep = "full/full"

for index in range(len(canvases)):

    canvas = canvases[index]
    image_url = canvas["images"][0]["resource"]["@id"].replace("full/full", rep)
    print(image_url)

    if index < th:
        continue

    file = "data/target_all/"+num+"/input/"+str(index+1).zfill(3)+".jpg"

    if not os.path.exists(file):
        download_img(image_url, file)

reason: import numpy


text:
from typing import Any, Tuple, Union, TypeVar, Optional

from abc import ABC

from anndata import AnnData

import numpy as np
from scipy.sparse import spmatrix

KernelExpression = TypeVar("KernelExpression")

class KernelMixin(ABC):
    """Mixin that exposes various properties of :class:`cellrank.kernels.KernelExpression`."""

    def __init__(self, kernel: "KernelExpression", **kwargs: Any):
        super().__init__(**kwargs)
        self._kernel = kernel
        self._n_obs = self.kernel.adata.n_obs

    @property
    def kernel(self) -> "KernelExpression":
        """Underlying kernel expression."""
        return self._kernel

    @property
    def adata(self) -> AnnData:
        """Annotated data object."""
        return self.kernel.adata

    @property
    def shape(self) -> Tuple[int, int]:
        """Shape of the kernel."""
        return self.kernel.shape

    @adata.setter
    def adata(self, adata: Optional[AnnData]) -> None:
        self.kernel.adata = adata

    def __len__(self) -> int:
        return self._n_obs

    @property
    def transition_matrix(self) -> Union[np.ndarray, spmatrix]:
        """Transition matrix of :attr:`kernel`."""
        return self.kernel.transition_matrix

    @property
    def backward(self) -> bool:
        """Direction of :attr:`kernel`."""
        return self.kernel.backward

max_stars_repo_name: WeilerP/cellrank

max_issues_repo_name: WeilerP/cellrank

max_forks_repo_name: WeilerP/cellrank

reason: import numpy,from scipy


text:
import datetime
import inspect
import numpy.testing as npt
import os.path
import pandas as pd
import pandas.util.testing as pdt
import sys
from tabulate import tabulate
import unittest

# #find parent directory and import model
# parentddir = os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir))
# sys.path.append(parentddir)
from ..screenip_exe import Screenip

test = {}

class TestScreenip(unittest.TestCase):
    """
    Unit tests for screenip.
    """
    print("screenip unittests conducted at " + str(datetime.datetime.today()))

    def setUp(self):
        """
        Setup routine for screenip unittest.
        :return:
        """

        pass
        # screenip2 = screenip_model.screenip(0, pd_obj_inputs, pd_obj_exp_out)
        # setup the test as needed
        # e.g. pandas to open screenip qaqc csv
        #  Read qaqc csv and create pandas DataFrames for inputs and expected outputs

    def tearDown(self):
        """
        Teardown routine for screenip unittest.
        :return:
        """
        pass
        # teardown called after each test
        # e.g. maybe write test results to some text file

    def create_screenip_object(self):
        # create empty pandas dataframes to create empty object for testing
        df_empty = pd.DataFrame()
        # create an empty screenip object
        screenip_empty = Screenip(df_empty, df_empty)
        return screenip_empty

    def test_screenip_unit_fw_bird(self):
        """
        unittest for function screenip.fw_bird:
        :return:
        """
        expected_results = pd.Series([0.0162, 0.0162, 0.0162], dtype='float')
        result = pd.Series([], dtype='float')

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        try:
            # for i in range(0,3):
            #     result[i] = screenip_empty.fw_bird()
            screenip_empty.no_of_runs = len(expected_results)
            screenip_empty.fw_bird()
            result = screenip_empty.out_fw_bird
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_fw_mamm(self):
        """
        unittest for function screenip.fw_mamm:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([0.172, 0.172, 0.172], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            screenip_empty.no_of_runs = len(expected_results)
            result = screenip_empty.fw_mamm()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_dose_bird(self):
        """
        unittest for function screenip.dose_bird:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([1000000., 4805.50175, 849727.21122], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            #(self.out_fw_bird * self.solubility)/(self.bodyweight_assessed_bird / 1000.)
            screenip_empty.out_fw_bird = pd.Series([10., 0.329, 1.8349], dtype='float')
            screenip_empty.solubility = pd.Series([100., 34.9823, 453.83], dtype='float')
            screenip_empty.bodyweight_assessed_bird = pd.Series([1.0, 2.395, 0.98], dtype='float')
            result = screenip_empty.dose_bird()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_dose_mamm(self):
        """
        unittest for function screenip.dose_mamm:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([8000000., 48205.7595, 3808036.37889], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            #(self.out_fw_mamm * self.solubility)/(self.bodyweight_assessed_mammal / 1000)
            screenip_empty.out_fw_mamm = pd.Series([20., 12.843, 6.998], dtype='float')
            screenip_empty.solubility = pd.Series([400., 34.9823, 453.83], dtype='float')
            screenip_empty.bodyweight_assessed_mammal = pd.Series([1., 9.32, 0.834], dtype='float')
            result = screenip_empty.dose_mamm()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_at_bird(self):
        """
        unittest for function screenip.at_bird:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([1000., 687.9231, 109.3361], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            #(self.ld50_avian_water) * ((self.bodyweight_assessed_bird / self.bodyweight_tested_bird)**(self.mineau_scaling_factor - 1.))
            screenip_empty.ld50_avian_water = pd.Series([2000., 938.34, 345.83], dtype='float')
            screenip_empty.bodyweight_assessed_bird = pd.Series([100., 39.49, 183.54], dtype='float')
            screenip_empty.ld50_bodyweight_tested_bird = pd.Series([200., 73.473, 395.485], dtype='float')
            screenip_empty.mineau_scaling_factor = pd.Series([2., 1.5, 2.5], dtype='float')
            result = screenip_empty.at_bird()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_at_mamm(self):
        """
        unittest for function screenip.at_mamm:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([11.89207, 214.0572, 412.6864], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            #(self.ld50_mammal_water) * ((self.bodyweight_tested_mammal / self.bodyweight_assessed_mammal)**0.25)
            screenip_empty.ld50_mammal_water = pd.Series([10., 250., 500.], dtype='float')
            screenip_empty.ld50_bodyweight_tested_mammal = pd.Series([200., 39.49, 183.54], dtype='float')
            screenip_empty.bodyweight_assessed_mammal = pd.Series([100., 73.473, 395.485], dtype='float')
            result = screenip_empty.at_mamm()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_fi_bird(self):
        """
        unittest for function screenip.fi_bird:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([0.012999, 0.026578, 0.020412], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            #0.0582 * ((bw_grams / 1000.)**0.651)
            bw_grams = pd.Series([100., 300., 200.], dtype='float')
            result = screenip_empty.fi_bird(bw_grams)
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_act(self):
        """
        unittest for function screenip.test_act:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([10.5737, 124.8032, 416.4873], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            #(self.noael_mammal_water) * ((self.bodyweight_tested_mammal / self.bodyweight_assessed_mammal)**0.25)
            screenip_empty.noael_mammal_water = pd.Series([10., 120., 400.], dtype='float')
            screenip_empty.noael_bodyweight_tested_mammal = pd.Series([500., 385.45, 673.854], dtype='float')
            screenip_empty.bodyweight_assessed_mammal = pd.Series([400., 329.45, 573.322], dtype='float')
            result = screenip_empty.act()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_det(self):
        """
        unittest for function screenip.det
        return:
        """
    #
    #     '''
    #     Dose Equiv. Toxicity:
    #
    #     The FI value (kg-diet) is multiplied by the reported NOAEC (mg/kg-diet) and then divided by
    #     the test animal's body weight to derive the dose-equivalent chronic toxicity value (mg/kg-bw):
    #
    #     Dose Equiv. Toxicity = (NOAEC * FI) / BW
    #
    #     NOTE: The user enters the lowest available NOAEC for the mallard duck, for the bobwhite quail,
    #     and for any other test species. The model calculates the dose equivalent toxicity values for
    #     all of the modeled values (Cells F20-24 and results worksheet) and then selects the lowest dose
    #     equivalent toxicity value to represent the chronic toxicity of the chemical to birds.
    #     '''
    #     try:
    #         # result =
    #         # self.assertEquals(result, )
    #         pass
    #     finally:
    #         pass
    #    return
    #
    #
    # def test_det_duck(self):
    #     """
    #     unittest for function screenip.det_duck:
    #     :return:
    #     """
    #     try:
    #         # det_duck = (self.noaec_duck * self.fi_bird(1580.)) / (1580. / 1000.)
    #         screenip_empty.noaec_duck = pd.Series([1.], dtype='int')
    #         screenip_empty.fi_bird = pd.Series([1.], dtype='int')
    #         result = screenip_empty.det_duck()
    #         npt.assert_array_almost_equal(result, 1000., 4, '', True)
    #     finally:
    #         pass
    #     return
    #
    # def test_det_quail(self):
    #     """
    #     unittest for function screenip.det_quail:
    #     :return:
    #     """
    #     try:
    #         # det_quail = (self.noaec_quail * self.fi_bird(178.)) / (178. / 1000.)
    #         screenip_empty.noaec_quail = pd.Series([1.], dtype='int')
    #         screenip_empty.fi_bird = pd.Series([1.], dtype='int')
    #         result = screenip_empty.det_quail()
    #         npt.assert_array_almost_equal(result, 1000., 4, '', True)
    #     finally:
    #         pass
    #     return
    #
    # def test_det_other_1(self):
    #     """
    #     unittest for function screenip.det_other_1:
    #     :return:
    #     """
    #     try:
    #         #det_other_1 = (self.noaec_bird_other_1 * self.fi_bird(self.bodyweight_bird_other_1)) / (self.bodyweight_bird_other_1 / 1000.)
    #         #det_other_2 = (self.noaec_bird_other_2 * self.fi_bird(self.bodyweight_bird_other_1)) / (self.bodyweight_bird_other_1 / 1000.)
    #         screenip_empty.noaec_bird_other_1 = pd.Series([400.]) # mg/kg-diet
    #         screenip_empty.bodyweight_bird_other_1 = pd.Series([100]) # grams
    #         result = screenip_empty.det_other_1()
    #         npt.assert_array_almost_equal(result, 4666, 4)
    #     finally:
    #         pass
    #     return
    #
    #   The following tests are configured such that:
        #   1. four values are provided for each needed input
        #   2. the four input values generate four values of out_det_* per bird type
        #   3. the inputs per bird type are set so that calculations of out_det_* will result in
        #      each bird type having one minimum among the bird types;
        #      thus all four calculations result in one minimum per bird type

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([4.2174, 4.96125, 7.97237, 10.664648], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            screenip_empty.bodyweight_bobwhite_quail = 178.
            screenip_empty.bodyweight_mallard_duck = 1580.
            screenip_empty.noaec_quail = pd.Series([100., 300., 75., 150.], dtype='float')
            screenip_empty.noaec_duck = pd.Series([400., 100., 200., 350.], dtype='float')
            screenip_empty.noaec_bird_other_1 = pd.Series([50., 200., 300., 250.], dtype='float')
            screenip_empty.noaec_bird_other_2 = pd.Series([350., 400., 250., 100.], dtype='float')
            screenip_empty.noaec_bodyweight_bird_other_1 = pd.Series([345.34, 453.54, 649.29, 294.56], dtype='float')
            screenip_empty.noaec_bodyweight_bird_other_2 = pd.Series([123.84, 85.743, 127.884, 176.34], dtype='float')
            screenip_empty.no_of_runs = len(expected_results)
            result = screenip_empty.det()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_acute_bird(self):
        """
        unittest for function screenip.acute_bird:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([10., 5.22093, 0.479639], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            # self.out_acute_bird = self.out_dose_bird / self.out_at_bird
            screenip_empty.out_dose_bird = pd.Series([100., 121.23, 43.994], dtype='float')
            screenip_empty.out_at_bird = pd.Series([10., 23.22, 91.723], dtype='float')
            result = screenip_empty.acute_bird()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_acuconb(self):
        """
        unittest for function screenip.acuconb:
        Message stating whether or not a risk is present
        :return:
        """
        # if self.out_acuconb == -1:
        #     if self.out_acute_bird == None:
        #         raise ValueError\
        #         ('acute_bird variable equals None and therefor this function cannot be run.')
        # if self.out_acute_bird < 0.1:
        #     self.out_acuconb = ('Drinking water exposure alone is NOT a potential concern for birds')
        # else:
        #     self.out_acuconb = ('Exposure through drinking water alone is a potential concern for birds')

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series(["Exposure through drinking water alone is a potential concern "
                            "for birds", "Drinking water exposure alone is NOT a potential "
                            "concern for birds", "Exposure through drinking water alone is a "
                            "potential concern for birds"], dtype='object')
        result = pd.Series([], dtype='object')

        try:
            screenip_empty.out_acute_bird = pd.Series([0.2, 0.09, 0.1], dtype='float')
            result = screenip_empty.acuconb()
            pdt.assert_series_equal(result, expected_results,  True)
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_acute_mamm(self):
        """
        unittest for function screenip.acute_mamm:
        :return:
        """
        # self.out_acute_mamm = self.out_dose_mamm / self.out_at_mamm

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([10., 14.68657, 2.124852], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            screenip_empty.out_dose_mamm = pd.Series([100., 34.44, 159.349], dtype='float')
            screenip_empty.out_at_mamm = pd.Series([10., 2.345, 74.993], dtype='float')
            result = screenip_empty.acute_mamm()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_acuconm(self):
        """
        unittest for function screenip.acuconm:
        Message stating whether or not a risk is present
        :return:
        """
        # if self.out_acuconm == -1:
        #     if self.out_acute_mamm == None:
        #         raise ValueError\
        #         ('acute_mamm variable equals None and therefor this function cannot be run.')
        #     if self.out_acute_mamm < 0.1:
        #         self.out_acuconm = ('Drinking water exposure alone is NOT a potential concern for mammals')
        #     else:
        #         self.out_acuconm = ('Exposure through drinking water alone is a potential concern for mammals')
        #     return self.out_acuconm

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series(["Drinking water exposure alone is NOT a potential concern "
                                      "for mammals", "Exposure through drinking water alone is a "
                                      "potential concern for mammals", "Drinking water exposure "
                                      "alone is NOT a potential concern for mammals"], dtype='object')
        result = pd.Series([], dtype='object')

        try:
            screenip_empty.out_acute_mamm = pd.Series([0.09, 0.2, 0.002], dtype='float')
            result = screenip_empty.acuconm()
            pdt.assert_series_equal(result, expected_results, True)
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_chron_bird(self):
        """
        unittest for function screenip.chron_bird:
        :return:
        """
        #self.out_chron_bird = self.out_dose_bird / self.out_det

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([0.5, 0.10891, 2.39857], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            screenip_empty.out_dose_bird = pd.Series([5., 1.32, 19.191], dtype='float')
            screenip_empty.out_det = pd.Series([10., 12.12, 8.001], dtype='float')
            result = screenip_empty.chron_bird()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_chronconb(self):
        """
        unittest for function screenip.chronconb:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series(["Drinking water exposure alone is NOT "
                                      "a potential concern for birds", "Exposure through "
                                      "drinking water alone is a potential concern for "
                                      "birds", "Drinking water exposure alone is NOT a "
                                      "potential concern for birds"], dtype='object')
        result = pd.Series([], dtype='object')

        try:
            screenip_empty.out_chron_bird = pd.Series([0.12, 3., 0.97], dtype='float')
            result = screenip_empty.chronconb()
            pdt.assert_series_equal(result, expected_results, True)
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_chron_mamm(self):
        """
        unittest for function screenip.chron_mamm:
        :return:
        """
        # self.out_chron_mamm = self.out_dose_mamm / self.out_act

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series([2.0, 14.1333, 244.7245], dtype='float')
        result = pd.Series([], dtype='float')

        try:
            screenip_empty.out_dose_mamm = pd.Series([8., 34.344, 23.983], dtype='float')
            screenip_empty.out_act = pd.Series([4., 2.43, 0.098], dtype='float')
            result = screenip_empty.chron_mamm()
            npt.assert_allclose(result, expected_results, rtol=1e-4, atol=0, err_msg='', verbose=True )
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

    def test_screenip_unit_chronconm(self):
        """
        unittest for function screenip.chronconm:
        :return:
        """

        # create empty pandas dataframes to create empty object for this unittest
        screenip_empty = self.create_screenip_object()

        expected_results = pd.Series(["Drinking water exposure alone is NOT a potential "
                            "concern for mammals", "Exposure through drinking water alone "
                            "is a potential concern for mammals", "Drinking water exposure "
                            "alone is NOT a potential concern for mammals"], dtype='object')
        result = pd.Series([], dtype='object')

        try:
            screenip_empty.out_chron_mamm = pd.Series([0.5, 1.0, 0.09], dtype='float')
            result = screenip_empty.chronconm()
            pdt.assert_series_equal(result, expected_results, True)
        finally:
            tab = [result, expected_results]
            print("\n")
            print(inspect.currentframe().f_code.co_name)
            print(tabulate(tab, headers='keys', tablefmt='rst'))
        return

# unittest will
# 1) call the setup method,
# 2) then call every method starting with "test",
# 3) then the teardown method
if __name__ == '__main__':
    unittest.main()
    #pass

max_stars_repo_name: qed-uber/ubertool

max_issues_repo_name: qed-uber/ubertool

max_forks_repo_name: quanted/ubertool


text:
# coding=utf-8
import time
from pathlib import Path
from typing import List, Dict, Union, Tuple

import pandas as pd
import numpy as np
from matplotlib import rcParams

rcParams["font.family"] = "monospace"
from matplotlib import pyplot as plt
import argparse
import scipy.stats

# plt.xkcd()
# set std params here because lazy
color_xa = "xkcd:pastel orange"
color_xb = "xkcd:pastel pink"
color_xc = "xkcd:forest green"

# lib functions
def nice_string_output(
    names: List[str], values: List[str], extra_spacing: int = 0,
):
    max_values = len(max(values, key=len))
    max_names = len(max(names, key=len))
    string = ""
    for name, value in zip(names, values):
        string += "{0:s} {1:>{spacing}} \n".format(
            name, value, spacing=extra_spacing + max_values + max_names - len(name),
        )
    return string[:-2]

def calculate_joint_rates_and_prob(*rates: Tuple[float]) -> Tuple[float, float]:
    """
    Returns joint rate and probability
    """
    K = np.sum(rates)[0]
    P = 1.0 / K

    return K, P

def run(args: Dict):
    TEST = args["test"]
    max_time = args["maxtime"]
    rate_production = args["alpha"]
    volume = args["volume"]
    rate_degradation = args["gamma"]

    plot_range = (0, 20)

    curr_time = 0
    curr_n_mrna = 0

    storage_mrna = [0]
    storage_time = [0]
    storage_steps = []

    while curr_time < max_time:
        curr_r_prod = rate_production * volume
        curr_r_deg = rate_degradation * curr_n_mrna

        # 1. calculate K
        K = sum((curr_r_deg, curr_r_prod))

        # 2. draw a random number a and find duration until next event bu -ln a /K
        a = np.random.uniform(0, 1, size=1)[0]

        tau = -np.log(a) / K

        # 3. Draw a new number b
        b = np.random.uniform(0, 1, size=1)[0]

        # TODO: make this argmin based for n r
        if b < curr_r_prod / K:  # produce
            curr_n_mrna += 1
            if TEST:
                print(f"Produced one MRNA! Count: {curr_n_mrna}")
        else:
            curr_n_mrna -= 1

            if TEST:
                print(f"Degraded one MRNA! Count: {curr_n_mrna}")

        curr_time += tau
        if TEST:
            print(f"Incremented time by: {tau:.2f} --> {curr_time:.2f}")

        storage_mrna.append(curr_n_mrna)
        storage_time.append(curr_time)
        storage_steps.append(tau)

    storage_mrna.pop()
    storage_time.pop()

    df = pd.DataFrame({"Time": storage_time, "N_MRNA": storage_mrna, "Steps":storage_steps})
    print(df.tail())
    print(f"N steps:{len(df)}")

    # plotting

    n_mean = sum(df["N_MRNA"] * df["Steps"]) / max_time

    param_str = nice_string_output(
        names=["alpha", "gamma", "volume", "max_time"],
        values=[
            f"{v:.2f}"
            for v in (args["alpha"], args["gamma"], args["volume"], args["maxtime"])
        ],
    )
    poisson_lambda = args["alpha"] * args["volume"] / args["gamma"]
    result_str = nice_string_output(
        names=[r"Poisson param", "Mean", "Var", "Fano", "N steps", "Mean Step"],
        values=[
            f"{v:.2f}"
            for v in (
                poisson_lambda,
                n_mean,
                # df["N_MRNA"].mean(),
                df["N_MRNA"].var(),
                df["N_MRNA"].var() / n_mean,
                len(df),
                df["Time"].diff().mean(),
            )
        ],
        extra_spacing=2,
    )

    print(param_str)
    print(result_str)

    ax_dist: plt.Axes
    ax_traj: plt.Axes
    fig: plt.Figure

    fig, (ax_traj, ax_dist) = plt.subplots(nrows=2, figsize=(6, 8))

    ax_traj.set_title("Time plot of Gillespie Algorithm")

    ax_traj.scatter(
        df["Time"], df["N_MRNA"], label="Simulated", c=color_xa, alpha=0.7, marker="2"
    )
    ax_traj.axhline(poisson_lambda, ls="-.", label="Analytical Mean")
    ax_traj.axhspan(
        poisson_lambda - poisson_lambda ** 0.5,
        poisson_lambda + poisson_lambda ** 0.5,
        ls="-.",
        label="Analytical STD",
        alpha=0.2,
        color=color_xc,
    )

    ax_traj.set_ylabel("Number of MRNA")
    ax_traj.set_xlabel("Time")

    ax_traj.legend(loc="upper left")

    ax_dist.hist(
        df["N_MRNA"],
        histtype="step",
        bins=max(plot_range),
        range=plot_range,
        label="MRNA Number",
        color=color_xa,
        density=True,
    )

    ax_dist.set_xlabel("Number of MRNA")
    ax_dist.set_ylabel("Density")
    ax_dist.set_title(
        rf"Poisson with $\lambda = {poisson_lambda:.3f}$ vs. Empirical Distribution"
    )

    analytical_poisson = scipy.stats.poisson(poisson_lambda)
    analytical_poisson_x = np.arange(*plot_range)
    analytical_poisson_y = analytical_poisson.pmf(analytical_poisson_x)

    ax_dist.vlines(
        analytical_poisson_x + 0.5,
        0,
        analytical_poisson_y,
        colors=color_xc,
        label="Analytical Poisson PMF",
        alpha=0.5,
    )

    ax_dist.legend()

    ax_traj.text(
        s="PARAMS:\n" + param_str, x=max_time * 0.7, y=9,
    )
    ax_dist.text(
        s="RESULTS:\n" + result_str, y=0.05, x=12,
    )

    fig.tight_layout()

    basename = args["out_name"] + "_simulation." + args["filetype"]
    fig.savefig(args["outdir"] / basename)
    fig.clf()

def parse_arguments() -> Dict[str, Union[int, float, str, Path]]:
    parser = argparse.ArgumentParser(description="Chapter 7 exercise 8 code")

    parser.add_argument("-a", "--alpha", type=float, default=3.0)
    parser.add_argument("-g", "--gamma", type=float, default=0.5)
    parser.add_argument("-v", "--volume", type=float, default=1.0)
    parser.add_argument("-t", "--maxtime", type=int, default=1000)

    parser.add_argument("-o", "--outdir", type=str, default="./figs")
    parser.add_argument("--filetype", type=str, choices=["png", "pdf"], default="pdf")

    parser.add_argument("--test", type=bool, default=False)  # change to false once done

    args = parser.parse_args()
    argdict = vars(args)  # returns a dict, easier to deal with

    if argdict["test"] == True:
        argdict["maxtime"] = 100

    po = Path(argdict["outdir"])
    if not po.exists():
        po.mkdir()
    print("Set output dir to: " + str(po.absolute()))
    argdict["outdir"] = po
    # Set output name here
    timestr = timestr = time.strftime("%Y%m%d_%H%M%S")
    out_str = f"ch7_e8_" + timestr

    argdict["out_name"] = out_str

    return argdict

if __name__ == "__main__":
    args = parse_arguments()
    run(args)

reason: import numpy,import scipy


text:
#  Authors: Zhaoshuo Li, Xingtong Liu, Francis X. Creighton, Russell H. Taylor, and Mathias Unberath
#
#  Copyright (c) 2020. Johns Hopkins University - All rights reserved.

import os

import numpy as np
import torch.utils.data as data
from PIL import Image
from natsort import natsorted

from dataset.preprocess import augment, normalization
from dataset.stereo_albumentation import random_crop, horizontal_flip
from utilities.python_pfm import readPFM

class MiddleburyBaseDataset(data.Dataset):
    def __init__(self, datadir, split='train'):
        super(MiddleburyBaseDataset, self).__init__()

        self.datadir = datadir
        self.split = split

        self.left_fname = None
        self.right_fname = None
        self.disp_left_fname = None
        self.disp_right_fname = None
        self.occ_left_fname = None
        self.occ_right_fname = None

        self._augmentation()

    def _read_data(self):
        self.left_data = [os.path.join(obj, self.left_fname) for obj in os.listdir(os.path.join(self.datadir))]
        self.right_data = [os.path.join(obj, self.right_fname) for obj in os.listdir(os.path.join(self.datadir))]
        self.disp_left_data = [os.path.join(obj, self.disp_left_fname) for obj in
                               os.listdir(os.path.join(self.datadir))]
        self.disp_right_data = [os.path.join(obj, self.disp_right_fname) for obj in
                                os.listdir(os.path.join(self.datadir))]
        self.occ_left_data = [os.path.join(obj, self.occ_left_fname) for obj in
                              os.listdir(os.path.join(self.datadir))]
        self.occ_right_data = [os.path.join(obj, self.occ_right_fname) for obj in
                               os.listdir(os.path.join(self.datadir))]
        self.left_data = natsorted(self.left_data)
        self.right_data = natsorted(self.right_data)
        self.disp_left_data = natsorted(self.disp_left_data)
        self.disp_right_data = natsorted(self.disp_right_data)
        self.occ_left_data = natsorted(self.occ_left_data)
        self.occ_right_data = natsorted(self.occ_right_data)

    def _augmentation(self):
        self.transformation = None

    def __len__(self):
        return len(self.left_data)

    def __getitem__(self, idx):
        input_data = {}

        # left
        left_fname = os.path.join(self.datadir, self.left_data[idx])
        left = np.array(Image.open(left_fname)).astype(np.uint8)
        input_data['left'] = left

        # right
        right_fname = os.path.join(self.datadir, self.right_data[idx])
        right = np.array(Image.open(right_fname)).astype(np.uint8)
        input_data['right'] = right

        if not self.split == 'test':  # no disp for test files
            # occ
            occ_left_fname = os.path.join(self.datadir, self.occ_left_data[idx])
            occ_right_fname = os.path.join(self.datadir, self.occ_right_data[idx])
            occ_left = np.array(Image.open(occ_left_fname)) == 128
            occ_right = np.array(Image.open(occ_right_fname)) == 128

            # disp
            disp_left_fname = os.path.join(self.datadir, self.disp_left_data[idx])
            disp_right_fname = os.path.join(self.datadir, self.disp_right_data[idx])

            disp_left, _ = readPFM(disp_left_fname)
            disp_right, _ = readPFM(disp_right_fname)

            if self.split == 'train':
                # horizontal flip
                input_data['left'], input_data['right'], input_data['occ_mask'], input_data['occ_mask_right'], \
                input_data['disp'], \
                input_data['disp_right'] = horizontal_flip(input_data['left'], input_data['right'], occ_left, occ_right,
                                                           disp_left,
                                                           disp_right, self.split)
                # random crop
                input_data = random_crop(360, 640, input_data, self.split)
            else:
                input_data['occ_mask'] = occ_left
                input_data['occ_mask_right'] = occ_right
                input_data['disp'] = disp_left
                input_data['disp_right'] = disp_right
            input_data = augment(input_data, self.transformation)
        else:
            input_data = normalization(**input_data)

        return input_data

class Middlebury2014Dataset(MiddleburyBaseDataset):
    def __init__(self, datadir, split='train'):
        super(Middlebury2014Dataset, self).__init__(datadir, split)

        self.left_fname = 'im0.png'
        self.right_fname = 'im1.png'
        self.disp_left_fname = 'disp0GT.pfm'
        self.disp_right_fname = 'disp1GT.pfm'
        self.occ_left_fname = 'mask0nocc.png'
        self.occ_right_fname = 'mask1nocc.png'

        self._read_data()

max_stars_repo_name: lebionick/stereo-transformer


text:
import numpy as np
import torch
import scipy.io as sio
import h5py
from sklearn.model_selection import StratifiedKFold, ShuffleSplit, GroupShuffleSplit
import importlib

class GraphDataset(torch.utils.data.Dataset):
    ''' Sample graphs and nodes in graph
    '''
    def __init__(self, files):
        self.files = files

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):

        var = ['static_R', 'label', 'subject_id']
        feature, label, subject_id = Eloadmat(self.files[idx], var)
        label[label==-1]=0
        feature = triu2mat(feature)
        return feature, label, subject_id

class fcDataset(torch.utils.data.Dataset):

    def __init__(self, files):
        self.files = files

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        var = ['static_R', 'label', 'subject_id']
        res = Eloadmat(self.files[idx], var)
        return res

def triu2mat(trilfeat):

    N, N2 = trilfeat.shape
    dim = int((1+np.sqrt(1+8*N2))/2)
    index = np.ones((dim,dim))
    index = np.triu(index,1)

    matrix = np.zeros((N,dim,dim),dtype=trilfeat.dtype)
    for i in range(N):
        temp = np.zeros((dim,dim),dtype=trilfeat.dtype)
        temp[index==1]=trilfeat[i,:]
        temp = temp + temp.T
        matrix[i,:,:] = temp
    return matrix

def Eloadmat(datafile, var):
    if isinstance(datafile, list):
        res = [[]] * len(var)
        for i, file in enumerate(datafile):
            tmpres = Eloadmat(file, var)
            for j, tmp in enumerate(tmpres):
                res[j].extend(tmp)

        for i, tmp in enumerate(res):
            res[i] = np.stack(tmp, 0)
        return res

    res = []
    try:
        ## read normal .mat
        data = sio.loadmat(datafile)
        for tmp in var:
            if tmp not in data.keys():
                res.append([])
                continue
            res.append(data[tmp])
    except:
        ## read -v7.3 .mat
        data = h5py.File(datafile, 'r')
        for tmp in var:
            if tmp not in data.keys():
                res.append([])
                continue
            X = np.transpose(data[tmp])
            if X.dtype == np.dtype('uint16'):
                X = bytes(X[:]).decode('utf-16')
            res.append(X)
        data.close()
    return res

def split_kfold(n_splits=10, index=None, label=None, shuffle=False):
    if label is None:
        label = np.ones(len(index))
    skf = StratifiedKFold(n_splits=n_splits, shuffle=shuffle)
    return list(skf.split(index, label))

def split_train_val_test(n_splits=10, index=None, label=None, shuffle=False):
    if label is None:
        label = np.ones(len(index))
    if not isinstance(n_splits, list):
        n_splits = [n_splits] * 2
    skf = StratifiedKFold(n_splits=n_splits[0], shuffle=shuffle)
    skf2 = StratifiedKFold(n_splits=n_splits[1], shuffle=False)
    all_idx = []
    for train_idx, test_idx in skf.split(index, label):
        tmpidx1, tmpidx2 = list(skf2.split(train_idx, np.array(label)[train_idx]))[0]
        all_idx.append((train_idx[tmpidx1], train_idx[tmpidx2], test_idx))

    return all_idx

def split_train_val_test_balance(val_test_num=100, index=None, label=None, shuffle=False):
    if label is None:
        label = np.ones(len(index))
    if not isinstance(val_test_num, list):
        val_test_num = [val_test_num] * 2

    mlb = int(max(label))
    sf = ShuffleSplit(n_splits=1, test_size=val_test_num[0], random_state=7)
    sf2 = ShuffleSplit(n_splits=1, test_size=val_test_num[1], random_state=7)

    train_idxs = []
    val_idxs = []
    test_idxs = []
    for i in range(mlb + 1):
        idx = np.nonzero(label == i)[0]
        tmpsub = index[idx]
        train_idx_tmp, test_idx = list(sf.split(tmpsub))[0]
        tmpidx1, tmpidx2 = list(sf2.split(tmpsub[train_idx_tmp]))[0]
        train_idx = idx[train_idx_tmp[tmpidx1]]
        val_idx = idx[train_idx_tmp[tmpidx2]]
        test_idx = idx[test_idx]
        train_idxs.append(train_idx)
        val_idxs.append(val_idx)
        test_idxs.append(test_idx)
    train_idxs = np.concatenate(train_idxs)
    val_idxs = np.concatenate(val_idxs)
    test_idxs = np.concatenate(test_idxs)
    return (train_idxs, val_idxs, test_idxs)

def apply_splits(files, index, subject=None, uniq_subject=None):
    if not (isinstance(index[0], list) or isinstance(index[0], np.ndarray)):
        index = [index]
    out_files = []
    for idx in index:
        if subject != None:
            sub = np.array(uniq_subject)[idx]
            idx = np.isin(np.array(subject), sub)
        test_files = np.array(files)[idx].tolist()
        out_files.append(test_files)

    return out_files

def file2loader(files, batch_size, shuffle=True, num_workers=0, sampler=None, labels=None):
    if not (isinstance(files[0], list) or isinstance(files[0], np.ndarray)):
        files = [files]
    if not isinstance(batch_size, list):
        batch_size = [batch_size] * len(files)
    if not isinstance(shuffle, list):
        if len(files) == 3:
            shuffle = [shuffle, False, False]
        else:
            shuffle = [shuffle, False]

    dataset_loader = []
    for i, file in enumerate(files):
        if i == 0 and sampler:
            flag = False
            sampler_dic = get_sampler(sampler)
            sampler = sampler_dic['sampler'](labels, **sampler_dic['params'])
        else:
            sampler = None
            flag = True
        dataset_sampler = GraphDataset(file)
        tmp_loader = torch.utils.data.DataLoader(
            dataset_sampler, batch_size=batch_size[i],
            shuffle=shuffle[i] & flag,
            sampler=sampler,
            num_workers=num_workers)
        dataset_loader.append(tmp_loader)

    return dataset_loader

def cross_val_loader(files, subject, uniq_subject, batch_size, ith_fold, label=None, n_splits=10, sampler=None):
    if subject != None:
        index = uniq_subject
    else:
        index = files
    split_index = split_train_val_test(n_splits=n_splits, index=index, label=label)
    split_files = apply_splits(files=files, index=split_index[ith_fold], subject=subject, uniq_subject=uniq_subject)
    if sampler:
        labels = Eloadmat(split_files[0], ['label'])[0]
    else:
        labels = None
    split_loader = file2loader(files=split_files, batch_size=batch_size, shuffle=True, num_workers=8, sampler=sampler,
                               labels=labels)

    return split_loader

def cross_train_test_loader(files, subject, uniq_subject, batch_size, ith_fold, label=None, n_splits=10, sampler=None):
    if subject != None:
        index = uniq_subject
    else:
        index = files
    split_index = split_kfold(n_splits=n_splits, index=index, label=label)
    split_files = apply_splits(files=files, index=split_index[ith_fold], subject=subject, uniq_subject=uniq_subject)
    if sampler:
        labels = Eloadmat(split_files[0], ['label'])[0]
    else:
        labels = None
    split_loader = file2loader(files=split_files, batch_size=batch_size, shuffle=True, num_workers=8, sampler=sampler,
                               labels=labels)

    return split_loader

def train_val_test_loader(files, subject, uniq_subject, batch_size, label=None, val_test_num=100, sampler=None):
    if subject != None:
        index = uniq_subject
    else:
        index = files
    split_index = split_train_val_test_balance(val_test_num=val_test_num, index=index, label=label)
    split_files = apply_splits(files=files, index=split_index, subject=subject, uniq_subject=uniq_subject)
    if sampler:
        labels = Eloadmat(split_files[0], ['label'])[0]
    else:
        labels = None
    split_loader = file2loader(files=split_files, batch_size=batch_size, shuffle=True, num_workers=8, sampler=sampler,
                               labels=labels)

    return split_loader

def source_import(file_path):
    """This function imports python module directly from source code using importlib"""
    spec = importlib.util.spec_from_file_location('', file_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module

def get_sampler(types):
    if types == 'ClassAwareSampler':
        sampler_defs = {'def_file': 'ClassAwareSampler.py', 'num_samples_cls': 4, 'type': 'ClassAwareSampler'}
    else:
        sampler_defs = {'alpha': 1.0, 'cycle': 0, 'decay_gap': 30, 'def_file': 'MixedPrioritizedSampler.py',
                        'epochs': 90, 'fixed_scale': 1, 'lam': 1.0, 'manual_only': True, 'nroot': 2.0, 'ptype': 'score',
                        'rescale': False, 'root_decay': None, 'type': 'MixedPrioritizedSampler'}

    if sampler_defs:
        if sampler_defs['type'] == 'ClassAwareSampler':
            sampler_dic = {
                'sampler': source_import(sampler_defs['def_file']).get_sampler(),
                'params': {'num_samples_cls': sampler_defs['num_samples_cls']}
            }
        elif sampler_defs['type'] in ['MixedPrioritizedSampler',
                                      'ClassPrioritySampler']:
            sampler_dic = {
                'sampler': source_import(sampler_defs['def_file']).get_sampler(),
                'params': {k: v for k, v in sampler_defs.items() \
                           if k not in ['type', 'def_file']}
            }
    else:
        sampler_dic = None
    return sampler_dic


text:
# Copyright (c) 2018, NVIDIA CORPORATION.

import numpy as np
import pytest

from cudf.core import DataFrame, Series
from cudf.tests.utils import assert_eq

@pytest.mark.parametrize("ncats,nelem", [(2, 2), (2, 10), (10, 100)])
def test_factorize(ncats, nelem):
    df = DataFrame()
    np.random.seed(0)

    # initialize data frame
    df["cats"] = arr = np.random.randint(2, size=10, dtype=np.int32)

    uvals, labels = df["cats"].factorize()
    np.testing.assert_array_equal(labels.to_array(), sorted(set(arr)))
    assert isinstance(uvals, Series)
    assert isinstance(labels, Series)

    encoder = dict((labels[idx], idx) for idx in range(len(labels)))
    handcoded = [encoder[v] for v in arr]
    np.testing.assert_array_equal(uvals.to_array(), handcoded)

def test_factorize_index():
    df = DataFrame()
    df["col1"] = ["C", "H", "C", "W", "W", "W", "W", "W", "C", "W"]
    df["col2"] = [
        2992443.0,
        2992447.0,
        2992466.0,
        2992440.0,
        2992441.0,
        2992442.0,
        2992444.0,
        2992445.0,
        2992446.0,
        2992448.0,
    ]

    assert_eq(
        df.col1.factorize()[0].to_array(), df.to_pandas().col1.factorize()[0]
    )
    assert_eq(
        df.col1.factorize()[1].to_pandas().values,
        df.to_pandas().col1.factorize()[1].values,
    )

    df = df.set_index("col2")

    assert_eq(
        df.col1.factorize()[0].to_array(), df.to_pandas().col1.factorize()[0]
    )
    assert_eq(
        df.col1.factorize()[1].to_pandas().values,
        df.to_pandas().col1.factorize()[1].values,
    )

max_stars_repo_name: isVoid/cudf

max_issues_repo_name: isVoid/cudf

max_forks_repo_name: isVoid/cudf


text:
"""
Example 1. Drawing a teapot from multiple viewpoints.
"""
import os
import argparse

import torch
import numpy as np
import tqdm
import imageio

import neural_renderer as nr

current_dir = os.path.dirname(os.path.realpath(__file__))
data_dir = os.path.join(current_dir, 'data')

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('-i', '--filename_input', type=str, default=os.path.join(data_dir, 'teapot.obj'))
  parser.add_argument('-o', '--filename_output', type=str, default=os.path.join(data_dir, 'example1.gif'))
  parser.add_argument('-g', '--gpu', type=int, default=0)
  args = parser.parse_args()

  # other settings
  camera_distance = 2.732
  elevation = 30
  texture_size = 2

  # load .obj
  vertices, faces = nr.load_obj(args.filename_input)
  vertices = vertices[None, :, :]  # [num_vertices, XYZ] -> [batch_size=1, num_vertices, XYZ]
  faces = faces[None, :, :]  # [num_faces, 3] -> [batch_size=1, num_faces, 3]

  # create texture [batch_size=1, num_faces, texture_size, texture_size, texture_size, RGB]
  textures = torch.ones(1, faces.shape[1], texture_size, texture_size, texture_size, 3, dtype=torch.float32).cuda()

  # to gpu

  # create renderer
  renderer = nr.Renderer(camera_mode='look_at')

  # draw object
  loop = tqdm.tqdm(range(0, 360, 4))
  writer = imageio.get_writer(args.filename_output, mode='I')
  for num, azimuth in enumerate(loop):
    loop.set_description('Drawing')
    renderer.eye = nr.get_points_from_angles(camera_distance, elevation, azimuth)
    images, _, _ = renderer(vertices, faces, textures)  # [batch_size, RGB, image_size, image_size]
    image = images.detach().cpu().numpy()[0].transpose((1, 2, 0))  # [image_size, image_size, RGB]
    writer.append_data((255 * image).astype(np.uint8))
  writer.close()

if __name__ == '__main__':
  main()