text:
#
# For licensing see accompanying LICENSE.txt file.
# Copyright (C) 2020 Apple Inc. All Rights Reserved.
#

from pylab import *

import argparse
import h5py
import inspect
import os
import pandas as pd
import scipy.linalg
import sklearn.preprocessing
import sklearn.metrics

import path_utils
path_utils.add_path_to_sys_path("../lib", mode="relative_to_current_source_dir", frame=inspect.currentframe())
import embree_utils
import octomap_utils
import random_walk_utils

parser = argparse.ArgumentParser()
parser.add_argument("--scene_dir", required=True)
parser.add_argument("--use_python_reference_implementation", action="store_true")
args = parser.parse_args()

assert os.path.exists(args.scene_dir)

path_utils.add_path_to_sys_path(os.path.join(args.scene_dir, "..", ".."), mode="relative_to_cwd", frame=inspect.currentframe())
import _dataset_config

print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK] Begin...")

scene_name       = os.path.basename(args.scene_dir)
asset_export_dir = os.path.join(args.scene_dir, "_asset_export")
detail_dir       = os.path.join(args.scene_dir, "_detail")
mesh_dir         = os.path.join(detail_dir, "mesh")
octomap_dir      = os.path.join(detail_dir, "octomap")
tmp_dir          = os.path.join(args.scene_dir, "_tmp")

metadata_cameras_asset_export_csv_file = os.path.join(asset_export_dir, "metadata_cameras_asset_export.csv")
metadata_cameras_csv_file              = os.path.join(detail_dir, "metadata_cameras.csv")
metadata_scene_csv_file                = os.path.join(detail_dir, "metadata_scene.csv")
mesh_vertices_hdf5_file                = os.path.join(mesh_dir, "mesh_vertices.hdf5")
mesh_faces_vi_hdf5_file                = os.path.join(mesh_dir, "mesh_faces_vi.hdf5")
mesh_faces_oi_hdf5_file                = os.path.join(mesh_dir, "mesh_faces_oi.hdf5")
octomap_bt_file                        = os.path.join(octomap_dir, "octomap.bt")
octomap_free_space_min_hdf5_file       = os.path.join(octomap_dir, "octomap_free_space_min.hdf5")
octomap_free_space_max_hdf5_file       = os.path.join(octomap_dir, "octomap_free_space_max.hdf5")

scene = [ s for s in _dataset_config.scenes if s["name"] == scene_name ][0]
with h5py.File(mesh_vertices_hdf5_file,          "r") as f: mesh_vertices          = f["dataset"][:]
with h5py.File(mesh_faces_vi_hdf5_file,          "r") as f: mesh_faces_vi          = f["dataset"][:]
with h5py.File(mesh_faces_oi_hdf5_file,          "r") as f: mesh_faces_oi          = f["dataset"][:]
with h5py.File(octomap_free_space_min_hdf5_file, "r") as f: octomap_free_space_min = f["dataset"][:]
with h5py.File(octomap_free_space_max_hdf5_file, "r") as f: octomap_free_space_max = f["dataset"][:]

obj_ids_unique    = unique(mesh_faces_oi)
color_vals_unique = arange(obj_ids_unique.shape[0])
np.random.seed(0)
np.random.shuffle(color_vals_unique)

df_scene              = pd.read_csv(metadata_scene_csv_file, index_col="parameter_name")
meters_per_asset_unit = df_scene.loc["meters_per_asset_unit"][0]
asset_units_per_meter = 1.0 / meters_per_asset_unit

# parameters

# when sampling views, generate images using these image and camera parameters
n_width_pixels  = 256
n_height_pixels = 192
n_fov_x         = pi/3

n_samples_random_walk            = 100  # each final generated camera trajectory consists of this many views
n_samples_octomap_query          = 1000 # generate this many preliminary candidates from a local neighborhood, test if they're in free space
n_samples_camera_pose_candidates = 20   # generate this many candidates from free space, compute view scores

n_voxel_size = scene["voxel_extent_meters"]*asset_units_per_meter

# local neighborhood bounding boxes for random walk
n_query_half_extent_relative_to_start   = array([np.inf,np.inf,0.25])*asset_units_per_meter
n_query_half_extent_relative_to_current = array([1.5,1.5,0.25])*asset_units_per_meter

if args.use_python_reference_implementation:

    height_pixels = n_height_pixels
    width_pixels  = n_width_pixels
    fov_x         = n_fov_x

    # when randomly sampling an up vector, perturb it according to these parameters
    n_camera_up_hint_noise_std_dev = 0.1
    n_camera_up_hint_nominal       = array([0,0,1])

    fov_y = 2.0 * arctan(height_pixels * tan(fov_x/2) / width_pixels)

    uv_min  = -1.0
    uv_max  = 1.0
    half_du = 0.5 * (uv_max - uv_min) / width_pixels
    half_dv = 0.5 * (uv_max - uv_min) / height_pixels

    u, v = meshgrid(linspace(uv_min+half_du, uv_max-half_du, width_pixels),
                    linspace(uv_min+half_dv, uv_max-half_dv, height_pixels)[::-1])

    ray_offset_x = u*tan(fov_x/2.0)
    ray_offset_y = v*tan(fov_y/2.0)
    ray_offset_z = -ones_like(ray_offset_x)

    rays_cam = dstack((ray_offset_x,ray_offset_y,ray_offset_z))
    V_cam    = matrix(rays_cam.reshape(-1,3)).T

    # margin parameter for line-of-sight queries: point A must be closer than the scene geometry to point B,
    # along the ray from B to A, by a margin of eps percent, in order for A to be considered visible from B
    eps = 0.01

    # when attempting to find the initial look-at position, sample the occupancy map slightly closer to the
    # initial look-from position than the point of mesh intersection, but at least delta units away from the
    # initial look-from position to avoid a degenerate (look-from, look-at) pair
    delta = 0.0001

    # constant term added to the view saliency score; as lamb goes to infinty, the distribution of view saliency
    # scores will approach a uniform distribution
    lamb = 0.0

    np.random.seed(0)

# get cameras from the original asset file
df_cameras_asset_export = pd.read_csv(metadata_cameras_asset_export_csv_file)
df_cameras = pd.DataFrame(columns=["camera_name"])
i = 0

