text:
# @author lucasmiranda42
# encoding: utf-8
# module deepof

"""

Functions and general utilities for the deepof package. See documentation for details

"""

import argparse
import gc
import multiprocessing
import os
from copy import deepcopy
from itertools import combinations, product
from typing import Tuple, Any, List, Union, NewType

import cv2
import networkx as nx
import numpy as np
import pandas as pd
import regex as re
import ruptures as rpt
import tensorflow as tf
from joblib import Parallel, delayed
from scipy.signal import savgol_filter
from sklearn import mixture
from sklearn.feature_selection import VarianceThreshold
from tqdm import tqdm

# DEFINE CUSTOM ANNOTATED TYPES #
project = NewType("deepof_project", Any)
coordinates = NewType("deepof_coordinates", Any)
table_dict = NewType("deepof_table_dict", Any)

# CONNECTIVITY FOR DLC MODELS

def connect_mouse_topview(animal_id=None) -> nx.Graph:
    """

    Creates a nx.Graph object with the connectivity of the bodyparts in the
    DLC topview model for a single mouse. Used later for angle computing, among others.

    Args:
        animal_id (str): if more than one animal is tagged, specify the animal identyfier as a string.

    Returns:
        connectivity (nx.Graph)

    """

    connectivity = {
        "Nose": ["Left_ear", "Right_ear", "Spine_1"],
        "Left_ear": ["Right_ear", "Spine_1"],
        "Right_ear": ["Spine_1"],
        "Spine_1": ["Center", "Left_fhip", "Right_fhip"],
        "Center": ["Left_fhip", "Right_fhip", "Spine_2", "Left_bhip", "Right_bhip"],
        "Spine_2": ["Left_bhip", "Right_bhip", "Tail_base"],
        "Tail_base": ["Tail_1", "Left_bhip", "Right_bhip"],
        "Tail_1": ["Tail_2"],
        "Tail_2": ["Tail_tip"],
    }

    connectivity = nx.Graph(connectivity)

    if animal_id:
        mapping = {
            node: "{}_{}".format(animal_id, node) for node in connectivity.nodes()
        }
        nx.relabel_nodes(connectivity, mapping, copy=False)

    return connectivity

# QUALITY CONTROL AND PREPROCESSING #

def str2bool(v: str) -> bool:
    """

    Returns the passed string as a boolean.

    Args:
        v (str): String to transform to boolean value.

    Returns:
        bool. If conversion is not possible, it raises an error

    """

    if isinstance(v, bool):
        return v  # pragma: no cover
    elif v.lower() in ("yes", "true", "t", "y", "1"):
        return True
    elif v.lower() in ("no", "false", "f", "n", "0"):
        return False
    raise argparse.ArgumentTypeError("Boolean compatible value expected.")

def likelihood_qc(dframe: pd.DataFrame, threshold: float = 0.9) -> np.array:
    """

    Returns a DataFrame filtered dataframe, keeping only the rows entirely above the threshold.

    Args:
        dframe (pandas.DataFrame): DeepLabCut output, with positions over time and associated likelihood.
        threshold (float): Minimum acceptable confidence.

    Returns:
        filt_mask (np.array): mask on the rows of dframe

    """

    Likes = np.array([dframe[i]["likelihood"] for i in list(dframe.columns.levels[0])])
    Likes = np.nan_to_num(Likes, nan=1.0)
    filt_mask = np.all(Likes > threshold, axis=0)

    return filt_mask

def bp2polar(tab: pd.DataFrame) -> pd.DataFrame:
    """

    Returns the DataFrame in polar coordinates.

    Args:
        tab (pandas.DataFrame): Table with cartesian coordinates.

    Returns:
        polar (pandas.DataFrame): Equivalent to input, but with values in polar coordinates.

    """

    tab_ = np.array(tab)
    complex_ = tab_[:, 0] + 1j * tab_[:, 1]
    polar = pd.DataFrame(np.array([abs(complex_), np.angle(complex_)]).T)
    polar.rename(columns={0: "rho", 1: "phi"}, inplace=True)
    return polar

def tab2polar(cartesian_df: pd.DataFrame) -> pd.DataFrame:
    """

    Returns a pandas.DataFrame in which all the coordinates are polar.

    Args:
        cartesian_df (pandas.DataFrame): DataFrame containing tables with cartesian coordinates.

    Returns:
        result (pandas.DataFrame): Equivalent to input, but with values in polar coordinates.

    """

    result = []
    for df in list(cartesian_df.columns.levels[0]):
        result.append(bp2polar(cartesian_df[df]))
    result = pd.concat(result, axis=1)
    idx = pd.MultiIndex.from_product(
        [list(cartesian_df.columns.levels[0]), ["rho", "phi"]],
    )
    result.columns = idx
    result.index = cartesian_df.index
    return result

def compute_dist(
    pair_array: np.array, arena_abs: int = 1, arena_rel: int = 1
) -> pd.DataFrame:
    """

    Returns a pandas.DataFrame with the scaled distances between a pair of body parts.

    Args:
        pair_array (numpy.array): np.array of shape N * 4 containing X, y positions.
        over time for a given pair of body parts.
        arena_abs (int): Diameter of the real arena in cm.
        arena_rel (int): Diameter of the captured arena in pixels.

    Returns:
        result (pd.DataFrame): pandas.DataFrame with the absolute distances between a pair of body parts.

    """

    lim = 2 if pair_array.shape[1] == 4 else 1
    a, b = pair_array[:, :lim], pair_array[:, lim:]
    ab = a - b

    dist = np.sqrt(np.einsum("...i,...i", ab, ab))
    return pd.DataFrame(dist * arena_abs / arena_rel)

def bpart_distance(
    dataframe: pd.DataFrame, arena_abs: int = 1, arena_rel: int = 1
) -> pd.DataFrame:
    """

    Returns a pandas.DataFrame with the scaled distances between all pairs of body parts.

    Args:
        dataframe (pandas.DataFrame): pd.DataFrame of shape N*(2*bp) containing X,y positions
        over time for a given set of bp body parts.
        arena_abs (int): Diameter of the real arena in cm.
        arena_rel (int): Diameter of the captured arena in pixels.

    Returns:
        result (pd.DataFrame): pandas.DataFrame with the absolute distances between all pairs of body parts.

    """

    indexes = combinations(dataframe.columns.levels[0], 2)
    dists = []
    for idx in indexes:
        dist = compute_dist(np.array(dataframe.loc[:, list(idx)]), arena_abs, arena_rel)
        dist.columns = [idx]
        dists.append(dist)

    return pd.concat(dists, axis=1)

def angle(a: np.array, b: np.array, c: np.array) -> np.array:
    """

    Returns a numpy.array with the angles between the provided instances.

    Args:
        a (np.array): 2D positions over time for a body part.
        b (np.array): 2D positions over time for a body part.
        c (np.array): 2D positions over time for a body part.

    Returns:
        ang (np.array): 1D angles between the three-point-instances.

    """

    ba = a - b
    bc = c - b

    cosine_angle = np.einsum("...i,...i", ba, bc) / (
        np.linalg.norm(ba, axis=1) * np.linalg.norm(bc, axis=1)
    )
    ang = np.arccos(cosine_angle)

    return ang

def angle_trio(bpart_array: np.array) -> np.array:
    """

    Returns a numpy.array with all three possible angles between the provided instances.

    Args:
        bpart_array (numpy.array): 2D positions over time for a bodypart.

    Returns:
        ang_trio (numpy.array): 2D all-three angles between the three-point-instances.

    """

    a, b, c = bpart_array
    ang_trio = np.array([angle(a, b, c), angle(a, c, b), angle(b, a, c)])

    return ang_trio

def rotate(
    p: np.array, angles: np.array, origin: np.array = np.array([0, 0])
) -> np.array:
    """

    Returns a 2D numpy.ndarray with the initial values rotated by angles radians.

    Args:
        p (numpy.ndarray): 2D Array containing positions of bodyparts over time.
        angles (numpy.ndarray): Set of angles (in radians) to rotate p with.
        origin (numpy.ndarray): Rotation axis (zero vector by default).

    Returns:
        - rotated (numpy.ndarray): rotated positions over time

    """

    R = np.array([[np.cos(angles), -np.sin(angles)], [np.sin(angles), np.cos(angles)]])

    o = np.atleast_2d(origin)
    p = np.atleast_2d(p)

    rotated = np.squeeze((R @ (p.T - o.T) + o.T).T)

    return rotated

def align_trajectories(data: np.array, mode: str = "all") -> np.array:
    """

    Returns a numpy.array with the positions rotated in a way that the center (0 vector)
    and the body part in the first column of data are aligned with the y axis.

    Args:
        data (numpy.ndarray): 3D array containing positions of body parts over time, where
        shape is N (sliding window instances) * m (sliding window size) * l (features)
        mode (string): Specifies if *all* instances of each sliding window get
        aligned, or only the *center*

    Returns:
        aligned_trajs (np.ndarray): 2D aligned positions over time.

    """

    angles = np.zeros(data.shape[0])
    data = deepcopy(data)
    dshape = data.shape

    if mode == "center":
        center_time = (data.shape[1] - 1) // 2
        angles = np.arctan2(data[:, center_time, 0], data[:, center_time, 1])
    elif mode == "all":
        data = data.reshape(-1, dshape[-1], order="C")
        angles = np.arctan2(data[:, 0], data[:, 1])
    elif mode == "none":
        data = data.reshape(-1, dshape[-1], order="C")
        angles = np.zeros(data.shape[0])

    aligned_trajs = np.zeros(data.shape)

    for frame in range(data.shape[0]):
        aligned_trajs[frame] = rotate(
            data[frame].reshape([-1, 2], order="C"),
            angles[frame],
        ).reshape(data.shape[1:], order="C")

    if mode == "all" or mode == "none":
        aligned_trajs = aligned_trajs.reshape(dshape, order="C")

    return aligned_trajs

def smooth_boolean_array(a: np.array) -> np.array:
    """

    Returns a boolean array in which isolated appearances of a feature are smoothed.

    Args:
        a (numpy.ndarray): Boolean instances.

    Returns:
        a (numpy.ndarray): Smoothened boolean instances.

    """

    for i in range(1, len(a) - 1):
        if a[i - 1] == a[i + 1]:
            a[i] = a[i - 1]
    return a == 1

def split_with_breakpoints(a: np.ndarray, breakpoints: list) -> np.ndarray:
    """

    Split a numpy.ndarray at the given breakpoints.

    Args:
        a (np.ndarray): N (instances) * m (features) shape
        breakpoints (list): list of breakpoints obtained with ruptures

    Returns:
        split_a (np.ndarray): N (instances) * l (maximum break length) * m (features) shape

    """
    rpt_lengths = list(np.array(breakpoints)[1:] - np.array(breakpoints)[:-1])
    max_rpt_length = np.max([breakpoints[0], np.max(rpt_lengths)])

    # Reshape experiment data according to extracted ruptures
    split_a = np.split(np.expand_dims(a, axis=0), breakpoints[:-1], axis=1)

    split_a = [
        np.pad(
            i,
            ((0, 0), (0, max_rpt_length - i.shape[1]), (0, 0)),
            constant_values=0.0,
        )
        for i in split_a
    ]
    split_a = np.concatenate(split_a, axis=0)

    return split_a

