text:
import sys
import os
import math
import numpy as np
import matplotlib.pyplot as plt
import gc
import gym
from gym import error, spaces
from gym.utils import seeding
sys.path.append(sys.path[0] + '/../..')
sys.path.append(sys.path[0] + '/..')
import matplotlib as mpl
mpl.rcParams['pdf.fonttype'] = 42

try:
    from envs.pyplanner2d import EMExplorer
    from envs.utils import load_config, plot_virtual_map, plot_virtual_map_cov, plot_path
except ImportError as e:
    raise error.DependencyNotInstalled('{}. Build em_exploration and export PYTHONPATH=build_dir'.format(e))

class ExplorationEnv(gym.Env):
    metadata = {'render.modes': ['human', 'state'],
                'render.pause': 0.001}

    def __init__(self,
                 map_size,
                 env_index,
                 test):
        dir_path = os.path.dirname(os.path.realpath(__file__))
        config = dir_path + '/exploration_env.ini'
        self._config = load_config(config)
        self.env_index = env_index
        self.map_size = map_size
        self.dist = 0.0
        self.test = test
        self._one_nearest_frontier = False
        self.nearest_frontier_point = 0
        self.seed()
        self._obs = self.reset()
        self._viewer = None
        self.reg_out = self._sim._planner_params.reg_out
        self._max_steps = self._sim._environment_params.max_steps
        self.max_step = self._max_steps

        num_actions = self._sim._planner_params.num_actions
        self._step_length = self._sim._planner_params.max_edge_length
        self._rotation_set = np.arange(0, np.pi * 2, np.pi * 2 / num_actions) - np.pi
        self._action_set = [np.array([np.cos(t) * self._step_length,
                                      np.sin(t) * self._step_length,
                                      t])
                            for t in self._rotation_set]
        self.action_space = spaces.Discrete(n=num_actions)
        assert (len(self._action_set) == num_actions)
        self._done = False
        self.loop_clo = False
        self._frontier = []
        self._frontier_index = []

        self.map_resolution = self._sim._virtual_map_params.resolution
        rows, cols = self._sim._virtual_map.to_array().shape
        self.leng_i_map = rows
        self.leng_j_map = cols
        self._max_sigma = self._sim._virtual_map.get_parameter().sigma0
        min_x, max_x = self._sim._map_params.min_x, self._sim._map_params.max_x
        min_y, max_y = self._sim._map_params.min_y, self._sim._map_params.max_y
        self._pose = spaces.Box(low=np.array([min_x, min_y, -math.pi]),
                                high=np.array([max_x, max_y, math.pi]), dtype=np.float32)
        self._vm_cov_sigma = spaces.Box(low=0, high=self._max_sigma, dtype=np.float32, shape=(rows, cols))
        self._vm_cov_angle = spaces.Box(low=-math.pi, high=math.pi, dtype=np.float32, shape=(rows, cols))
        self._vm_prob = spaces.Box(low=0.0, high=1.0, shape=(rows, cols), dtype=np.float32)
        self.observation_space = spaces.Tuple([self._pose,
                                               self._vm_prob,
                                               self._vm_cov_sigma,
                                               self._vm_cov_angle])
        self.ext = 20.0

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return seed

    def _get_obs(self):
        cov_array = self._sim._virtual_map.to_cov_array()
        self._obs_show = np.array(
            [self._sim.vehicle_position.x, self._sim.vehicle_position.y, self._sim.vehicle_position.theta]), \
                         self._sim._virtual_map.to_array(), cov_array[0], cov_array[1]
        self._obs = self._sim._virtual_map.to_array()
        return self._obs

    def _get_utility(self, action=None):
        if action is None:
            distance = 0.0
        else:
            angle_weight = self._config.getfloat('Planner', 'angle_weight')
            distance = math.sqrt(action.x ** 2 + action.y ** 2 + angle_weight * action.theta ** 2)
        return self._sim.calculate_utility(distance)

    def step(self, action):
        if self._sim._planner_params.reg_out:
            action = self._action_set[action]
        u1 = self._get_utility()
        self._sim.simulate([action.x, action.y, action.theta])
        self.dist = self.dist + math.sqrt(action.x ** 2 + action.y ** 2)
        u2 = self._get_utility(action)
        return self._get_obs(), self.done(), {}

    def plan(self):
        if not self._sim.plan():
            self._done = True
            return []

        actions = []
        for edge in self._sim._planner.iter_solution():
            if self._sim._planner_params.reg_out:
                actions.insert(0, (np.abs(np.asarray(self._rotation_set) - edge.get_odoms()[0].theta)).argmin())
            else:
                actions.insert(0, edge.get_odoms()[0].theta)
        return actions

    def rrt_plan(self, goal_key):
        if not self._sim.rrt_plan(goal_key, self._frontier):
            self._done = True
            return []

        actions = []
        for edge in self._sim._planner.iter_solution():
            actions.insert(0, edge.get_odoms()[0])
        return actions

    def line_plan(self, goal_key, fro=[0, 0]):
        actions = self._sim.line_plan(goal_key, fro)
        return actions

    def actions_all_goals(self):
        key_size = self._sim._slam.key_size()
        land_size = self.get_landmark_size()
        fro_size = len(self._frontier)
        all_actions = [[]] * (key_size + fro_size)
        # actions for frontiers
        for i, vi in enumerate(self._frontier):
            all_actions[i + key_size] = self.line_plan(key_size, vi)

        return all_actions

    def rewards_all_goals(self, all_actions):
        key_size = self._sim._slam.key_size()
        land_size = self.get_landmark_size()
        fro_size = len(self._frontier)
        rewards = [np.nan] * (key_size + fro_size)

        # calculating rewards for each actions
        for i, _ in enumerate(self._frontier):
            rewards[i + key_size] = self._sim.simulations_reward(all_actions[i + key_size])
        act_max = np.nanargmax(rewards)
        if self.is_nf(act_max):
            self.loop_clo = False
            rewards = np.interp(rewards, (np.nanmin(rewards), np.nanmax(rewards)), (-1.0, 0.0))
        else:
            self.loop_clo = True
            rewards = np.interp(rewards, (np.nanmin(rewards), np.nanmax(rewards)), (-1.0, 1.0))
        rewards[np.isnan(rewards)] = 0
        return rewards

    def status(self):
        return self._sim._virtual_map.explored()

    def done(self):
        return self._done or self._sim.step > self._max_steps or self.status() > 0.85

    def get_landmark_error(self, sigma0=1.0):
        error = 0.0
        for key, predicted in self._sim.map.iter_landmarks():
            landmark = self._sim.environment.get_landmark(key)
            error += np.sqrt((landmark.point.x - predicted.point.x) ** 2 + (landmark.point.y - predicted.point.y) ** 2)
        error += sigma0 * (self._sim.environment.get_landmark_size() - self._sim.map.get_landmark_size())
        return error / self._sim.environment.get_landmark_size()

    def get_dist(self):
        return self.dist

    def get_landmark_size(self):
        return self._sim._slam.map.get_landmark_size()

    def get_key_size(self):
        return self._sim._slam.key_size()

    def print_graph(self):
        self._sim._slam.print_graph()

    def max_uncertainty_of_trajectory(self):
        land_size = self.get_landmark_size()
        self._sim._slam.adjacency_degree_get()
        features = np.array(self._sim._slam.features_out())
        return np.amax(features[land_size:])

    def graph_matrix(self):
        self.frontier()
        trace_map = self._sim._virtual_map.to_cov_trace()
        key_size = self._sim._slam.key_size()
        land_size = self.get_landmark_size()
        fro_size = len(self._frontier)

        self._sim._slam.adjacency_degree_get()
        adjacency = np.array(self._sim._slam.adjacency_out())
        features = np.array(self._sim._slam.features_out())
        adjacency = np.pad(adjacency, ((0, fro_size), (0, fro_size)), 'constant')
        features = np.pad(features, ((0, fro_size), (0, 0)), 'constant')

        robot_location = [self._sim.vehicle_position.x, self._sim.vehicle_position.y]

        # add frontiers to adjacency matrix
        for i in range(fro_size):
            frontier_point = self._frontier[i]
            for j in range(len(self._frontier_index[i])):
                index_node = self._frontier_index[i][j]
                if index_node == 0:
                    self.nearest_frontier_point = i+key_size
                    dist = self.points2dist(frontier_point, robot_location)
                    adjacency[key_size - 1][i + key_size] = dist
                    adjacency[i + key_size][key_size - 1] = dist
                else:
                    dist = self.points2dist(frontier_point, self._sim._slam.get_key_points(index_node - 1))
                    adjacency[index_node - 1][i + key_size] = dist
                    adjacency[i + key_size][index_node - 1] = dist

        # add frontiers to features matrix col 1: trace of cov
        for i in range(fro_size):
            indx = self.coor2index(self._frontier[i][0], self._frontier[i][1])
            f = trace_map[indx[0]][indx[1]]
            features[key_size + i][0] = f

        # add frontiers to features matrix col 2: distance to the robot
        features_2 = np.zeros(np.shape(features))
        for i in range(key_size):
            key_point = self._sim._slam.get_key_points(i)
            dist = self.points2dist(key_point, robot_location)
            features_2[i][0] = dist
        for i in range(fro_size):
            frontier_point = self._frontier[i]
            dist = self.points2dist(frontier_point, robot_location)
            features_2[key_size + i][0] = dist

        # add frontiers to features matrix col 5: direction to the robot
        features_5 = np.zeros(np.shape(features))
        root_theta = self._sim.vehicle_position.theta
        for i in range(key_size):
            key_point = self._sim._slam.get_key_points(i)
            tdiff = self.diff_theta(key_point, robot_location, root_theta)
            features_5[i][0] = tdiff
        for i in range(fro_size):
            frontier_point = self._frontier[i]
            tdiff = self.diff_theta(frontier_point, robot_location, root_theta)
            features_5[key_size + i][0] = tdiff

        # add frontiers to features matrix col 3: probability
        features_3 = np.zeros(np.shape(features))
        for i in range(key_size):
            key_point = self._sim._slam.get_key_points(i)
            indx = self.coor2index(key_point[0], key_point[1])
            probobility = self._obs[indx[0]][indx[1]]
            features_3[i][0] = probobility
        for i in range(fro_size):
            frontier_point = self._frontier[i]
            indx = self.coor2index(frontier_point[0], frontier_point[1])
            probobility = self._obs[indx[0]][indx[1]]
            features_3[key_size + i][0] = probobility

        # add frontiers to features matrix clo 4: index of locations
        features_4 = np.zeros(np.shape(features))
        for i in range(key_size - 1):
            features_4[i][0] = -1
        features_4[key_size - 1][0] = 0
        for i in range(fro_size):
            features_4[key_size + i][0] = 1

        features = np.concatenate((features, features_2, features_5, features_3, features_4), axis=1)

        # create global features
        avg_landmarks_error = np.mean(features[1:land_size + 1][:, 0])
        global_features = np.array([avg_landmarks_error])
        return adjacency, features, global_features, fro_size

    def is_nf(self, id):
        if self.nearest_frontier_point == id:
            return True
        else:
            return False

    def frontier(self):
        vehicle_location = [self._sim.vehicle_position.x, self._sim.vehicle_position.y]

        a = self._obs < 0.45
        free_index_i, free_index_j = np.nonzero(a)

        all_landmarks = []
        all_frontiers = []
        landmark_keys = range(self.get_landmark_size())

        self._frontier = []
        self._frontier_index = []

        for land_key in landmark_keys:
            points = list(self._sim._slam.get_key_points(land_key))
            all_landmarks.append(points)

        for ptr in range(len(free_index_i)):
            cur_i = free_index_i[ptr]
            cur_j = free_index_j[ptr]
            count = 0
            cur_i_min = free_index_i[ptr] - 1 if free_index_i[ptr] - 1 >= 0 else 0
            cur_i_max = free_index_i[ptr] + 1 if free_index_i[ptr] + 1 < self.leng_i_map else self.leng_i_map - 1
            cur_j_min = free_index_j[ptr] - 1 if free_index_j[ptr] - 1 >= 0 else 0
            cur_j_max = free_index_j[ptr] + 1 if free_index_j[ptr] + 1 < self.leng_j_map else self.leng_j_map - 1

            for ne_i in range(cur_i_min, cur_i_max + 1):
                for ne_j in range(cur_j_min, cur_j_max + 1):
                    if 0.49 < self._obs[ne_i][ne_j] < 0.51:
                        count += 1

            if count >= 2:
                ind2co = self.index2coor(cur_i, cur_j)
                if self._sim._map_params.min_x + self.ext <= ind2co[0] <= self._sim._map_params.max_x - self.ext \
                        and self._sim._map_params.min_y + self.ext <= ind2co[
                    1] <= self._sim._map_params.max_y - self.ext:
                    all_frontiers.append(ind2co)

        cur_fro = all_frontiers[self.nearest_frontier(vehicle_location, all_frontiers)]
        self._frontier.append(cur_fro)
        self._frontier_index.append([0])

        if not self._one_nearest_frontier:
            for ip, p in enumerate(all_landmarks):
                cur_fro = all_frontiers[self.nearest_frontier(p, all_frontiers)]
                try:
                    self._frontier_index[self._frontier.index(cur_fro)].append(ip + 1)
                except ValueError:
                    self._frontier.append(cur_fro)
                    self._frontier_index.append([ip + 1])

        not_go = []
        for i, vi in enumerate(self._frontier):
            temp_list = []
            for j, vj in enumerate(self._frontier[i:]):
                if vi == vj and i not in not_go:
                    temp_list.append(i + j)
                    if i + j != i:
                        not_go.append(i + j)
            self._frontier_index.append(temp_list)

    def nearest_frontier(self, point, all_frontiers):
        min_dist = float("Inf")
        min_index = None
        for index, fro_points in enumerate(all_frontiers):
            dist = self.points2dist(point, fro_points)
            if dist < min_dist:
                min_dist = dist
                min_index = index
        return min_index

    def show_frontier(self, index):
        plt.plot(np.transpose(self._frontier)[0], np.transpose(self._frontier)[1], 'mo')
        plt.plot(np.transpose(self._frontier)[0][index], np.transpose(self._frontier)[1][index], 'ro')

    def index2coor(self, matrix_i, matrix_j):
        x = (matrix_j + 0.5) * self.map_resolution + self._sim._map_params.min_x
        y = (matrix_i + 0.5) * self.map_resolution + self._sim._map_params.min_y
        return [x, y]

    def coor2index(self, x, y):
        map_j = int(round((x - self._sim._map_params.min_x) / self.map_resolution - 0.5))
        map_i = int(round((y - self._sim._map_params.min_y) / self.map_resolution - 0.5))
        return [map_i, map_j]

    def points2dist(self, point1, point2):
        dist = np.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2)
        return dist

    def diff_theta(self, point1, point2, root_theta):
        goal_theta = math.atan2(point1[1] - point2[1], point1[0] - point2[0])
        if goal_theta < 0:
            goal_theta = math.pi * 2 + goal_theta
        if root_theta < 0:
            root_theta = math.pi * 2 + root_theta
        diff = goal_theta - root_theta
        if diff < 0:
            diff = math.pi * 2 + diff
        return diff

    def reset(self):
        self._done = False
        while True:
            # Reset seed in configuration
            if not self.test:
                seed1 = self.np_random.randint(0, np.iinfo(np.int32).max, dtype=np.int32)
                seed2 = self.np_random.randint(0, np.iinfo(np.int32).max, dtype=np.int32)
            else:
                seed1 = self.env_index
                seed2 = self.env_index
            landmark_size = int(self.map_size ** 2 * 0.005)
            self._config.set('Simulator', 'seed', str(seed1))
            self._config.set('Planner', 'seed', str(seed1))
            self._config.set('Simulator', 'lo', str(seed2))
            self._config.set('Environment', 'min_x', str(-self.map_size / 2))
            self._config.set('Environment', 'min_y', str(-self.map_size / 2))
            self._config.set('Environment', 'max_x', str(self.map_size / 2))
            self._config.set('Environment', 'max_y', str(self.map_size / 2))
            self._config.set('Simulator', 'num', str(landmark_size))

            # Initialize new instance and perfrom a 360 degree scan of the surrounding
            self._sim = EMExplorer(self._config)
            for step in range(4):
                odom = 1, 1, math.pi / 2.0
                u1 = self._get_utility()
                self._sim.simulate(odom)

            if self._sim._slam.map.get_landmark_size() < 1:
                print("regenerate a environment")
                self.env_index = self.env_index + 50
                continue

            # Return initial observation
            return self._get_obs()

    def render(self, mode='human', close=False, action_index=-1):
        if close:
            return
        if mode == 'human':
            if self._viewer is None:
                self._sim.plot()
                self._viewer = plt.gcf()
                plt.ion()
                plt.tight_layout()
                plt.xlim((self._sim._map_params.min_x + 14, self._sim._map_params.max_x - 14))
                plt.ylim((self._sim._map_params.min_y + 14, self._sim._map_params.max_y - 14))
                plt.show()
            else:
                self._viewer.clf()
                self._sim.plot()
                # plot_path(self._sim._planner, dubins=False)
                plt.xlim((self._sim._map_params.min_x + 14, self._sim._map_params.max_x - 14))
                plt.ylim((self._sim._map_params.min_y + 14, self._sim._map_params.max_y - 14))
                plt.draw()
            if action_index != -1:
                self.show_frontier(action_index)
            plt.pause(ExplorationEnv.metadata['render.pause'])
        elif mode == 'state':
            # print self._obs
            # assert(len(self._obs_show) == 3)
            print (self._viewer is None)
            if self._viewer is None:
                self._viewer = plt.subplots(1, 3, figsize=(18, 6))
                plot_virtual_map(self._sim._virtual_map, self._sim._map_params,
                                 alpha=1.0, ax=self._viewer[1][0])
                plot_virtual_map_cov(self._obs_show[1:], self._max_sigma,
                                     self._sim._map_params,
                                     alpha=1.0, ax=self._viewer[1][1])
                plt.sca(self._viewer[1][2])
                self._sim.plot()
                plt.ion()
                plt.tight_layout()
                plt.show()
            else:
                self._viewer[1][0].clear()
                plot_virtual_map(self._sim._virtual_map, self._sim._map_params,
                                 alpha=1.0, ax=self._viewer[1][0])
                self._viewer[1][1].clear()
                plot_virtual_map_cov(self._obs_show[1:], self._max_sigma,
                                     self._sim._map_params,
                                     alpha=1.0, ax=self._viewer[1][1])
                self._viewer[1][2].clear()
                plt.sca(self._viewer[1][2])
                self._sim.plot()
                plt.draw()
            plt.pause(ExplorationEnv.metadata['render.pause'])
        else:
            super(ExplorationEnv, self).render(mode=mode)