for c in df_cameras_asset_export.to_records():

    in_camera_name                            = c["camera_name"]
    in_camera_dir                             = os.path.join(asset_export_dir, in_camera_name)
    in_camera_keyframe_positions_hdf5_file    = os.path.join(in_camera_dir, "camera_keyframe_positions.hdf5")
    in_camera_keyframe_orientations_hdf5_file = os.path.join(in_camera_dir, "camera_keyframe_orientations.hdf5")

    assert len(in_camera_name.lstrip("cam_")) != 2 or not in_camera_name.lstrip("cam_").isdigit()

    out_camera_name                                 = "cam_%02d" % i
    out_camera_dir                                  = os.path.join(args.scene_dir, "_detail", out_camera_name)
    out_camera_preview_dir                          = os.path.join(args.scene_dir, "_detail", out_camera_name, "preview")
    out_camera_keyframe_frame_indices_hdf5_file     = os.path.join(out_camera_dir, "camera_keyframe_frame_indices.hdf5")
    out_camera_keyframe_positions_hdf5_file         = os.path.join(out_camera_dir, "camera_keyframe_positions.hdf5")
    out_camera_keyframe_look_at_positions_hdf5_file = os.path.join(out_camera_dir, "camera_keyframe_look_at_positions.hdf5")
    out_camera_keyframe_orientations_hdf5_file      = os.path.join(out_camera_dir, "camera_keyframe_orientations.hdf5")
    out_metadata_camera_csv_file                    = os.path.join(out_camera_dir, "metadata_camera.csv")

    if not os.path.exists(out_camera_dir): os.makedirs(out_camera_dir)
    if not os.path.exists(out_camera_preview_dir): os.makedirs(out_camera_preview_dir)

    with h5py.File(in_camera_keyframe_positions_hdf5_file,    "r") as f: in_camera_keyframe_positions    = f["dataset"][:]
    with h5py.File(in_camera_keyframe_orientations_hdf5_file, "r") as f: in_camera_keyframe_orientations = f["dataset"][:]

    print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK] Input camera " + in_camera_name + ", output camera " + out_camera_name + "...")

    if not args.use_python_reference_implementation:

        camera_look_from_positions, camera_look_at_positions, camera_orientations, intersection_distances, prim_ids = \
            random_walk_utils.generate_camera_trajectory_random_walk(
                mesh_vertices,
                mesh_faces_vi,
                octomap_bt_file,
                octomap_free_space_min,
                octomap_free_space_max,
                in_camera_keyframe_positions[0],
                in_camera_keyframe_orientations[0],
                n_width_pixels,
                n_height_pixels,
                n_fov_x,
                n_samples_random_walk,
                n_samples_octomap_query,
                n_samples_camera_pose_candidates,
                n_voxel_size,
                n_query_half_extent_relative_to_current,
                n_query_half_extent_relative_to_start,
                tmp_dir=tmp_dir)

        if camera_look_from_positions is None or camera_look_at_positions is None or camera_orientations is None or intersection_distances is None or prim_ids is None:
            print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK] WARNING: random_walk_utils.generate_camera_trajectory_random_walk DID NOT EXECUTE SUCCESSFULLY. SKIPPING.")
            if not os.listdir(out_camera_preview_dir): os.rmdir(out_camera_preview_dir)
            if not os.listdir(out_camera_dir): os.rmdir(out_camera_dir)
            i = i+1
            continue

        #
        # save preview images
        #

        for j in range(n_samples_random_walk):

            prim_ids_curr = prim_ids[j]
            obj_ids_curr  = mesh_faces_oi[prim_ids_curr]

            color_vals = ones_like(obj_ids_curr)*np.nan
            for obj_id,color_val in zip(obj_ids_unique,color_vals_unique):
                color_vals[obj_id == obj_ids_curr] = color_val
            color_vals[prim_ids_curr == -1] = np.nan

            out_camera_preview_jpg_file = os.path.join(out_camera_preview_dir, "frame.%04d.jpg" % j)
            imsave(out_camera_preview_jpg_file, color_vals, vmin=np.min(color_vals_unique), vmax=np.max(color_vals_unique))

        #
        # final output values
        #

        num_keyframes                         = n_samples_random_walk
        out_camera_keyframe_frame_indices     = arange(num_keyframes)
        out_camera_keyframe_positions         = camera_look_from_positions
        out_camera_keyframe_look_at_positions = camera_look_at_positions
        out_camera_keyframe_orientations      = camera_orientations
        out_camera_frame_time_seconds         = 1.0

    if args.use_python_reference_implementation:

        # # note that skipping a camera trajectory like this will mean that metadata_cameras.csv will be
        # # incorrectly exported, so this if statement should be uncommented for debugging purposes only
        # if out_camera_name == "cam_00":
        #     i = i+1
        #     continue

        #
        # start_camera_look_from_position
        #

        start_camera_look_from_position = in_camera_keyframe_positions[0]

        #
        # start_camera_look_at_position
        #

        start_camera_R_world_from_cam = in_camera_keyframe_orientations[0]
        start_camera_look_at_dir      = -start_camera_R_world_from_cam[:,2]

        # compute intersection distance for center ray
        intersection_distances, intersection_normals, prim_ids = embree_utils.generate_ray_intersections(mesh_vertices, mesh_faces_vi, matrix(start_camera_look_from_position).A, matrix(start_camera_look_at_dir).A, tmp_dir=tmp_dir)
        intersection_distance = max(intersection_distances[0] - 1.75*n_voxel_size, delta)
        query_position        = start_camera_look_from_position + intersection_distance*start_camera_look_at_dir
        octomap_sample        = octomap_utils.generate_octomap_samples(octomap_bt_file, array([query_position]), tmp_dir)[0]

        if isfinite(intersection_distance) and octomap_sample == 0:
            start_camera_look_at_position = start_camera_look_from_position + intersection_distance*start_camera_look_at_dir
            computed_intersection_distance_image = False

        else:

            # try to find an unoccupied cell for the initial look-at position by shooting camera rays,
            # testing for intersections with the scene mesh, and testing the occupancy map cells slightly
            # before the intersections (when proceeding from the optical center of the camera to the
            # mesh intersection point); if no unoccupied cells can be found, try perturbing the camera
            # forwards along the camera's look-at vector slightly and trying again
            camera_z_axis                          = start_camera_R_world_from_cam[:,2]
            camera_perturb_attempts                = 8
            camera_perturb_length                  = 0.25*n_voxel_size
            all_intersection_distances_at_infinity = False
            encountered_unoccupied_cell            = False

            for p in range(camera_perturb_attempts):

                V_world              = start_camera_R_world_from_cam*V_cam
                ray_directions_world = V_world.T.A
                ray_positions_world  = ones_like(ray_directions_world) * (start_camera_look_from_position + p*camera_perturb_length*camera_z_axis)

                intersection_distances, intersection_normals, prim_ids = embree_utils.generate_ray_intersections(mesh_vertices, mesh_faces_vi, ray_positions_world, ray_directions_world, tmp_dir=tmp_dir)

                if not any(isfinite(intersection_distances)):
                    all_intersection_distances_at_infinity = True
                    break

                # clip rays against octomap bounding box, see https://tavianator.com/fast-branchless-raybounding-box-intersections/
                ray_directions_world = sklearn.preprocessing.normalize(ray_directions_world)
                t_min                = ones_like(intersection_distances)*-np.inf
                t_max                = ones_like(intersection_distances)*np.inf

                t_x0 = (octomap_free_space_min[0] - ray_positions_world[:,0]) / ray_directions_world[:,0]
                t_x1 = (octomap_free_space_max[0] - ray_positions_world[:,0]) / ray_directions_world[:,0]
                mask = logical_not(isclose(ray_directions_world[:,0], 0))
                t_min[mask] = np.maximum(t_min[mask], np.minimum(t_x0[mask], t_x1[mask]))
                t_max[mask] = np.minimum(t_max[mask], np.maximum(t_x0[mask], t_x1[mask]))

                t_y0 = (octomap_free_space_min[1] - ray_positions_world[:,1]) / ray_directions_world[:,1]
                t_y1 = (octomap_free_space_max[1] - ray_positions_world[:,1]) / ray_directions_world[:,1]
                mask = logical_not(isclose(ray_directions_world[:,1], 0))
                t_min[mask] = np.maximum(t_min[mask], np.minimum(t_y0[mask], t_y1[mask]))
                t_max[mask] = np.minimum(t_max[mask], np.maximum(t_y0[mask], t_y1[mask]))

                t_z0 = (octomap_free_space_min[2] - ray_positions_world[mask][:,2]) / ray_directions_world[:,2]
                t_z1 = (octomap_free_space_max[2] - ray_positions_world[mask][:,2]) / ray_directions_world[:,2]
                mask = logical_not(isclose(ray_directions_world[:,2], 0))
                t_min[mask] = np.maximum(t_min[mask], np.minimum(t_z0[mask], t_z1[mask]))
                t_max[mask] = np.minimum(t_max[mask], np.maximum(t_z0[mask], t_z1[mask]))

                assert all(isfinite(t_min))
                assert all(isfinite(t_max))
                assert all(t_max > t_min) # assert all rays intersect bounding box
                assert all(t_min < 0.5*1.75*n_voxel_size) # assert all rays start from inside bounding box (with a bit of slack because bounding box min and max might be off by a half voxel)

                # import mayavi_utils
                # import mayavi.mlab
                # tmp = ray_positions_world + t_max[:,newaxis]*ray_directions_world
                # mayavi_utils.points3d_color_by_scalar(tmp, scalars=zeros_like(t_max), scale_factor=1.0, opacity=1.0)
                # mayavi.mlab.triangular_mesh(mesh_vertices[:,0], mesh_vertices[:,1], mesh_vertices[:,2], mesh_faces_vi, representation="surface", color=(0.75,0.75,0.75), opacity=0.25)
                # mayavi.mlab.show()

                t_max[logical_not(isfinite(intersection_distances))] = np.inf
                intersection_distances = np.maximum(np.minimum(intersection_distances, t_max) - 1.75*n_voxel_size,  delta)

                query_positions = (start_camera_look_from_position + p*camera_perturb_length*camera_z_axis) + intersection_distances[:,newaxis]*ray_directions_world
                octomap_samples = octomap_utils.generate_octomap_samples(octomap_bt_file, query_positions, tmp_dir)

                # import mayavi_utils
                # import mayavi.mlab
                # tmp = query_positions
                # mayavi_utils.points3d_color_by_scalar(tmp, scalars=octomap_samples, scale_factor=1.0, opacity=1.0)
                # mayavi.mlab.triangular_mesh(mesh_vertices[:,0], mesh_vertices[:,1], mesh_vertices[:,2], mesh_faces_vi, representation="surface", color=(0.75,0.75,0.75), opacity=0.25)
                # mayavi.mlab.show()

                if any(octomap_samples == 0):
                    encountered_unoccupied_cell = True
                    break

            if all_intersection_distances_at_infinity:
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK] WARNING: CAMERA DOESN'T OBSERVE ANY PART OF THE SCENE. ALL INTERSECTION DISTANCES ARE AT INFINITY. GIVING UP.")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                continue

            if not encountered_unoccupied_cell:
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK] WARNING: CAMERA DOESN'T OBSERVE ANY PART OF THE SCENE. ALL OBSERVED OCTOMAP SAMPLES ARE UNKNOWN OR OCCUPIED. GIVING UP.")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK]")
                continue

            computed_intersection_distance_image = True

            half_dx = 0.5
            half_dy = 0.5
            pixel_center = array([(width_pixels+1)/2.0, (height_pixels+1)/2.0])
            pixel_x, pixel_y = meshgrid(linspace(half_dx, width_pixels+half_dx,  width_pixels),
                                        linspace(half_dy, height_pixels+half_dy, height_pixels))

            pixels                     = dstack((pixel_x,pixel_y)).reshape(-1,2)
            center_to_pixels           = pixels - pixel_center
            center_to_pixels_distances = linalg.norm(center_to_pixels, axis=1)
            center_to_pixels_distances[logical_not(isfinite(intersection_distances))] = np.inf
            center_to_pixels_distances[octomap_samples != 0]                          = np.inf
            assert any(isfinite(center_to_pixels_distances))

            selected_index                = argsort(center_to_pixels_distances)[0]
            ray_direction_world           = ray_directions_world[selected_index]
            intersection_distance         = intersection_distances[selected_index]
            start_camera_look_at_position = start_camera_look_from_position + intersection_distance*ray_direction_world

        #
        # save preview image
        #

        if not computed_intersection_distance_image:
            V_world = start_camera_R_world_from_cam*V_cam
            ray_directions_world = V_world.T.A
            ray_positions_world  = ones_like(ray_directions_world) * start_camera_look_from_position
            intersection_distances, intersection_normals, prim_ids = \
                embree_utils.generate_ray_intersections(mesh_vertices, mesh_faces_vi, ray_positions_world, ray_directions_world, tmp_dir=tmp_dir)

        prim_ids      = prim_ids.reshape(height_pixels, width_pixels)
        prim_ids_curr = prim_ids
        obj_ids_curr  = mesh_faces_oi[prim_ids_curr]

        color_vals = ones_like(obj_ids_curr)*np.nan
        for obj_id,color_val in zip(obj_ids_unique,color_vals_unique):
            color_vals[obj_id == obj_ids_curr] = color_val
        color_vals[prim_ids_curr == -1] = np.nan

        out_camera_preview_jpg_file = os.path.join(out_camera_preview_dir, "frame.%04d.jpg" % 0)
        imsave(out_camera_preview_jpg_file, color_vals, vmin=np.min(color_vals_unique), vmax=np.max(color_vals_unique))

        #
        # camera_up_hints
        #

        camera_up_hint_noise = np.random.randn(n_samples_random_walk,3)*n_camera_up_hint_noise_std_dev
        camera_up_hints      = n_camera_up_hint_nominal + camera_up_hint_noise

        #
        # compute random walk
        #

        current_camera_look_from_position = start_camera_look_from_position
        current_camera_look_at_position   = start_camera_look_at_position
        current_camera_orientation        = start_camera_R_world_from_cam

        random_walk_camera_look_from_positions = [current_camera_look_from_position]
        random_walk_camera_look_at_positions   = [current_camera_look_at_position]
        random_walk_camera_orientations        = [current_camera_orientation]

        for j in range(1,n_samples_random_walk):

            print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK] j = " + str(j))

            #
            # look_from_position_candidates
            #

            current_position = current_camera_look_from_position
            start_position   = start_camera_look_from_position

            query_extent    = 2*n_query_half_extent_relative_to_current
            query_offsets   = np.random.rand(n_samples_octomap_query, 3)*query_extent - n_query_half_extent_relative_to_current
            query_positions = current_position + query_offsets

            # filter according to global min and max boundaries
            query_positions = query_positions[all(query_positions >= start_position - n_query_half_extent_relative_to_start, axis=1)]
            query_positions = query_positions[all(query_positions <= start_position + n_query_half_extent_relative_to_start, axis=1)]
            assert query_positions.shape[0] > 0

            # filter according to occupancy
            octomap_samples = octomap_utils.generate_octomap_samples(octomap_bt_file, query_positions, tmp_dir)
            query_positions = query_positions[octomap_samples == 0]
            assert query_positions.shape[0] > 0

            # filter according to line of sight
            current_to_query           = query_positions - current_position
            current_to_query_distances = linalg.norm(current_to_query, axis=1)
            query_ray_directions       = current_to_query / current_to_query_distances[:,newaxis]
            query_ray_positions        = ones_like(query_ray_directions)*current_position
            intersection_distances, intersection_normals, prim_ids = embree_utils.generate_ray_intersections(mesh_vertices, mesh_faces_vi, query_ray_positions, query_ray_directions, tmp_dir=tmp_dir)
            query_positions = query_positions[(1-eps)*intersection_distances >= current_to_query_distances]
            assert query_positions.shape[0] > 0

            look_from_position_candidates = query_positions

            #
            # look_at_position_candidates
            #

            current_position = current_camera_look_at_position
            start_position   = start_camera_look_at_position

            query_extent    = 2*n_query_half_extent_relative_to_current
            query_offsets   = np.random.rand(n_samples_octomap_query, 3)*query_extent - n_query_half_extent_relative_to_current
            query_positions = current_position + query_offsets

            # filter according to global min and max boundaries
            query_positions = query_positions[all(query_positions >= start_position - n_query_half_extent_relative_to_start, axis=1)]
            query_positions = query_positions[all(query_positions <= start_position + n_query_half_extent_relative_to_start, axis=1)]
            assert query_positions.shape[0] > 0

            # filter according to occupancy
            octomap_samples = octomap_utils.generate_octomap_samples(octomap_bt_file, query_positions, tmp_dir)
            query_positions = query_positions[octomap_samples == 0]
            assert query_positions.shape[0] > 0

            # filter according to line of sight
            current_to_query           = query_positions - current_position
            current_to_query_distances = linalg.norm(current_to_query, axis=1)
            query_ray_directions       = current_to_query
            query_ray_positions        = ones_like(query_ray_directions)*current_position
            intersection_distances, intersection_normals, prim_ids = embree_utils.generate_ray_intersections(mesh_vertices, mesh_faces_vi, query_ray_positions, query_ray_directions, tmp_dir=tmp_dir)
            query_positions = query_positions[(1-eps)*intersection_distances >= current_to_query_distances]
            assert query_positions.shape[0] > 0

            look_at_position_candidates = query_positions

            # import mayavi_utils
            # import mayavi.mlab
            # tmp = array([current_camera_look_from_position, current_camera_look_at_position])
            # mayavi.mlab.plot3d(tmp[:,0], tmp[:,1], tmp[:,2], tube_radius=0.2, opacity=1.0, color=(0.75,0.0,0.0))
            # mayavi_utils.points3d_color_by_scalar(tmp, scalars=[0,0], scale_factor=0.4, opacity=1.0)
            # # mayavi_utils.points3d_color_by_scalar(query_positions, scalars=octomap_samples, scale_factor=0.2, opacity=1.0)
            # mayavi.mlab.triangular_mesh(mesh_vertices[:,0], mesh_vertices[:,1], mesh_vertices[:,2], mesh_faces_vi, representation="surface", color=(0.75,0.75,0.75), opacity=0.25)
            # mayavi.mlab.show()
            # assert False

            #
            # raycast against triangle mesh for each camera pose candidate
            #

            look_from_position_candidates = look_from_position_candidates[np.random.choice(look_from_position_candidates.shape[0], size=n_samples_camera_pose_candidates)]
            look_at_position_candidates   = look_at_position_candidates[np.random.choice(look_at_position_candidates.shape[0],     size=n_samples_camera_pose_candidates)]
            camera_up_hint                = camera_up_hints[j]

            camera_orientation_candidates = zeros((n_samples_camera_pose_candidates,3,3))

            for k in range(n_samples_camera_pose_candidates):

                camera_look_from_position = look_from_position_candidates[k]
                camera_look_at_position   = look_at_position_candidates[k]
                camera_look_at_dir        = sklearn.preprocessing.normalize(array([camera_look_at_position - camera_look_from_position]))[0]

                # The convention here is that the camera's positive x axis points right, the positive y
                # axis points up, and the positive z axis points away from where the camera is looking.
                camera_z_axis = -sklearn.preprocessing.normalize(array([camera_look_at_dir]))
                camera_x_axis = -sklearn.preprocessing.normalize(cross(camera_z_axis, camera_up_hint))
                camera_y_axis = sklearn.preprocessing.normalize(cross(camera_z_axis, camera_x_axis))

                R_world_from_cam = c_[ matrix(camera_x_axis).T, matrix(camera_y_axis).T, matrix(camera_z_axis).T ]
                V_world = R_world_from_cam*V_cam

                camera_orientation_candidates[k] = R_world_from_cam

                if k == 0:
                    ray_directions_world = V_world.T.A
                    ray_positions_world  = ones_like(V_world.T.A) * camera_look_from_position
                else:
                    ray_directions_world = r_[ray_directions_world, V_world.T.A]
                    ray_positions_world  = r_[ray_positions_world,  ones_like(V_world.T.A) * camera_look_from_position]

            intersection_distances, intersection_normals, prim_ids = \
                embree_utils.generate_ray_intersections(mesh_vertices, mesh_faces_vi, ray_positions_world, ray_directions_world, tmp_dir=tmp_dir)

            #
            # select a camera pose candidate with probability proportional to the view score
            #

            prim_ids    = prim_ids.reshape(n_samples_camera_pose_candidates, height_pixels, width_pixels)
            view_scores = zeros((n_samples_camera_pose_candidates))
            for k in range(n_samples_camera_pose_candidates):
                prim_ids_curr    = prim_ids[k]
                prim_ids_unique  = unique(prim_ids_curr)
                valid_fraction   = float(count_nonzero(prim_ids_curr != -1)) / float(height_pixels*width_pixels)
                num_unique_prims = prim_ids_unique.shape[0]
                view_scores[k]   = (valid_fraction**2)*num_unique_prims + lamb

            selected_index = np.random.choice(n_samples_camera_pose_candidates, p=view_scores/sum(view_scores))
            # selected_index = np.argmax(view_scores)

            #
            # save preview image
            #

            prim_ids_curr = prim_ids[selected_index]
            obj_ids_curr  = mesh_faces_oi[prim_ids_curr]

            color_vals = ones_like(obj_ids_curr)*np.nan
            for obj_id,color_val in zip(obj_ids_unique,color_vals_unique):
                color_vals[obj_id == obj_ids_curr] = color_val
            color_vals[prim_ids_curr == -1] = np.nan

            out_camera_preview_jpg_file = os.path.join(out_camera_preview_dir, "frame.%04d.jpg" % j)
            imsave(out_camera_preview_jpg_file, color_vals, vmin=np.min(color_vals_unique), vmax=np.max(color_vals_unique))

            #
            # append selected camera pose to random walk
            #

            current_camera_look_from_position = look_from_position_candidates[selected_index]
            current_camera_look_at_position   = look_at_position_candidates[selected_index]
            current_camera_orientation        = camera_orientation_candidates[selected_index]

            random_walk_camera_look_from_positions.append(current_camera_look_from_position)
            random_walk_camera_look_at_positions.append(current_camera_look_at_position)
            random_walk_camera_orientations.append(current_camera_orientation)

        random_walk_camera_look_from_positions = array(random_walk_camera_look_from_positions)
        random_walk_camera_look_at_positions   = array(random_walk_camera_look_at_positions)
        random_walk_camera_orientations        = array(random_walk_camera_orientations)

        num_keyframes                         = n_samples_random_walk
        out_camera_keyframe_frame_indices     = arange(num_keyframes)
        out_camera_keyframe_positions         = random_walk_camera_look_from_positions
        out_camera_keyframe_look_at_positions = random_walk_camera_look_at_positions
        out_camera_keyframe_orientations      = random_walk_camera_orientations
        out_camera_frame_time_seconds         = 1.0

    df_cameras = df_cameras.append({"camera_name":out_camera_name}, ignore_index=True)

    with h5py.File(out_camera_keyframe_frame_indices_hdf5_file,     "w") as f: f.create_dataset("dataset", data=out_camera_keyframe_frame_indices)
    with h5py.File(out_camera_keyframe_positions_hdf5_file,         "w") as f: f.create_dataset("dataset", data=out_camera_keyframe_positions)
    with h5py.File(out_camera_keyframe_look_at_positions_hdf5_file, "w") as f: f.create_dataset("dataset", data=out_camera_keyframe_look_at_positions)
    with h5py.File(out_camera_keyframe_orientations_hdf5_file,      "w") as f: f.create_dataset("dataset", data=out_camera_keyframe_orientations)

    df_camera = pd.DataFrame(columns=["parameter_name", "parameter_value"], data={"parameter_name": ["frame_time_seconds"], "parameter_value": [out_camera_frame_time_seconds]})
    df_camera.to_csv(out_metadata_camera_csv_file, index=False)

    i = i+1