def rolling_window(
    a: np.array, window_size: int, window_step: int, automatic_changepoints: str = False
) -> np.ndarray:
    """

    Returns a 3D numpy.array with a sliding-window extra dimension.

    Args:
        a (np.ndarray): N (instances) * m (features) shape
        window_size (int): Size of the window to apply
        window_step (int): Step of the window to apply
        automatic_changepoints (str): Changepoint detection algorithm to apply.
        If False, applies a fixed sliding window.

    Returns:
        rolled_a (np.ndarray): N (sliding window instances) * l (sliding window size) * m (features)

    """

    breakpoints = None

    if automatic_changepoints:
        # Define change point detection model using ruptures
        # Remove dimensions with low variance (occurring when aligning the animals with the y axis)
        rpt_model = rpt.KernelCPD(
            kernel=automatic_changepoints, min_size=window_size, jump=window_step
        ).fit(VarianceThreshold(threshold=1e-3).fit_transform(a))

        # Extract change points from current experiment
        breakpoints = rpt_model.predict(pen=4.0)
        rolled_a = split_with_breakpoints(a, breakpoints)

    else:
        shape = (a.shape[0] - window_size + 1, window_size) + a.shape[1:]
        strides = (a.strides[0],) + a.strides
        rolled_a = np.lib.stride_tricks.as_strided(
            a, shape=shape, strides=strides, writeable=True
        )[:window_step]

    return rolled_a, breakpoints

def rupture_per_experiment(
    table_dict: table_dict,
    to_rupture: np.ndarray,
    rupture_indices: list,
    automatic_changepoints: str,
    window_size: int,
    window_step: int,
) -> np.ndarray:
    """

    Apply the rupture method independently to each experiment, and concatenate into a single dataset
    at the end. Returns a dataset and the rupture indices, adapted to be used in a concatenated version
    of the labels.

    Args:
        table_dict (deepof.data.table_dict): table_dict with all experiments.
        to_rupture (np.ndarray): Array with dataset to rupture.
        rupture_indices (list): Indices of tables to rupture. Useful to select training and test sets.
        automatic_changepoints (str): Rupture method to apply.
        If false, a sliding window of window_length * window_size is obtained.
        If one of "l1", "l2" or "rbf", different automatic change point detection algorithms are applied
        on each independent experiment.
        window_size (int): If automatic_changepoints is False, specifies the length of the sliding window.
        If not, it determines the minimum size of the obtained time series breaks.
        window_step (int): If automatic_changepoints is False, specifies the stride of the sliding window.
        If not, it determines the minimum step size of the obtained time series breaks.

    Returns:
        ruptured_dataset (np.ndarray): Dataset with all ruptures concatenated across the first axis.
        rupture_indices (list): Indices of ruptures.

    """

    # Generate a base ruptured training set and a set of breaks
    ruptured_dataset, break_indices = None, None
    cumulative_shape = 0
    # Iterate over all experiments and populate them
    for i, tab in enumerate(table_dict.values()):
        if i in rupture_indices:
            current_size = tab.shape[0]
            current_train, current_breaks = rolling_window(
                to_rupture[cumulative_shape : cumulative_shape + current_size],
                window_size,
                window_step,
                automatic_changepoints,
            )
            # Add shape of the current tab as the last breakpoint,
            # to avoid skipping breakpoints between experiments
            if current_breaks is not None:
                current_breaks = np.array(current_breaks) + cumulative_shape
                cumulative_shape += current_size

            try:  # pragma: no cover
                # To concatenate the current ruptures with the ones obtained
                # until now, pad the smallest to the length of the largest
                # alongside axis 1 (temporal dimension) with zeros.
                if ruptured_dataset.shape[1] >= current_train.shape[1]:
                    current_train = np.pad(
                        current_train,
                        (
                            (0, 0),
                            (0, ruptured_dataset.shape[1] - current_train.shape[1]),
                            (0, 0),
                        ),
                    )
                elif ruptured_dataset.shape[1] < current_train.shape[1]:
                    ruptured_dataset = np.pad(
                        ruptured_dataset,
                        (
                            (0, 0),
                            (0, current_train.shape[1] - ruptured_dataset.shape[1]),
                            (0, 0),
                        ),
                    )

                # Once that's taken care of, concatenate ruptures alongside axis 0
                ruptured_dataset = np.concatenate([ruptured_dataset, current_train])
                if current_breaks is not None:
                    break_indices = np.concatenate([break_indices, current_breaks])
            except (ValueError, AttributeError):
                ruptured_dataset = current_train
                if current_breaks is not None:
                    break_indices = current_breaks

    return ruptured_dataset, break_indices

def smooth_mult_trajectory(
    series: np.array, alpha: int = 0, w_length: int = 11
) -> np.ndarray:
    """

    Returns a smoothed a trajectory using a Savitzky-Golay 1D filter.

    Args:
        series (numpy.ndarray): 1D trajectory array with N (instances)
        alpha (int): 0 <= alpha < w_length; indicates the difference between the degree of the polynomial and the window
        length for the Savitzky-Golay filter used for smoothing. Higher values produce a worse fit, hence more smoothing.
        w_length (int): Length of the sliding window to which the filter fit. Higher values yield a coarser fit,
        hence more smoothing.

    Returns:
        smoothed_series (np.ndarray): smoothed version of the input, with equal shape

    """

    if alpha is None:
        return series

    smoothed_series = savgol_filter(
        series, polyorder=(w_length - alpha), window_length=w_length, axis=0
    )

    assert smoothed_series.shape == series.shape

    return smoothed_series

def moving_average(time_series: pd.Series, lag: int = 5) -> pd.Series:
    """

    Fast implementation of a moving average function.

    Args:
        time_series (pd.Series): Uni-variate time series to take the moving average of.
        lag (int): size of the convolution window used to compute the moving average.

    Returns:
        moving_avg (pd.Series): Uni-variate moving average over time_series.

    """

    moving_avg = np.convolve(time_series, np.ones(lag) / lag, mode="same")

    return moving_avg

def mask_outliers(
    time_series: pd.DataFrame,
    likelihood: pd.DataFrame,
    likelihood_tolerance: float,
    lag: int,
    n_std: int,
    mode: str,
) -> pd.DataFrame:
    """

    Returns a mask over the bivariate trajectory of a body part, identifying as True all detected outliers.
    An outlier can be marked with one of two criteria: 1) the likelihood reported by DLC is below likelihood_tolerance,
    and/or 2) the deviation from a moving average model is greater than n_std.

    Args:
        time_series (pd.DataFrame): Bi-variate time series representing the x, y positions of a single body part
        likelihood (pd.DataFrame): Data frame with likelihood data per body part as extracted from deeplabcut
        likelihood_tolerance (float): Minimum tolerated likelihood, below which an outlier is called
        lag (int): Size of the convolution window used to compute the moving average
        n_std (int): Number of standard deviations over the moving average to be considered an outlier
        mode (str): If "and" (default) both x and y have to be marked in order to call an outlier. If "or", one is enough.

    Returns
        mask (pd.DataFrame): Bi-variate mask over time_series. True indicates an outlier.

    """

    moving_avg_x = moving_average(time_series["x"], lag)
    moving_avg_y = moving_average(time_series["y"], lag)

    residuals_x = time_series["x"] - moving_avg_x
    residuals_y = time_series["y"] - moving_avg_y

    outlier_mask_x = np.abs(residuals_x) > np.mean(
        residuals_x[lag:-lag]
    ) + n_std * np.std(residuals_x[lag:-lag])
    outlier_mask_y = np.abs(residuals_y) > np.mean(
        residuals_y[lag:-lag]
    ) + n_std * np.std(residuals_y[lag:-lag])
    outlier_mask_l = likelihood < likelihood_tolerance
    mask = None

    if mode == "and":
        mask = (outlier_mask_x & outlier_mask_y) | outlier_mask_l
    elif mode == "or":
        mask = (outlier_mask_x | outlier_mask_y) | outlier_mask_l

    return mask

def full_outlier_mask(
    experiment: pd.DataFrame,
    likelihood: pd.DataFrame,
    likelihood_tolerance: float,
    exclude: str,
    lag: int,
    n_std: int,
    mode: str,
) -> pd.DataFrame:
    """

    Iterates over all body parts of experiment, and outputs a dataframe where all x, y positions are
    replaced by a boolean mask, where True indicates an outlier.

    Args:
        experiment (pd.DataFrame): Data frame with time series representing the x, y positions of a every body part
        likelihood (pd.DataFrame): Data frame with likelihood data per body part as extracted from deeplabcut
        likelihood_tolerance (float): Minimum tolerated likelihood, below which an outlier is called
        exclude (str): Body part to exclude from the analysis (to concatenate with bpart alignment)
        lag (int): Size of the convolution window used to compute the moving average
        n_std (int): Number of standard deviations over the moving average to be considered an outlier
        mode (str): If "and" (default) both x and y have to be marked in order to call an outlier. If "or", one is enough.

    Returns:
        full_mask (pd.DataFrame): Mask over all body parts in experiment. True indicates an outlier

    """

    body_parts = experiment.columns.levels[0]
    full_mask = experiment.copy()

    if exclude:
        full_mask.drop(exclude, axis=1, inplace=True)

    for bpart in body_parts:
        if bpart != exclude:

            mask = mask_outliers(
                experiment[bpart],
                likelihood[bpart],
                likelihood_tolerance,
                lag,
                n_std,
                mode,
            )

            full_mask.loc[:, (bpart, "x")] = mask
            full_mask.loc[:, (bpart, "y")] = mask
            continue

    return full_mask

def interpolate_outliers(
    experiment: pd.DataFrame,
    likelihood: pd.DataFrame,
    likelihood_tolerance: float,
    exclude: str = "",
    lag: int = 5,
    n_std: int = 3,
    mode: str = "or",
    limit: int = 10,
) -> pd.DataFrame:
    """

    Marks all outliers in experiment and replaces them using a uni-variate linear interpolation approach.
    Note that this approach only works for equally spaced data (constant camera acquisition rates).

    Args:
        experiment (pd.DataFrame): Data frame with time series representing the x, y positions of a every body part.
        likelihood (pd.DataFrame): Data frame with likelihood data per body part as extracted from deeplabcut.
        likelihood_tolerance (float): Minimum tolerated likelihood, below which an outlier is called.
        exclude (str): Body part to exclude from the analysis (to concatenate with bpart alignment).
        lag (int): Size of the convolution window used to compute the moving average.
        n_std (int): Number of standard deviations over the moving average to be considered an outlier.
        mode (str): If "and" both x and y have to be marked in order to call an outlier. If "or" (default), one is enough.
        limit (int): Maximum of consecutive outliers to interpolate. Defaults to 10.

    Returns:
        interpolated_exp (pd.DataFrame): Interpolated version of experiment.

    """

    interpolated_exp = experiment.copy()

    mask = full_outlier_mask(
        experiment, likelihood, likelihood_tolerance, exclude, lag, n_std, mode
    )

    interpolated_exp[mask] = np.nan
    interpolated_exp.interpolate(
        method="linear", limit=limit, limit_direction="both", inplace=True
    )
    # Add original frames to what happens before lag
    interpolated_exp = pd.concat(
        [experiment.iloc[:lag, :], interpolated_exp.iloc[lag:, :]]
    )

    return interpolated_exp

def filter_columns(columns: list, selected_id: str) -> list:
    """

    Given a set of TableDict columns, returns those that correspond to a given animal, specified in selected_id.

    Args:
        columns (list): List of columns to filter.
        selected_id (str): Animal ID to filter for.

    Returns:
        filtered_columns (list): List of filtered columns.

    """

    columns_to_keep = []
    for column in columns:
        # Speed transformed columns
        if type(column) == str and column.startswith(selected_id):
            columns_to_keep.append(column)
        # Raw coordinate columns
        if column[0].startswith(selected_id) and column[1] in [
            "x",
            "y",
            "rho",
            "phi",
        ]:
            columns_to_keep.append(column)
        # Raw distance and angle columns
        elif len(column) in [2, 3] and all([i.startswith(selected_id) for i in column]):
            columns_to_keep.append(column)

    return columns_to_keep