if __name__ == '__main__':
    import sys

    ExplorationEnv.metadata['render.pause'] = 0.001
    lo_num = 7
    map_size = 40
    total_reward = np.empty([0, 0])
    TEST = False

    mode = 'human'
    env = ExplorationEnv(map_size, lo_num, TEST)
    t = 0
    epoch = 0
    done = False
    flag = False
    actions = []
    env.render(mode=mode)

    for i in range(1000):
        adjacency, featrues, global_features, fro_size = env.graph_matrix()

        key_size = env.get_key_size()
        land_size = env.get_landmark_size()
        all_actions = env.actions_all_goals()
        rewards = env.rewards_all_goals(all_actions)
        # print "rewards: ", rewards, "\n"
        rewards[0:key_size] = np.nan
        act_index = np.nanargmax(rewards)
        max_reward = rewards[act_index]

        actions = all_actions[act_index]
        print("###############################")
        temp_reward = 0
        print("max_reward: ", str(max_reward))

        for a in actions:
            obs, reward, done, _ = env.step(a)
            temp_reward += reward
            env.render(mode=mode, action_index=act_index-key_size)
            print("step: ", str(t), "reward: ", str(reward))

            # print "landmark error: ", env.get_landmark_error()
            t = t + 1
            if done:
                break

        ls = env.get_landmark_size()
        epoch += 1
        adjacency, featrues, global_features, fro_size = env.graph_matrix()
        print('done: ', done, 'explored: ', env.status())

        if done:
            print("done")
            print("total steps: ", t)
            print("epoch is: ", epoch)
            print("error: ", env.get_landmark_error())
            input("Press Enter to continue...")
            del env
            gc.collect()
            env = ExplorationEnv(map_size, lo_num, TEST)
            env.render(mode=mode)
            t = 0
            print("new one")
        flag = False
    plt.pause(1e10)


text:
# Data ETL

import pandas as pd
import numpy as np