df_cameras.to_csv(metadata_cameras_csv_file, index=False)

print("[HYPERSIM: SCENE_GENERATE_CAMERA_TRAJECTORIES_RANDOM_WALK] Finished.")

max_stars_repo_name: americast/ml-hypersim


text:
from copy import deepcopy
import matplotlib.pyplot as plt
import numpy as np
from astropy import units, wcs
from astropy.io import fits
from scipy.interpolate import interp1d, UnivariateSpline
import scipy as sp
from scipy import linalg

class IFUcube:
    """
    A class for applying Principal Component Analyses (PCA) Tomography in
    Integral Field Unit (IFU) spectroscopy data cubes. See a description
    of the method in Steiner et al. 2009 (https://arxiv.org/pdf/0901.2266).
    """
    def __init__(self, fname, scidata='SCI',
                 primary='PRIMARY', spectral_dimension=3,
                 emission_line_mask=None):
        """
        Instantiates the class IFUcube.

        Parameters
        ----------
        fname : string
            Name of the FITS file containing the IFU datacube. This
            should be the standard output from the FITS file.

        scidata: integer or string
            Extension of the FITS file containing the scientific data.

        primary: integer or string
            Extension of the FITS file containing the basic header.

        Returns
        -------
        Nothing.
        """

        self.fitsfile = fname
        self.tomograms = None
        self.eigenvectors = None
        self.eigenvalues = None

        with fits.open(fname) as hdu:
            self.data = hdu[scidata].data
            self.header = hdu[primary].header
            self.header_data = hdu[scidata].header
            self.wcs = wcs.WCS(self.header_data)

            labels = ['TOMOGRAMS', 'EIGENVECTORS', 'EIGENVALUES']

            try:
                self.tomograms = hdu[labels[0]].data
                self.eigenvectors = hdu[labels[1]].data
                self.eigenvalues = hdu[labels[2]].data
            except:
                pass

        self.wl = self.wcs.sub((spectral_dimension,)).wcs_pix2world(np.arange(len(self.data)), 0)[0]

        if self.wcs.wcs.cunit[2] == units.m:
            self.wl *= 1e+10

        if emission_line_mask is not None:
            self.em_mask = 1 * (emission_line_mask)
        else:
            self.em_mask = None

    def run_pca(self, components='all', emission_line_mask=None,
                emission_line_treatment='interpolate', write_fits=False,
                out_ext='_PCA.fits', normalization=False,
                norm_win=[None, None]):
        """
        Runs the PCA Tomography algoritm in the IFUcube
        (see detail in Steiner et al. 2009, https://arxiv.org/pdf/0901.2266.

        Parameters
        --------
        components: integer or string
            How many components to be saved. The standard is 'all',
            but it can become too large for large datacubes (e.g. MUSE datacubes).

            The other option is to select the number (integer) of
            components to be saved, e.g. components = 100.

        emission_line_mask: list of array
            The emission lines to be masked. It has to be a Boolean array
            (False value for the emission lines) with the same size
            as the spectral dimension of the datacube.

        emission_line_treatment: string
            If a 'emission_line_mask' is passed, 'emission_line_treatment' designs the treatment of the emission lines,
            wheter they will interpolated by 'best-fitting' curve (emission_line_treatment = 'interpolate')
            or they will be assign as nan's (emission_line_treatment = 'nan'). the standart is 'interpolate'.

        write_fits: Boolean
            Whether to save the results in a fits file. It can be read after with 'red_pca()' method.

        out_ext: string
            The extension to the fits file to be saved. Only applies if 'write_fits = True'.

        normalization: Boolean
            Whether to normalize the spectra before applying the PCA Tomography (NOT WORKING, TO BE IMPLEMENTED!).

        Returns
        -------
        Nothing.
        """

        if emission_line_mask is not None:
            self.em_mask = 1 * (emission_line_mask)
        else:
            self.em_mask = None

        print('--------------------------------------------')
        print('IT CAN TAKE SOME TIME AND SOME MEMORY!')
        print('--------------------------------------------')

        self.new_data = deepcopy(self.data)

        # Masking Emission lines
        if self.em_mask is not None:
            if np.shape(np.array(self.em_mask))[0] == self.new_data.shape[0]:
                self.em_mask[0] = True
                self.em_mask[-1] = True
                if(emission_line_treatment == 'interpolate'):
                    begin = []
                    end = []

                    for i in range(len(self.em_mask) - 1):
                        if (self.em_mask[i + 1] - self.em_mask[i]) == -1:
                            begin.append(i)
                        if (self.em_mask[i + 1] - self.em_mask[i]) == 1:
                            end.append(i)

                    for j in range(len(begin)):
                        for m in range(self.new_data.shape[1]):
                            for n in range(self.new_data.shape[2]):
                                xs = [self.wl[begin[j]], self.wl[end[j]]]
                                ys = [self.data[begin[j], m, n], self.data[end[j], m, n]]
                                interp = interp1d(xs, ys)
                                self.new_data[begin[j]:(end[j] + 1), m, n] = interp(self.wl[begin[j]:(end[j] + 1)])
                elif(emission_line_treatment == 'nan'):
                    self.new_data[:, :, self.em_mask] = 0
                else:
                    raise RuntimeError("emission_line_treatment needs tobe either 'interpolate' or 'nan'")
            else:
                raise RuntimeError("'emission_line_mask' {}\
                    needs to have the same shape as the spectral dimension\
                    of the IFUcube {}".format(self.new_data.shape[0],
                                              np.shape(np.array(self.em_mask))))

        # applying nan mask
        self.nan_mask = ~np.isfinite(self.data)
        self.new_data[self.nan_mask] = 0.0

        self.zones = self._gen_zones()
        n = len(self.zones)

        # Median spectra
        Q_lambda = ((self.new_data.sum(axis=1)).sum(axis=1)) / n

        # subtractin mean spectra
        I_ij_lambda = (self.new_data.T) - (Q_lambda)
        I_ij_lambda = I_ij_lambda.T

        Q_lambda = None

        # from 3D to 2D
        I_B_lambda = np.zeros((n, I_ij_lambda.shape[0]))
        n = 0
        for i in self.zones:
            I_B_lambda[n, :] = I_ij_lambda[:, i[1], i[0]]
            n = n + 1

        I_ij_lambda = None
        # covariance Matrix
        C_cov = (sp.dot(I_B_lambda.transpose(), I_B_lambda)) / (n - 1)

        # solving eigen-problem
        w_k, E_k = linalg.eigh(C_cov)

        # Organazing by eigenvalue
        S = sp.argsort(w_k)[::-1]
        self.eigenvalues = w_k[S]
        self.eigenvectors = E_k[:, S]

        # New coordenats system
        T_B_k = sp.dot(I_B_lambda, self.eigenvectors)

        if (components != 'all'):
            try:
                n_comp = int(components)
            except ValueError:
                print("'components' needs to be 'all' or an integer")
            self.tomograms = np.zeros((n_comp,
                                       self.new_data.shape[1],
                                       self.new_data.shape[2]))

            for n, i in enumerate(self.zones):
                self.tomograms[:, i[1], i[0]] = T_B_k[n, :(n_comp)]
            T_B_k[:] = None

            self.eigenvectors = self.eigenvectors[:, :n_comp]
            self.eigenvalues = self.eigenvalues[:n_comp]

        elif(components == 'all'):

            self.tomograms = np.zeros(np.shape(self.new_data))
            for n, i in enumerate(self.zones):
                self.tomograms[:, i[1], i[0]] = T_B_k[n, :]
            T_B_k[:] = None

        if write_fits:
            hdu = fits.PrimaryHDU([], header=self.header_data)
            data = fits.ImageHDU(self.data, name='SCI', header=self.header_data)
            tom = fits.ImageHDU(self.tomograms, name='Tomograms')
            eigenvectors = fits.ImageHDU(self.eigenvectors, name='Eigenvectors')
            eigenvalues = fits.ImageHDU(self.eigenvalues, name='Eigenvalues')
            hdul = fits.HDUList([hdu, data, tom, eigenvectors, eigenvalues])
            hdul.writeto(self.fitsfile[:-5] + out_ext)

        print('DONE!')

    def _gen_zones(self):
        return np.array([[i, j] for i in range(self.data.shape[2]) for j in range(self.data.shape[1])])

    @staticmethod
    def read_pca(fname, scidata='SCI', primary='PRIMARY', emission_line_mask=None, spectral_dimension=3):
        """
        Method for reading fits file created with 'run_pca()' method.

        """
        temp_obj = IFUcube(fname, emission_line_mask=emission_line_mask)
        return temp_obj

    def plot_tomogram(self, component, contrast=1):
        """
        Plots the tomogram and its respective Eigenvector for a given component.

        Parameters
        --------
        component: integer
            Which component to be plotted. The component numbers starts in 1.

        contrast: integer
            The colour contrast for the matplotlib pcolormesh plot.

        """

        data = self.tomograms[component - 1, :, ]
        eigh = self.eigenvectors[:, component - 1]
        if self.em_mask is not None:
            eigh[~self.em_mask] = np.nan

        fig = plt.figure(figsize=(12, 4))
        ax1 = plt.subplot2grid((3, 12), (0, 0), rowspan=3, colspan=3)
        y, x = np.indices(np.shape(data))
        y, x = y - np.max(y) / 2, x - np.max(x) / 2

        if (eigh[(abs(eigh) == np.nanmax(abs(eigh)))] < 0):
            eigh = eigh * - 1.0
            data = data * - 1.0
        img = ax1.pcolormesh(x, y, data,
                             vmin=np.percentile(data, contrast),
                             vmax=np.percentile(data, 100 - contrast))
        ax1.set_title('Tomogram ' + str(component))

        ax2 = plt.subplot2grid((3, 12), (0, 4), rowspan=3, colspan=8)
        ax2.plot(self.wl, eigh, label="eigen-spectra " + str(component))
        ax2.set_ylabel(r"E$_{} $".format(str(component)) + ' [Arbitrary Units]')
        ax2.set_xlabel(r'Wavelenght [$\AA$]')
        plt.legend()
        plt.show()

    def clean_data(self, components, fit_spline=True, smoothing_factor=0.1, write_fits=False):
        """
        Removes choosen components of the datacube, it is particularlly useful
        in 'Instrumental Fingerprint' removal from IFU spectroscopy data cubes
        (see the details in Menezes et al. 2014, https://arxiv.org/abs/1401.7078).

        Parameters
        --------
        components: integer or list
            Which components to be remove. It can be one component
            (e.g. components = 5) or a list (e.g. components = [22,30,45])

        fit_spline: Boolean
            Wheter to fit or not a spline to the Eigenvector before removing it.
            The standart it True.

        smoothing_factor: float
            The smoothing factor of the spline. Only applies if 'fit_spline = True'.

        write_fits: Boolean
            Whether to save the results in a fits file. It can be read after with 'red_pca()' method.

        out_ext: string
            The extension to the fits file to be saved. Only applies if 'write_fits = True'.
            See the documentation in scipy manual:
            https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.UnivariateSpline.html.

        write_fits: Boolean
            Whether to add a 'CLEAN DATA' extension in the IFU data cube.

        Returns
        -------
        clean_data: array
            The array containing the data with the component removed.
        """

        clean_data = self.data

        if isinstance(components, int):
            data = self.tomograms[components - 1, :, ]
            eigh = self.eigenvectors[:, components - 1]
            if (eigh[(abs(eigh) == np.nanmax(abs(eigh)))] < 0):
                eigh = eigh * -1.0
                data = data * -1.0

            if self.em_mask is not None:
                eigh[~self.em_mask] = np.nan

            if fit_spline:
                spline = UnivariateSpline(self.wl, eigh, s=smoothing_factor)
                plt.plot(self.wl, eigh, label="eigen-spectra " + str(components))
                plt.plot(self.wl, spline(self.wl), label="spline fit")
                plt.show()

                eigen_cube = data * spline(self.wl)[:, np.newaxis, np.newaxis]
                clean_data = clean_data - eigen_cube
            elif(~fit_spline):
                eigen_cube = data * eigh[:, np.newaxis, np.newaxis]
                clean_data = clean_data - eigen_cube
        if isinstance(components, list):
            for i in components:
                data = self.tomograms[i - 1, :, ]
                eigh = self.eigenvectors[:, i - 1]
                if (eigh[(abs(eigh) == np.nanmax(abs(eigh)))] < 0):
                    eigh = eigh * -1.0
                    data = data * -1.0

                if fit_spline:
                    spline = UnivariateSpline(self.wl, eigh, s=smoothing_factor)
                    plt.plot(self.wl, eigh, label="eigen-spectra " + str(i))
                    plt.plot(self.wl, spline(self.wl), lab="spline fit")
                    plt.show()
                    eigen_cube = data * spline(self.wl)[:, np.newaxis, np.newaxis]
                    clean_data = clean_data - eigen_cube
                elif(~fit_spline):
                    eigen_cube = data * eigh[:, np.newaxis, np.newaxis]
                    clean_data = clean_data - eigen_cube
        if write_fits:
            hdul = fits.open(self.fitsfile)
            hdul.append(fits.ImageHDU(clean_data, name='CLEAN SCI', header=self.header_data))
            hdul.writeto(self.fitsfile, overwrite=True)

        return clean_data

    def plot_scree_test(self, components=15, log_scale=True, scree_line=0.1):
        """
        Plots the scree test for the PCA Tomography (see details in Steiner et al. 2009,
         https://academic.oup.com/mnras/article/395/1/64/1078341).

        Parameters
        --------
        components: integer
            Which components to be ploted.

        log_scale: Boolean
            Whether to show the Covariance axis (y) in log scale.

        scree_line: float
            Where to draw the reference line for the scree test.

        """

        if components > len(self.eigenvalues):
            components = int(len(self.eigenvalues))

        var = (self.eigenvalues / np.sum(self.eigenvalues))[:(components)] * 100

        component = np.arange(1, components + 1, 1)
        plt.xlim(0, components + 1)

        plt.scatter(component, var)
        plt.ylabel(r'Variance (%)')
        plt.xlabel(r'Eigenvector number')
        plt.xticks(component)

        scale = var[0] - var[-1]
        if log_scale:
            plt.yscale('log')

            plt.hlines(scree_line,
                       xmin=0, xmax=components + 1,
                       color='red',
                       linestyle='--',
                       linewidth=0.5,
                       label='Variance = ' + str(scree_line) + ' %')
        else:
            plt.hlines(scree_line,
                       xmin=0,
                       xmax=components + 1,
                       color='red',
                       linestyle='--',
                       linewidth=0.5,
                       label='Variance = ' + str(scree_line) + ' %')
            plt.hlines(0, xmin=0, xmax=components + 1,
                       color='black',
                       linewidth=0.5,
                       label='Variance = 0.0 %')

        plt.legend()
        plt.show()

    def get_tomogram(self, component):
        tomogram = self.tomograms[component - 1, :, ]
        return tomogram

    def get_eigenvector(self, component):
        eigh = self.eigenvectors[:, component - 1]
        return eigh