# noinspection PyUnboundLocalVariable
def recognize_arena(
    videos: list,
    vid_index: int,
    path: str = ".",
    tables: dict = None,
    recoglimit: int = 500,
    high_fidelity: bool = False,
    arena_type: str = "circular",
    detection_mode: str = "rule-based",
    cnn_model: tf.keras.models.Model = None,
) -> Tuple[np.array, int, int]:
    """

    Returns numpy.array with information about the arena recognised from the first frames
    of the video. WARNING: estimates won't be reliable if the camera moves along the video.

    Args:
        videos (list): Relative paths of the videos to analise.
        vid_index (int): Element of videos list to use.
        path (str): Full path of the directory where the videos are.
        tables (dict): Dictionary with DLC time series in DataFrames as values.
        recoglimit (int): Number of frames to use for position estimates.
        high_fidelity (bool): If True, runs arena recognition on the whole video. Slow, but
        potentially more accurate in poor lighting conditions.
        arena_type (string): Arena type; must be one of ['circular'].
        detection_mode (str): Algorithm to use to detect the arena. "cnn" uses a
        pretrained model based on ResNet50 to predict the ellipse parameters from
        the image. "rule-based" (default) uses a simpler (and faster) image segmentation approach.
        cnn_model (tf.keras.models.Model): Model to use if detection_mode=="cnn".

    Returns:
        arena (np.ndarray): 1D-array containing information about the arena.
        "circular" (3-element-array) -> x-y position of the center and the radius.
        h (int): Height of the video in pixels.
        w (int): Width of the video in pixels.

    """

    cap = cv2.VideoCapture(os.path.join(path, videos[vid_index]))

    if high_fidelity:
        recoglimit = int(1e10)  # set recoglimit to a very big value

    if tables is not None:
        # Select relevant table to check animal positions; if animals are close to the arena, do not take those frames
        # into account
        centers = tables[list(tables.keys())[vid_index]].iloc[:recoglimit, :]

        # Fix the edge case where there are less frames than the minimum specified for recognition
        recoglimit = np.min([recoglimit, centers.shape[0]])

        # Select animal centers
        centers = centers.loc[
            :, [bpart for bpart in centers.columns if "Tail" not in bpart[0]]
        ]
        centers_shape = centers.shape

    # Loop over the first frames in the video to get resolution and center of the arena
    arena, fnum, h, w = None, 0, None, None

    while cap.isOpened() and fnum < recoglimit:
        ret, frame = cap.read()
        # if frame is read correctly ret is True
        if not ret:  # pragma: no cover
            print("Can't receive frame (stream end?). Exiting ...")
            break

        if arena_type == "circular":

            # Detect arena and extract positions
            temp_center, temp_axes, temp_angle = circular_arena_recognition(
                frame, detection_mode=detection_mode, cnn_model=cnn_model
            )
            temp_arena = np.array([[*temp_center, *temp_axes, temp_angle]])

            # Set if not assigned, else concat and return the median
            if arena is None:
                arena = temp_arena
            else:
                arena = np.concatenate([arena, temp_arena], axis=0)

            if h is None and w is None:
                w, h = frame.shape[0], frame.shape[1]

        fnum += 1

    cap.release()
    cv2.destroyAllWindows()

    # Compute the distance between animal centers and the center of the video, for
    # the arena to be based on frames which minimize obstruction of its borders
    if tables is not None:
        center_distances = np.nanmax(
            np.linalg.norm(
                centers.to_numpy().reshape(-1, 2) - (w / 2, h / 2), axis=1
            ).reshape(-1, centers_shape[1] // 2),
            axis=1,
        )
        # Within the frame recognition limit, only the 1% less obstructed will contribute to the arena
        # fitting
        center_quantile = np.quantile(center_distances, 0.05)
        arena = arena[center_distances < center_quantile]
        weights = 1 / center_distances[center_distances < center_quantile]

    # Compute the median across frames and return to tuple format for downstream compatibility
    arena = np.average(arena[~np.any(np.isnan(arena), axis=1)], axis=0)
    arena = (tuple(arena[:2].astype(int)), tuple(arena[2:4].astype(int)), arena[4])

    return arena, h, w

def circular_arena_recognition(
    frame: np.array,
    detection_mode: str = "rule-based",
    cnn_model: tf.keras.models.Model = None,
) -> np.array:
    """

    Returns x,y position of the center, the lengths of the major and minor axes,
    and the angle of the recognised arena.

    Args:
        frame (np.ndarray): numpy.ndarray representing an individual frame of a video
        detection_mode (str): Algorithm to use to detect the arena. "cnn" uses a
        pretrained model based on ResNet50 to predict the ellipse parameters from
        the image. "rule-based" uses a simpler (and faster) image segmentation approach.
        cnn_model (tf.keras.models.Model): Model to use if detection_mode=="cnn".

    Returns:
        circles (np.ndarray): 3-element-array containing x,y positions of the center
        of the arena, and a third value indicating the radius.

    """

    if detection_mode == "rule-based":

        # Convert image to grayscale, threshold it and close it with a 5x5 kernel
        kernel = np.ones((5, 5))
        gray_image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        ret, thresh = cv2.threshold(gray_image, 255 // 4, 255, 0)
        for _ in range(5):
            thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)

        # Obtain contours from the image, and retain the largest one
        cnts, _ = cv2.findContours(
            thresh.astype(np.int64), cv2.RETR_FLOODFILL, cv2.CHAIN_APPROX_TC89_KCOS
        )
        main_cnt = np.argmax([len(c) for c in cnts])

        # Detect the main ellipse containing the arena
        ellipse_params = cv2.fitEllipse(cnts[main_cnt])

        # Parameters to return
        center_coordinates = tuple([int(i) for i in ellipse_params[0]])
        axes_length = tuple([int(i) // 2 for i in ellipse_params[1]])
        ellipse_angle = ellipse_params[2]

    elif detection_mode == "cnn":

        input_shape = tuple(cnn_model.input.shape[1:-1])
        image_temp = cv2.resize(frame, input_shape)
        image_temp = image_temp / 255

        # Detect the main ellipse containing the arena
        predicted_arena = cnn_model.predict(image_temp[np.newaxis, :])[0]

        # Parameters to return
        center_coordinates = tuple(
            (predicted_arena[:2] * frame.shape[:2][::-1] / input_shape).astype(int)
        )
        axes_length = tuple(
            (predicted_arena[2:4] * frame.shape[:2][::-1] / input_shape).astype(int)
        )
        ellipse_angle = predicted_arena[4]

    else:
        raise ValueError(
            "Invalid detection mode. Select between 'cnn' and 'rule-based'"
        )

    # noinspection PyUnboundLocalVariable
    return center_coordinates, axes_length, ellipse_angle

def rolling_speed(
    dframe: pd.DatetimeIndex,
    window: int = 3,
    rounds: int = 3,
    deriv: int = 1,
    center: str = None,
    shift: int = 2,
    typ: str = "coords",
) -> pd.DataFrame:
    """

    Returns the average speed over n frames in pixels per frame.

    Args:
        dframe (pandas.DataFrame): Position over time dataframe.
        window (int): Number of frames to average over.
        rounds (int): Float rounding decimals.
        deriv (int): Position derivative order; 1 for speed, 2 for acceleration, 3 for jerk, etc.
        center (str): For internal usage only; solves an issue with pandas.MultiIndex that arises when centering frames
        to a specific body part.
        shift (int): Window shift for rolling speed calculation.
        typ (str): Type of dataset. Intended for internal usage only.

    Returns:
        speeds (pd.DataFrame): Data frame containing 2D speeds for each body part in the original data or their
        consequent derivatives.

    """

    original_shape = dframe.shape
    if center:
        body_parts = [bp for bp in dframe.columns.levels[0] if bp != center]
    else:
        try:
            body_parts = dframe.columns.levels[0]
        except AttributeError:
            body_parts = dframe.columns

    speeds = pd.DataFrame

    for der in range(deriv):
        features = 2 if der == 0 and typ == "coords" else 1

        distances = (
            np.concatenate(
                [
                    np.array(dframe).reshape([-1, features], order="C"),
                    np.array(dframe.shift(shift)).reshape([-1, features], order="C"),
                ],
                axis=1,
            )
            / shift
        )

        distances = np.array(compute_dist(distances))
        distances = distances.reshape(
            [
                original_shape[0],
                (original_shape[1] // 2 if typ == "coords" else original_shape[1]),
            ],
            order="C",
        )
        distances = pd.DataFrame(distances, index=dframe.index)
        speeds = np.round(distances.rolling(window).mean(), rounds)
        dframe = speeds

    speeds.columns = body_parts

    return speeds.fillna(0.0)

# MACHINE LEARNING FUNCTIONS #

def gmm_compute(x: np.array, n_components: int, cv_type: str) -> list:
    """

    Fits a Gaussian Mixture Model to the provided data and returns evaluation metrics.

    Args:
        x (numpy.ndarray): Data matrix to train the model
        n_components (int): Number of Gaussian components to use
        cv_type (str): Covariance matrix type to use.
        Must be one of "spherical", "tied", "diag", "full".

    Returns:
        - gmm_eval (list): model and associated BIC for downstream selection.

    """

    gmm = mixture.GaussianMixture(
        n_components=n_components,
        covariance_type=cv_type,
        max_iter=100000,
        init_params="kmeans",
    )
    gmm.fit(x)
    gmm_eval = [gmm, gmm.bic(x)]

    return gmm_eval

def gmm_model_selection(
    x: pd.DataFrame,
    n_components_range: range,
    part_size: int,
    n_runs: int = 100,
    n_cores: int = False,
    cv_types: Tuple = ("spherical", "tied", "diag", "full"),
) -> Tuple[List[list], List[np.ndarray], Union[int, Any]]:
    """

    Runs GMM clustering model selection on the specified X dataframe, outputs the bic distribution per model,
    a vector with the median BICs and an object with the overall best model.

    Args:
        x (pandas.DataFrame): Data matrix to train the models
        n_components_range (range): Generator with numbers of components to evaluate
        n_runs (int): Number of bootstraps for each model
        part_size (int): Size of bootstrap samples for each model
        n_cores (int): Number of cores to use for computation
        cv_types (tuple): Covariance Matrices to try. All four available by default

    Returns:
        - bic (list): All recorded BIC values for all attempted parameter combinations
        (useful for plotting)
        - m_bic(list): All minimum BIC values recorded throughout the process
        (useful for plottinh)
        - best_bic_gmm (sklearn.GMM): Unfitted version of the best found model

    """

    # Set the default of n_cores to the most efficient value
    if not n_cores:
        n_cores = min(multiprocessing.cpu_count(), n_runs)

    bic = []
    m_bic = []
    lowest_bic = np.inf
    best_bic_gmm = 0

    pbar = tqdm(total=len(cv_types) * len(n_components_range))

    for cv_type in cv_types:

        for n_components in n_components_range:

            res = Parallel(n_jobs=n_cores, prefer="threads")(
                delayed(gmm_compute)(
                    x.sample(part_size, replace=True), n_components, cv_type
                )
                for _ in range(n_runs)
            )
            bic.append([i[1] for i in res])

            pbar.update(1)
            m_bic.append(np.median([i[1] for i in res]))
            if m_bic[-1] < lowest_bic:
                lowest_bic = m_bic[-1]
                best_bic_gmm = res[0][0]

    return bic, m_bic, best_bic_gmm

# RESULT ANALYSIS FUNCTIONS #

def cluster_transition_matrix(
    cluster_sequence: np.array,
    nclusts: int,
    autocorrelation: bool = True,
    return_graph: bool = False,
) -> Tuple[Union[nx.Graph, Any], np.ndarray]:
    """Computes the transition matrix between clusters and the autocorrelation in the sequence.

    Args:
        cluster_sequence (numpy.array): Sequence of cluster assignments.
        nclusts (int): Number of clusters in the sequence.
        autocorrelation (bool): Whether to compute the autocorrelation of the sequence.
        return_graph (bool): Whether to return the transition matrix as an networkx.DiGraph object.

    Returns:
        trans_normed (numpy.ndarray / networkx.Graph): Transition matrix as numpy.ndarray or networkx.DiGraph.
        autocorr (numpy.array): If autocorrelation is True, returns a numpy.ndarray with all autocorrelation
        values on cluster assignment.
    """

    # Stores all possible transitions between clusters
    clusters = [str(i) for i in range(nclusts)]
    cluster_sequence = cluster_sequence.astype(str)

    trans = {t: 0 for t in product(clusters, clusters)}
    k = len(clusters)

    # Stores the cluster sequence as a string
    transtr = "".join(list(cluster_sequence))

    # Assigns to each transition the number of times it occurs in the sequence
    for t in trans.keys():
        trans[t] = len(re.findall("".join(t), transtr, overlapped=True))

    # Normalizes the counts to add up to 1 for each departing cluster
    trans_normed = np.zeros([k, k]) + 1e-5
    for t in trans.keys():
        trans_normed[int(t[0]), int(t[1])] = np.round(
            trans[t]
            / (sum({i: j for i, j in trans.items() if i[0] == t[0]}.values()) + 1e-5),
            3,
        )

    # If specified, returns the transition matrix as an nx.Graph object
    if return_graph:
        trans_normed = nx.Graph(trans_normed)

    if autocorrelation:
        cluster_sequence = list(map(int, cluster_sequence))
        autocorr = np.corrcoef(cluster_sequence[:-1], cluster_sequence[1:])
        return trans_normed, autocorr

    return trans_normed

# TODO:
#    - Add center / time in zone to supervised_tagging

reason: import numpy,from scipy,import networkx


text:
#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
=================================================
Classification Accuracy as a Substantive
Quantity of Interest: Measuring Polarization
in Westminster Systems

Andrew Peterson & Arthur Spirling 2017
=================================================

Generate classifier accuracy scores and predictions,
using pre-generated matrices from gen_mats.py.

"""
# Authors:
# Andrew Peterson <ajp502 at nyu dot edu>
# Arthur Spirling < at nyu dot edu>
# License: BSD 3 clause

# run with:
# python sess_estimates.py
# or uncomment sys.argv and use python sess_estimates.py 1 '/foo/datain/' '/foo/input_mats/' '/foo/output/' 'runid'
# arguments: (1) normalize (2) input directory (3) output mat directory (4) run id

import os
import pickle as pickle
import sys
import logging

import pandas as pd
import numpy as np
import re
import string
import itertools
import os.path
import time

import scipy
from scipy.io import mmread
from scipy.sparse import csr_matrix
from sklearn import preprocessing
from sklearn.preprocessing import maxabs_scale
from sklearn.feature_extraction.text import CountVectorizer

from sklearn.linear_model import SGDClassifier
from sklearn.linear_model import PassiveAggressiveClassifier
from sklearn.linear_model import Perceptron
from sklearn.linear_model import LogisticRegression

from sklearn.model_selection import StratifiedKFold
from sklearn.model_selection import train_test_split

from utils import has_regex, prep_year_data, fit_pred_offline_classifiers  # Functions for preparing data, etc.

#---------------------------------------

#---------------------------------------
def randomize_by_member(df):
	""" Randomize labels by member,
	keeping the same proportion, a la
	Gentzkow, et al 2015,
	Measuring Polarization in High-dimensional Data"""
	#mp_ids = list(set(df.memref))
	uniq = df.drop_duplicates('memref')
	mp_ids = uniq.memref
	proportion_cons = np.mean(uniq.y_binary) # (cons = 1)
	cons = [x>(1-proportion_cons) for x in np.random.rand(1,len(mp_ids))]
	cons = [int(x) for x in cons[0] ]
	random_cons = dict(zip(mp_ids, cons))
	df['rand_cons']= df['memref'].map(random_cons)
	return(np.array(df.rand_cons))
#---------------------------------------

#---------------------------------------
def load_X_y(data_in, mats_dir, indx, yrmth, randomize, normalized=1):
	errors = []
	try:
		df = prep_year_data(data_in, yrmth, minlen=40)
		logging.info("length df: %d" % len(df))
	except:
		logging.error("failed getting year-month %s" % yrmth)
		errors.append(yrmth)
	df.index = range(len(df))
	testsize = len(df)/10

	if randomize:
		y = randomize_by_member(df)
	else:
		y = df.y_binary
	if (np.mean(y)==0 or np.mean(y)==1 or len(df)==0):
		logging.warning("no variation in year: %s" % yrmth)
		errors.append(yrmth)
		return(np.NaN, np.NaN)
	if normalized:
		X = mmread(mats_dir + 'topic_aug_mat_normalized_j5_' + str(indx) + '.mtx')
		X = csr_matrix(X)
	else:
		X = mmread(mats_dir + 'topic_aug_mat_' + str(indx) + '.mtx')
	logging.info("Num errors: %d" % len(errors))
	return(X, y)
#---------------------------------------
#
#---------------------------------------
def run_estimates(data_in, mats_dir, results_dir, run_id, sess_indx, randomize=0, normalized=1):
	""" Run classifiers for each Parliamentary session."""
	yearly_stats = {}
	yearly_preds = {}

	setup_time = []
	train_time = []
	starttime = time.time()
	errors = []
	for indx, yrmth in sess_indx.items():
		tick = time.time()
		logging.info("currently running: %s" % yrmth)
		X, y = load_X_y(data_in, mats_dir, indx, yrmth, randomize, normalized)
		st_time = time.time() - tick
		setup_time.append(st_time)
		logging.info("setup time: %d s"  % st_time)

		tick = time.time()
		skf = StratifiedKFold(n_splits=10, shuffle=True, random_state=1234)
		cls_stats = {}
		preds= {}
		foldid = 0
		for train_index, test_index in skf.split(X, y):
			#logging.info("fold: %d" % foldid)
			#logging.info("TRAIN: %s"  train_index)#, "TEST:", test_index)
			X_train, X_test = X[train_index], X[test_index]
			y_train, y_test = y[train_index], y[test_index]
			cls_stats[foldid], preds[foldid] = fit_pred_offline_classifiers(X_train, y_train, X_test, y_test, X)
			foldid += 1
		fit_time = time.time() - tick
		train_time.append(fit_time)
		yearly_stats[indx] = cls_stats
		yearly_preds[indx] = preds

	all_yearstime = time.time() - starttime
	logging.info("%ds required to vectorize and to fit 4 classifiers \t" % all_yearstime)
	pickle.dump(yearly_stats, open(results_dir + "/yearly_stats_rand"+ str(randomize) + "_run" + str(run_id) +".pkl", "wb" ) )
	pickle.dump(yearly_preds, open(results_dir + "/yearly_predictions_rand"+ str(randomize)+ "_run" + str(run_id)  +".pkl", "wb" ) )

#---------------------------------------
# Stats from estimates
#---------------------------------------
def avergage_per_classifier(cls_stats, classifier_names):
	accuracies = {}
	median = {}
	vs = {}
	for classif in classifier_names:
		accs = []
		for fold, stats in cls_stats.items():
			relevant = stats[classif]
			accs.append(relevant['accuracy'])
		accuracies[classif] = np.mean(accs)
		vs[classif] = np.var(accs)
		median[classif] = np.median(accs)
	return(accuracies, median, vs)
#---------------------------------------
#
#---------------------------------------
def stats_from_estimates(yearly_stats, sess_indx, randomize, run_id, results_dir):
	""" """
	classifier_names = ['SAG', 'SGD', 'Perceptron','Passive-Aggressive'] #classifiers.keys()

	rows = []
	for indx, yr in sess_indx.items():
		#logging.info(str(yr))
		try:
			curr = yearly_stats[indx]
			mns, meds, vs = avergage_per_classifier(curr, classifier_names )
			rows.append([indx, yr, mns['SAG'], mns['SGD'], mns['Perceptron'], mns['Passive-Aggressive'],
				meds['SAG'], meds['SGD'], meds['Perceptron'], meds['Passive-Aggressive'],
				vs['SAG'], vs['SGD'], vs['Perceptron'], vs['Passive-Aggressive'] ])
		except:
			logging.error("Error getting stats for: ", str(yr))

	res = pd.DataFrame(data=rows, columns = ['index', 'yrmth',
		'mn_sag','mn_sgd','mn_pcpt','mn_passAgr',
		'md_sag','md_sgd','md_pcpt','md_passAgr',
		'var_sag','var_sgd','var_pcpt','var_passAgr'	])

	res.to_csv(results_dir + '/acc_allmembers_rand' + str(randomize)+ "_run_" + str(run_id)  +".csv", index=False)
#---------------------------------------
#
#---------------------------------------
def main():
	curr_dir = os.getcwd()
	curr_dir = re.sub('/UK_data', '', curr_dir)

	logging.basicConfig(filename= curr_dir + '/log_files/sess_estimates.log',level=logging.INFO,format='%(asctime)s %(lineno)s: %(message)s')
	logging.info('Start.')

	#sess_indx_file = (provide path to full index for full data)
	#sess_indx = pickle.load(open(sess_indx_file, 'rb'))
	sess_indx = {9: '1944-11', 74: '2008-12'}

	#normalize = sys.argv[1] # normalize_X = True
	randomize = 0 # sys.argv[1]
	data_in = curr_dir + "/data/" #sys.argv[2]  # data_in = '/nas/tz/uk/wtopics/'
	mats_dir = curr_dir + "/" #sys.argv[3]
	run_id = "replication" #sys.argv[4]
	results_dir = curr_dir

	run_estimates(data_in, mats_dir, results_dir, run_id, sess_indx, randomize=randomize, normalized=1)
	logging.info("Calculating accuracy stats from estimates")

	yearly_stats = pickle.load(open(results_dir + "/yearly_stats_rand"+ str(randomize) + "_run" + str(run_id) +".pkl", 'rb'))
	stats_from_estimates(yearly_stats, sess_indx, randomize, run_id, results_dir)

if __name__ == "__main__":
	main()

ext: py

reason: import numpy,import scipy,from scipy


text:
import numpy as np
import csv
import os

def CIFAR100sparse2coarse(targets):
    """Convert Pytorch CIFAR100 sparse targets to coarse targets.
    Usage:
        trainset = torchvision.datasets.CIFAR100(path)
        trainset.targets = sparse2coarse(trainset.targets)
    """
    coarse_labels = np.array([4, 1, 14, 8, 0, 6, 7, 7, 18, 3,
                              3, 14, 9, 18, 7, 11, 3, 9, 7, 11,
                              6, 11, 5, 10, 7, 6, 13, 15, 3, 15,
                              0, 11, 1, 10, 12, 14, 16, 9, 11, 5,
                              5, 19, 8, 8, 15, 13, 14, 17, 18, 10,
                              16, 4, 17, 4, 2, 0, 17, 4, 18, 17,
                              10, 3, 2, 12, 12, 16, 12, 1, 9, 19,
                              2, 10, 0, 1, 16, 12, 9, 13, 15, 13,
                              16, 19, 2, 4, 6, 19, 5, 5, 8, 19,
                              18, 1, 2, 15, 6, 0, 17, 8, 14, 13])
    return coarse_labels[targets]

def write_csv(data, filename, write_header=False, mode='a'):
    fieldnames = ['class', 'comfort_level', 'correct']
    with open(filename, mode, encoding='UTF8', newline='') as f:
        writer = csv.writer(f)
        if write_header:
            writer.writerow(fieldnames)
        writer.writerows(data)
        f.flush()
        os.fsync(f.fileno())

def get_nap_params(cfg, model_name, dataset_name):
    layers_to_monitor = []
    for l in cfg[model_name][dataset_name]["layers_to_monitor"]:
        layers_to_monitor.append(tuple(l.values()))

max_stars_repo_name: BartlomiejOlber/od-test

max_issues_repo_name: BartlomiejOlber/od-test

max_forks_repo_name: BartlomiejOlber/od-test

reason: import numpy


text:
from .gsh_functions import hex_eval
from .gsh_functions import cub_eval
from .gsh_functions import tri_eval
from .gsh_functions import hex_basis_info
from .gsh_functions import cub_basis_info
from .gsh_functions import tri_basis_info
from .imag_ffts import _ImagFFTBasis
import numpy as np

class GSHBasis(_ImagFFTBasis):

    r"""
    Discretize a continuous field made up three Euler angles (in radians) used
    to represent crystal orientation into continuous local states using the
    Generalized Spherical Harmonic (GSH) basis. This basis uses the following
    equation to discretize the orientation field.

    .. math::

       \frac{1}{\Delta x} \int_s m(g, x) dx =
       \sum_{l, m, n} m[l, \tilde{m}, n, s] T_l^{\tilde{m}n}(g)

    where the :math:`T_l^{\tilde{m}n}` are GSH basis functions and the
    local state space :math:`H` is mapped into the orthogonal, periodic
    domain of the GSH functions

    The mapping of :math:`H` into some desired periodic domain is done
    automatically in PyMKS by using the `domain` key work argument to
    select the desired crystal symmetry.

    >>> X = np.array([[0.1, 0.2, 0.3],
    ...               [6.5, 2.3, 3.4]])
    >>> gsh_basis = GSHBasis(n_states = [3], domain='hexagonal')
    >>> def test_gsh(x):
    ...     phi = x[:, 1]
    ...     t915 = np.cos(phi)
    ...     return 0.15e2 / 0.2e1 * t915 ** 2 - 0.5e1 / 0.2e1

    >>> assert(np.allclose(np.squeeze(gsh_basis.discretize(X)), test_gsh(X)))

    If you select an invalid crystal symmetry PyMKS will give an error

    >>> gsh_basis = GSHBasis(n_states=[3], domain='squishy')
    ... # doctest: +ELLIPSIS
    Traceback (most recent call last):
    ...
    RuntimeError: invalid crystal symmetry

    >>> gsh_basis = GSHBasis(n_states=[3], domain='hex') # doctest: +ELLIPSIS
    Traceback (most recent call last):
    ...
    RuntimeError: invalid crystal symmetry
    """

    def __init__(self, n_states=15, domain=None):
        """
        Instantiate a `Basis`

        Args:
            n_states (int, array): An array of local states to be used. states
                requested. If an integer is provided, all local states up
                to that number will be used.
            domain (list, optional): indicate the desired crystal symmetry for
                the GSH. Valid choices for symmetry are "hexagonal", "cubic" or
                "triclinic" if no symmetry is desired (not specifying any
                symmetry has the same effect)
        """
        if isinstance(n_states, int):
            n_states = np.arange(n_states)
        if domain is None or domain == 'triclinic':
            self.domain = 'triclinic'
        elif domain == 'hexagonal':
            self.domain = 'hexagonal'
        elif domain == 'cubic':
            self.domain = 'cubic'
        else:
            raise RuntimeError("invalid crystal symmetry")
        full_indx = self._gsh_basis_info()
        super(GSHBasis, self).__init__(n_states=n_states, domain=self.domain)
        self.basis_indices = full_indx[self.n_states, :]

    def check(self, X):
        """Warns the user if Euler angles apear to be defined in degrees
        instead of radians"""
        if (np.min(X) < -90.) or (np.max(X) > 90.):
            Warning("X may be defined in degrees instead of radians")
        if X.shape[-1] != 3:
            raise RuntimeError('X must have 3 angles (in radians) in the ' +
                               'last dimention')

    def _check_shape(self, X_shape, y_shape):
        """
        Checks the shape of the microstructure and response data to
        ensure that they are correct.

        Firstly, the response data "y" must have a dimension to index the
        microstructure instantiation and at least one dimension to index the
        local microstructural information.

        Second, the shape of X and y must me equal except for the last
        dimension of X.

        Finally, the length of the final dimension of X must be 3.
        This is because we assume that Bunge Euler angles are assigned for
        each location in the microstructure
        """
        if not len(y_shape) > 1:
            raise RuntimeError("The shape of y is incorrect.")
        if X_shape[-1] != 3:
            raise RuntimeError("X must have 3 continuous local states " +
                               "(euler angles in radians) in the last axis.")
        if y_shape != X_shape[:-1]:
            raise RuntimeError("The X and y must have the same number of " +
                               "samples and microstructure shape.")

    def _pred_shape(self, X):
        """
        Function to describe the expected output shape of a given
        microstructure X.
        """
        _shape = X.shape[:-1]  # X has Euler angles, while output is scalar
        if len(_shape) < 2:
            _shape = (X.shape[0],) + (X.shape[1] / 3,)
        return _shape

    def discretize(self, X):
        """
        Discretize `X`.

        Args:
            X (ND array): The microstructure, an `(n_samples, n_x, ..., 3)`
                shaped array where `n_samples` is the number of samples,
                `n_x` is the spatial discretization and the last dimension
                contains the Bunge Euler angles in radians.
        Returns:
            Float valued field of of Generalized Spherical Harmonics
            coefficients.

        >>> X = np.array([[0.1, 0.2, 0.3],
        ...               [6.5, 2.3, 3.4]])
        >>> gsh_basis = GSHBasis(n_states = [1])
        >>> def q(x):
        ...     phi1 = x[:, 0]
        ...     phi = x[:, 1]
        ...     phi2 = x[:, 2]
        ...     x_GSH = ((0.3e1 / 0.2e1) * (0.1e1 + np.cos(phi)) *
        ...              np.exp((-1*1j) * (phi1 + phi2)))
        ...     return x_GSH
        >>> assert(np.allclose(np.squeeze(gsh_basis.discretize(X)), q(X)))
        """
        self.check(X)
        self._select_axes(X[..., 0])
        return self._gsh_eval(X)

    def _reshape_feature(self, X, size):
        """
        Helper function used to check the shape of the microstructure,
        and change to appropriate shape.

        Args:
            X: The microstructure, an `(n_samples, n_x, ...)` shaped array
                where `n_samples` is the number of samples and `n_x` is thes
                patial discretization.
            size: the new size of the array

        Returns:
            microstructure with shape (n_samples, size)
        """
        _shape = (X.shape[0],) + size + (3,)
        return X.reshape(_shape)

    def _gsh_basis_info(self):
        """
        Returns a an array with the indices used in the GSH functions for a
        a specific crystal symmetry.
        """
        if self.domain == 'triclinic':
            return tri_basis_info()
        elif self.domain == 'hexagonal':
            return hex_basis_info()
        elif self.domain == 'cubic':
            return cub_basis_info()

    def _gsh_eval(self, X):
        """
        Discretize `X` based on crystal symmetry.

        Args:
            X (ND array): The microstructure, an `(n_samples, n_x, ..., 3)`
                shaped array where `n_samples` is the number of samples,
                `n_x` is the spatial discretization and the last dimension
                contains the Bunge Euler angles in radians.
        Returns:
            Float valued field of of Generalized Spherical Harmonics
            coefficients.
        """
        if self.domain == 'triclinic':
            return tri_eval(X, self.n_states)
        elif self.domain == 'hexagonal':
            return hex_eval(X, self.n_states)
        elif self.domain == 'cubic':
            return cub_eval(X, self.n_states)

max_stars_repo_name: edavgaun/pymks

max_issues_repo_name: edavgaun/pymks

max_forks_repo_name: edavgaun/pymks


text:
"""
follows Gary's IterControl.m in
iris_mt_scratch/egbert_codes-20210121T193218Z-001/egbert_codes/matlabPrototype_10-13-20/TF/classes
"""
import numpy as np

from aurora.transfer_function.regression.helper_functions import rme_beta

class IterControl(object):
    """ """

    def __init__(
        self,
        max_number_of_iterations=10,
        max_number_of_redescending_iterations=2,
        **kwargs,
    ):
        """

        Parameters
        ----------
        max_number_of_iterations: int
            Set to zero for OLS, otherwise, this is how many times the RME
            will refine the estimate.
        max_number_of_redescending_iterations : int
            1 or 2 is fine at most.  If set to zero we ignore the redescend code block.
        tolerance : float
            minimum fractional change in any term in b.  Any smaller change
            than this will trigger convergence
        epsilon : float
            NOT USED: REMOVE
        kwargs

        Class Variables:
        <Specific to Egbert's Robust Regression>
        r0: float
            Effectively infinty for OLS, this controls the point at which residuals
            transition from being penalized by a squared vs a linear function.  The
            units are standard deviations of the residual distribution.  Normally
            1.4, 1.5 or thereabouts
        u0: float
            u0 is a parameter for the redescending portion of the robust regression.
            It is controlled by a double exponential formula (REFERENCE NEEDED).  It
            makes for severe downweighting about u0.  The function is continuous
            "math friendly" (all derivates exist etc) so you can prove theorems about it
            etc.
        </Specific to Egbert's Robust Regression>
        """
        self.number_of_iterations = 0
        self.number_of_redescending_iterations = 0

        self.tolerance = 0.005
        self.epsilon = 1000
        self.max_number_of_iterations = max_number_of_iterations
        self.max_number_of_redescending_iterations = (
            max_number_of_redescending_iterations
        )

        # <regression-M estimator params>
        self.r0 = 1.5
        self.u0 = 2.8
        # </regression-M estimator params>

        # <Additional properties>
        # used sometimes to control one or another of the iterative algorithms
        # These were translated from the matlab code and may be moved in future
        self.return_covariance = True
        self.save_cleaned = False
        self.robust_diagonalize = False
        # </Additional properties>

    def converged(self, b, b0):
        """
        Parameters
        ----------
        b : complex-valued numpy array
            the most recent regression estimate
        b0 : complex-valued numpy array
            The previous regression estimate

        Returns
        -------
        converged: bool
            True of the regression has terminated, False otherwise

        Notes:
        The variable maximum_change finds the maximum amplitude component of the vector
        1-b/b0.  Looking at the formula, one might want to cast this instead as
        1 - abs(b/b0), however, that will be insensitive to phase changes in b,
        which is complex valued.  The way it is coded np.max(np.abs(1 - b / b0)) is
        correct as it stands.
        """

        converged = False
        maximum_change = np.max(np.abs(1 - b / b0))
        tolerance_cond = maximum_change <= self.tolerance
        iteration_cond = self.number_of_iterations >= self.max_number_of_iterations
        if tolerance_cond or iteration_cond:
            converged = True
            # These print statments are not very clear and
            # Should be reworded.
            # if tolerance_cond:
            #    print(
            #        f"Converged Due to MaxChange < Tolerance after "
            #        f" {self.number_of_iterations} of "
            #        f" {self.max_number_of_iterations} iterations"
            #    )
            # elif iteration_cond:
            #    print(
            #        f"Converged Due to maximum number_of_iterations "
            #        f" {self.max_number_of_iterations}"
            #    )
        else:
            converged = False

        return converged

    @property
    def continue_redescending(self):
        maxxed_out = (
            self.number_of_redescending_iterations
            >= self.max_number_of_redescending_iterations
        )
        if maxxed_out:
            return False
        else:
            return True

    @property
    def correction_factor(self):
        """
        TODO: This is an RME specific property.  Suggest move r0, u0 and this method
        into an RME-config class.

        See notes on usage in
        transfer_function.regression.helper_functions.rme_beta

        Returns
        -------
        correction_factor : float
            correction factor used for scaling the residual error_variance
        """
        correction_factor = 1.0 / rme_beta(self.r0)
        return correction_factor

max_stars_repo_name: simpeg/aurora

max_issues_repo_name: simpeg/aurora

max_forks_repo_name: simpeg/aurora


text:
"""
Test the grid masking functions
"""
import numpy as np
import numpy.testing as npt
import xarray as xr
import pytest

from ..mask import distance_mask
from ..coordinates import grid_coordinates

def test_distance_mask():
    "Check that the mask works for basic input"
    region = (0, 5, -10, -4)
    coords = grid_coordinates(region, spacing=1)
    mask = distance_mask((2.5, -7.5), maxdist=2, coordinates=coords)
    true = [
        [False, False, False, False, False, False],
        [False, False, True, True, False, False],
        [False, True, True, True, True, False],
        [False, True, True, True, True, False],
        [False, False, True, True, False, False],
        [False, False, False, False, False, False],
        [False, False, False, False, False, False],
    ]
    assert mask.tolist() == true

def test_distance_mask_projection():
    "Check that the mask works when given a projection"
    region = (0, 5, -10, -4)
    coords = grid_coordinates(region, spacing=1)
    mask = distance_mask(
        (2.5, -7.5),
        maxdist=20,
        coordinates=coords,
        projection=lambda e, n: (10 * e, 10 * n),
    )
    true = [
        [False, False, False, False, False, False],
        [False, False, True, True, False, False],
        [False, True, True, True, True, False],
        [False, True, True, True, True, False],
        [False, False, True, True, False, False],
        [False, False, False, False, False, False],
        [False, False, False, False, False, False],
    ]
    assert mask.tolist() == true

def test_distance_mask_grid():
    "Check that the mask works for grid input"
    region = (0, 5, -10, -4)
    shape = (7, 6)
    east, north = grid_coordinates(region, shape=shape)
    coords = {"easting": east[0, :], "northing": north[:, 0]}
    data_vars = {"scalars": (["northing", "easting"], np.ones(shape))}
    grid = xr.Dataset(data_vars, coords=coords)
    masked = distance_mask((2.5, -7.5), maxdist=2, grid=grid)
    true = [
        [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, 1, 1, np.nan, np.nan],
        [np.nan, 1, 1, 1, 1, np.nan],
        [np.nan, 1, 1, 1, 1, np.nan],
        [np.nan, np.nan, 1, 1, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
        [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    ]
    npt.assert_array_equal(true, masked.scalars.values)

def test_distance_mask_missing_args():
    "Check that the function fails if no coordinates or grid are given"
    with pytest.raises(ValueError):
        distance_mask((2.5, -7.5), maxdist=2)

def test_distance_mask_wrong_shapes():
    "Check that the function fails if coordinates have different shapes"
    coords = np.ones((10, 3)), np.zeros((2, 4))
    with pytest.raises(ValueError):
        distance_mask((2.5, -7.5), maxdist=2, coordinates=coords)

max_issues_repo_name: lheagy/verde

max_forks_repo_name: lheagy/verde


text:
"""
"""
# Installed
from scipy.optimize import linear_sum_assignment as linear_assignment
from scipy.spatial.transform import Rotation as R
from shapely.geometry import Polygon
from numba import jit
import numpy as np

DISTANCES = {}
ASSIGNMENTS = {}

@jit
def S_cov(H, P, R):
	return np.linalg.inv(H @ P @ H.T + R)

def euclidean(trackers, frame, match_idx, **kargs):
	@jit
	def do_jit(a, b):
		N, M = len(a), len(b)
		d = np.empty((N,M))
		for n in range(N):
			d[n] = np.sum((a[n] - b)**2, axis=-1)
		return d

	if match_idx is None:
		M = frame.shape[-1]
		match_idx = slice(M)
	else:
		M = len(match_idx)

	N = len(trackers)
	a = np.empty((N,M))
	b = frame.data[:,match_idx]
	trk_idx = np.empty(N)

	for i, (idx, tracker) in enumerate(trackers.items()):
		a[i] = tracker.x[match_idx,].flatten()
		trk_idx[i] = idx

	return do_jit(a, b), trk_idx

def mahalanobis(trackers, frame, match_idx, **kargs):
	@jit
	def do_jit(a, b, S):
		N, M = len(a), len(b)
		d = np.empty((N,M))
		for n in range(N):
			for m in range(M):
				c = a[n] - b[m]
				d[n,m] = c @ S[n] @ c.T
		return d

	if match_idx is None:
		M = frame.shape[-1]
		match_idx = slice(M)
	else:
		M = len(match_idx)

	N = len(trackers)
	a = np.empty((N,M))
	b = frame.data[:,match_idx]
	S = np.empty((N,M,M))
	trk_idx = np.empty(N)

	for i, (idx, tracker) in enumerate(trackers.items()):
		a[i] = tracker.x[match_idx,].flatten()
		S[i] = spatial.S_cov(tracker.H, tracker.P, tracker.R)[match_idx, match_idx] #tracker.SI[match_idx, match_idx]
		trk_idx[i] = idx

	return do_jit(a, b, S), trk_idx

def bbox_to_poly_2d(bbox):
	if bbox.shape[-1] == 5:
		x, y, l, w, yaw = bbox.T
		r = R.from_euler('zyx', (yaw, 0, 0))
	elif bbox.shape[-1] == 6:
		x, y, l, w, xr, yr = bbox.T
		yaw = vec_to_yaw(xr, yr)
		r = R.from_euler('zyx', (yaw, 0, 0))
	elif bbox.shape[-1] == 8:
		x, y, l, w, xq, yq, zq, wq = bbox.T
		r = R.from_quat((xq, yq, zq, wq))
	else:
		raise ValueError("BBox with dimension {} not understood!".format(bbox.shape[-1]))

	l = l * 0.5
	w = w * 0.5
	corners = ((l, w, 0.0), (l, -w, 0.0), (-l, -w, 0.0), (-l, w, 0.0))
	corners = r.apply(corners)
	corners += (x, y, 0.0)
	return Polygon(corners[:,:2])

def iou_2d(trackers, frame, match_idx, **kargs):
	"""
	"""
	N = len(trackers)
	M = len(frame)
	iou = np.empty((N,M))
	trk_idx = np.empty(N)

	for n, (trk_id, tracker) in enumerate(trackers.items()):
		a = bbox_to_poly_2d(tracker.x[match_idx,].flatten())
		trk_idx[n] = trk_id
		for m, data in enumerate(frame.data):
			b = bbox_to_poly_2d(data[match_idx,])
			inter = a.intersection(b).area
			iou[n,m] = inter / (a.area + b.area - inter)
	return 1.0-iou, trk_idx

@jit
def yaw_to_vec(yaw):
	x = np.cos(yaw)
	y = np.sin(yaw)
	z = 0.0
	return x,y,z

def vec_to_yaw(x, y, z=0):
	pi = np.where(x > 0.0, np.pi, -np.pi)
	with np.errstate(divide='ignore', over='ignore'):
		if x == 0:
			a = np.inf
		else:
			a = y / x
		yaw = np.arctan(a) + (y < 0) * pi
	return yaw

def quat_to_vec(x, y, z, w, order='zyx'):
	r = R.from_quat((x, y, z, w))
	yaw, pitch, roll = r.as_euler(order)
	x = np.cos(yaw)
	y = np.sin(yaw)
	z = np.sin(pitch)
	return x,y,z

def hungarian_match(cost, **kargs):
	indices = linear_assignment(cost)
	am = np.in1d(range(cost.shape[0]), indices[0])
	bm = np.in1d(range(cost.shape[1]), indices[1])
	return indices, am, bm

def greedy_match(cost, **kargs):
	"""
	"""
	matched_indices = []
	N, M = cost.shape
	rank = np.argsort(cost.flatten())
	aidx, bidx = np.unravel_index(rank, cost.shape)
	am = np.zeros(N, dtype=bool)
	bm = np.zeros(M, dtype=bool)
	taken = np.zeros(N*M, dtype=bool)
	for i, (a, b) in enumerate(zip(aidx, bidx)):
		if np.all(am): #all trackers already taken?
			break
		elif np.all(bm): #all detections already taken?
			break
		elif ~am[a] and ~bm[b]:
			taken[i] = True
			am[a] = True #mark as taken
			bm[b] = True
	return (aidx[taken], bidx[taken]), am, bm

DISTANCES['mahalanobis'] = mahalanobis
DISTANCES['euclidean'] = euclidean
DISTANCES['iou_2d'] = iou_2d
ASSIGNMENTS['hungarian'] = hungarian_match
ASSIGNMENTS['greedy'] = greedy_match

# Test
if __name__ == '__main__':
	np.random.seed(0)
	a = np.random.rand(10,3)
	mask = euclidean(a, a) < 0.25
	print(mask)
	x, y = np.nonzero(mask)
	merge = np.zeros(a.shape)
	merge[x] += a[y]
	merge, n = np.unique(merge, return_counts=True, axis=0)
	merge /= n.reshape(-1,1)
	print(merge)

reason: import numpy,from scipy,from numba


text:
# -*- coding: utf-8 -*-
"""sliceoptim module for interaction with Slic3r software and generation of samples batches.
"""

__author__ = "Nils Artiges"
__copyright__ = "Nils Artiges"
__license__ = "apache 2.0"

import logging
from pathlib import Path
from subprocess import CalledProcessError, run
import pandas as pd
import numpy as np
import os

from pandas.core.frame import DataFrame

from sliceoptim import __version__

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from sliceoptim.core import Printer, Filament

log = logging.getLogger("sliceoptim" + __name__)

class Sample:
    """Class to define samples to test slicing on. Connects with Slic3r software with cli.
    """

    def __init__(
        self,
        input_file: str,
        is_first_layer: bool,
        printer: "Printer",
        filament: "Filament",
        config_file: Path = None,
        params: dict = None,
    ) -> None:
        """Instantiates a test Sample, holding a stl file and a set of slicing parameters.
           The generated object can produce corresponding gcode used to build test grids.
           Uses Slic3r as slicing engine.

        Args:
            input_file (str): stl file to slice.
            is_first_layer (bool): Declares if the Sample targets first-layer tests.
            printer (Printer): Printer object holding printer's characteristics.
            filament (Filament): Filament object holding filament's characteristics.
            params ([type], optional): Additional parameters to pass to SLic3r. Defaults to None.
            config_file (str, pathlike): file path for a Slic3r .ini config file.
        """
        # params is a dict of slic3r printing options with corresponding values
        self.input_file = input_file
        self.config_file = config_file
        self.printer = printer
        self.filament = filament
        self.is_first_layer = is_first_layer
        self.temp_output_dir = "/tmp/printer_optim"
        self.quality = None
        self.cost = None
        self.__print_time = None
        self.__params = {}
        if is_first_layer:
            self.set_first_layer_presets()
        else:
            self.set_standard_print_presets()
        if (printer is not None) and (filament is not None):
            self.set_main_presets(
                printer=printer, filament=filament
            )  # extract and set first layer temperature params
        if config_file is not None:
            with open(config_file, "r") as conf_file:
                conf = conf_file.read()
                self.set_param(
                    "first-layer-bed-temperature",
                    int(conf.split("first_layer_bed_temperature = ")[1].split("\n")[0]),
                )
                self.set_param(
                    "first-layer-temperature",
                    int(conf.split("first_layer_temperature = ")[1].split("\n")[0]),
                )
        if params is not None:
            for p, v in params.items():
                self.set_param(p, v)
        pass

    @property
    def output_file(self) -> str:
        """Getter for output file.

        Returns:
            str: Returns the path for the test Sample output file.
        """
        if self.get_param("output") == "default":
            filepath = self.input_file.replace(".stl", ".gcode")
        else:
            filepath = self.get_param("output")
        return filepath

    @output_file.setter
    def output_file(self, value: str) -> None:
        """Setter for output file.

        Args:
            value (str): output file path.
        """
        self.set_param("output", value)
        log.debug("Sample output file set to {}".format(value))
        pass

    @property
    def print_time(self) -> float:
        """Getter for Sample print time estimate.

        Returns:
            [float]: Returns estimated print time.
        """
        if self.__print_time is None:
            gcode = self.gcode  # print time is updated at gcode generation
        log.debug("Estimated print time for Sample is {} min".format(self.__print_time))
        return self.__print_time

    @print_time.setter
    def print_time(self, value):
        """[summary]

        Args:
            value ([type]): No print_time value can be set (it is estimated internally).

        Raises:
            AttributeError: Print time can not be set directly.
        """
        raise AttributeError("Print time can not be set (it is computed form gcode).")

    def __estimate_print_time_from_gcode(self, gcode: str) -> float:
        """Estimates print time in minutes from given GCODE.

        Args:
            gcode (str): GCODE whose print time is estimated from.

        Returns:
            float: A print time estimated from simple movements decomposition.
        """

        def parse_line(line: str) -> dict:
            """Parses a line of GCODE towards the resulting coordinates.

            Args:
                line (str): GCODE instruction line.

            Returns:
                dict: New coordinates as a dict.
            """
            move = {"G": None, "X": None, "Y": None, "Z": None, "F": None}
            for key, val in move.items():
                if key in line:
                    d = line.split(key)[1].split(" ")[0]
                    try:
                        if key == "G":
                            move[key] = int(d)
                        else:
                            move[key] = float(d)
                    except ValueError:
                        continue
            return move

        lines = gcode.splitlines()

        minutes = 0
        feedrate = 0
        last_state = {"G": 0, "X": 0, "Y": 0, "Z": 0, "F": 0}
        for l in lines:
            move = parse_line(l)
            distance = 0
            if move["G"] == 1:
                if move["F"] is not None:
                    feedrate = move["F"]
                    last_state["F"] = move["F"]
                for c in ["X", "Y", "Z"]:
                    if move[c] is not None and move[c] != last_state[c]:
                        distance = np.abs(move[c] - last_state[c])
                        last_state[c] = move[c]
            if feedrate != 0 and distance != 0:
                minutes += distance / feedrate
        return minutes

    @property
    def gcode(self) -> str:
        """Getter for generated GCODE.

        Returns:
            str: generated GCODE as a string.
        """
        temp_output_file = self.temp_output_dir + "/sample.gcode"
        if not os.path.isdir(self.temp_output_dir):
            os.makedirs(self.temp_output_dir)
        self.write_gcode(temp_output_file)
        with open(temp_output_file) as file:
            gcode = file.read()
            file.close()
        os.remove(temp_output_file)
        self.__print_time = self.__estimate_print_time_from_gcode(gcode=gcode)
        return gcode

    @gcode.setter
    def gcode(self, value):
        """Setter for GCODE. Raises an error because GCODE is internally computed.

        Args:
            value ([type]): gcode cannot be set.

        Raises:
            AttributeError: GCODE cannot be set (only computed).
        """
        raise AttributeError("GCODE cannot be set (only computed).")

    def erase_start_end_gcode(self) -> None:
        """Erases starting en ending GCODE in Slic3r parameters."""
        self.set_param("start-gcode", "")
        self.set_param("end-gcode", "")
        pass

    def set_param(self, name: str, value) -> None:
        """Set a Slic3r parameter for the current Sample.

        Args:
            name (str): Name of the corresponding Slic3r parameter.
            value ([type]): Desired value for the parameter.
        """
        # remove param if none or default
        if (value is None) or (value == "default"):
            if name in self.__params.keys():
                self.__params.pop(name)
        # convert type if value is a float
        elif isinstance(value, float):
            if ("temperature" in name) or ("speed" in name):
                value = int(value)
            self.__params[name] = value
        else:
            self.__params[name] = value
        # reset computed print times and quality rates when a param is changed
        self.__print_time = None
        self.quality = None
        pass

    def get_param(self, name: str) -> str:
        """Get set value for a SLic3r parameter.

        Args:
            name (str): Name of queried parameter.

        Returns:
            str: Parameter value.
        """
        value = "default"
        if name in self.__params.keys():
            value = self.__params[name]
        return value

    def reset_params(self) -> None:
        """Reset all Slic3r parameters."""
        self.__params = {}
        pass

    def write_gcode(self, output_file=None) -> None:
        """Write GCODE to given / default output file

        Args:
            output_file (str, optional): File path to write test sample GCODE. Defaults to None.

        Raises:
            ValueError: Retrived error from Slic3r call.
        """
        cmd = [
            "slic3r",
        ]
        params = self.__params.copy()
        if self.config_file is not None:
            params["load"] = self.config_file
        if output_file is not None:
            params["output"] = output_file
        for key, val in params.items():
            cmd.append("--" + key)
            if val is not None:
                cmd.append(str(val))
        cmd.append(self.input_file)
        res = run(cmd, capture_output=True)
        try:
            res.check_returncode()
        except CalledProcessError:
            raise ValueError(str(res.stderr))
        log.debug("GCODE generated for Sample")
        pass

    def get_info(self) -> dict:
        """Get Sample information from Slic3r parsing.

        Raises:
            ValueError: Error from Slic3r call.

        Returns:
            dict: Sample information from Slic3r parsing.
                Example of returned info:
                {'filename': 'calicat.stl',
                'size_x': 28.535534,
                'size_y': 28.5,
                'size_z': 35.0,
                'min_x': 0.464466,
                'min_y': 0.0,
                'min_z': 0.0,
                'max_x': 29.0,
                'max_y': 28.5,
                'max_z': 35.0,
                'number_of_facets': 876.0,
                'manifold': 'yes',
                'number_of_parts': 1.0,
                'volume': 12501.378906}
        """
        infos = {}
        cmd = ["slic3r", self.input_file, "--info"]
        res = run(cmd, capture_output=True)
        try:
            res.check_returncode()
        except CalledProcessError:
            raise ValueError(str(res.stderr))
        for inf in res.stdout.decode("utf8").split("\n"):
            if "[" in inf:
                infos["filename"] = inf[1:-1]
            elif " = " in inf:
                inf = inf.split(" = ")
                try:
                    infos[inf[0]] = float(inf[1])
                except ValueError:
                    infos[inf[0]] = inf[1]
        return infos

    def set_first_layer_presets(self) -> None:
        """Set Slic3r parameters for first layer case."""
        self.set_param("top-solid-layers", 0)
        self.set_param("bottom-solid-layers", 1)
        self.set_param("fill-density", 0)
        self.set_param("skirts", 0)
        self.set_param("perimeters", 1)
        log.debug("Set first layer presets for Sample.")
        pass

    def set_standard_print_presets(self) -> None:
        """Set Slic3r presets for a standard sample print."""
        self.set_param("top-solid-layers", 2)
        self.set_param("bottom-solid-layers", 2)
        self.set_param("fill-density", 20)
        self.set_param("skirts", 0)
        self.set_param("perimeters", 2)
        log.debug("Set first standard presets for Sample.")
        pass

    def set_main_presets(self, printer: "Printer", filament: "Filament") -> None:
        """Set Slic3r presets derived from pinter and filament parameters.

        Args:
            printer (Printer): Printer object to derive presets from.
            filament (Filament): Filament object to derive presets from.
        """
        self.set_param(
            "bed-shape",
            "0x0,{x}x0,{x}x{y},0x{y}".format(
                x=printer.bed_size[0], y=printer.bed_size[1]
            ),
        )
        self.set_param("min-print-speed", printer.min_speed)
        self.set_param("max-print-speed", printer.max_speed)
        self.set_param("nozzle-diameter", printer.nozzle_diameter)
        # filament related parameters
        self.set_param("filament-diameter", filament.diameter)
        log.debug("Set main presets for Sample.")
        pass

    def get_param_from_gcode(self, gcode: str, param_name: str) -> float:
        """Get a parameter from a GCODE file sliced with Slic3r.

        Args:
            gcode (str): GCODE file to parse, as a string.
            param_name (str): name of the parameter we look for.

        Raises:
            ValueError: Raises an error if the parameter name is not fond in GCODE.

        Returns:
            float: Returns the parameter value, as a float if possible.
        """
        start_str = "; " + param_name + " = "
        end_str = "\n" if param_name != "filament used" else "mm"
        try:
            val = gcode.split(start_str)[1].split(end_str)[0]
        except IndexError:
            raise ValueError("Parameter name could not be found in GCODE.")
        if "%" in val:
            val = float(val.split("%")[0])
            # special cases for parameters in percent
            if param_name == "first_layer_extrusion_width":
                val = self.get_param_from_gcode(gcode, "first_layer_height") * val / 100
        else:
            val = float(val)
        if val == 0:
            # special cases for default values
            if param_name == "extrusion_width":
                val = 1.25 * self.get_param_from_gcode(gcode, "nozzle_diameter")
        return val

class SampleGrid:
    """Class for generation of test sample batches.
    """

    def __init__(
        self,
        sample_input_file: str,
        is_first_layer: bool,
        printer: "Printer",
        filament: "Filament",
        designs: DataFrame,
        output_path: str,
        spacing=0.0,
        sample_default_params={},
        config_file=None,
    ) -> None:
        """Instantiates a SampleGrid object, holding and generating a batch of test samples.

        Args:
            sample_input_file (str): stl file path for the test sample.
            is_first_layer (bool): True if the test targets a first layer optimization.
            printer (Printer): Printer object
            filament (Filament): Filament object.
            designs (DataFrame): DataFrame of design parameters for samples (each row corresponds to a Sample).
                                 This argument is optional (can be generated by SampleGrid later).
            output_path (str): File path to write generated GCODE.
            spacing (float, optional): Absolute spacing between samples in mm. Defaults to 0.0.
            sample_default_params (dict, optional): Default parameters dict for Samples.
            config_file (str, pathlike): file path for a Slic3r .ini config file.
        """
        self.help_path = Path("src/sliceoptim/data/slic3r_help_file.txt")
        self.sample_input_file = sample_input_file
        self.is_first_layer = is_first_layer
        self.spacing = spacing
        self.printer = printer
        self.filament = filament
        self.output_path = output_path
        self.sample_default_params = sample_default_params
        self.config_file = config_file
        self.__default_sample = self.gen_test_sample()
        self.__default_sample_info = self.__default_sample.get_info()
        self.check_designs_input(designs)
        self.__samples_list = []
        self.__designs = designs
        self.init_gcode = (
            "G28 ; home all axes\n"
            "G1 Z5 F5000 ; lift nozzle\n"
            "G29 A; Activate UBL\n"
            "G29 L1; Load the mesh stored in slot 1 (from G29 S1)\n"
            "G29 J; Probe 3 points\n"
            "G1 X2.0 Y2.0 Z0.0 F5000.0 ; go to edge of print area\n"
            "G92 E0 ; Reset extruder length to zero\n"
            "M190 S40; Set and wait for bed temperature\n"
            "M109 S190; Set and wait for hotend temperature\n"
            "G1 Z0.200 F1000.0 ; Go to Start Z position\n"
            "G1 Y62.0 E9.0 F1000.0 ; intro line\n"
            "G1 Y102.0 E21.5 F1000.0 ; intro line\n"
            "G1 E20; small retract\n"
            "G92 E0.0 ; reset extruder distance position\n"
        )
        self.end_gcode = (
            "M104 S0 ; turn off temperature\n"
            "G28 X0 ; home X axis\n"
            "M84 ; disable motors\n"
        )
        log.info("Created new SampleGrid")

    @property
    def designs(self) -> pd.DataFrame:
        """Getter for Sample designs DataFrame.

        Returns:
            pd.DataFrame: DataFrame of design parameters for samples
                          Each row corresponds to a Sample, and column to a parameter.
        """
        return self.__designs

    @designs.setter
    def designs(self, val: pd.DataFrame):
        """Setter for Sample designs DataFrame.

        Args:
            val (pd.DataFrame): DataFrame of design parameters for samples
                                Each row corresponds to a Sample, and column to a parameter.
        """
        self.check_designs_input(val)
        self.__samples_list = []
        self.__designs = val
        pass

    def check_designs_input(self, designs: pd.DataFrame) -> bool:
        """Utility method to check if designs are properly formatted.

        Args:
            designs (pd.DataFrame): designs DataFrame to test.

        Raises:
            TypeError: Not a DataFrame.
            ValueError: Wrong length case.

        Returns:
            bool: True if designs DataFrame is correct.
        """
        max_s = self.grid_shape()["max_samples_count"]
        if not isinstance(designs, pd.DataFrame):
            raise TypeError('"designs" entry must be a DataFrame')
        elif len(designs) > max_s:
            raise ValueError(
                "length of 'designs' DataFrame can not be greater "
                "than maximum number of samples on heat bed ({})".format(max_s)
            )
        log.debug("SampleGrid designs are OK")
        pass

    @property
    def samples_list(self) -> list:
        """List of Samples in SampleGrid.

        Returns:
            list: All samples as a list.
        """
        if not self.__samples_list:
            samples_list = []
            df_coords = self.gen_sample_coordinates()
            for s in df_coords.iterrows():
                sample_params = {}
                s_index = s[0]
                s_x = s[1]["x"]
                s_y = s[1]["y"]
                coords = str(s_x) + "," + str(s_y)
                sample_params["print-center"] = coords
                sample_params["start-gcode"] = ""
                sample_params["end-gcode"] = ""
                for param_name in self.designs.columns:
                    val = self.designs.iloc[s_index][param_name]
                    sample_params[param_name] = val
                samples_list.append(self.gen_test_sample(sample_params=sample_params))
                if s_index < len(df_coords) - 1:
                    next_x = df_coords.iloc[s_index + 1]["x"]
                    next_y = df_coords.iloc[s_index + 1]["y"]
                    temp = samples_list[-1].get_param("first-layer-temperature")
                    bed_temp = samples_list[-1].get_param("first-layer-bed-temperature")
                    inter_gcode = self.inter_gcode(
                        temperature=temp, bed_temperature=bed_temp, x=next_x, y=next_y
                    )
                    samples_list[-1].set_param("end-gcode", inter_gcode)
            self.__samples_list = samples_list
        return self.__samples_list

    @property
    def quality_list(self) -> list:
        """Returns quality of samples as a list.

        Returns:
            list: quality values of samples.
        """
        return [s.quality for s in self.samples_list]

    @quality_list.setter
    def quality_list(self, val: list) -> None:
        """Sets quality values of samples from a list

        Args:
            val (list): list of samples quality values
        """
        if len(val) != len(self.samples_list):
            raise ValueError(
                "Dimension mismatch between quality values list and samples number."
            )
        for i, s in enumerate(self.samples_list):
            s.quality = val[i]
        pass

    # TODO : implement cost list and interface with dataframe exports/imports
    @property
    def cost_list(self) -> list:
        """Returns quality of samples as a list.

        Returns:
            list: quality values of samples.
        """
        return [s.cost for s in self.samples_list]

    @quality_list.setter
    def cost_list(self, val: list) -> None:
        """Sets quality values of samples from a list

        Args:
            val (list): list of samples quality values
        """
        if len(val) != len(self.samples_list):
            raise ValueError(
                "Dimension mismatch between quality values list and samples number."
            )
        for i, s in enumerate(self.samples_list):
            s.quality = val[i]
        pass

    def gen_test_sample(self, sample_params: dict = {}) -> Sample:
        """Generates and returns a new test Sample (not added to the SampleGrid).

        Args:
            sample_params (dict, optional): Parameters dict for the new sample. Defaults to {}.

        Returns:
            Sample: Generated Sample object.
        """
        params = self.sample_default_params.copy()
        params.update(sample_params)
        sample = Sample(
            input_file=self.sample_input_file,
            is_first_layer=self.is_first_layer,
            printer=self.printer,
            filament=self.filament,
            params=params,
            config_file=self.config_file,
        )
        log.info("New test Sample generated.")
        return sample

    def inter_gcode(self, temperature, bed_temperature, x, y):
        """Generates intermediate GCODE to be inserted between two Samples.

        Args:
            temperature (float): Temperature of the next Sample.
            bed_temperature (float): Bed temperature of the next Sample.
            x (float): X coordinate of the next Sample.
            y (float): Y coordinate of the next Sample.

        Returns:
            str: Generated intermediate GCODE.
        """
        gcode = "; ### Inter-samples GCODE ### \n"
        height = self.__default_sample_info["size_z"] + 0.2
        gcode += "G1 Y{y} Z{z}\n".format(y=y, z=height)
        gcode += "G1 X{}\n".format(x)
        gcode += "G1 Z0.0\n"
        if isinstance(bed_temperature, int):
            gcode += "M190 R{}\n".format(int(bed_temperature))
        if isinstance(temperature, int):
            gcode += "M109 R{}\n".format(int(temperature))
        log.info("Intermediate GCODE generated.")
        gcode += "; ### End of inter-samples GCODE ### \n"
        return gcode

    def grid_shape(self) -> dict:
        """Compute and returns SampleGrid shape data.

        Returns:
            dict: {"max_samples_count": maximum number of Samples,
                   "n_cols": number of columns (X axis),
                   "n_rows": number of rows (Y axis),
                   "size_x": size of X axis in mm,
                   "size_y": size of Y axis in mm}
        """
        max_x = self.printer.bed_size[0] * 0.85
        max_y = self.printer.bed_size[1] * 0.85
        s_info = self.__default_sample_info
        s_x = s_info["size_x"]
        s_y = s_info["size_y"]
        n_cols = np.floor((max_x + self.spacing) / (s_x + self.spacing))
        n_rows = np.floor((max_y + self.spacing) / (s_y + self.spacing))
        n_samples = 0
        if self.is_first_layer:
            n_samples = n_cols * n_rows
        else:
            n_samples = min(n_cols, n_rows)
        size_x = n_cols * (s_x + self.spacing) - self.spacing
        size_y = n_rows * (s_y + self.spacing) - self.spacing
        return {
            "max_samples_count": int(n_samples),
            "n_cols": int(n_cols),
            "n_rows": int(n_rows),
            "size_x": size_x,
            "size_y": size_y,
        }

    def compute_first_sample_coords(self, grid_shape) -> dict:
        """Computes and returns the coordinates for the first Sample.

        Args:
            grid_shape (dict): SampleGrid shape data as returned by "grid_shape" method.

        Returns:
            dict: "x" and "y" coordinates
        """
        sample_info = self.__default_sample_info
        x = 0.5 * (
            self.printer.bed_size[0] - grid_shape["size_x"] + sample_info["size_x"]
        )
        y = self.printer.bed_size[1] - 0.5 * (
            self.printer.bed_size[1] - grid_shape["size_y"] + sample_info["size_y"]
        )
        coords = {"x": x, "y": y}
        log.info(
            "First Sample coordinates on current SampleGrid are X: {x} and Y: {y}.".format(
                **coords
            )
        )
        return coords

    def gen_sample_coordinates(self) -> pd.DataFrame:
        """Generate all sample coordinates and return them as a DataFrame.

        Returns:
            pd.DataFrame: DataFrame of coordinates with columns "x" and "y".
        """
        # occidental reading disposition
        # length of dataframe
        n_samples = len(self.designs)
        # compute grid shape or diag shape : ToDo !
        grid_shape = self.grid_shape()
        # generate coordinates
        sample_info = self.__default_sample_info
        first_coords = self.compute_first_sample_coords(grid_shape)
        x = np.arange(0, grid_shape["n_cols"], 1)
        y = np.arange(0, grid_shape["n_rows"], 1)
        if self.is_first_layer:
            xx, yy = np.meshgrid(x, y)
        else:  # diagonal disposition
            xx = x
            yy = y
        xx = first_coords["x"] + xx * (sample_info["size_x"] + self.spacing)
        yy = first_coords["y"] - yy * (sample_info["size_y"] + self.spacing)
        # add to dataframe
        coords = {"x": xx.flatten()[0:n_samples], "y": yy.flatten()[0:n_samples]}
        df = pd.DataFrame(coords)
        log.info("Samples coordinates generated.")
        return df

    def write_gcode(self, output_path=None):
        """Writes GCODE to the (predefined) output path.

        Args:
            output_path (str, optional): Output file path to write GCODE on. Defaults to None.
                                         If set to None, writes to default output path.
        """
        if output_path is None:
            output_path = self.output_path
        if os.path.exists(output_path):
            os.remove(output_path)
        output = open(output_path, "a+")  # append mode
        output.write(self.init_gcode)
        for s in self.samples_list:
            output.write(s.gcode)
        output.write(self.end_gcode)
        output.close()
        log.info("SampleGrid GCODE written.")
        pass

max_stars_repo_name: oiesauvage/sliceoptim

max_issues_repo_name: oiesauvage/sliceoptim

max_forks_repo_name: oiesauvage/sliceoptim


text:
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 26 15:03:56 2019

@author: jone
"""
#%%
import numpy as np
import cv2

# 모두 0으로 되어 있는 빈 Canvas(검정색)
img = np.zeros((512, 512, 3), np.uint8)

# Line 그리기
img = cv2.line(img, (0, 0), (511, 511), (255, 0, 0), 5)

# 사각형 그리기
cv2.rectangle(img, (384, 0), (510, 128), (0, 255, 0), 3)

# 원 그리기
cv2.circle(img, (447, 63), 63, (0, 0, 255), -1)

# 타원 그리기
cv2.ellipse(img, (256, 256), (100, 50), 0, 0, 180, 255, -1)

# Polygon 그리기
# 각 꼭지점은 2차원 행렬로 선언
pts = np.array([[10, 5], [20, 30], [70, 20], [50, 10]], np.int32)
# 이미지에 표현하기 위해 3차원 행렬로 변환. 변환이전과 이후의 행렬 갯수는 동일해야 함
# -1은 원본에 해당하는 값을 그대로 유지
pts1 = pts.reshape((-1, 1, 2))
img = cv2.polylines(img, [pts1], True, (0, 255, 255))

# 이미지에 Text 추가
cv2.putText(img, 'OpenCV', (10, 500), cv2.FONT_HERSHEY_SIMPLEX, 4, (255, 255, 255), 2)

cv2.imshow('image', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

max_stars_repo_name: zeroam/TIL

max_issues_repo_name: zeroam/TIL

max_forks_repo_name: zeroam/TIL