def load_data():

    # for reproducibility
    np.random.seed(42)

    id_col = "unit_number"
    time_col = "time"
    feature_cols = ["op_setting_1", "op_setting_2", "op_setting_3"] + [
        "sensor_measurement_{}".format(x) for x in range(1, 22)
    ]
    column_names = [id_col, time_col] + feature_cols

    train_x_orig = pd.read_csv(
        "../DataSets/CMAPSS/train_FD001.csv", header=None, sep="\s+", decimal="."
    )
    train_x_orig.columns = column_names

    test_x_orig = pd.read_csv(
        "../DataSets/CMAPSS/test_FD001.csv", header=None, sep="\s+", decimal="."
    )
    test_x_orig.columns = column_names

    test_y_orig = pd.read_csv(
        "../DataSets/CMAPSS/RUL_FD001.csv", header=None, names=["T"]
    )

    # Make engine numbers and days zero-indexed
    train_x_orig.iloc[:, 0:2] -= 1
    test_x_orig.iloc[:, 0:2] -= 1

    print(train_x_orig.shape)
    print(test_x_orig.shape)

    train_idx = np.random.choice(
        range(train_x_orig.unit_number.unique().max() + 1), replace=False, size=80
    )  # selecting 80% units for training
    train_idx.sort()

    vld_idx = np.array(
        [
            x
            for x in range(train_x_orig.unit_number.unique().max() + 1)
            if x not in train_idx
        ]
    )  # remaining are validation indices

    train = train_x_orig[train_x_orig.unit_number.isin(train_idx)]  # training data
    train.reset_index(drop=True, inplace=True)
    vld = train_x_orig[train_x_orig.unit_number.isin(vld_idx)]  # validation data
    vld.reset_index(drop=True, inplace=True)

    # Truncating the validation data randomly 5 times each
    vld_trunc = []
    test_y = []
    max_cycle = []
    test_index = []  # for debugging purposes
    temp_or_test_cycles = []
    counter = -1

    for i in set(vld.unit_number.unique()):
        # print(f'unit number is {i}')
        for j in range(1, 6):  # 5 truncations per instance
            counter += 1
            np.random.seed(i * j)
            temp_df = vld[vld.unit_number == i]
            temp_df.reset_index(drop=True, inplace=True)  # important
            length = temp_df.shape[0]
            temp_or_test_cycles.append(length)
            level = np.random.choice(np.arange(5, 96), 1)[0]
            r = np.int(length * (1 - level / 100))
            temp_df = temp_df.truncate(after=r)
            temp_df["unit_number"] = np.repeat(counter, temp_df.shape[0])
            vld_trunc.append(temp_df)
            max_cycle.append(length)

    vld_trunc = pd.concat(vld_trunc)
    vld_trunc.reset_index(drop=True, inplace=True)

    return (
        train,
        feature_cols,
        vld_trunc,
        vld,
        np.array(max_cycle),
        test_x_orig,
        test_y_orig,
    )

ext: py

max_stars_repo_name: MariosKef/RULe

max_issues_repo_name: MariosKef/RULe

max_forks_repo_name: MariosKef/RULe

reason: import numpy


text:
import argparse
import torch
from torchvision.utils import make_grid

import railrl.torch.pytorch_util as ptu
import pickle
import uuid
import cv2
from PIL import Image
import numpy as np
import os.path as op

filename = str(uuid.uuid4())

def simulate_policy(args):
    ptu.set_gpu_mode(True)
    model = pickle.load(open(args.file, "rb")) # joblib.load(args.file)
    model.to(ptu.device)
    imgs = np.load(args.imgfile)
    import ipdb; ipdb.set_trace()
    z = model.encode(ptu.np_to_var(imgs))
    samples = model.decode(z).cpu()

    recon_imgs = samples.data.view(64, model.input_channels, model.imsize,
                              model.imsize)
    recon_imgs = recon_imgs.cpu()
    grid = make_grid(recon_imgs, nrow=8)
    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()
    im = Image.fromarray(ndarr)
    im.show()
    # cv2.imshow('img', im)
    # cv2.waitKey(1)
    # for sample in samples:
    #     tensor = tensor.cpu()
    #     img = ptu.get_numpy(tensor)
    comparison = torch.cat([
        recon_imgs,
        imgs,
    ])
    save_dir = osp.join(logger.get_snapshot_dir(), 'r%d.png' % epoch)
    save_image(comparison.data.cpu(), save_dir, nrow=n)

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('file', type=str,
                        help='path to the snapshot file')
    parser.add_argument('--imgfile', type=str,
                        help='path to the img (npy) file')
    parser.add_argument('--H', type=int, default=300,
                        help='Max length of rollout')
    parser.add_argument('--speedup', type=float, default=10,
                        help='Speedup')
    parser.add_argument('--gpu', action='store_true')
    parser.add_argument('--pause', action='store_true')
    parser.add_argument('--hide', action='store_true')
    parser.add_argument('--enable_render', action='store_true')
    parser.add_argument('--multitaskpause', action='store_true')
    args = parser.parse_args()

    simulate_policy(args)


text:
"""
gimbal/mcmc.py
"""

import contextlib
import numpy as onp
from tqdm.auto import trange

import jax.config
import jax.numpy as jnp
import jax.random as jr
from jax import lax, jit, partial, vmap
from jax.scipy.special import logsumexp

import tensorflow_probability.substrates.jax as tfp
import tensorflow_probability.substrates.jax.distributions as tfd

from ssm.messages import hmm_sample

from .util import (triangulate, project,
                  xyz_to_uv, uv_to_xyz, signed_angular_difference,
                  Rxy_mat, cartesian_to_polar,
                  children_of, tree_graph_laplacian, hvmfg_natural_parameter,)
from . import util_io

def initialize_parameters(params,
                          pos_location_0=0., pos_variance_0=1e8,
                          state_transition_count=10.,
                          regularizer=1e-6):
    num_keypoints = len(params['parents'])
    num_cameras =  len(params['camera_matrices'])
    dim = params['camera_matrices'].shape[-1] - 1
    dim_obs = params['camera_matrices'].shape[-2] - 1
    num_states = len(params['state_probability'])

    # Must explicitly cast arrays (that are not results of calculations) to
    # desired dtype. For example, if we are in x64 mode and we have
    #   arr = jnp.array([1,2,3], dtype=jnp.float32)
    # the following operations preserve original dtype (undesirable)
    #   jnp.broadcast_to(arr, arr.shape) -> jnp.float32
    #   jnp.asarray(arr) -> jnp.float32
    #   jnp.array(arr) -> jnp.float32
    # The following operation changes dtype
    #   arr.asdtype(jnp.float64) -> jnp.float64
    dtype = jnp.float64 if jax.config.read('jax_enable_x64') else jnp.float32

    params['children'] = children_of(params['parents'])

    # -----------------------------------
    # Canonical reference frame
    # -----------------------------------
    x_axis, z_axis = params['crf_abscissa'], params['crf_normal']
    y_axis = jnp.cross(z_axis, x_axis)
    params['crf_axes'] = jnp.stack([x_axis, y_axis, z_axis], axis=0)

    # -----------------------------------
    # Skeletal and positional parameters
    # -----------------------------------
    params['pos_radius'] \
        = jnp.broadcast_to(params['pos_radius'], (num_keypoints,))\
             .astype(dtype)

    params['pos_radial_variance'] \
        = jnp.broadcast_to(params['pos_radial_variance'], (num_keypoints,))\
             .astype(dtype)
    params['pos_radial_precision'] \
        = tree_graph_laplacian(
            params['parents'],
            1 / params['pos_radial_variance'][...,None,None] * jnp.eye(dim)
            ) + regularizer * jnp.eye(num_keypoints*dim)
    params['pos_radial_covariance'] \
        = jnp.linalg.inv(params['pos_radial_precision'])

    # NB: Using jnp.linalg.eigvals throws the NotImplementedError with message "Nonsymmetric eigendecomposition is only implemented on the CPU backend. However, these matrices ARE symmetric...
    assert jnp.all(onp.linalg.eigvals(params['pos_radial_precision']) > 0.), \
        f"Expected positive definite radial precision matrix, but got eigenvalues of {jnp.linalg.eigvals(params['pos_radial_precision'])}.\nConsider adjusting value of regularizer."
    assert jnp.all(onp.linalg.eigvals(params['pos_radial_covariance']) > 0.), \
        f"Expected positive definite radial covariance matrix, but got eigenvalues of {jnp.linalg.eigvals(params['pos_radial_covariance'])}.\nConsider adjusting value of regularizer."

    params['pos_dt_variance'] \
        = jnp.broadcast_to(params['pos_dt_variance'], (num_keypoints,))\
             .astype(dtype)
    params['pos_dt_covariance'] \
        = jnp.kron(jnp.diag(params['pos_dt_variance']), jnp.eye(dim))
    params['pos_dt_precision'] \
        = jnp.kron(jnp.diag(1./params['pos_dt_variance']), jnp.eye(dim))

    params['pos_precision_t'] \
        = params['pos_radial_precision'] + params['pos_dt_precision']
    params['pos_covariance_t'] \
        = jnp.linalg.inv(params['pos_precision_t'])

    params['pos_variance_0'] \
        = jnp.broadcast_to(params.get('pos_variance_0', pos_variance_0),
                           (num_keypoints,))\
             .astype(dtype)
    params['pos_covariance_0'] \
        = jnp.kron(jnp.diag(params['pos_variance_0']), jnp.eye(dim))

    params['pos_location_0'] \
        = jnp.broadcast_to(params.get('pos_location_0', pos_location_0),
                           (num_keypoints, dim))\
             .astype(dtype)

    # -----------------------------
    # Observation error parameters
    # -----------------------------
    params['obs_outlier_probability'] \
        = jnp.broadcast_to(params['obs_outlier_probability'],
                           (num_cameras, num_keypoints))\
             .astype(dtype)

    params['obs_outlier_location'] \
        = jnp.broadcast_to(params['obs_outlier_location'],
                           (num_cameras, num_keypoints, dim_obs))\
             .astype(dtype)
    params['obs_outlier_variance'] \
        = jnp.broadcast_to(params['obs_outlier_variance'],
                          (num_cameras, num_keypoints))\
             .astype(dtype)
    params['obs_outlier_covariance'] \
        = jnp.kron(params['obs_outlier_variance'][..., None, None],
                   jnp.eye(dim_obs))

    params['obs_inlier_location'] \
        = jnp.broadcast_to(params['obs_inlier_location'],
                          (num_cameras, num_keypoints, dim_obs))\
             .astype(dtype)
    params['obs_inlier_variance'] \
        = jnp.broadcast_to(params['obs_inlier_variance'],
                          (num_cameras, num_keypoints))\
             .astype(dtype)
    params['obs_inlier_covariance'] \
        = jnp.kron(params['obs_inlier_variance'][..., None, None],
                   jnp.eye(dim_obs))

    # -----------------------------
    # Pose state parameters
    # -----------------------------
    params['state_transition_count'] \
        = jnp.broadcast_to(params.get('state_transition_count', state_transition_count),
                          (num_states,))\
             .astype(dtype)

    params['state_probability'] \
        = jnp.broadcast_to(params.get('state_probability', 1./num_states),
                           (num_states,))\
             .astype(dtype)

    params['state_directions'] \
        = jnp.broadcast_to(params.get('state_directions', jnp.array([0,0,1.])),
                           (num_states, num_keypoints, dim))\
             .astype(dtype)

    params['state_concentrations'] \
        = jnp.broadcast_to(params.get('state_concentrations', 0.),
                           (num_states, num_keypoints))\
             .astype(dtype)

    return params