if __name__ == "__main__":
    a = 1

ext: py

max_stars_repo_name: muryelgp/PCAtomography

max_issues_repo_name: muryelgp/PCAtomography

max_forks_repo_name: muryelgp/PCAtomography

reason: import numpy,import scipy,from scipy,from astropy


text:
"""
Test code plotting results.
"""

import os
import shutil

import numpy as np
import pytest
from bilby.core.grid import Grid
from bilby.core.result import Result, read_in_result
from cwinpy import PulsarParameters
from cwinpy.plot import DEFAULT_BOUNDS, Plot
from matplotlib.figure import Figure

class TestPlotting(object):
    """
    Test the plotting class.
    """

    @classmethod
    def setup_class(cls):
        # set results files containing data to plot
        cls.lppen = os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "data",
            "plotting_test_lppen.hdf",
        )

        cls.cwinpy = os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "data",
            "plotting_test_cwinpy.json.gz",
        )

        cls.grid = os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "data",
            "plotting_test_grid.json.gz",
        )

        # labels for each results set
        cls.allresults = {
            "lppen": cls.lppen,
            "CWInPy": cls.cwinpy,
            "Grid": cls.grid,
        }

        cls.pulsar = os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "data",
            "plotting_test.par",
        )

        cls.outputdir = "plots"
        os.makedirs(cls.outputdir, exist_ok=True)

    @classmethod
    def teardown_class(cls):
        """
        Remove test output plots.
        """

        shutil.rmtree(cls.outputdir)

    def test_plot_inputs(self):
        """
        Test inputs to plotting class.
        """

        with pytest.raises(TypeError):
            # fails with no positional arguments
            _ = Plot()

        with pytest.raises(TypeError):
            # fails with position argument of the wrong type
            _ = Plot(1)

        # test reading in different objects
        for res in [self.lppen, read_in_result(self.cwinpy), Grid.read(self.grid)]:
            plot = Plot(res)
            assert len(plot.results) == 1

        # test reading in dictionary - parameters is not specified, so should
        # fail with inconsistent parameter error
        with pytest.raises(ValueError):
            _ = Plot(
                {
                    "lppen": self.lppen,
                    "CWInPy": read_in_result(self.cwinpy),
                    "grid": Grid.read(self.grid),
                }
            )

        # try reading non-existent files
        with pytest.raises(IOError):
            _ = Plot({"one": "no_file.hdf", "two": "blah.json"})

        plot = Plot(self.allresults, parameters=["h0", "psi"])
        assert len(plot.results) == len(self.allresults)

        for label, restype in zip(self.allresults.keys(), [Result, Result, Grid]):
            assert isinstance(plot.results[label], restype)

    def test_parameters(self):
        """
        Test input parameters for plotting.
        """

        # pass parameter with wrong type
        with pytest.raises(TypeError):
            _ = Plot(self.cwinpy, parameters=1)

        # pass unavailable parameter
        with pytest.raises(ValueError):
            _ = Plot(self.cwinpy, parameters="blah")

        # pass real parameter - should fail for default plot type "corner",
        # which cannot be used for a single parameter
        with pytest.raises(TypeError):
            _ = Plot(self.cwinpy, parameters="h0")

        plot = Plot(self.cwinpy, parameters="h0", plottype="hist")
        assert plot.parameters == ["h0"]
        assert plot.plottype == "hist"
        assert plot.fig is None

        # pass two parameters with plottype "kde" - should fail as this can
        # only be used for 1 parameter
        with pytest.raises(TypeError):
            _ = Plot(self.cwinpy, parameters=["h0", "cosiota"], plottype="kde")

        plot = Plot(self.cwinpy, parameters=["h0", "cosiota"])
        assert plot.parameters == ["h0", "cosiota"]
        assert plot.plottype == "corner"
        assert len(plot.results) == 1

        # pass three parameters with plot type "hist" - should fail for as this
        # can only be used for 1 parameter
        with pytest.raises(TypeError):
            _ = Plot(self.cwinpy, parameters=["h0", "cosiota", "psi"], plottype="hist")

        plot = Plot(self.cwinpy, parameters=["h0", "cosiota", "psi"], plottype="corner")
        assert plot.parameters == ["h0", "cosiota", "psi"]
        assert plot.plottype == "corner"
        assert len(plot.results) == 1

    def test_plottype(self):
        """
        Test plot type is valid.
        """

        with pytest.raises(TypeError):
            _ = Plot(self.cwinpy, parameters="h0", plottype="skdfjkfda")

    def test_untrig(self):
        """
        Test inverting trigonometric function.
        """

        # fail as cannot invert value for Grid
        with pytest.raises(TypeError):
            _ = Plot(self.grid, parameters="cosiota", plottype="hist", untrig="cosiota")

        plot = Plot(self.cwinpy, parameters="iota", untrig="cosiota", plottype="hist")
        assert plot.parameters == ["iota"]
        assert len(plot.results) == 1
        assert "iota" in plot.results["result"].posterior.columns
        assert all(
            (DEFAULT_BOUNDS["iota"]["low"] <= plot.results["result"].posterior["iota"])
            & (
                DEFAULT_BOUNDS["iota"]["high"]
                > plot.results["result"].posterior["iota"]
            )
        )

    def test_pulsar_parameters(self):
        """
        Test reading in pulsar parameters.
        """

        # fail if parameter file is not right
        with pytest.raises(IOError):
            _ = Plot(
                self.cwinpy,
                parameters="iota",
                untrig="cosiota",
                plottype="hist",
                pulsar=985348,
            )

        plot = Plot(
            self.cwinpy,
            parameters="iota",
            untrig="cosiota",
            plottype="hist",
            pulsar=self.pulsar,
        )

        assert list(plot.injection_parameters.keys()) == ["iota"]
        assert plot.injection_parameters["iota"] == np.arccos(
            PulsarParameters(self.pulsar)["COSIOTA"]
        )

    def test_plot_1d(self):
        """
        Test 1D plots.
        """

        # histogram
        plot = Plot(self.allresults, parameters="h0", plottype="hist")
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "hist"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)

        plot.save("oned_hist.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

        # histogram with KDE
        plot = Plot(self.allresults, parameters="h0", plottype="hist", kde=True)
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "hist"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)

        plot.save("oned_hist_with_kde.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

        # KDE
        plot = Plot(self.allresults, parameters="h0", plottype="kde")
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "kde"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)

        plot.save("oned_kde.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

    def test_plot_2d(self):
        """
        Test 2D plots.
        """

        # corner
        plot = Plot(self.allresults, parameters=["h0", "cosiota"], plottype="corner")
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "corner"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)
        assert len(plot.fig.axes) == 4

        plot.save("twod_corner.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

        # triangle
        plot = Plot(self.allresults, parameters=["h0", "cosiota"], plottype="triangle")
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "triangle"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)
        assert len(plot.fig.axes) == 4

        plot.save("twod_triangle.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

        # reverse triangle
        plot = Plot(
            self.allresults, parameters=["h0", "cosiota"], plottype="reverse_triangle"
        )
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "reverse_triangle"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)
        assert len(plot.fig.axes) == 4

        plot.save("twod_reverse_triangle.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

        # contour
        plot = Plot(self.allresults, parameters=["h0", "cosiota"], plottype="contour")
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "contour"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)
        assert len(plot.fig.axes) == 1

        plot.save("twod_contour.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

    def test_plot_nd(self):
        """
        Test mulitple dimension plots.
        """

        # corner
        plot = Plot(
            self.allresults,
            parameters=["h0", "cosiota", "psi", "phi0"],
            plottype="corner",
        )
        assert len(plot.results) == len(self.allresults)
        assert plot.plottype == "corner"

        fig = plot.plot()  # create plot
        assert isinstance(fig, Figure) and isinstance(plot.fig, Figure)
        assert len(plot.fig.axes) == 16

        plot.save("nd_corner.png", dpi=150)
        fig.clf()
        fig.close()
        del plot

    def test_credible_interval(self):
        """
        Test credible interval calculation.
        """

        plot = Plot(
            self.allresults["CWInPy"],
            parameters=["h0", "cosiota", "psi", "phi0"],
            plottype="corner",
        )

        with pytest.raises(ValueError):
            # invalid parameter
            _ = plot.credible_interval(parameter="skdg")

        interval = plot.credible_interval("psi")
        assert len(interval) == 2
        assert interval[0] < interval[1]
        assert (
            interval[0] > DEFAULT_BOUNDS["psi"]["low"]
            and interval[1] < DEFAULT_BOUNDS["psi"]["high"]
        )

        # test upper limit
        ul = plot.upper_limit("h0")

        assert isinstance(ul, float)
        assert ul > 0.0

        plot = Plot(
            self.allresults,
            parameters=["h0", "cosiota", "psi", "phi0"],
            plottype="corner",
        )
        interval = plot.credible_interval(parameter="cosiota")
        assert len(interval) == len(self.allresults)

        for label in self.allresults:
            credint = interval[label]
            assert (
                credint[0] > DEFAULT_BOUNDS["cosiota"]["low"]
                and credint[1] < DEFAULT_BOUNDS["cosiota"]["high"]
            )

        uls = plot.upper_limit("h0")
        for label in self.allresults:
            assert isinstance(uls[label], float)
            assert uls[label] > 0.0

reason: import numpy


text:
"""Functions to help generate Barrier Certificates (Extra).

Written by: The Robotarium Team
Modified by: Zahi Kakish (zmk5)

"""
import time
from typing import Union
from typing import Callable

from cvxopt import matrix
from cvxopt import solvers

import numpy as np

import quadprog as solver2

solvers.options['show_progress'] = False
solvers.options['reltol'] = 1e-2 # was e-2
solvers.options['feastol'] = 1e-2 # was e-4
solvers.options['maxiters'] = 50 # default is 100

def create_robust_barriers(
        max_num_obstacles: int = 100,
        max_num_robots: int = 30,
        d: int = 5,
        wheel_vel_limit: float = 12.5,
        base_length: float = 0.105,
        wheel_radius: float = 0.016,
        projection_distance: float = 0.05,
        gamma: Union[int, float] = 150,
        safety_radius: float = 0.12) -> Callable:
    """Create a robust barrier certificate function."""
    D = np.matrix([
        [wheel_radius/2, wheel_radius/2],
        [-wheel_radius/base_length, wheel_radius/base_length]])
    L = np.matrix([[1, 0], [0, projection_distance]]) * D
    disturb = np.matrix([
        [-d, -d, d, d],
        [-d, d, d, -d]])
    num_disturbs = np.size(disturb[1, :])

    # TODO: Take out 4*max_num_robots?
    max_num_constraints = (max_num_robots**2-max_num_robots)//2 + max_num_robots*max_num_obstacles + 4*max_num_robots
    A = np.matrix(np.zeros([max_num_constraints, 2*max_num_robots]))
    b = np.matrix(np.zeros([max_num_constraints, 1]))
    Os = np.matrix(np.zeros([2, max_num_robots]))
    ps = np.matrix(np.zeros([2, max_num_robots]))
    Ms = np.matrix(np.zeros([2, 2*max_num_robots]))

    def robust_barriers(dxu, x, obstacles):
        """Output robust barrier certificate array."""
        num_robots = np.size(dxu[0,:])

        if obstacles.size != 0:
            num_obstacles = np.size(obstacles[0,:])
        else:
            num_obstacles = 0

        if num_robots < 2:
            temp = 0
        else:
            temp = (num_robots**2-num_robots)//2

        if num_robots == 0:
            return []

        # Generate constraints for barrier certificates based on the size of the safety radius
        num_constraints = temp + num_robots*num_obstacles
        A[0:num_constraints, 0:2*num_robots] = 0
        Os[0, 0:num_robots] = np.cos(x[2, :])
        Os[1, 0:num_robots] = np.sin(x[2, :])
        ps[:, 0:num_robots] = x[0:2, :] + projection_distance*Os[:, 0:num_robots]
        # Ms Real Form
        # Ms[0, 0:2*num_robots:2] = Os[0, 0:num_robots]
        # Ms[0, 1:2*num_robots:2] = -projection_distance*Os[1, 0:num_robots]
        # Ms[1, 1:2*num_robots:2] = projection_distance*Os[0, 0:num_robots]
        # Ms[1, 0:2*num_robots:2] = Os[1, 0:num_robots]
        # Flipped Ms to be able to perform desired matrix multiplication
        Ms[0, 0:2*num_robots:2] = Os[0, 0:num_robots]
        Ms[0, 1:2*num_robots:2] = Os[1, 0:num_robots]
        Ms[1, 1:2*num_robots:2] = projection_distance*Os[0, 0:num_robots]
        Ms[1, 0:2*num_robots:2] = -projection_distance*Os[1, 0:num_robots]
        MDs  = (Ms.T * D).T
        temp = np.copy(MDs[1, 0:2*num_robots:2])
        MDs[1, 0:2*num_robots:2] =  MDs[0, 1:2*num_robots:2]
        MDs[0, 1:2*num_robots:2] = temp

        count = 0

        for i in range(num_robots-1):
            diffs = ps[:,i] - ps[:, i+1:num_robots]
            hs = np.sum(np.square(diffs),0) - safety_radius**2 # 1 by N
            h_dot_is = 2*diffs.T*MDs[:,2*i:2*i+2] # N by 2
            h_dot_js = np.matrix(np.zeros((2,num_robots - (i+1))))
            h_dot_js[0, :] = -np.sum(2*np.multiply(diffs, MDs[:,2*(i+1):2*num_robots:2]), 0)
            h_dot_js[1, :] = -np.sum(2*np.multiply(diffs, MDs[:,2*(i+1)+1:2*num_robots:2]), 0)
            new_constraints = num_robots - i - 1
            A[count:count+new_constraints, (2*i):(2*i+2)] = h_dot_is
            A[range(count, count+new_constraints), range(2*(i+1), 2*num_robots, 2)] = h_dot_js[0, :]
            A[range(count, count+new_constraints), range(2*(i+1)+1, 2*num_robots, 2)] = h_dot_js[1, :]
            b[count:count+new_constraints] = -gamma*(np.power(hs, 3)).T - np.min(h_dot_is*disturb, 1) - np.min(h_dot_js.T*disturb, 1)
            count += new_constraints

        if obstacles.size != 0:
            # Do obstacles
            for i in range(num_robots):
                diffs = (ps[:, i] - obstacles)
                h = np.sum(np.square(diffs),0) - safety_radius**2
                h_dot_i = 2*diffs.T*MDs[:,2*i:2*i+2]
                A[count:count+num_obstacles, (2*i):(2*i+2)] = h_dot_i
                b[count:count+num_obstacles] = -gamma*(np.power(h, 3)).T  - np.min(h_dot_i*disturb, 1)
                count = count + num_obstacles

        # Adding Upper Bounds On Wheels
        A[count:count+2*num_robots, 0:2*num_robots] = -np.eye(2*num_robots)
        b[count:count+2*num_robots] = -wheel_vel_limit
        count += 2*num_robots
        # Adding Lower Bounds on Wheels
        A[count:count+2*num_robots, 0:2*num_robots] = np.eye(2*num_robots)
        b[count:count+2*num_robots] = -wheel_vel_limit
        count += 2*num_robots

        # Solve QP program generated earlier
        L_all = np.kron(np.eye(num_robots), L)
        dxu = np.linalg.inv(D)*dxu # Convert user input to differential drive
        vhat = np.matrix(np.reshape(dxu, (2*num_robots, 1), order='F'))
        H = 2*L_all.T*L_all
        f = np.transpose(-2*np.transpose(vhat)*np.transpose(L_all)*L_all)

        # Alternative Solver
        # start = time.time()
        # vnew2 = solvers.qp(matrix(H), matrix(f), -matrix(A[0:count,0:2*num_robots]), -matrix( b[0:count]))['x'] # , A, b) Omit last 2 arguments since our QP has no equality constraints
        # print('Time Taken by cvxOpt: {} s'.format(time.time() - start))

        start = time.time()
        vnew = solver2.solve_qp(H, -np.squeeze(np.array(f)), A[0:count, 0:2*num_robots].T, np.squeeze(np.array(b[0:count])))[0]
        print('Time Taken by quadprog: {} s'.format(time.time() - start))
        # Initial Guess for Solver at the Next Iteration
        # vnew = quadprog(H, double(f), -A(1:num_constraints,1:2*num_robots), -b(1:num_constraints), [], [], -wheel_vel_limit*ones(2*num_robots,1), wheel_vel_limit*ones(2*num_robots,1), [], opts);
        # Set robot velocities to new velocities
        dxu = np.reshape(vnew, (2, num_robots), order='F')
        dxu = D * dxu

        return dxu

    return robust_barriers


text:
# -*- coding: utf-8 -*-
"""
Created on Thu Apr 29 17:39:01 2021

@author: Amirali
"""

import matplotlib.pyplot as plt
import numpy as np

filePath = "../Data/experiment2.txt"
f = open(filePath, 'r')
try:
    data = f.readlines()
    splitted = data[1].split("[")
    splitted = splitted[1].split("]")
    splitted = splitted[0]
    splitted = splitted.split(", ")
    intSplitted = np.zeros((len(data)-1, len(splitted)), dtype=int)
    for i in range(1, len(data)):
        splitted = data[i].split("[")
        splitted = splitted[1].split("]")
        splitted = splitted[0]
        splitted = splitted.split(", ")
        cntr = 0
        for num in splitted:
            intSplitted[i-1][cntr] = int(num)
            cntr = cntr + 1
    print(data)
    print(data[1])
    print(splitted[0])
    print(type(splitted[0]))
    print(intSplitted)
finally:
    f.close()

MG_NOPP = intSplitted[0]
MG_OPP = intSplitted[1]
NL_NOPP = intSplitted[2]
NL_OPP = intSplitted[3]
QP_NOPP = intSplitted[4]
QP_OPP = intSplitted[5]
p = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]
improvement_MG = np.divide(np.subtract(MG_NOPP, MG_OPP), MG_NOPP)
improvement_NL = np.divide(np.subtract(NL_NOPP, NL_OPP), NL_NOPP)
improvement_QP = np.divide(np.subtract(QP_NOPP, QP_OPP), QP_NOPP)

plt.plot(p, MG_NOPP, "-*", label='MG_NOPP')
plt.plot(p, MG_OPP, "-o", label='MG_OPP')
plt.plot(p, NL_NOPP, "-*", label='NL_NOPP')
plt.plot(p, NL_OPP, "-o", label='NL_OPP')
plt.plot(p, QP_NOPP, "-*", label='QP_NOPP')
plt.plot(p, QP_OPP, "-o", label='QP_OPP')
plt.legend()
plt.savefig('M5N20L6PSW0.8-algos.pdf', format='pdf')
plt.show()
#plt.subplot(1,2,2)
plt.plot(p, improvement_MG, "-*", label='MG')
plt.plot(p, improvement_NL, "-o", label='NL')
plt.plot(p, improvement_QP, "-o", label='QP')
plt.legend()
plt.savefig('M5N20L6PSW0.8-impros.pdf', format='pdf')

BIGGER_SIZE = 10
MARKER_SIZE = 4
LINE_WIDTH = 0.9

plt.rcParams['pdf.fonttype'] = 42
plt.rcParams['ps.fonttype'] = 42
plt.rcParams["font.family"] = "Times New Roman"
plt.rc('font', size=BIGGER_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=BIGGER_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=BIGGER_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=BIGGER_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=BIGGER_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=BIGGER_SIZE)