@jit
def log_joint_probability(params, observations,
                          outliers, positions, directions,
                          heading, pose_state, transition_matrix):
    """Compute the log joint probabiltiy of sampled values under the
    prior parameters.

    Parameters
    ----------
        params: dict
        observations: ndarray, shape (N, C, K, D_obs).
        outliers: ndarray, shape (N, C, K).
        positions: ndarray, shape (N, K, D).
        directions: ndarray, shape (N, K, D).
        heading: ndarray, shape (N,).
        pose_state: ndarray, shape (N,).
        transition_matrix: ndarray, shape (S,S).
    """

    C, S = observations.shape[1], transition_matrix.shape[0]

    Z_prior = tfd.Bernoulli(probs=params['obs_outlier_probability'])

    # Y_ins = tfd.MultivariateNormalFullCovariance(params['obs_inlier_location'], params['obs_inlier_covariance'])  # batch shape (num_cams, num_joints), event_shape (dim_obs)
    # Y_outs = tfd.MultivariateNormalFullCovariance(params['obs_outlier_location'], params['obs_outlier_covariance'])
    Y_ins = [tfd.MultivariateNormalFullCovariance(
                        params['obs_inlier_location'][c],
                        params['obs_inlier_covariance'][c]) \
            for c in range(C)]
    Y_outs = [tfd.MultivariateNormalFullCovariance(
                        params['obs_outlier_location'][c],
                        params['obs_outlier_covariance'][c]) \
            for c in range(C)]
    X_t0 = tfd.MultivariateNormalFullCovariance(
                params['pos_location_0'].ravel(),
                params['pos_covariance_0'])
    # U_given_S = VMF(params['pstate']['mus'], params['pstate']['kappas'])  # batch shape (num_states, num_joints), event_shape (dim)

    # =================================================

    def log_likelihood(xt, yts, zts, params):
        """ Compute log likelihood of observations for a single time step.

        log p(y[t] | x[t], z[t]) =
                z[t] log N(y[t] | proj(x[t]; P), omega_in)
                + (1-z[t]) log N(y[t] | proj(x[t]; P), omega_out)
        """

        # Does not work :( ... Error message
        # Shapes must be 1D sequences of concrete values of integer type, got [Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/2)>].
        # If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.
        # obs_err = np.stack([yts[c] - project(xt, params["Ps"][c]) \
        #                     for c in range(yts.shape[0])], axis=0)
        # lp = (1-zts) * Y_ins.log_prob(obs_err) + zts * Y_outs.log_prob(obs_err)

        lp = 0
        # for c in range(yts.shape[0]):
        for c in range(C):
            obs_err = yts[c] - project(params["camera_matrices"][c], xt)
            lp += (1-zts[c]) * Y_ins[c].log_prob(obs_err)
            lp += zts[c] * Y_outs[c].log_prob(obs_err)
            # Y_in = tfd.MultivariateNormalFullCovariance(params['obs_inlier_location'][c], params['obs_inlier_covariance'][c])
            # Y_out = tfd.MultivariateNormalFullCovariance(params['obs_outlier_location'][c], params['obs_outlier_covariance'][c])
            # lp += (1-zts[c]) * Y_in.log_prob(obs_err)
            # lp += zts[c] * Y_out.log_prob(obs_err)
        return jnp.sum(lp)

    def log_pos_given_dir(xt, ut, params):
        """ Compute log likelihood of positions given directions for a
        single time step.

        log p(x[t] | u[t]) = prod_j N(x[t,parj] + r[j] u[t,j], sigma^2 I)
        """
        ht = hvmfg_natural_parameter(
                    params['children'], params['pos_radius'],
                    params['pos_radial_variance'],
                    ut)
        Xt_given_Ut = tfd.MultivariateNormalFullCovariance(
                            params['pos_radial_covariance'] @ ht.ravel(),
                            params['pos_radial_covariance'])
        return Xt_given_Ut.log_prob(xt.ravel())

    def log_pos_dynamics(xtp1, xt, params):
        """ Compute log p(x[t+1] | x[t]) = log N(x[t], sigma^2 I) """
        Xtp1_given_Xt = tfd.MultivariateNormalFullCovariance(
                                xt.ravel(), params['pos_dt_covariance'])
        return Xtp1_given_Xt.log_prob(xtp1.ravel())

    def log_dir_given_state(ut, ht, st, params):
        """ Compute log p(R(-h[t]) u[t] | h[t], s[t]) = log vMF(nu_s[t], kappa_s[t]) """
        canon_ut = ut @ Rxy_mat(-ht).T
        u_given_st = tfd.VonMisesFisher(
                        params['state_directions'][st],
                        params['state_concentrations'][st])
        return jnp.sum(u_given_st.log_prob(canon_ut))

    # =================================================

    # p(y | x, z)
    lp = jnp.sum(vmap(log_likelihood, in_axes=(0,0,0,None))
                     (positions, observations, outliers, params))

    # p(z; rho)
    lp += jnp.sum(Z_prior.log_prob(outliers))

    # p(x | u)
    lp += jnp.sum(vmap(log_pos_given_dir, in_axes=(0,0,None))
                      (positions, directions, params))

    # p(x[t] | x[t-1])
    lp += jnp.sum(vmap(log_pos_dynamics, in_axes=(0,0,None))
                      (positions[1:], positions[:-1], params))
    lp += X_t0.log_prob(positions[0].ravel())

    # p(u | s)
    # canon_u = jnp.einsum('tmn, tjn-> tjm', R_mat(-h), u)
    # lp = jnp.sum(U_given_S[s].log_prob(canon_u))
    lp += jnp.sum(vmap(log_dir_given_state, in_axes=(0,0,0,None))
                      (directions, heading, pose_state, params))

    # p(h) = VonMises(0, 0), so lp = max entropy over circle, which is constant

    # p(s[0]) + sum p(s[t+1] | s[t])
    lp += jnp.sum(jnp.log(transition_matrix[pose_state[1:], pose_state[:-1]]))
    lp += jnp.sum(jnp.log(params['state_probability']))

    return lp

def initialize(seed, params, observations, init_positions=None):
    """Initialize latent variables of model.

    Parameters
    ----------
        seed: jax.random.PRNGKey
        params: dict
        observations: ndarray, shape (N, C, K, D_obs)
        init_positions: ndarray, shape (N, K, D), optional.
            Initial guess of 3D positions. If None (default), positions
            are triangulated using direct linear triangulation.

    Returns
    -------
        samples: dict
    """

    seed = iter(jr.split(seed, 3))

    N, C, K, D_obs = observations.shape

    # ---------------------
    # Initialize positions
    # ---------------------
    if init_positions is None:
        if D_obs == 2:
            obs = jnp.moveaxis(observations, 1, 0)
            positions = triangulate(params['camera_matrices'], obs)
        elif D_obs == 3:
            positions = observations.copy()
        else:
            raise ValueError(f"Expected observation data dimension {{2,3}}, but receieved {D_obs}.")
    else:
        positions = jnp.asarray(init_positions)

    # Initialize HMC results
    hmc_log_accept_ratio = jnp.array(0.)
    hmc_proposed_gradients = jnp.empty_like(positions)

    # --------------------------
    # Derive initial directions
    # --------------------------
    directions = positions - positions[:,params['parents']]
    directions /= jnp.linalg.norm(directions, axis=-1, keepdims=True)

    # Use placeholder value for undefined root direction
    directions = directions.at[:,0].set(0.)

    # -------------------------------
    # Derive initial heading vectors
    # -------------------------------
    k_base, k_tip = params['crf_keypoints']
    u_axis, v_axis, n_axis = params['crf_axes']

    # Unnormalized direction vector
    heading = positions[:,k_tip] - positions[:,k_base]

    # Project 3D direction vectors onto 2D plane of coordinate reference
    # frame (CRF). Then, re-embed the 2D plane in the ambient 3D frame.
    # In the default case, the CRF is defined as the standard Cartesian
    # reference (i.e. poses aligned to "x"-axis, and headings considered
    # in the "x-y" plane, rotating about normal "z"-axis.) The following
    # two operations then effectively zero-out the z-axis coordinates.
    heading = xyz_to_uv(heading, u_axis, v_axis)
    heading = uv_to_xyz(heading, u_axis, v_axis)

    # Normalize vector to unit length
    heading /= jnp.linalg.norm(heading, axis=-1, keepdims=True)

    # Calculate angle from canonical reference direction
    heading = signed_angular_difference(
                jnp.broadcast_to(u_axis, heading.shape), heading, n_axis)

    # ---------------------------
    # Sample outliers from prior
    # ---------------------------
    outliers = jr.uniform(next(seed), (N,C,K)) < params['obs_outlier_probability']

    # Consider any NaN observations as outliers
    outliers = jnp.where(jnp.isnan(observations).any(axis=-1),
                         True, outliers)

    # ------------------------------------------------------
    # Sample pose states and transition matrix from uniform
    # ------------------------------------------------------
    # TODO sample from prior parameters, state_probability and state_counts
    num_states = len(params['state_probability'])
    pose_state = jr.randint(next(seed), (N,), 0, num_states)
    transition_matrix = jnp.ones((num_states, num_states)) / num_states

    log_prob = log_joint_probability(
                    params, observations,
                    outliers, positions, directions,
                    heading, pose_state, transition_matrix,
                    )

    return dict(
        outliers=outliers,
        positions=positions,
        directions=directions,
        heading=heading,
        pose_state=pose_state,
        transition_matrix=transition_matrix,
        log_probability=log_prob,
        hmc_log_accept_ratio=hmc_log_accept_ratio,
        hmc_proposed_gradients=hmc_proposed_gradients,
    )

@jit
def sample_positions(seed, params, observations, samples,
                     step_size=1e-1, num_leapfrog_steps=1):
    """Sample positions by taking one Hamiltonian Monte Carlo step."""

    N, C, K, D_obs = observations.shape

    def objective(positions):
        return log_joint_probability(
            params,
            observations,
            samples['outliers'],
            positions,
            samples['directions'],
            samples['heading'],
            samples['pose_state'],
            samples['transition_matrix'],
            )

    last_positions = samples['positions']   # shape (N, K, D)

    hmc = tfp.mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=objective,
                num_leapfrog_steps=num_leapfrog_steps,
                step_size=step_size
            )

    positions, kernel_results = hmc.one_step(
                last_positions,
                hmc.bootstrap_results(last_positions),
                seed=seed)

    return positions, kernel_results

@jit
def sample_outliers(seed, params, observations, samples):
    """Sample outliers

    TODO define inlier/outlier distributions beforehand. These are static.
    """

    predicted_observations = jnp.stack(
        [project(P, samples['positions']) for P in params['camera_matrices']],
        axis=1)
    error = observations - predicted_observations

    # Log probability of predicted observation being inlier
    Y_inlier = tfd.MultivariateNormalFullCovariance(
                                    params['obs_inlier_location'],
                                    params['obs_inlier_covariance']
                                    )
    lp_k0 = jnp.log(1-params['obs_outlier_probability'])
    lp_k0 += Y_inlier.log_prob(error)

    # Log probability of predicted observation being outlier
    Y_outlier = tfd.MultivariateNormalFullCovariance(
                                    params['obs_outlier_location'],
                                    params['obs_outlier_covariance']
                                    )
    lp_k1 = jnp.log(params['obs_outlier_probability'])
    lp_k1 += Y_outlier.log_prob(error)

    # Update posterior
    lognorm = logsumexp(jnp.stack([lp_k0, lp_k1], axis=-1), axis=-1)
    p_isoutlier = jnp.exp(lp_k1 - lognorm)

    # Draw new samples
    outliers = jr.uniform(seed, observations.shape[:-1]) < p_isoutlier

    # Any NaN observations are obviously drawn from outlier distribution
    outliers = jnp.where(jnp.isnan(observations[...,0]), True, outliers)

    return outliers

@jit
def sample_directions(seed, params, samples):
    """Sample directions from 3D von Mises-Fisher distribution."""

    positions = samples['positions']    # (N,K,3)
    heading = samples['heading']        # (N,)
    state = samples['pose_state']       # (N,)
    D = positions.shape[-1]

    # Rotate mean directional prior (defined in the canonical reference
    # frame) into the absolute/ambient reference frame
    mu_priors = jnp.einsum('tmn, tjn-> tjm',
                           Rxy_mat(heading), params['state_directions'][state])
    kappa_priors = params['state_concentrations'][state]

    # Keypoint position contributions to direction vectors
    dx = positions - positions[:, params['parents']]

    # Calculate posterior parameters
    mu_tilde  = kappa_priors[...,None] * mu_priors
    mu_tilde += dx * (params['pos_radius']/params['pos_radial_variance'])[:,None]

    mu_post = mu_tilde / jnp.linalg.norm(mu_tilde, axis=-1, keepdims=True)
    kappa_post = jnp.linalg.norm(mu_tilde, axis=-1)

    # Sample from posterior
    directions = tfd.VonMisesFisher(mu_post, kappa_post).sample(seed=seed)
    directions = directions.at[:,0,:].set(jnp.zeros(D))

    return directions

@jit
def sample_headings(seed, params, samples):
    """Sample headings from uniform circular distribution."""

    state = samples['pose_state']   # (N,)

    # Polar representation of 3D vectors, array shapes (N,K)
    mu_thetas, mu_phis = cartesian_to_polar(params['state_directions'][state][:,1:])
    us_thetas, us_phis = cartesian_to_polar(samples['directions'][:,1:])

    # Update parameters, which we find by solving
    #   k_likelihood sin(theta_likelihood) = sum w_i sin(theta_i)
    #   k_likelihood cos(theta_likelihood) = sum w_i cos(theta_i)
    # Recall: Our prior for headings is vMF with concentration 0
    # so, theta_post = theta_likelihood, k_post = k_likelihood
    azim_weights = jnp.sin(mu_phis) * jnp.sin(us_phis)
    k_sin_thetas = jnp.sum(azim_weights * jnp.sin(us_thetas - mu_thetas), axis=-1)
    k_cos_thetas = jnp.sum(azim_weights * jnp.cos(us_thetas - mu_thetas), axis=-1)

    theta_post = jnp.arctan2(k_sin_thetas, k_cos_thetas)
    kappa_post = jnp.sqrt(k_sin_thetas**2 + k_cos_thetas**2)

    return tfd.VonMises(theta_post, kappa_post).sample(seed=seed)

@jit
def U_given_S_log_likelihoods(params, samples):
    """Calculate log likelihood of sampled directions under conditional
    prior distribution.

    Function seperated from sample_state so that this portion can be jitted.
    """
    directions = samples['directions']  # (N, K, 3)
    heading = samples['heading']        # (N,)

    # Rotate direction in absolute frame to canonical reference frame
    canonical_directions = jnp.einsum('tmn, tjn-> tjm',
                                      Rxy_mat(-heading), directions)

    # Calculate log likelihood of each set of directions, for each state
    # TODO Initialize
    # shape: (N, S)
    U_given_S = tfd.VonMisesFisher(params['state_directions'],
                                   params['state_concentrations'])
    return jnp.sum(U_given_S.log_prob(canonical_directions[:,None,...]),
                   axis=-1)

def sample_state(seed, params, samples):

    transition_matrix = samples['transition_matrix']    # shape (S, S)

    lls = U_given_S_log_likelihoods(params, samples)

    # Run forward filter, then backward sampler to draw states
    # Must move to CPU to sample, then move back to GPU
    states = hmm_sample(onp.asarray(params['state_probability'], dtype=onp.float64),
                        onp.asarray(transition_matrix[None, :, :], dtype=onp.float64),
                        onp.asarray(lls, dtype=onp.float64))
    return jnp.asarray(states)

@jit
def sample_transition_matrix(seed, params, samples):

    state = samples['pose_state']   # shape (N,)

    N, S = len(state), len(params['state_transition_count'])

    def count_all(t, counts):
        return counts.at[state[t], state[t+1]].add(1)

    counts = lax.fori_loop(0, N-1, count_all, jnp.zeros((S, S)))

    return tfd.Dirichlet(params['state_transition_count'] + counts).sample(seed=seed)

# Cannot jit because smaple_state requires moving arrays to CPU
# @jit
def step(seed, params, observations, samples):
    """Execute a single iteration of MCMC sampling."""
    seeds = jr.split(seed, 6)

    positions, kernel_results = \
                    sample_positions(seeds[0], params, observations, samples)
    samples['positions'] = positions
    samples['hmc_log_accept_ratio'] = kernel_results.log_accept_ratio
    samples['hmc_proposed_gradients'] = \
                    kernel_results.proposed_results.grads_target_log_prob[0]

    samples['outliers'] = \
                    sample_outliers(seeds[1], params, observations, samples)

    samples['directions'] = \
                    sample_directions(seeds[2], params, samples)

    samples['heading'] = \
                    sample_headings(seeds[3], params, samples)

    samples['pose_state'] = \
                    sample_state(seeds[4], params, samples)

    samples['transition_matrix'] = \
                    sample_transition_matrix(seeds[5], params, samples)

    samples['log_probability'] = log_joint_probability(
                        params, observations,
                        samples['outliers'],
                        samples['positions'], samples['directions'],
                        samples['heading'], samples['pose_state'],
                        samples['transition_matrix'],
                        )
    return samples

def predict(seed, params, observations, init_positions=None,
            num_mcmc_iterations=1000,
            enable_x64=False,
            hmc_options={'init_step_size':1e-1, 'num_leapfrog_steps':1},
            out_options={},
            ):
    """Predict latent pose variables from observations using MCMC sampling.

    Parameters
    ----------
        seed: jax.random.PRNGKey
        params: dict
        observations: ndarray, shape (N, C, K, D_obs)
        init_positions: None or ndarray, shape (N, K, D), optional
            Initial guess of 3D positions. If None (default), initial
            guess made from observations in mcmc.initialize
        num_mcmc_iterations: int
        enable_x64: bool, optional. default: False
            If True, enable jax computations with double precision (float64).
            Else, perform computations with single precision (float32).
        hmc_options: dict, optional.
            step_size: int. Step size per leapfrog integration.
                Larger step sizes lead to faster progress, but too-large step sizes make rejection exponentially more likely. When possible, it's often helpful to match per-variable step sizes to the standard deviations of the target distribution in each variable.
            num_leapfrog_steps: int. Number of leapfrog steps to take.
        out_options: dict, optional. If empty (default), samples are not saved.
            path: str. Path to which to save file.
            chunk_size: int, optional. MCMC iteration interval at which to save.
                Large chunk size reduces I/O overhead of writing to memory, but requires larger local memory to store sample outputs.

    Returns
    -------
        buffer, dict
            Contains last chunk_size of samples
    """

    # Enable double-precision if specified
    jax.config.update("jax_enable_x64", enable_x64)

    # Initialize
    seed, init_seed = jr.split(seed, 2)

    params = initialize_parameters(params)
    samples = initialize(init_seed, params,
                              observations, init_positions)

    # Initialize buffer and output file (if specified)
    chunk_size = out_options.get('chunk_size', num_mcmc_iterations)
    burnin = out_options.get('burnin', 0)

    hdf5_output = None
    if isinstance(out_options.get('path', None), str):
        hdf5_output = util_io.SavePredictionsToHDF(
                                        out_options['path'], samples,
                                        max_iter = num_mcmc_iterations-burnin,
                                        chunk_size=chunk_size,
                                        **out_options.get('hdf_kwargs', {}))

    buffer = {k: jnp.empty((chunk_size, *v.shape), dtype=v.dtype)
              for k, v in samples.items()}

    # Setup progress bar
    pbar = trange(num_mcmc_iterations)
    pbar.set_description("lp={:.2f}".format(samples['log_probability']))

    for itr in pbar:
        samples = \
            step(jr.fold_in(seed, itr), params, observations, samples)

        # ------------------------------------------------------------------
        # Update the progress bar
        pbar.set_description("lp={:.2f}".format(samples['log_probability']))
        pbar.update(1)

        # Update buffer
        for k, v in buffer.items():
            buffer[k] = buffer[k].at[itr % chunk_size].set(samples[k])

        # Save chunk
        if (hdf5_output is not None):
            if (itr >= burnin) and (itr+1) % chunk_size == 0:
                hdf5_output.update(buffer)

    # Store final chunk of samples, if needed
    if (hdf5_output is not None) and (itr+1) % chunk_size != 0:
        # Truncate buffer to only contain not-yet-written samples
        for k, v in buffer.items():
            buffer[k] = buffer[k][:itr % chunk_size + 1]

        hdf5_output.update(buffer)

    return buffer