f, ax = plt.subplots(1,2, figsize=(6,2.5), squeeze=False)

plt.xlabel('Generation probability')
ax[0][0].plot(p, MG_NOPP, "->", label='MG_NOPP', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][0].plot(p, MG_OPP, "-o", label='MG_OPP', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][0].plot(p, NL_NOPP, "->", label='NL_NOPP', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][0].plot(p, NL_OPP, "-o", label='NL_OPP', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][0].plot(p, QP_NOPP, "-*", label='QP_NOPP', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][0].plot(p, QP_OPP, "-o", label='QP_OPP', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][0].set_ylabel('Average total waiting time (slots)')
ax[0][0].set_xlabel('Generation probability')
ax[0][0].legend(loc='best', fancybox=True,frameon=False,framealpha=0.8)
plt.ylabel('Improvement ratio')
plt.xlabel('Generation probability')
ax[0][1].plot(p, improvement_MG, "->", label='MG', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][1].plot(p, improvement_NL, "-o", label='NL', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][1].plot(p, improvement_QP, "-o", label='QP', markersize=MARKER_SIZE, linewidth=LINE_WIDTH)
ax[0][1].legend(loc='best', fancybox=True,frameon=False,framealpha=0.8)
plt.subplots_adjust(wspace=0.35)
plt.subplots_adjust(hspace=0.1)
plt.tight_layout()
plt.savefig('M5N20L6PSW0.8-sideBySide.pdf', format='pdf')
plt.show()

max_stars_repo_name: QuantumLite/QNS-Lite

max_issues_repo_name: QuantumLite/QNS-Lite

max_forks_repo_name: QuantumLite/QNS-Lite


text:
from __future__ import print_function
import os
import sys
print(sys.path)
sys.path.append(os.path.dirname(os.getcwd()))
print(sys.path)

from show3d_balls import showpoints
import argparse
import numpy as np
import torch
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable
from pointnet.dataset import ShapeNetDataset
from pointnet.model import PointNetDenseCls
import matplotlib.pyplot as plt

#showpoints(np.random.randn(2500,3), c1 = np.random.uniform(0,1,size = (2500)))

parser = argparse.ArgumentParser()

parser.add_argument('--model', type=str, default='', help='model path')
parser.add_argument('--idx', type=int, default=0, help='model index')
#parser.add_argument('--dataset', type=str, default='', help='dataset path')
parser.add_argument('--dataset', type=str, default='/home/djq/pointnet.pytorch/shapenetcore_partanno_segmentation_benchmark_v0', help="dataset path")
parser.add_argument('--class_choice', type=str, default='', help='class choice')

# 输出一行状态栏参数如下:
# 				Namespace(class_choice='Airplane', dataset='',
# 							idx=2, model='seg/seg_model_Chair_1.pth')

opt = parser.parse_args()
print(opt)

d = ShapeNetDataset(
    root=opt.dataset,
    class_choice=[opt.class_choice],
    split='test',
    data_augmentation=False)

idx = opt.idx

# print("d:{}".format(d))
print("model %d/%d" % (idx, len(d))) # d代表全部的飞机数量
# model 2/341

point, seg = d[idx] # 模型里的第idx的点云
print(point.size(), seg.size()) # seg代表每一个点的标签
# torch.Size([2500, 3]) torch.Size([2500])

point_np = point.numpy() #将torch转为numpy

# 可视化
cmap = plt.cm.get_cmap("hsv", 10)
cmap = np.array([cmap(i) for i in range(10)])[:, :3]
gt = cmap[seg.numpy() - 1, :]

# 载入模型
state_dict = torch.load(opt.model)
classifier = PointNetDenseCls(k= state_dict['conv4.weight'].size()[0])
classifier.load_state_dict(state_dict)
classifier.eval() #评估

# 点云转置
point = point.transpose(1, 0).contiguous()
print('point.transpose(1, 0).shape: ',point.shape)

point = Variable(point.view(1, point.size()[0], point.size()[1]))
pred, _, _ = classifier(point) #分割
# print('\npred.shape:',pred[0],'\n')
pred_choice = pred.data.max(2)[1]
print(pred_choice) #输出每一个点的预测类别

#print(pred_choice.size())
print(pred_choice.numpy()[0])  #[1 1 1 ... 1 1 1]
pred_color = cmap[pred_choice.numpy()[0], :] #根据分类结果显示颜色
print('\npred_color: ',pred_color.shape,'\n')

#print(pred_color.shape)
showpoints(point_np, gt, pred_color) #pred_colord的为(2500, 3)的矩阵
print(point_np.shape)


text:
import numpy as np
from pytheas import Scatt3D

# import numpy.testing as npt
from testutils import *

def model(verbose=False):
    fem = Scatt3D()
    fem.rm_tmp_dir()
    fem.eps_des = 2 - 0 * 1j
    fem.hx_des = 3
    fem.hy_des = 3
    fem.hz_des = 0.5
    fem.hx_box = fem.hx_des * 1.1
    fem.hy_box = fem.hy_des * 1.1
    fem.hz_box = fem.hz_des * 3
    fem.R_sph = 0.05
    fem.z_sph = -fem.hz_des / 2 * 1.2 - fem.R_sph
    fem.parmesh = 2
    fem.parmesh_des = 2
    fem.parmesh_pml = fem.parmesh * 2 / 3
    fem.parmesh_host = 3
    fem.recomb_des = False
    fem.matprop_pattern = [1.4, 2 - 0.02 * 1j, 3 - 0.01j]  # refractive index values
    if verbose:
        fem.getdp_verbose = 4
        fem.gmsh_verbose = 4
        fem.python_verbose = 1
    fem.initialize()
    fem.make_mesh()

    return fem

def test_scatt3D(verbose=False):
    fem = model(verbose=verbose)
    mat = pattern(xsym=True, ysym=True)
    fem.register_pattern(mat.pattern, mat._threshold_val)
    fem.compute_solution()
    return fem

max_stars_repo_name: benvial/pytheas

max_issues_repo_name: benvial/pytheas

max_forks_repo_name: benvial/pytheas


text:
import os
import sys
import unittest
import matplotlib as mpl
import scikit_posthocs._posthocs as sp
import scikit_posthocs._omnibus as som
import scikit_posthocs._outliers as so
import scikit_posthocs._plotting as splt
import scikit_posthocs._global as spg
import seaborn as sb
import numpy as np
import matplotlib.axes as ma
from pandas import DataFrame
if os.environ.get('DISPLAY', '') == '':
    print('No display found. Using non-interactive Agg backend')
    mpl.use('Agg')
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