reason: import numpy,import jax,from jax


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

from FaceID import face_class

import os
import cv2
import numpy as np
import sqlite3
import shutil

# get current project location path
currentPath = os.getcwd()

# SQL arguments settings
SQLite_path = os.path.join(currentPath, 'facelabeldb', str("face_"+face_class))
SQLite_name = "facelabel.db3"
SQLite_source = os.path.join(
    currentPath, 'facelabeldb', 'db_sources', SQLite_name)

dbconn = sqlite3.connect(os.path.join(SQLite_path, SQLite_name))
curs = dbconn.cursor()
db_TableName = 'facelabel'
column_faceindex = 'faceindex'
column_facename = 'facename'

def load_faceimg(faceimgpath, imgtype, img_resize):

    load_index = 0
    images = []  # face image
    labels = []  # face label
    label_class = []

    print("Start Scanning faceimg folder image class.\n")

    # run SQL DB Table clearn
    SQL_delete_syntax = '''
    DELETE FROM {}
    '''.format(db_TableName)
    SQL_run = curs.execute(SQL_delete_syntax)

    if SQL_run:

        dbconn.commit()

    else:

        print('Run SQL_delete_syntax Faild.')

    # scan faceimg folder
    for dir_item in os.listdir(faceimgpath):

        #print("No."+str(load_index)+" [ "+dir_item+" ]")
        print("\nNo.{} [ {} ]".format(str(load_index), dir_item))

        load_index += 1
        label_class.append(dir_item)

        # join root directory and class folder name
        full_path = os.path.abspath(os.path.join(faceimgpath) + dir_item)
        print("Full Detail Path = {}".format(full_path))

        class_name = dir_item
        num_img = 0

        # read image
        for dir_item in os.listdir(full_path):

            # print(dir_item)
            num_img += 1
            img_path = ""

            if dir_item.endswith(imgtype)and os.path.exists(str(full_path)):

                img_path = os.path.abspath(
                    os.path.join(full_path) + "\\" + dir_item)
                # print(img_path)
                image = cv2.imread(img_path)
                image = cv2.resize(image, (img_resize, img_resize))

                # get image size
                img_width, img_height, channels = image.shape

                # if num_img < 10: # 1~9

                #    dir_item = str(dir_item).rstrip("_000" + str(num_img) + ".jpg").rstrip()

                # elif num_img >= 10 and num_img < 100: # 10~99

                #    dir_item = str(dir_item).rstrip("_00" + str(num_img) + ".jpg").rstrip()

                # elif num_img >= 100 and num_img < 1000: # 100~999

                #    dir_item = str(dir_item).rstrip("_0" + str(num_img) + ".jpg").rstrip()

                # else: # 10~99

                #    dir_item = str(dir_item).rstrip("_" + str(num_img) + ".jpg").rstrip()

                # add image to array
                images.append(image)
                labels.append(class_name)

        print("Image Quantity: {}".format(num_img))

    return images, labels, img_width, img_height, channels, load_index, label_class

def char_to_code(labels, label_class, n_class):

    # for index,item in label_class:

    #    # run SQL_Insert_syntax
    #    SQL_Insert_syntax = '''
    #    INSERT INTO {}('{}','{}')
    #    VALUES('{}','{}')
    #    '''.format(db_TableName,column_faceindex,column_facename,str(index),str(item))
    #    SQL_run = curs.execute(SQL_Insert_syntax)

    #    if SQL_run:
    #        dbconn.commit()
    #    else:
    #        print ('Run SQL_Insert_syntax Faild.')

    for index, classname in label_class:

        labels = np.array(
            [int(index) if label == classname else label for label in labels])

        # run SQL_Insert_syntax
        SQL_Insert_syntax = '''
        INSERT INTO {}('{}','{}')
        VALUES('{}','{}')
        '''.format(db_TableName, column_faceindex, column_facename, str(index), str(classname))
        SQL_run = curs.execute(SQL_Insert_syntax)

        if SQL_run:
            dbconn.commit()
        else:
            print('Run SQL_Insert_syntax Faild.')

    return labels


text:
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sat Jun 10 10:54:42 2017

@author: Filippo Broggini
"""

# %%

#%load_ext autoreload
#%autoreload 2

# %%

import sys

#sys.dont_write_bytecode = True

import numpy as np
import scipy.misc as sp

import matplotlib.pyplot as plt

import seispy.GeologicalModelling as GM
import seispy.SeismicModelling2D  as SM

# %% Layered model - Deterministic

ints = np.array([20, 30, 40, 20, 10, 20, 18, 35])
v = np.array([1500, 2000, 1500, 1800, 2000, 2200, 2500, 2100])
rho = v*0.6
Layers0 = GM.LayeredModel({'dims': [np.sum(ints), 1000], 'type': 'layer'})
Layers0.Deterministic(ints,v,rho)
Layers0.Apply()

plt.imshow(Layers0.V)
plt.colorbar()

# %% Fault model - Deterministic

ints = np.array([20, 30, 60, 20, 20, 18, 35])
v = np.array([1500, 2000, 1500, 1800, 2000, 2500, 2100])
rho = v*0.6
faultlim = [400,600]
offset = 50
Fault0 = GM.FaultModel({'dims': [np.sum(ints), 1000], 'type': 'layer'})
Fault0.Deterministic(ints,v,rho,faultlim=faultlim,offset=offset)
Fault0.Apply()

plt.imshow(Fault0.V)
plt.colorbar()

# %% Trap model - Deterministic
perc = 0
center_x = 50
center_z = np.array([80, 90])
v = np.array([0, 600, 0])
rho = v*0.6
Trap0 = GM.TrapModel({'dims': [100, 100], 'type': 'trap'})
Trap0.Deterministic(center_x, center_z, v, rho, perc)
Trap0.Apply()

plt.imshow(Trap0.V)
plt.colorbar()

# %% synthetic_model

model = 1

if model==0:
    Fake = np.concatenate((Layers0.V, Fault0.V, Layers0.V), axis=0)
    Fake[270:370, 600:700] = Fake[270:370, 600:700] + Trap0.V
    Fake[400:500, 650:750] = Fake[400:500, 650:750] + Trap0.V
    Fake[100:200, 650:750] = Fake[100:200, 650:750] + Trap0.V
    Fake[200:300, 150:250] = Fake[200:300, 150:250] + Trap0.V
    Fake[350:450, 250:350] = Fake[350:450, 250:350] + Trap0.V
    Fake[300:400, 750:850] = Fake[300:400, 750:850] + Trap0.V
if model==1:
    Fake = np.concatenate((Layers0.V, Fault0.V, Layers0.V), axis=0)
#    Fake[270:370, 600:700] = Fake[270:370, 600:700] + Trap0.V
    Fake[400:500, 650:750] = Fake[400:500, 650:750] + Trap0.V
    Fake[100:200, 650:750] = Fake[100:200, 650:750] + Trap0.V
    Fake[200:300, 150:250] = Fake[200:300, 150:250] + Trap0.V
#    Fake[350:450, 250:350] = Fake[350:450, 250:350] + Trap0.V
#    Fake[300:400, 750:850] = Fake[300:400, 750:850] + Trap0.V
#Fake[300:400, 600:700] = Trap0.V
if model==2:
    Fake = np.concatenate((Layers0.V, Fault0.V, Layers0.V), axis=0)
#    Fake[270:370, 600:700] = Fake[270:370, 600:700] + Trap0.V
    Fake[400:500, 650:750] = Fake[400:500, 650:750] + Trap0.V
    Fake[100:200, 650:750] = Fake[100:200, 650:750] + Trap0.V
    Fake[200:300, 150:250] = Fake[200:300, 150:250] + Trap0.V

plt.figure(figsize=(12, 7))
plt.imshow(Fake)
plt.colorbar()
plt.show()

# %%

Seismod = SM.SeismicModelling2D({'V'      : Fake,
                                 'Rho'    : Fake*0.6,
                                 'dt'     : 0.004,
                                 'ot'     : 0,
                                 'ntrick' : 31,
                                 'f0'     : 10})

Seismod.Apply()
Seismod.Visualize(cbarlims=[-4e12,4e12])

sp.imsave('../datasets/seismic/synthetics/targets/synthetic_model' + str(model) + '.png',Seismod.stack)

# %%
#filepath='/home/bfilippo/Work/EAGE_Hackatoon_2017/datasets/faults/'
#filename='fault'
#
#Layers.Save(filepath=filepath, filename=filename)
#
#plt.show()

reason: import numpy,import scipy


text:
from __future__ import print_function
import sys
import os
import numpy as np
import json
import chainer
import chainer.links as L
import chainer.functions as F
from chainer import cuda
from chainer import training, Variable
from chainer.training import extensions
from chainer.optimizer import WeightDecay, GradientClipping
from collections import defaultdict, OrderedDict

from py_utils import read_pretrained_embeddings, read_model_defs
from biaffine import Biaffine
from param import Param

UNK = "*UNKNOWN*"
START = "*START*"
END = "*END*"
IGNORE = -1

class FeatureExtractor(object):
    def __init__(self, model_path):
        self.model_path = model_path
        self.words = read_model_defs(model_path / 'words.txt')
        self.chars = read_model_defs(model_path / 'chars.txt')
        self.unk_word = self.words[UNK]
        self.start_word = self.words[START]
        self.end_word = self.words[END]
        self.unk_char = self.chars[UNK]
        self.start_char = self.chars[START]
        self.end_char = self.chars[END]

    def process(self, words):
        """
        words: list of unicode tokens
        """
        w = np.array([self.start_word] + [self.words.get(
            x, self.unk_word) for x in words] + [self.end_word], 'i')
        l = max(len(x) for x in words)
        c = -np.ones((len(words) + 2, l), 'i')
        c[0, 0] = self.start_char
        c[-1, 0] = self.end_char
        for i, word in enumerate(words, 1):
            for j in range(len(word)):
                c[i, j] = self.chars.get(word[j], self.unk_char)
        return w, c, np.array([l], 'i')

class BiaffineJaLSTMParser(chainer.Chain):
    def __init__(self, model_path):
        Param.load(self, model_path / 'tagger_defs.txt')
        self.extractor = FeatureExtractor(model_path)
        self.in_dim = self.word_dim + self.char_dim
        super(BiaffineJaLSTMParser, self).__init__(
                emb_word=L.EmbedID(self.n_words, self.word_dim),
                emb_char=L.EmbedID(self.n_chars, 50, ignore_label=IGNORE),
                conv_char=L.Convolution2D(1, self.char_dim,
                    (3, 50), stride=1, pad=(1, 0)),
                lstm_f=L.NStepLSTM(self.nlayers, self.in_dim,
                    self.hidden_dim, 0.32),
                lstm_b=L.NStepLSTM(self.nlayers, self.in_dim,
                    self.hidden_dim, 0.32),
                arc_dep=L.Linear(2 * self.hidden_dim, self.dep_dim),
                arc_head=L.Linear(2 * self.hidden_dim, self.dep_dim),
                rel_dep=L.Linear(2 * self.hidden_dim, self.dep_dim),
                rel_head=L.Linear(2 * self.hidden_dim, self.dep_dim),
                biaffine_arc=Biaffine(self.dep_dim),
                biaffine_tag=L.Bilinear(self.dep_dim, self.dep_dim, len(self.targets)))

    def forward(self, ws, cs, ls, dep_ts=None):
        batchsize = len(ws)
        xp = chainer.cuda.get_array_module(ws[0])
        ws = map(self.emb_word, ws)
        cs = [F.squeeze(
            F.max_pooling_2d(
                self.conv_char(
                    F.expand_dims(
                        self.emb_char(c), 1)), (int(l[0]), 1)))
                    for c, l in zip(cs, ls)]
        xs_f = [F.dropout(F.concat([w, c]), 0.5) for w, c in zip(ws, cs)]
        xs_b = [x[::-1] for x in xs_f]
        cx_f, hx_f, cx_b, hx_b = self._init_state(xp, batchsize)
        _, _, hs_f = self.lstm_f(hx_f, cx_f, xs_f)
        _, _, hs_b = self.lstm_b(hx_b, cx_b, xs_b)
        hs_b = [x[::-1] for x in hs_b]
        hs = [F.concat([h_f, h_b]) for h_f, h_b in zip(hs_f, hs_b)]

        dep_ys = [self.biaffine_arc(
            F.elu(F.dropout(self.arc_dep(h), 0.32)),
            F.elu(F.dropout(self.arc_head(h), 0.32))) for h in hs]

        if dep_ts is not None:
            heads = dep_ts
        else:
            heads = [F.argmax(y, axis=1) for y in dep_ys]

        cat_ys = [
                self.biaffine_tag(
            F.elu(F.dropout(self.rel_dep(h), 0.32)),
            F.elu(F.dropout(self.rel_head(
                F.embed_id(t, h, ignore_label=IGNORE)), 0.32))) \
                        for h, t in zip(hs, heads)]

        return cat_ys, dep_ys

    def predict(self, xs):
        xs = [self.extractor.process(x) for x in xs]
        ws, ss, ps = zip(*xs)
        with chainer.no_backprop_mode(), chainer.using_config('train', False):
            cat_ys, dep_ys = self.forward(ws, ss, ps)
        return zip([F.log_softmax(y[1:-1]).data for y in cat_ys],
                [F.log_softmax(y[1:-1, :-1]).data for y in dep_ys])

    def predict_doc(self, doc, batchsize=16):
        res = []
        for i in range(0, len(doc), batchsize):
            res.extend([(i + j, 0, y)
                for j, y in enumerate(self.predict(doc[i:i + batchsize]))])
        return res

    def _init_state(self, xp, batchsize):
        res = [Variable(xp.zeros(( # forward cx, hx, backward cx, hx
                self.nlayers, batchsize, self.hidden_dim), 'f')) for _ in range(4)]
        return res

    @property
    def cats(self):
        return list(zip(*sorted(self.targets.items(), key=lambda x: x[1])))[0]

max_issues_repo_name: masashi-y/Camelthorn

max_forks_repo_name: masashi-y/Camelthorn


text:
"""generate a solid file from an irregular shaped mask

"""
import logging
import os
import subprocess
from shutil import which
import numpy as np

def make_solid_file(clipped_mask, out_name, dx=1000, dz=1000):
    """Make a solid file and vtk file from a clipped mask, write to out_name, requires pfmask-to-sol from pftools

    Parameters
    ----------
    clipped_mask : ndarray
        the mask array cropped to the inner shape, with 1's in the mask, 0's outside
    out_name : str
        name for the output .vtk and .pfsol files
    dx : int, optional
        horizontal cell size (Default value = 1000)
    dz : int, optional
        vertical cell size (Default value = 1000)

    Returns
    -------
    batches : list
        list of batches (patches?) located
    """
    # TODO Add ability to buid features as patches on top layer of mask
    pf_mask_to_sol_path = find_mask_to_sol_exe()
    if pf_mask_to_sol_path is None:
        msg = 'Could not locate pfmask-to-pfsol utility needed to generate solid file (.pfsol)' \
              ' ensure PARFLOW_DIR environment variable is set'
        logging.exception(msg)
        raise Exception(msg)
    # TODO: Should not have to flip this mask. Find why it is upside down
    mask_mat = np.flip(clipped_mask, axis=1)
    #

    if len(mask_mat.shape) == 3:
        mask_mat = np.squeeze(mask_mat, axis=0)
    # create back borders
    # Back borders occur where mask[y+1]-mask[y] is negative
    # (i.e. the cell above is a zero and the cell is inside the mask, i.e. a 1)
    back_mat = np.zeros(mask_mat.shape)

    # create front borders
    # Front borders occur where mask[y-1]-mask[y] is negative
    # (i.e. the cell above is a zero and the cell is inside the mask, i.e. a 1)
    front_mat = np.zeros(mask_mat.shape)

    # create left borders
    # Left borders occur where mask[x-1]-mask[x] is negative
    left_mat = np.zeros(mask_mat.shape)

    # create right borders
    # Right borders occur where mask[x+1]-mask[x] is negative
    right_mat = np.zeros(mask_mat.shape)

    # deal with top and bottom patches
    # 3 = regular overland boundary
    # 4 = Lake
    # 5 = Sink
    # 6 = bottom
    # 8 = Stream
    # 9 = Reservoir

    top_mat = mask_mat * 3

    bottom_mat = mask_mat * 6

    # create header and write ascii files
    header = 'ncols ' + str(mask_mat.shape[1]) + '\n'
    header += 'nrows ' + str(mask_mat.shape[0]) + '\n'
    header += 'xllcorner 0.0\n'
    header += 'yllcorner 0.0\n'
    header += 'cellsize ' + str(dx) + '\n'
    header += 'NODATA_value 0.0\n'

    patches = {os.path.join(os.path.dirname(out_name), 'Back_Border.asc'): back_mat,
               os.path.join(os.path.dirname(out_name), 'Front_Border.asc'): front_mat,
               os.path.join(os.path.dirname(out_name), 'Right_Border.asc'): right_mat,
               os.path.join(os.path.dirname(out_name), 'Left_Border.asc'): left_mat,
               os.path.join(os.path.dirname(out_name), 'Bottom_Border.asc'): bottom_mat,
               os.path.join(os.path.dirname(out_name), 'Top_Border.asc'): top_mat}

    list_patches = list(patches.keys())
    for patch in patches:
        with open(patch, 'w') as fo:
            fo.write(header)
            np.savetxt(fo, patches[patch].reshape([-1, 1]), '%.3f', ' ')

    # Create solid domain file
    # This part assumes that you have installed pf-mask-utilities or ParFlow with PARFLOW_DIR set
    out_vtk = out_name + '.vtk'
    out_pfsol = out_name + '.pfsol'

    if os.path.isfile(out_vtk):
        os.remove(out_vtk)

    if os.path.isfile(out_pfsol):
        os.remove(out_pfsol)

    sub_cmd_str = [pf_mask_to_sol_path[0],
                   '--mask-back', list_patches[0],
                   '--mask-front', list_patches[1],
                   '--mask-right', list_patches[2],
                   '--mask-left ' + list_patches[3],
                   '--mask-bottom ' + list_patches[4],
                   '--mask-top ' + list_patches[5],
                   '--vtk', out_vtk,
                   '--pfsol', out_pfsol,
                   pf_mask_to_sol_path[1], str(dz)]
    logging.info(f'begin mask_to_sol subprocess, command executed: {sub_cmd_str}')
    create_sub = subprocess.run(sub_cmd_str, stdout=subprocess.PIPE)
    temp_list = create_sub.stdout.decode('utf-8').split('\n')
    batches = ''
    for line in temp_list:
        if 'Number of triangles in patch' in line:
            line = line.strip()
            batches += line.split()[-3] + ' '
    logging.info(f'identified batches in domain {batches}')
    return batches

def find_mask_to_sol_exe():
    """find the pf_mask_to_pfsol or mask-to-pfsol utility on the system

    Returns
    -------
    pf_mask_to_sol_path : tuple
        tuple of (path to executable, depth/z-bottom flag for argument) or None if no executable was found
    """
    pf_mask_to_sol_path = None
    possible_paths = {('mask-to-pfsol', '--depth'): [os.environ.get('PFMASKUTILS'), which('mask-to-pfsol')],
                      ('pfmask-to-pfsol', '--z-bottom'):
                          [f'{os.environ.get("PARFLOW_DIR")}/bin', which('pfmask-to-pfsol')]}
    for executable, possible_paths in possible_paths.items():
        executable_path = next((path for path in possible_paths if path is not None), None)
        if executable_path is not None:
            if os.path.isfile(os.path.join(executable_path, executable[0])):
                pf_mask_to_sol_path = (os.path.join(executable_path, executable[0]), executable[1])
                break
    if pf_mask_to_sol_path is None:
        logging.exception('Could not locate mask to solid utility!')
    else:
        logging.info(f'searching for mask_to_sol executable resulted in: {pf_mask_to_sol_path[0]}')
    return pf_mask_to_sol_path

max_stars_repo_name: arezaii/subsetter

max_issues_repo_name: arezaii/subsetter

max_forks_repo_name: arezaii/subsetter


text:
# Copyright 2020 Google LLC
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Test get_weight_scale function with auto and auto_po2 modes of quantizers.py."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import logging
from numpy.testing import assert_allclose
import pytest
from tensorflow.keras import backend as K
from qkeras import binary
from qkeras import get_weight_scale
from qkeras import ternary

# expected value if input is uniform distribution is:
#   - alpha = m/2.0 for binary
#   - alpha = (m+d)/2.0 for ternary

def test_binary_auto():
  """Test binary auto scale quantizer."""

  np.random.seed(42)
  N = 1000000
  m_list = [1.0, 0.1, 0.01, 0.001]

  for m in m_list:
    x = np.random.uniform(-m, m, (N, 10)).astype(K.floatx())
    x = K.constant(x)

    quantizer = binary(alpha="auto")
    q = K.eval(quantizer(x))

    result = get_weight_scale(quantizer, q)
    expected = m / 2.0
    logging.info("expect %s", expected)
    logging.info("result %s", result)
    assert_allclose(result, expected, rtol=0.02)

def test_binary_auto_po2():
  """Test binary auto_po2 scale quantizer."""

  np.random.seed(42)
  N = 1000000
  m_list = [1.0, 0.1, 0.01, 0.001]

  for m in m_list:
    x = np.random.uniform(-m, m, (N, 10)).astype(K.floatx())
    x = K.constant(x)

    quantizer_ref = binary(alpha="auto")
    quantizer = binary(alpha="auto_po2")

    q_ref = K.eval(quantizer_ref(x))
    q = K.eval(quantizer(x))

    ref = get_weight_scale(quantizer_ref, q_ref)

    expected = np.power(2.0, np.round(np.log2(ref)))
    result = get_weight_scale(quantizer, q)

    assert_allclose(result, expected, rtol=0.0001)

def test_ternary_auto():
  """Test ternary auto scale quantizer."""

  np.random.seed(42)
  N = 1000000
  m_list = [1.0, 0.1, 0.01, 0.001]

  for m in m_list:
    x = np.random.uniform(-m, m, (N, 10)).astype(K.floatx())
    x = K.constant(x)

    quantizer = ternary(alpha="auto")
    q = K.eval(quantizer(x))

    d = m/3.0
    result = np.mean(get_weight_scale(quantizer, q))
    expected = (m + d) / 2.0
    assert_allclose(result, expected, rtol=0.02)

def test_ternary_auto_po2():
  """Test ternary auto_po2 scale quantizer."""

  np.random.seed(42)
  N = 1000000
  m_list = [1.0, 0.1, 0.01, 0.001]

  for m in m_list:
    x = np.random.uniform(-m, m, (N, 10)).astype(K.floatx())
    x = K.constant(x)

    quantizer_ref = ternary(alpha="auto")
    quantizer = ternary(alpha="auto_po2")

    q_ref = K.eval(quantizer_ref(x))
    q = K.eval(quantizer(x))

    ref = get_weight_scale(quantizer_ref, q_ref)

    expected = np.power(2.0, np.round(np.log2(ref)))
    result = get_weight_scale(quantizer, q)

    assert_allclose(result, expected, rtol=0.0001)

if __name__ == "__main__":
  pytest.main([__file__])

max_stars_repo_name: kshithijiyer/qkeras

max_issues_repo_name: kshithijiyer/qkeras

max_forks_repo_name: kshithijiyer/qkeras

reason: import numpy,from numpy


text:
from collections import defaultdict

from scipy import spatial
import numpy as np
import torch

from collections import defaultdict
from src.modules.lcfcn import lcfcn_loss
from scipy import spatial
import numpy as np
import torch

from collections import defaultdict

from scipy import spatial
import numpy as np
import torch

class SegMeter:
    def __init__(self, split):
        self.cf = None
        self.n_samples = 0
        self.split = split
        self.ae = 0
        self.game = 0

    def val_on_batch(self, model, batch):
        masks = batch["masks"].squeeze()
        self.n_samples += batch['images'].shape[0]
        pred_mask = model.predict_on_batch(batch).squeeze()

        # counts
        blobs = lcfcn_loss.get_blobs(pred_mask)
        points = lcfcn_loss.blobs2points(blobs)
        pred_counts = float(points.sum())
        self.ae += np.abs(float((batch['points']==1).sum()) - pred_counts)

        gt_points = batch['points'].squeeze().clone()
        gt_points[gt_points!=1] = 0
        self.game += lcfcn_loss.compute_game(pred_points=points.squeeze(),
                                             gt_points=gt_points, L=3)
        #
        # print(masks.sum())
        ind = masks != 255
        masks = masks[ind]
        pred_mask = pred_mask[ind]

        labels = np.arange(model.n_classes)
        cf = confusion_multi_class(torch.as_tensor(pred_mask).float().cuda(), masks.cuda().float(),
                                    labels=labels)

        if self.cf is None:
            self.cf = cf
        else:
            self.cf += cf

    def get_avg_score(self):
        # return -1
        Inter = np.diag(self.cf)
        G = self.cf.sum(axis=1)
        P = self.cf.sum(axis=0)
        union = G + P - Inter

        nz = union != 0
        iou = Inter / np.maximum(union, 1)
        mIoU = np.mean(iou[nz])
        iou[~nz] = np.nan
        val_dict = {'%s_score' % self.split: mIoU}
        for c in range(self.cf.shape[1]):
            val_dict['%s_class%d' % (self.split, c)] = iou[c]
        val_dict['%s_mae' % (self.split)] = self.ae / self.n_samples
        val_dict['%s_game' % (self.split)] = self.game / self.n_samples
        return val_dict

def confusion_multi_class(prediction, truth, labels):
    """
    cf = confusion_matrix(y_true=prediction.cpu().numpy().ravel(),
            y_pred=truth.cpu().numpy().ravel(),
                    labels=labels)
    """
    nclasses = labels.max() + 1
    cf2 = torch.zeros(nclasses, nclasses, dtype=torch.float, device=prediction.device)
    prediction = prediction.view(-1).long()
    truth = truth.view(-1)
    to_one_hot = torch.eye(int(nclasses), dtype=cf2.dtype, device=prediction.device)
    for c in range(nclasses):
        true_mask = (truth == c)
        pred_one_hot = to_one_hot[prediction[true_mask]].sum(0)
        cf2[:, c] = pred_one_hot

    return cf2.cpu().numpy()

def confusion_binary_class(prediction, truth):
    confusion_vector = prediction / truth

    tp = torch.sum(confusion_vector == 1).item()
    fp = torch.sum(confusion_vector == float('inf')).item()
    tn = torch.sum(torch.isnan(confusion_vector)).item()
    fn = torch.sum(confusion_vector == 0).item()
    cm = np.array([[tn,fp],
                   [fn,tp]])
    return cm

class SegMeterBinary:
    def __init__(self, split):
        self.cf = None
        self.struct_list = []
        self.split = split

    def val_on_batch(self, model, batch):
        masks_org = batch["masks"]

        pred_mask_org = model.predict_on_batch(batch)
        ind = masks_org != 255
        masks = masks_org[ind]
        pred_mask = pred_mask_org[ind]
        self.n_classes = model.n_classes
        if model.n_classes == 1:
            cf = confusion_binary_class(torch.as_tensor(pred_mask).float().cuda(), masks.cuda().float())
        else:
            labels = np.arange(model.n_classes)
            cf = confusion_multi_class(torch.as_tensor(pred_mask).float().cuda(), masks.cuda().float(),
                                    labels=labels)

        if self.cf is None:
            self.cf = cf
        else:
            self.cf += cf

        # structure
        struct_score = float(struct_metric.compute_struct_metric(pred_mask_org, masks_org))
        self.struct_list += [struct_score]

    def get_avg_score(self):
        TP = np.diag(self.cf)
        TP_FP = self.cf.sum(axis=1)
        TP_FN = self.cf.sum(axis=0)
        TN = TP[::-1]

        FP = TP_FP - TP
        FN = TP_FN - TP

        iou = TP / (TP + FP + FN)
        dice = 2*TP / (FP + FN + 2*TP)

        iou[np.isnan(iou)] = -1
        dice[np.isnan(dice)] = -1

        mDice = np.mean(dice)
        mIoU = np.mean(iou)

        prec = TP / (TP + FP)
        recall = TP / (TP + FN)
        spec = TN/(TN+FP)
        fscore = (( 2.0 * prec * recall ) / (prec + recall))

        val_dict = {}
        if self.n_classes == 1:
            val_dict['%s_dice' % self.split] = dice[0]
            val_dict['%s_iou' % self.split] = iou[0]

            val_dict['%s_prec' % self.split] = prec[0]
            val_dict['%s_recall' % self.split] = recall[0]
            val_dict['%s_spec' % self.split] = spec[0]
            val_dict['%s_fscore' % self.split] = fscore[0]

            val_dict['%s_score' % self.split] = dice[0]
            val_dict['%s_struct' % self.split] = np.mean(self.struct_list)
        return val_dict

# def confusion_multi_class(prediction, truth, labels):
#     """
#     cf = confusion_matrix(y_true=prediction.cpu().numpy().ravel(),
#             y_pred=truth.cpu().numpy().ravel(),
#                     labels=labels)
#     """
#     nclasses = labels.max() + 1
#     cf2 = torch.zeros(nclasses, nclasses, dtype=torch.float,
#                       device=prediction.device)
#     prediction = prediction.view(-1).long()
#     truth = truth.view(-1)
#     to_one_hot = torch.eye(int(nclasses), dtype=cf2.dtype,
#                            device=prediction.device)
#     for c in range(nclasses):
#         true_mask = (truth == c)
#         pred_one_hot = to_one_hot[prediction[true_mask]].sum(0)
#         cf2[:, c] = pred_one_hot

#     return cf2.cpu().numpy()

def confusion_binary_class(pred_mask, gt_mask):
    intersect = pred_mask.bool() & gt_mask.bool()

    fp_tp = (pred_mask ==1).sum().item()
    fn_tp = gt_mask.sum().item()
    tn_fn = (pred_mask ==0).sum().item()

    tp = (intersect == 1).sum().item()
    fp = fp_tp - tp
    fn = fn_tp - tp
    tn = tn_fn - fn

    cm = np.array([[tp, fp],
                   [fn, tn]])
    return cm

reason: import numpy,from scipy