class TestPosthocs(unittest.TestCase):

    # Global tests
    def test_global_simes_test(self):
        a = np.array([0.9, 0.1, 0.01, 0.99, 1., 0.02, 0.04])
        result = spg.global_simes_test(a)
        self.assertAlmostEqual(result, 0.07)

    def test_global_f_test(self):
        a = np.array([0.9, 0.1, 0.01, 0.99, 1., 0.02, 0.04])
        result, _ = spg.global_f_test(a)
        self.assertAlmostEqual(result, 0.01294562)

    # Plotting tests
    def test_sign_array(self):

        p_values = np.array([[0., 0.00119517, 0.00278329],
                             [0.00119517, 0., 0.18672227],
                             [0.00278329, 0.18672227, 0.]])
        test_results = splt.sign_array(p_values)
        correct_results = np.array([[1, 1, 1], [1, 1, 0], [1, 0, 1]])
        self.assertTrue(np.all(test_results == correct_results))

    def test_sign_table(self):

        p_values = np.array([[1., 0.00119517, 0.00278329],
                             [0.00119517, 1., 0.18672227],
                             [0.00278329, 0.18672227, 1.]])

        correct_results = np.array([['-', '**', '**'],
                                    ['**', '-', 'NS'],
                                    ['**', 'NS', '-']], dtype=object)
        correct_resultsl = np.array([['-', '', ''],
                                    ['**', '-', ''],
                                    ['**', 'NS', '-']], dtype=object)
        correct_resultsu = np.array([['-', '**', '**'],
                                    ['', '-', 'NS'],
                                    ['', '', '-']], dtype=object)

        with self.assertRaises(ValueError):
            splt.sign_table(p_values, lower=False, upper=False)

        self.assertTrue(np.all(splt.sign_table(p_values, lower=False, upper=True) == correct_resultsu))
        self.assertTrue(np.all(splt.sign_table(p_values, lower=True, upper=False) == correct_resultsl))
        self.assertTrue(np.all(splt.sign_table(p_values, lower=True, upper=True) == correct_results))

    def test_sign_plot(self):

        x = np.array([[1, 1, 1],
                      [1, 1, 0],
                      [1, 0, 1]])
        a = splt.sign_plot(x, flat=True, labels=False)
        with self.assertRaises(ValueError):
            splt.sign_plot(x.astype(float), flat=True, labels=False)
        self.assertTrue(isinstance(a, ma._subplots.Axes))

    def test_sign_plot_nonflat(self):

        x = np.array([[1., 0.00119517, 0.00278329],
                      [0.00119517, 1., 0.18672227],
                      [0.00278329, 0.18672227, 1.]])
        a, cbar = splt.sign_plot(x, cbar=True, labels=False)

        with self.assertRaises(ValueError):
            splt.sign_plot(x, cmap=[1, 1], labels=False)
        with self.assertRaises(ValueError):
            splt.sign_plot(x.astype(np.int64), labels=False)

        self.assertTrue(isinstance(a, ma._subplots.Axes) and isinstance(cbar, mpl.colorbar.ColorbarBase))

    # Outliers tests
    def test_outliers_iqr(self):

        x = np.array([4, 5, 6, 10, 12, 4, 3, 1, 2, 3, 23, 5, 3])

        x_filtered = np.array([4, 5, 6, 10, 4, 3, 1, 2, 3, 5, 3])
        indices = np.delete(np.arange(13), [4, 10])
        outliers_indices = np.array([4, 10])
        outliers = np.array([12, 23])

        test_outliers = so.outliers_iqr(x, ret='outliers')
        test_outliers_indices = so.outliers_iqr(x, ret='outliers_indices')
        test_indices = so.outliers_iqr(x, ret='indices')
        test_filtered = so.outliers_iqr(x, ret='filtered')

        self.assertTrue(np.all(test_outliers == outliers)
                        and np.all(test_outliers_indices == outliers_indices)
                        and np.all(test_indices == indices)
                        and np.all(test_filtered == x_filtered))

    def test_outliers_grubbs(self):

        x = np.array([199.31, 199.53, 200.19, 200.82,
                      201.92, 201.95, 202.18, 245.57])
        test_results = so.outliers_grubbs(x)
        correct_results = np.array(
            [199.31, 199.53, 200.19, 200.82, 201.92, 201.95, 202.18])
        self.assertTrue(so.outliers_grubbs(x, hypo=True))
        self.assertTrue(np.all(test_results == correct_results))

    def test_outliers_tietjen(self):

        x = np.array([-1.40, -0.44, -0.30, -0.24, -0.22, -0.13, -0.05, 0.06, 0.10,
                      0.18, 0.20, 0.39, 0.48, 0.63, 1.01])
        test_results = so.outliers_tietjen(x, 2)
        correct_results = np.array([-0.44, -0.3, -0.24, -0.22, -0.13, -0.05, 0.06,
                                    0.1, 0.18, 0.2, 0.39, 0.48, 0.63])
        self.assertTrue(so.outliers_tietjen(x, 2, hypo=True))
        self.assertTrue(np.all(test_results == correct_results))

    def test_outliers_gesd(self):

        x = np.array([-0.25, 0.68, 0.94, 1.15, 1.2, 1.26, 1.26, 1.34, 1.38,
            1.43, 1.49, 1.49, 1.55, 1.56, 1.58, 1.65, 1.69, 1.7, 1.76, 1.77,
            1.81, 1.91, 1.94, 1.96, 1.99, 2.06, 2.09, 2.1, 2.14, 2.15, 2.23,
            2.24, 2.26, 2.35, 2.37, 2.4, 2.47, 2.54, 2.62, 2.64, 2.9, 2.92,
            2.92, 2.93, 3.21, 3.26, 3.3, 3.59, 3.68, 4.3, 4.64, 5.34, 5.42,
            6.01])
        correct_mask = np.zeros_like(x, dtype=bool)
        correct_mask[-3:] = True
        test_results = so.outliers_gesd(x, 5)
        test_mask_results = so.outliers_gesd(x, 5, hypo=True)
        correct_results = np.array([-0.25, 0.68, 0.94, 1.15, 1.2, 1.26,
            1.26,  1.34,  1.38,  1.43,  1.49,  1.49,  1.55,  1.56,  1.58,
            1.65,  1.69,  1.7 ,  1.76,  1.77,  1.81,  1.91,  1.94,  1.96,
            1.99,  2.06,  2.09,  2.1 ,  2.14,  2.15,  2.23,  2.24,  2.26,
            2.35,  2.37,  2.4 ,  2.47,  2.54,  2.62,  2.64,  2.9 ,  2.92,
            2.92,  2.93,  3.21,  3.26,  3.3 ,  3.59,  3.68,  4.3 ,  4.64])
        self.assertTrue(isinstance(so.outliers_gesd(x, 5, report=True), str))
        self.assertTrue(np.all(test_results == correct_results))
        self.assertTrue(np.all(test_mask_results == correct_mask))

    # Statistical tests
    df = sb.load_dataset("exercise")
    df[df.columns[df.dtypes == 'category']] = df[df.columns[df.dtypes == 'category']].astype(object)
    df_bn = np.array([[4,3,4,4,5,6,3],
                      [1,2,3,5,6,7,7],
                      [1,2,6,4,1,5,1]])

    # DataFrame conversion tests
    def test_convert_to_block_df(self):
        a = np.array([[0, 0, 4],
                       [1, 0, 1],
                       [2, 0, 1],
                       [0, 1, 3],
                       [1, 1, 2],
                       [2, 1, 2],
                       [0, 2, 4],
                       [1, 2, 3],
                       [2, 2, 6],
                       [0, 3, 4],
                       [1, 3, 5],
                       [2, 3, 4],
                       [0, 4, 5],
                       [1, 4, 6],
                       [2, 4, 1],
                       [0, 5, 6],
                       [1, 5, 7],
                       [2, 5, 5],
                       [0, 6, 3],
                       [1, 6, 7],
                       [2, 6, 1]], dtype=float)
        df_a = DataFrame(a, columns=['blk_col', 'grp_col', 'y_col'])

        result = sp.posthoc_nemenyi_friedman(a, y_col=2, group_col=1, block_col=0, melted=True)[0].values
        result2 = sp.posthoc_nemenyi_friedman(self.df_bn)[0].values
        result3 = sp.posthoc_nemenyi_friedman(df_a, y_col='y_col', group_col='grp_col', block_col='blk_col', melted=True)[0].values
        self.assertTrue(np.allclose(result, result2))
        self.assertTrue(np.allclose(result, result3))
        self.assertTrue(np.allclose(result2, result3))

    # Omnibox tests
    def test_osrt(self):
        df = DataFrame(dict(zip(['a', 'b', 'c'], self.df_bn.tolist()))).melt()
        p,_,_ = som.test_osrt(df, val_col='value', group_col='variable')
        result = 0.3157646
        self.assertTrue(np.allclose(p, result, atol=1.e-3))

    def test_durbin(self):
        r_result = np.array([0.205758, 8.468354, 6])
        result = som.test_durbin(self.df_bn)
        self.assertTrue(np.allclose(result, r_result))

    def test_mackwolfe(self):
        x = [[22, 23, 35], [60, 59, 54], [98, 78, 50], [60, 82, 59], [22, 44, 33], [23, 21, 25]]
        result, _ = som.test_mackwolfe(x, p=2)
        self.assertFalse(som.test_mackwolfe(x, p=20))
        self.assertFalse(som.test_mackwolfe(x, p=0))
        self.assertTrue(np.allclose(result, 0.0006812725))

    def test_mackwolfe_nperm(self):
        x = [[22, 23, 35], [60, 59, 54], [98, 78, 50], [60, 82, 59], [22, 44, 33], [23, 21, 25]]
        _, stat = som.test_mackwolfe(x, n_perm=50)
        self.assertTrue(np.allclose(stat, 3.2024699769846983))

    # Post hoc tests
    def test_posthoc_anderson(self):

        r_results = np.array([[1, 1.35079e-02, 8.64418e-09],
                              [1.35079e-02, 1, 1.644534e-05],
                              [8.64418e-09, 1.644534e-05, 1]])

        results = sp.posthoc_anderson(self.df, val_col = 'pulse', group_col = 'kind', p_adjust = 'holm')
        self.assertTrue(np.allclose(results.values, r_results, atol=3.e-3))

    def test_posthoc_conover(self):

        r_results = np.array([[1, 9.354690e-11, 1.131263e-02],
                              [9.354690e-11, 1, 5.496288e-06],
                              [1.131263e-02, 5.496288e-06, 1]])

        results = sp.posthoc_conover(self.df, val_col = 'pulse', group_col = 'kind', p_adjust = 'holm').values
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_dunn(self):

        r_results = np.array([[1, 9.570998e-09, 4.390066e-02],
                              [9.570998e-09, 1, 1.873208e-04],
                              [4.390066e-02, 1.873208e-04, 1]])

        results = sp.posthoc_dunn(self.df, val_col = 'pulse', group_col = 'kind', p_adjust = 'holm').values
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_nemenyi(self):

        r_results = np.array([[1, 2.431833e-08, 1.313107e-01],
                              [2.431833e-08, 1, 4.855675e-04],
                              [1.313107e-01, 4.855675e-04, 1]])

        results = sp.posthoc_nemenyi(self.df, val_col = 'pulse', group_col = 'kind').values
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_nemenyi_tukey(self):

        r_results = np.array([[1, 9.793203e-09, 1.088785e-01],
                              [9.793203e-09, 1, 0.0002789016],
                              [1.088785e-01, 0.0002789016, 1]])

        results = sp.posthoc_nemenyi(self.df, val_col = 'pulse', group_col = 'kind', dist = 'tukey').values
        self.assertTrue(np.allclose(results, r_results, atol=1.e-3))

    def test_posthoc_nemenyi_friedman(self):

        p_results = np.array([[1., 0.9, 0.82163255, 0.9, 0.9, 0.21477876, 0.9],
                              [0.9, 1., 0.87719193, 0.9, 0.9, 0.25967965, 0.9],
                              [0.82163255, 0.87719193, 1., 0.9, 0.9, 0.9, 0.9],
                              [0.9, 0.9, 0.9, 1., 0.9, 0.87719193, 0.9],
                              [0.9, 0.9, 0.9, 0.9, 1., 0.87719193, 0.9],
                              [0.21477876, 0.25967965, 0.9, 0.87719193, 0.87719193, 1., 0.54381888],
                              [0.9, 0.9, 0.9, 0.9, 0.9, 0.54381888, 1.]])
        results = sp.posthoc_nemenyi_friedman(self.df_bn)
        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_conover_friedman(self):

        results = sp.posthoc_conover_friedman(self.df_bn)
        p_results = np.array([[1.000000, 0.9325836, 0.2497915, 0.3203414, 0.3203414, 0.0517417, 0.6136576],
                              [0.9325836, 1.000000, 0.28339156, 0.36072942, 0.36072942, 0.06033626, 0.67344846],
                              [0.2497915, 0.28339156, 1.000000, 0.8657092, 0.8657092, 0.3607294, 0.5026565],
                              [0.3203414, 0.36072942, 0.8657092, 1.000000, 1.000000, 0.2833916, 0.6136576],
                              [0.3203414, 0.36072942, 0.8657092, 1.000000, 1.000000, 0.2833916, 0.6136576],
                              [0.0517417, 0.06033626, 0.3607294, 0.2833916, 0.2833916, 1.000000, 0.1266542],
                              [0.6136576, 0.67344846, 0.5026565, 0.6136576, 0.6136576, 0.1266542, 1.000000]])
        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_conover_friedman_tukey(self):

        results = sp.posthoc_conover_friedman(self.df_bn, p_adjust='single-step')
        p_results = np.array([[1.0000000,    np.nan,    np.nan,    np.nan,    np.nan,     np.nan,   np.nan],
                              [1.0000000,  1.000000,    np.nan,    np.nan,    np.nan,     np.nan,   np.nan],
                              [0.8908193, 0.9212619,  1.000000,    np.nan,    np.nan,     np.nan,   np.nan],
                              [0.9455949, 0.9642182, 0.9999978,  1.000000,    np.nan,     np.nan,   np.nan],
                              [0.9455949, 0.9642182, 0.9999978, 1.0000000,  1.000000,     np.nan,   np.nan],
                              [0.3177477, 0.3686514, 0.9642182, 0.9212619, 0.9212619,   1.000000,   np.nan],
                              [0.9986057, 0.9995081, 0.9931275, 0.9986057, 0.9986057,  0.6553018, 1.000000]])
        p_results[p_results > 0.9] = 0.9
        tri_upper = np.triu_indices(p_results.shape[0], 1)
        p_results[tri_upper] = np.transpose(p_results)[tri_upper]
        np.fill_diagonal(p_results, 1)
        self.assertTrue(np.allclose(results, p_results, atol=3.e-2))

    def test_posthoc_conover_friedman_non_melted(self):

        df = DataFrame(self.df_bn)
        results = sp.posthoc_conover_friedman(df, melted=False)
        p_results = np.array([[1.000000, 0.9325836, 0.2497915, 0.3203414, 0.3203414, 0.0517417, 0.6136576],
                              [0.9325836, 1.000000, 0.28339156, 0.36072942, 0.36072942, 0.06033626, 0.67344846],
                              [0.2497915, 0.28339156, 1.000000, 0.8657092, 0.8657092, 0.3607294, 0.5026565],
                              [0.3203414, 0.36072942, 0.8657092, 1.000000, 1.000000, 0.2833916, 0.6136576],
                              [0.3203414, 0.36072942, 0.8657092, 1.000000, 1.000000, 0.2833916, 0.6136576],
                              [0.0517417, 0.06033626, 0.3607294, 0.2833916, 0.2833916, 1.000000, 0.1266542],
                              [0.6136576, 0.67344846, 0.5026565, 0.6136576, 0.6136576, 0.1266542, 1.000000]])
        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_miller_friedman(self):

        results = sp.posthoc_miller_friedman(self.df_bn)

        p_results = np.array([[1.0, 1.0, 0.9411963, 0.9724396000000001, 0.9724396000000001, 0.4717981, 0.9993864],
                              [1.0, 1.0, 0.9588993, 0.9823818000000001, 0.9823818000000001, 0.5256257, 0.9997869],
                              [0.9411963, 0.9588993, 1.0, 0.9999991, 0.9999991, 0.9823818000000001, 0.9968575999999999],
                              [0.9724396000000001, 0.9823818000000001, 0.9999991, 1.0, 1.0, 0.9588993, 0.9993864],
                              [0.9724396000000001, 0.9823818000000001, 0.9999991, 1.0, 1.0, 0.9588993, 0.9993864],
                              [0.4717981, 0.5256257, 0.9823818000000001, 0.9588993, 0.9588993, 1.0, 0.7803545999999999],
                              [0.9993864, 0.9997869, 0.9968575999999999, 0.9993864, 0.9993864, 0.7803545999999999, 1.0]])

        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_siegel_friedman(self):

        results = sp.posthoc_siegel_friedman(self.df_bn)

        p_results = np.array([[1.000000, 0.92471904, 0.18587673, 0.25683926, 0.25683926, 0.01816302, 0.57075039],
                              [0.92471904, 1.0000000, 0.2193026, 0.2986177, 0.2986177, 0.0233422, 0.6366016],
                              [0.18587673, 0.2193026, 1.0000000, 0.8501067, 0.8501067, 0.2986177, 0.4496918],
                              [0.25683926, 0.2986177, 0.8501067, 1.000000, 1.0000000, 0.2193026, 0.5707504],
                              [0.25683926, 0.2986177, 0.8501067, 1.0000000, 1.0000000, 0.2193026, 0.5707504],
                              [0.01816302, 0.0233422, 0.2986177, 0.2193026, 0.2193026, 1.000000, 0.07260094],
                              [0.57075039, 0.6366016, 0.4496918, 0.5707504, 0.5707504, 0.07260094, 1.000000]])

        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_durbin(self):
        results = sp.posthoc_durbin(self.df_bn, p_adjust = 'holm')

        p_results = np.array([[1.000000, 1.000000, 1.0, 1.0, 1.0, 0.381364, 1.0],
                              [1.000000, 1.000000, 1.0, 1.0, 1.0, 0.444549, 1.0],
                              [1.000000, 1.000000, 1.0, 1.0, 1.0, 1.000000, 1.0],
                              [1.000000, 1.000000, 1.0, 1.0, 1.0, 1.000000, 1.0],
                              [1.000000, 1.000000, 1.0, 1.0, 1.0, 1.000000, 1.0],
                              [0.381364, 0.444549, 1.0, 1.0, 1.0, 1.000000, 1.0],
                              [1.000000, 1.000000, 1.0, 1.0, 1.0, 1.000000, 1.0]])
        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_quade(self):
        results = sp.posthoc_quade(self.df_bn)

        p_results = np.array([[1.00000000, 0.67651326, 0.15432143, 0.17954686, 0.2081421 , 0.02267043, 0.2081421],
                              [ 0.67651326,1.00000000, 0.29595042, 0.33809987, 0.38443835, 0.0494024 , 0.38443835],
                              [ 0.15432143, 0.29595042,1.00000000, 0.92586499, 0.85245022, 0.29595042, 0.85245022],
                              [ 0.17954686, 0.33809987, 0.92586499,1.00000000, 0.92586499, 0.25789648, 0.92586499],
                              [ 0.2081421 , 0.38443835, 0.85245022, 0.92586499,1.00000000, 0.22378308, 1.00000000],
                              [ 0.02267043, 0.0494024 , 0.29595042, 0.25789648, 0.22378308,1.00000000, 0.22378308],
                              [ 0.2081421 , 0.38443835, 0.85245022, 0.92586499, 1.00000000, 0.22378308,1.00000000]])
        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_quade_norm(self):

        results = sp.posthoc_quade(self.df_bn, dist='normal')

        p_results = np.array([[1.00000000, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
                              [ 0.5693540320,1.00000000, np.nan, np.nan, np.nan, np.nan, np.nan],
                              [ 0.0430605548, 0.145913303,1.00000000, np.nan, np.nan, np.nan, np.nan],
                              [ 0.0578705783, 0.184285855, 0.8993796,1.00000000, np.nan, np.nan, np.nan],
                              [ 0.0766885196, 0.229662468, 0.8003530, 0.8993796,1.00000000, np.nan, np.nan],
                              [ 0.0005066018, 0.003634715, 0.1459133, 0.1139777, 0.08782032,1.00000000, np.nan],
                              [ 0.0766885196, 0.229662468, 0.8003530, 0.8993796, 1.00000000, 0.08782032,1.00000000]])
        tri_upper = np.triu_indices(p_results.shape[0], 1)
        p_results[tri_upper] = np.transpose(p_results)[tri_upper]
        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_npm_test(self):
        results = sp.posthoc_npm_test(self.df_bn)

        p_results = np.array([[1., .9, .9],
                              [.9, 1., .9],
                              [.9, .9, 1.]])

        self.assertTrue(np.allclose(results, p_results))

    def test_posthoc_vanwaerden(self):
        r_results = np.array([[1, 1.054709e-02, 6.476665e-11],
                              [1.054709e-02, 1, 4.433141e-06],
                              [6.476665e-11, 4.433141e-06, 1]])

        results = sp.posthoc_vanwaerden(self.df, val_col = 'pulse', group_col = 'kind', p_adjust='holm')
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_dscf(self):
        r_results = np.array([[1, 4.430682e-02, 9.828003e-08],
                              [4.430682e-02, 1, 5.655274e-05],
                              [9.828003e-08, 5.655274e-05, 1]])

        results = sp.posthoc_dscf(self.df, val_col = 'pulse', group_col = 'kind')
        self.assertTrue(np.allclose(results, r_results, atol=0.001))

    def test_posthoc_ttest(self):

        r_results = np.array([[1, 9.757069e-03, 4.100954e-07],
                              [9.757069e-03, 1, 1.556010e-05],
                              [4.100954e-07, 1.556010e-05, 1]])

        results = sp.posthoc_ttest(self.df, val_col = 'pulse', group_col = 'kind', equal_var = False, p_adjust='holm')
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_ttest_pooled(self):

        x = [[1,2,3,5,1], [12,31,54,50,40], [10,12,6,74,11]]
        r_results = np.array([[1, 0.04226866, 0.24706893],
                              [0.04226866, 1, 0.2482456],
                              [0.24706893, 0.2482456, 1]])

        results = sp.posthoc_ttest(x, equal_var = False, p_adjust='holm', pool_sd=True)
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_tukey_hsd(self):

        x = [[1,2,3,4,5], [35,31,75,40,21], [10,6,9,6,1]]
        g = [['a'] * 5, ['b'] * 5, ['c'] * 5]
        results = sp.posthoc_tukey_hsd(np.concatenate(x), np.concatenate(g))
        n_results = np.array([[1, 1, 0], [1, 1, 1], [0, 1, 1]])
        self.assertTrue(np.allclose(n_results, results))

    def test_posthoc_mannwhitney(self):

        r_results = np.array([[1, 3.420508e-08, 1.714393e-02],
                              [3.420508e-08, 1, 1.968352e-05],
                              [1.714393e-02, 1.968352e-05, 1]])

        results = sp.posthoc_mannwhitney(self.df, val_col = 'pulse', group_col = 'kind').values
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_mannwhitney_ndarray(self):

        _x = [[1,2,3,5,1], [12,31,54,50,40], [10,12,6,74,11]]
        x = np.array(_x)
        g = np.repeat([0,1,2], 5)
        nd = np.column_stack((x.ravel(), g))
        xdf = DataFrame(dict(zip(list('abc'), _x))).melt(var_name='groups', value_name='vals')
        results = sp.posthoc_mannwhitney(xdf, val_col = 'vals', group_col = 'groups').values
        nd_results = sp.posthoc_mannwhitney(nd, val_col=0, group_col=1).values
        self.assertTrue(np.allclose(nd_results, results))

    def test_posthoc_wilcoxon(self):

        r_results = np.array([[1, 2.337133e-03, 2.857818e-06],
                              [2.337133e-03, 1, 1.230888e-05],
                              [2.857818e-06, 1.230888e-05, 1]])

        results = sp.posthoc_wilcoxon(self.df.sort_index(), val_col = 'pulse', group_col = 'kind')
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_scheffe(self):

        r_results = np.array([[1., 3.378449e-01, 3.047472e-10],
                              [3.378449e-01, 1., 2.173209e-07],
                              [3.047472e-10, 2.173209e-07, 1.]])

        results = sp.posthoc_scheffe(self.df.sort_index(), val_col = 'pulse', group_col = 'kind')
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_tamhane(self):

        r_results = np.array([[1, 2.898653e-02, 4.100954e-07],
                              [2.898653e-02, 1, 2.333996e-05],
                              [4.100954e-07, 2.333996e-05, 1]])

        results = sp.posthoc_tamhane(self.df.sort_index(), val_col = 'pulse', group_col = 'kind')
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_tamhane_nw(self):

        r_results = np.array([[1, 2.883219e-02, 4.780682e-08],
                              [2.883219e-02, 1, 8.643683e-06],
                              [4.780682e-08, 8.643683e-06, 1]])

        results = sp.posthoc_tamhane(self.df.sort_index(), val_col = 'pulse', group_col = 'kind', welch = False)
        self.assertTrue(np.allclose(results, r_results))

    def test_posthoc_tukey(self):
        r_results = np.array([[1, 3.042955e-01, 4.308631e-10],
                              [3.042955e-01, 1, 9.946571e-08],
                              [4.308631e-10, 9.946571e-08, 1]])

        results = sp.posthoc_tukey(self.df.sort_index(), val_col = 'pulse', group_col = 'kind')
        self.assertTrue(np.allclose(results, r_results, atol=1.e-3))

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

max_stars_repo_name: maximtrp/scikit-posthocs