text:
"""Predictor implementations."""

import copy
import logging
from abc import ABC, abstractmethod
from time import time
from typing import Callable, List, Tuple, Union

import numpy as np
from scipy.stats import pearsonr

try:
    import sklearn.gaussian_process as skl_gp
    import sklearn.linear_model as skl_lm
    import sklearn.model_selection as skl_ms
    import sklearn.neural_network as skl_nn
except ImportError:
    skl_lm = skl_gp = skl_nn = skl_ms = None

logger = logging.getLogger("ABC.Predictor")

class Predictor(ABC):
    """Generic predictor model class.

    A predictor should define:

    * `fit(x, y, w=None)` to fit the model on a sample of data x and outputs y,
      where x has shape (n_sample, n_feature), and
      y has shape (n_sample, n_out).
      Further, gets as a third argument the sample weights
      if `weight_samples` is set. Not all predictors support this.
    * `predict(X)` to predict outputs of shape (n_out,), where X has shape
      (n_sample, n_feature).

    """

    @abstractmethod
    def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None:
        """Fit the predictor to labeled data.

        Parameters
        ----------
        x: Samples, shape (n_sample, n_feature).
        y: Targets, shape (n_sample, n_out).
        w: Weights, shape (n_sample,).
        """

    @abstractmethod
    def predict(self, x: np.ndarray, normalize: bool = False) -> np.ndarray:
        """Predict outputs using the model.

        Parameters
        ----------
        x:
            Samples, shape (n_sample, n_feature) or (n_feature,).
        normalize:
            Whether outputs should be normalized, or on the original scale.

        Returns
        -------
        y: Predicted targets, shape (n_sample, n_out).
        """

def wrap_fit_log(fit):
    """Wrapper for fit logging."""

    def wrapped_fun(self, x: np.ndarray, y: np.ndarray, w: np.ndarray):
        start_time = time()

        # actual fitting
        ret = fit(self, x, y, w)

        logger.info(f"Fitted {self} in {time() - start_time:.2f}s")
        if self.log_pearson:
            # shape: n_sample, n_out
            y_pred = self.predict(x)
            coeffs = [
                pearsonr(y[:, i], y_pred[:, i])[0]
                for i in range(y_pred.shape[1])
            ]
            logger.info(
                " ".join(
                    [
                        "Pearson correlations:",
                        *[f"{coeff:.3f}" for coeff in coeffs],
                    ]
                ),
            )

        return ret

    return wrapped_fun

class SimplePredictor(Predictor):
    """Wrapper around generic predictor routines."""

    def __init__(
        self,
        predictor,
        normalize_features: bool = True,
        normalize_labels: bool = True,
        joint: bool = True,
        weight_samples: bool = False,
        log_pearson: bool = True,
    ):
        """
        Parameters
        ----------
        predictor:
            Predictor model to use, fulfilling the predictor contract.
        normalize_features:
            Whether to apply z-score normalization to the input data.
        normalize_labels:
            Whether to apply z-score normalization to the parameters.
        joint:
            Whether the predictor learns one model for all targets, or
            separate models per target.
        weight_samples:
            Whether to use importance sampling weights. Not that not all
            predictors may support weighted samples.
        log_pearson:
            Whether to log Pearson correlation coefficients after fitting.
        """
        self.predictor = predictor
        # only used if not joint
        self.single_predictors: Union[List, None] = None

        self.normalize_features: bool = normalize_features
        self.normalize_labels: bool = normalize_labels

        self.joint: bool = joint
        self.weight_samples: bool = weight_samples

        self.log_pearson: bool = log_pearson

        # indices to use
        self.use_ixs: Union[np.ndarray, None] = None

        # z-score normalization coefficients
        self.mean_x: Union[np.ndarray, None] = None
        self.std_x: Union[np.ndarray, None] = None
        self.mean_y: Union[np.ndarray, None] = None
        self.std_y: Union[np.ndarray, None] = None

    @wrap_fit_log
    def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None:
        """Fit the predictor to labeled data.

        Parameters
        ----------
        x: Samples, shape (n_sample, n_feature).
        y: Targets, shape (n_sample, n_out).
        w: Weights, shape (n_sample,).
        """
        # remove trivial features
        self.set_use_ixs(x=x)
        x = x[:, self.use_ixs]

        # normalize features
        if self.normalize_features:
            self.mean_x = np.mean(x, axis=0)
            self.std_x = np.std(x, axis=0)
            x = (x - self.mean_x) / self.std_x

        # normalize labels
        if self.normalize_labels:
            self.mean_y = np.mean(y, axis=0)
            self.std_y = np.std(y, axis=0)
            y = (y - self.mean_y) / self.std_y

        if self.joint:
            # fit a model with all parameters as joint response variables
            if self.weight_samples:
                self.predictor.fit(x, y, w)
            else:
                self.predictor.fit(x, y)
        else:
            # set up predictors
            if self.single_predictors is None:
                n_par = y.shape[1]
                self.single_predictors: List = [
                    copy.deepcopy(self.predictor) for _ in range(n_par)
                ]
            # fit a model for each parameter separately
            for predictor, y_ in zip(self.single_predictors, y.T):
                if self.weight_samples:
                    predictor.fit(x, y_, w)
                else:
                    predictor.fit(x, y_)

    def predict(self, x: np.ndarray, normalize: bool = False) -> np.ndarray:
        """Predict outputs using the model.

        Parameters
        ----------
        x:
            Samples, shape (n_sample, n_feature) or (n_feature,).
        normalize:
            Whether outputs should be normalized, or on the original scale.

        Returns
        -------
        y: Predicted targets, shape (n_sample, n_out).
        """
        # to 2d matrix
        if x.ndim == 1:
            x = x.reshape(1, -1)

        # remove trivial features
        x = x[:, self.use_ixs]

        # normalize features
        if self.normalize_features:
            x = (x - self.mean_x) / self.std_x

        if self.joint:
            y = self.predictor.predict(x)
        else:
            y = [
                predictor.predict(x).flatten()
                for predictor in self.single_predictors
            ]
            y = np.array(y).T

        if not normalize and self.normalize_labels:
            y = y * self.std_y + self.mean_y

        # some predictors may return flat arrays if n_out == 1
        if y.ndim == 1:
            y = y.reshape(-1, 1)

        return y

    def set_use_ixs(self, x: np.ndarray, log: bool = True) -> None:
        """Set feature indices to use.

        Parameters
        ----------
        x: Feature matrix, shape (n_sample, n_feature).
        log: Whether to log.
        """
        # remove trivial features
        self.use_ixs = np.any(x != x[0], axis=0)

        # log omitted indices
        if log and not self.use_ixs.all():
            logger.info(
                "Ignore trivial features "
                f"{list(np.flatnonzero(~self.use_ixs))}"
            )

    def __repr__(self) -> str:
        rep = f"<{self.__class__.__name__} predictor={self.predictor}"
        # print everything that is customized
        if not self.normalize_features:
            rep += f" normalize_features={self.normalize_features}"
        if not self.normalize_labels:
            rep += f" normalize_labels={self.normalize_labels}"
        if not self.joint:
            rep += f" joint={self.joint}"
        if self.weight_samples:
            rep += f" weight_samples={self.weight_samples}"
        return rep + ">"

class LinearPredictor(SimplePredictor):
    """Linear predictor model.

    Based on [#fearnheadprangle2012]_.

    .. [#fearnheadprangle2012]
        Fearnhead, Paul, and Dennis Prangle.
        "Constructing summary statistics for approximate Bayesian computation:
        Semi‐automatic approximate Bayesian computation."
        Journal of the Royal Statistical Society: Series B
        (Statistical Methodology) 74.3 (2012): 419-474.
    """

    def __init__(
        self,
        normalize_features: bool = True,
        normalize_labels: bool = True,
        joint: bool = True,
        weight_samples: bool = False,
        log_pearson: bool = True,
        **kwargs,
    ):
        # check installation
        if skl_lm is None:
            raise ImportError(
                "This predictor requires an installation of scikit-learn. "
                "Install e.g. via `pip install pyabc[scikit-learn]`"
            )

        default_kwargs = {
            'fit_intercept': True,
            'normalize': True,
        }
        default_kwargs.update(kwargs)
        predictor = skl_lm.LinearRegression(**default_kwargs)

        super().__init__(
            predictor=predictor,
            normalize_features=normalize_features,
            normalize_labels=normalize_labels,
            joint=joint,
            weight_samples=weight_samples,
            log_pearson=log_pearson,
        )

    def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None:
        super().fit(x, y, w)
        # log
        if self.joint:
            logger.debug(
                "Linear regression coefficients (n_target, n_feature):\n"
                f"{self.predictor.coef_}"
            )
        else:
            for i_pred, predictor in enumerate(self.single_predictors):
                logger.debug(
                    "Linear regression coefficients (n_target, n_feature):\n"
                    f"for predictor {i_pred}: {predictor.coef_}"
                )

class LassoPredictor(SimplePredictor):
    """Lasso (least absolute shrinkage and selection) model.

    Linear model with l1 regularization.
    """

    def __init__(
        self,
        normalize_features: bool = True,
        normalize_labels: bool = True,
        joint: bool = True,
        log_pearson: bool = True,
        **kwargs,
    ):
        """Additional keyword arguments are passed on to the model."""
        # check installation
        if skl_lm is None:
            raise ImportError(
                "This predictor requires an installation of scikit-learn. "
                "Install e.g. via `pip install pyabc[scikit-learn]`"
            )

        # translate arguments
        all_args = {
            'fit_intercept': True,
            'normalize': True,
        }
        all_args.update(kwargs)
        predictor = skl_lm.Lasso(**all_args)

        super().__init__(
            predictor=predictor,
            normalize_features=normalize_features,
            normalize_labels=normalize_labels,
            joint=joint,
            weight_samples=False,
            log_pearson=log_pearson,
        )

class GPPredictor(SimplePredictor):
    """Gaussian process model.

    Similar to [#borowska2021]_.

    .. [#borowska2021]
        Borowska, Agnieszka, Diana Giurghita, and Dirk Husmeier.
        "Gaussian process enhanced semi-automatic approximate Bayesian
        computation: parameter inference in a stochastic differential equation
        system for chemotaxis."
        Journal of Computational Physics 429 (2021): 109999.
    """

    def __init__(
        self,
        kernel: Union[Callable, skl_gp.kernels.Kernel] = None,
        normalize_features: bool = True,
        normalize_labels: bool = True,
        joint: bool = True,
        log_pearson: bool = True,
        **kwargs,
    ):
        """
        Parameters
        ----------
        kernel:
            Covariance kernel. Can be either a kernel, o
        """
        # check installation
        if skl_gp is None:
            raise ImportError(
                "This predictor requires an installation of scikit-learn. "
                "Install e.g. via `pip install pyabc[scikit-learn]`"
            )

        # default kernel
        if kernel is None:
            kernel = GPKernelHandle()
        self.kernel: Union[Callable, skl_gp.kernels.Kernel] = kernel

        self.kwargs = kwargs

        super().__init__(
            predictor=None,
            normalize_features=normalize_features,
            normalize_labels=normalize_labels,
            joint=joint,
            weight_samples=False,
            log_pearson=log_pearson,
        )

    def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None:
        # need to recreate the model

        # set indices to keep
        self.set_use_ixs(x=x, log=False)

        # kernel
        kernel = self.kernel
        if not isinstance(kernel, skl_gp.kernels.Kernel):
            n_in = sum(self.use_ixs)
            kernel = kernel(n_in)

        self.predictor = skl_gp.GaussianProcessRegressor(
            kernel=kernel, **self.kwargs
        )

        super().fit(x=x, y=y, w=w)

class GPKernelHandle:
    """Convenience class for Gaussian process kernel construction.

    Allows to create kernels depending on problem dimensions.
    """

    # kernels supporting features specific length scales
    ARD_KERNELS = ["RBF", "Matern"]

    def __init__(
        self,
        kernels: List[str] = None,
        kernel_kwargs: List[dict] = None,
        ard: bool = True,
    ):
        """
        Parameters
        ----------
        kernels:
            Names of `sklearn.kernel` covariance kernels. Defaults to a
            radial basis function (a.k.a. squared exponential) kernel "RBF"
            and a "WhiteKernel" to explain noise in the data. The resulting
            kernel is the sum of all kernels.
        kernel_kwargs:
            Optional arguments passed to the kernel constructors.
        ard:
            Automatic relevance determination by assigning a separate length
            scale per input variable. Only supported by some kernels,
            currently "RBF" and "Matern".
            If set to True, the capable kernels are automatically informed.
            It the underlying scitki-learn toolbox extends support, this
            list needs to be updated.
        """
        if kernels is None:
            kernels = ["RBF", "WhiteKernel"]
        self.kernels: List[str] = kernels

        if kernel_kwargs is None:
            kernel_kwargs = [{} for _ in self.kernels]
        self.kernel_kwargs = kernel_kwargs

        self.ard: bool = ard

    def __call__(self, n_in: int) -> "skl_gp.kernels.Kernel":
        """
        Parameters
        ----------
        n_in: Input (feature) dimension.

        Returns
        -------
        kernel: Kernel created from inputs.
        """
        kernels = [
            getattr(skl_gp.kernels, kernel)(
                length_scale=np.ones(n_in), **kernel_kwargs
            )
            if self.ard and kernel in GPKernelHandle.ARD_KERNELS
            else getattr(skl_gp.kernels, kernel)(**kernel_kwargs)
            for kernel, kernel_kwargs in zip(self.kernels, self.kernel_kwargs)
        ]
        return sum(kernels)

class MLPPredictor(SimplePredictor):
    """Multi-layer perceptron regressor predictor.

    See e.g. [#jiang2017]_.

    .. [#jiang2017]
        Jiang, Bai, et al.
        "Learning summary statistic for approximate Bayesian computation via
        deep neural network."
        Statistica Sinica (2017): 1595-1618.
    """

    def __init__(
        self,
        normalize_features: bool = True,
        normalize_labels: bool = True,
        joint: bool = True,
        hidden_layer_sizes: Union[Tuple[int, ...], Callable] = None,
        log_pearson: bool = True,
        **kwargs,
    ):
        """Additional keyword arguments are passed on to the model.

        Parameters
        ----------
        hidden_layer_sizes:
            Network hidden layer sizes. Can be either a tuple of ints, or
            a callable taking input dimension, output dimension, and number
            of samples and returning a tuple of ints.
            The :class:`HiddenLayerSize` provides some useful defaults.
        """
        # check installation
        if skl_nn is None:
            raise ImportError(
                "This predictor requires an installation of scikit-learn. "
                "Install e.g. via `pip install pyabc[scikit-learn]`"
            )

        if hidden_layer_sizes is None:
            hidden_layer_sizes = HiddenLayerHandle()
        self.hidden_layer_sizes = hidden_layer_sizes

        self.kwargs = {
            'solver': 'lbfgs',
            'max_iter': 10000,
            'early_stopping': True,
        }
        self.kwargs.update(kwargs)

        super().__init__(
            predictor=None,
            normalize_features=normalize_features,
            normalize_labels=normalize_labels,
            joint=joint,
            weight_samples=False,
            log_pearson=log_pearson,
        )

    def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None:
        # need to recreate the model

        # set indices to keep
        self.set_use_ixs(x=x, log=False)

        # hidden layer sizes
        hidden_layer_sizes = self.hidden_layer_sizes
        if callable(hidden_layer_sizes):
            n_in, n_out, n_sample = sum(self.use_ixs), y.shape[1], x.shape[0]
            hidden_layer_sizes = hidden_layer_sizes(n_in, n_out, n_sample)

        self.predictor = skl_nn.MLPRegressor(
            hidden_layer_sizes=hidden_layer_sizes, **self.kwargs
        )

        super().fit(x=x, y=y, w=w)

class HiddenLayerHandle:
    """Convenience class for various layer size strategies.

    Allows to define sizes depending on problem dimensions.
    """

    HEURISTIC = "heuristic"
    MEAN = "mean"
    MAX = "max"
    METHODS = [HEURISTIC, MEAN, MAX]

    def __init__(
        self,
        method: Union[str, List[str]] = MEAN,
        n_layer: int = 1,
        max_size: int = np.inf,
        alpha: float = 1.0,
    ):
        """
        Parameters
        ----------
        method:
            Method to use. Can be any of:

            * "heuristic" bases the number of neurons on the number of samples
              to avoid overfitting. See
              https://stats.stackexchange.com/questions/181.
            * "mean" takes the mean of input and output dimension.
            * "max" takes the maximum of input and output dimension.

            Additionally, a list of methods can be passed, in which case
            the minimum over all is used.
        n_layer:
            Number of layers.
        max_size:
            Maximum layer size. Applied to all strategies.
        alpha:
            Factor used in "heuristic". The higher, the fewer neurons.
            Recommended is a value in the range 2-10.
        """
        if isinstance(method, str):
            method = [method]
        for m in method:
            if m not in HiddenLayerHandle.METHODS:
                raise ValueError(
                    f"Method {m} must be in {HiddenLayerHandle.METHODS}"
                )
        self.method = method
        self.n_layer = n_layer
        self.max_size = max_size
        self.alpha = alpha

    def __call__(
        self,
        n_in: int,
        n_out: int,
        n_sample: int,
    ) -> Tuple[int, ...]:
        """
        Parameters
        ----------
        n_in: Input (feature) dimension.
        n_out: Output (target) dimension.
        n_sample: Number of samples.

        Returns
        -------
        hidden_layer_sizes: Tuple of hidden layer sizes.
        """
        neurons_arr = []
        for method in self.method:
            if method == HiddenLayerHandle.HEURISTIC:
                # number of neurons
                neurons = n_sample / (self.alpha * (n_in + n_out))
                # divide by number of layers
                neurons /= self.n_layer
            elif method == HiddenLayerHandle.MEAN:
                neurons = 0.5 * (n_in + n_out)
            elif method == HiddenLayerHandle.MAX:
                neurons = max(n_in, n_out)
            else:
                raise ValueError(f"Did not recognize method {self.method}.")
            neurons_arr.append(neurons)

        # take minimum over proposed values
        neurons_per_layer = min(neurons_arr)

        # cap
        neurons_per_layer = min(neurons_per_layer, self.max_size)

        # only >=2 dim makes sense, round
        neurons_per_layer = int(max(2.0, neurons_per_layer))

        layer_sizes = tuple(neurons_per_layer for _ in range(self.n_layer))
        logger.info(f"Layer sizes: {layer_sizes}")

        return layer_sizes

class ModelSelectionPredictor(Predictor):
    """Model selection over a set of predictors.

    Picks the model with minimum k-fold cross valdation score and retrains on
    full data set.
    """

    CROSS_VALIDATION = "cross_validation"
    TRAIN_TEST_SPLIT = "train_test_split"
    SPLIT_METHODS = [CROSS_VALIDATION, TRAIN_TEST_SPLIT]

    def __init__(
        self,
        predictors: List[Predictor],
        split_method: str = TRAIN_TEST_SPLIT,
        n_splits: int = 5,
        test_size: float = 0.2,
        f_score: Callable = None,
    ):
        """
        Parameters
        ----------
        predictors:
            Set of predictors over which to perform model selection.
        split_method:
            Method how to split the data set into training and test data,
            can be "cross_validation" for a full `n_splits` fold cross
            validation, or "train_test_split" for a single separation of a
            test set of size `test_size`.
        n_splits:
            Number of splits to use in k-fold cross validation.
        test_size:
            Fraction of samples to randomly pick as test set, when using a
            single training and test set.
        f_score:
            Score function to assess prediction quality. Defaults to
            root mean square error normalized by target standard variation.
            Takes arguments y1, y2, std for prediction, ground truth, and
            standard variation, and returns the score as a float.
        """
        super().__init__()
        self.predictors: List[Predictor] = predictors

        if split_method not in ModelSelectionPredictor.SPLIT_METHODS:
            raise ValueError(
                f"Split method {split_method} must be in "
                f"{ModelSelectionPredictor.SPLIT_METHODS}",
            )
        self.split_method: str = split_method

        self.n_splits: int = n_splits
        self.test_size: float = test_size

        if f_score is None:
            self.f_score = root_mean_square_error

        # holds the chosen predictor model
        self.chosen_one: Union[Predictor, None] = None

    def fit(self, x: np.ndarray, y: np.ndarray, w: np.ndarray = None) -> None:
        # output normalization
        std_y = np.std(y, axis=0)

        n_sample = x.shape[0]
        if self.split_method == ModelSelectionPredictor.CROSS_VALIDATION:
            # k-fold cross validation
            kfold = skl_ms.KFold(n_splits=self.n_splits, shuffle=True)
            splits = kfold.split(np.arange(n_sample))
            n_splits = self.n_splits
        else:
            # a single training and test set
            splits = skl_ms.train_test_split(
                np.arange(n_sample), test_size=self.test_size
            )
            # as iterable
            splits = [splits]
            n_splits = 1

        # scores on training and test set
        scores_test = np.empty((len(self.predictors), n_splits))
        #  for debugging
        scores_train = np.empty((len(self.predictors), n_splits))

        # iterate over training and test sets
        for i_split, (train_ixs, test_ixs) in enumerate(splits):
            # training and test sets
            x_train, y_train = x[train_ixs], y[train_ixs]
            w_train = w
            if w is not None:
                w_train = w[train_ixs]
            x_test, y_test = x[test_ixs], y[test_ixs]

            # iterate over predictors
            for i_pred, predictor in enumerate(self.predictors):
                # fit predictor on training set
                predictor.fit(x=x_train, y=y_train, w=w_train)

                # test score of z-score normalized variables
                y_test_pred = predictor.predict(x=x_test)
                scores_test[i_pred, i_split] = self.f_score(
                    y1=y_test_pred,
                    y2=y_test,
                    sigma=std_y,
                )

                # for debugging, log training scores
                y_train_pred = predictor.predict(x=x_train)
                scores_train[i_pred, i_split] = self.f_score(
                    y1=y_train_pred,
                    y2=y_train,
                    sigma=std_y,
                )

        # the score of a predictor is the mean over all folds
        scores_test = np.mean(scores_test, axis=1)
        scores_train = np.mean(scores_train, axis=1)

        # logging
        for predictor, score_train, score_test in zip(
            self.predictors, scores_train, scores_test
        ):
            logger.info(
                f"Test score {score_test:.3e} (train {score_train:.3e}) for "
                f"{predictor}"
            )

        # best predictor has minimum score
        self.chosen_one = self.predictors[np.argmin(scores_test)]

        # retrain chosen model on full data set
        self.chosen_one.fit(x=x, y=y, w=w)

    def predict(self, x: np.ndarray, normalize: bool = False) -> np.ndarray:
        # predict from the chosen model
        return self.chosen_one.predict(x=x, normalize=normalize)

def root_mean_square_error(
    y1: np.ndarray,
    y2: np.ndarray,
    sigma: Union[np.ndarray, float],
) -> float:
    """Root mean square error of `y1 - y2 / sigma`.

    Parameters
    ----------
    y1: Model simulations, shape (n_sample, n_par).
    y2: Ground truth values, shape (n_sample, n_par).
    sigma: Normalizations, shape (n_sample,) or (1,).

    Returns
    -------
    val: The normalized root mean square error value.
    """
    return np.sqrt(
        np.sum(
            ((y1 - y2) / sigma) ** 2,
        )
        / y1.size,
    )

def root_mean_square_relative_error(
    y1: np.ndarray,
    y2: np.ndarray,
) -> float:
    """Root mean square relative error of `(y1 - y2) / y2`.

    Note that this may behave badly for ground truth parameters close to 0.

    Parameters
    ----------
    y1: Model simulations.
    y2: Ground truth values.

    Returns
    -------
    val: The normalized root mean square relative error value.
    """
    return np.sqrt(
        np.sum(
            ((y1 - y2) / y2) ** 2,
        )
        / y1.size,
    )

max_stars_repo_name: Gabriel-p/pyABC

max_issues_repo_name: Gabriel-p/pyABC

max_forks_repo_name: Gabriel-p/pyABC

reason: import numpy,from scipy


text:
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import contextlib
from collections import deque

import numpy as np
import paddle

from ..transformers import UnifiedTransformerLMHeadModel, UnifiedTransformerTokenizer
from ..datasets import load_dataset
from ..data import Pad
from .utils import dygraph_mode_guard
from .task import Task

usage = r"""
           from paddlenlp import Taskflow

           # 非交互模式
           dialogue = Taskflow("dialogue")
           dialogue(["吃饭了吗"])
           '''
           ['刚吃完饭,你在干什么呢?']
           '''
           dialogue(["你好", "吃饭了吗"], ["你是谁？"])
           '''
           ['吃过了,你呢', '我是李明啊']
           '''

           dialogue = Taskflow("dialogue")
           # 进入交互模式 (输入exit退出)
           dialogue.interactive_mode(max_turn=3)

           '''
           [Human]:你好
           [Bot]:你好,很高兴认识你,我想问你一下,你喜欢运动吗?
           [Human]:喜欢
           [Bot]:那你喜欢什么运动啊?
           [Human]:篮球,你喜欢篮球吗
           [Bot]:当然了,我很喜欢打篮球的。
           '''
         """

class DialogueTask(Task):
    """
    Task of Chinese open domain dialogue.
    Args:
        task(string): The name of task.
        model(string): The model name in the task.
        kwargs (dict, optional): Additional keyword arguments passed along to the specific task.
    """

    resource_files_names = {
        "model_state": "model_state.pdparams",
        "model_config": "model_config.json",
    }
    resource_files_urls = {
        "plato-mini": {
            "model_state": [
                "https://bj.bcebos.com/paddlenlp/taskflow/dialogue/plato-mini/model_state.pdparams",
                "450be85b9b7f0bc03b12252a75af04f3"
            ],
            "model_config": [
                "https://bj.bcebos.com/paddlenlp/taskflow/dialogue/plato-mini/model_config.json",
                "5e853fda9a9b573815ad112e494a65af"
            ],
        }
    }

    def __init__(self, task, model, batch_size=1, max_seq_len=512, **kwargs):
        super().__init__(task=task, model=model, **kwargs)
        self._static_mode = False
        self._usage = usage
        self._check_task_files()
        self._construct_tokenizer(model)
        self._batch_size = batch_size
        self._max_seq_len = max_seq_len
        self._interactive_mode = False
        if self._static_mode:
            self._get_inference_model()
        else:
            self._construct_model(model)

    def _construct_input_spec(self):
        """
        Construct the input spec for the convert dygraph model to static model.
        """
        self._input_spec = [
            paddle.static.InputSpec(
                shape=[None, None], dtype="int64", name='input_ids'),
            paddle.static.InputSpec(
                shape=[None], dtype="int64", name='token_type_ids'),
        ]

    def _construct_model(self, model):
        """
        Construct the inference model for the predictor.
        """
        model_instance = UnifiedTransformerLMHeadModel.from_pretrained(
            self._task_path)
        model_path = os.path.join(self._task_path, "model_state.pdparams")
        state_dict = paddle.load(model_path)
        model_instance.set_state_dict(state_dict)
        self._model = model_instance
        self._model.eval()

    def _construct_tokenizer(self, model):
        """
        Construct the tokenizer for the predictor.
        """
        self._tokenizer = UnifiedTransformerTokenizer.from_pretrained(model)

    def _batchify_fn(self, batch_examples):
        #padding = False if self._batch_size == 1 else True
        pad_func = Pad(pad_val=self._tokenizer.pad_token_id,
                       pad_right=False,
                       dtype='int64')

        def pad_mask(batch_attention_mask):
            batch_size = len(batch_attention_mask)
            max_len = max(map(len, batch_attention_mask))
            attention_mask = np.ones(
                (batch_size, max_len, max_len), dtype='float32') * -1e4
            for i, mask_data in enumerate(attention_mask):
                seq_len = len(batch_attention_mask[i])
                mask_data[-seq_len:, -seq_len:] = np.array(
                    batch_attention_mask[i], dtype='float32')
            # In order to ensure the correct broadcasting mechanism, expand one
            # dimension to the second dimension (n_head of Transformer).
            attention_mask = np.expand_dims(attention_mask, axis=1)
            return attention_mask

        input_ids = pad_func(
            [example['input_ids'] for example in batch_examples])
        token_type_ids = pad_func(
            [example['token_type_ids'] for example in batch_examples])
        position_ids = pad_func(
            [example['position_ids'] for example in batch_examples])
        attention_mask = pad_mask(
            [example['attention_mask'] for example in batch_examples])

        return input_ids, token_type_ids, position_ids, attention_mask

    def _check_input_text(self, inputs):
        if self._interactive_mode:
            if isinstance(inputs, str):
                self.context.append(inputs.strip())
                inputs = [list(self.context)]
                return inputs
            else:
                raise ValueError(
                    "In the interactive mode, the input data shold be a string")
        elif not isinstance(inputs[0], list):
            raise ValueError(
                "If not in the interactive mode, the input data should be a list."
            )
        return inputs

    def _batchify(self, data, max_seq_len, batch_size):
        """
        Generate input batches.
        """
        padding = False if batch_size == 1 else True
        pad_func = Pad(pad_val=self._tokenizer.pad_token_id,
                       pad_right=False,
                       dtype=np.int64)

        def pad_mask(batch_attention_mask):
            batch_size = len(batch_attention_mask)
            max_len = max(map(len, batch_attention_mask))
            attention_mask = np.ones(
                (batch_size, max_len, max_len), dtype='float32') * -1e4
            for i, mask_data in enumerate(attention_mask):
                seq_len = len(batch_attention_mask[i])
                mask_data[-seq_len:, -seq_len:] = np.array(
                    batch_attention_mask[i], dtype='float32')
            # In order to ensure the correct broadcasting mechanism, expand one
            # dimension to the second dimension (n_head of Transformer).
            attention_mask = np.expand_dims(attention_mask, axis=1)
            return attention_mask

        def _parse_batch(batch_examples):
            if padding:
                input_ids = pad_func(
                    [example['input_ids'] for example in batch_examples])
                token_type_ids = pad_func(
                    [example['token_type_ids'] for example in batch_examples])
                position_ids = pad_func(
                    [example['position_ids'] for example in batch_examples])
                attention_mask = pad_mask(
                    [example['attention_mask'] for example in batch_examples])
            else:
                input_ids = np.asarray(
                    [example['input_ids'] for example in batch_examples],
                    dtype=np.int64)
                token_type_ids = np.asarray(
                    [example['token_type_ids'] for example in batch_examples],
                    dtype=np.int64)
                position_ids = np.asarray(
                    [example['position_ids'] for example in batch_examples],
                    dtype=np.int64)
                attention_mask = np.asarray(
                    [example['attention_mask'] for example in batch_examples])
                attention_mask = np.expand_dims(attention_mask, 0)

            return input_ids, token_type_ids, position_ids, attention_mask

        examples = []
        for texts in data:
            examples.append(self._convert_text_to_input(texts, max_seq_len))

        # Seperates data into some batches.
        one_batch = []
        for example in examples:
            one_batch.append(example)
            if len(one_batch) == batch_size:
                yield _parse_batch(one_batch)
                one_batch = []
        if one_batch:
            yield _parse_batch(one_batch)

    def _convert_text_to_input(self, texts, max_seq_len):
        """
        Convert input strings to tokens.
        """
        return self._tokenizer.dialogue_encode(
            texts,
            max_seq_len=max_seq_len,
            add_start_token_as_response=True,
            is_split_into_words=False)

    def _preprocess(self, inputs):
        """
        Transform the raw text to the model inputs, two steps involved:
           1) Transform the raw text to token ids.
           2) Generate the other model inputs from the raw text and token ids.
        """
        inputs = self._check_input_text(inputs)
        # Get the config from the kwargs
        num_workers = self.kwargs[
            'num_workers'] if 'num_workers' in self.kwargs else 0
        lazy_load = self.kwargs[
            'lazy_load'] if 'lazy_load' in self.kwargs else False

        batches = self._batchify(inputs, self._max_seq_len, self._batch_size)

        outputs = {}
        outputs['batches'] = batches
        outputs['text'] = inputs
        return outputs

    def _run_model(self, inputs):
        """
        Run the task model from the outputs of the `_tokenize` function.
        """
        all_ids = []
        all_scores = []

        for batch in inputs["batches"]:
            input_ids, token_type_ids, position_ids, attention_mask = map(
                paddle.to_tensor, batch)
            ids, scores = self._model.generate(
                input_ids=input_ids,
                token_type_ids=token_type_ids,
                position_ids=position_ids,
                attention_mask=attention_mask,
                max_length=64,
                min_length=1,
                decode_strategy='sampling',
                temperature=1.0,
                top_k=5,
                top_p=1.0,
                num_beams=0,
                length_penalty=1.0,
                early_stopping=False,
                use_faster=False,
                num_return_sequences=1)
            all_ids.extend([ids])
            all_scores.extend([scores])
        inputs['ids'] = all_ids
        inputs['scores'] = all_scores
        return inputs

    def _post_process_response(self, token_ids, tokenizer):
        '''
        Post-process the decoded sequence. Truncate from the first <eos>.
        '''
        eos_pos = len(token_ids)
        for i, tok_id in enumerate(token_ids):
            if tok_id == tokenizer.sep_token_id:
                eos_pos = i
                break
        token_ids = token_ids[:eos_pos]
        tokens = tokenizer.convert_ids_to_tokens(token_ids)
        tokens = tokenizer.merge_subword(tokens)
        return token_ids, tokens

    @contextlib.contextmanager
    def interactive_mode(self, max_turn=3):
        """
        Enter the interactive mode.
        """
        self._interactive_mode = True
        self.max_turn = max_turn
        self.context = deque(maxlen=self.max_turn)
        yield
        self.context.clear()
        self._interactive_mode = False

    def _get_in_turn_repetition(self, pred, is_cn=False):
        '''
        Get in-turn repetition.
        '''
        if len(pred) == 0:
            return 1.0
        if isinstance(pred[0], str):
            pred = [tok.lower() for tok in pred]
            if is_cn:
                pred = "".join(pred)
        tri_grams = set()
        for i in range(len(pred) - 2):
            tri_gram = tuple(pred[i:i + 3])
            if tri_gram in tri_grams:
                return True
            tri_grams.add(tri_gram)
        return False

    def _select_response(self,
                         ids,
                         scores,
                         tokenizer,
                         max_dec_len=None,
                         num_return_sequences=1,
                         keep_space=True):
        '''
        Select response with the highest score.
        '''
        ids = ids.numpy().tolist()
        scores = scores.numpy()

        if len(ids) != len(scores) or (len(ids) % num_return_sequences) != 0:
            raise ValueError(
                "the length of `ids` is {}, but the `num_return_sequences` is {}".
                format(len(ids), num_return_sequences))

        group = []
        tmp = []
        for pred, score in zip(ids, scores):
            pred_token_ids, pred_tokens = self._post_process_response(pred,
                                                                      tokenizer)
            num_token = len(pred_token_ids)
            if keep_space:
                response = " ".join(pred_tokens)
            else:
                response = "".join(pred_tokens)

            in_turn_repetition = self._get_in_turn_repetition(pred_tokens, True) \
                or self._get_in_turn_repetition(pred_token_ids)
            # not ending
            if max_dec_len is not None and num_token >= max_dec_len:
                score -= 1e3
            elif in_turn_repetition:
                score -= 1e3

            tmp.append([response, score])
            if len(tmp) == num_return_sequences:
                group.append(tmp)
                tmp = []

        results = []
        for preds in group:
            preds = sorted(preds, key=lambda x: -x[1])
            results.append(preds[0][0])
        return results

    def _postprocess(self, inputs):
        all_ids = inputs['ids']
        all_scores = inputs['scores']
        texts = inputs['text']

        results = []
        for ids, scores, text in zip(all_ids, all_scores, texts):
            results.extend(
                self._select_response(
                    ids,
                    scores,
                    self._tokenizer,
                    num_return_sequences=1,
                    keep_space=False))

        if self._interactive_mode:
            self.context.append(results[0].strip())
        return results

ext: py

reason: import numpy


text:
from scipy.spatial import Delaunay
import prody as pr
import numpy as np
from .pointTriangleDistance import pointTriangleDistance as distance
from numba import jit
import copy

@jit("f8(f8[:],f8[:],f8[:],f8[:])", nopython=True, cache=True)
def vol(a, b, c, d):
    M = np.zeros((3, 3))
    M[0, :] = np.subtract(a, d)
    M[1, :] = np.subtract(b, d)
    M[2, :] = np.subtract(c, d)
    return np.abs(np.linalg.det(M)) / 6

@jit("f8(f8[:,:])", nopython=True, cache=True)
def get_radius(points):
    a = np.linalg.norm(points[0] - points[1])
    a1 = np.linalg.norm(points[2] - points[3])
    b = np.linalg.norm(points[0] - points[2])
    b1 = np.linalg.norm(points[1] - points[3])
    c = np.linalg.norm(points[0] - points[3])
    c1 = np.linalg.norm(points[1] - points[2])
    p = (a * a1 + b * b1 + c * c1) / 2
    V = vol(points[0], points[1], points[2], points[3])
    if V > 0:
        return 1 / (6 * V) * np.sqrt(p * (p - a * a1) * (p - b * b1) * (p - c * c1))
    else:
        return np.inf

@jit("i4[:,:](f8[:,:], i4[:,:], f8)", nopython=True, cache=True)
def _calc_alpha_simplex(C, S, a):
    M = S.shape[0]
    N = S.shape[1]
    Result = np.zeros((M, N))
    j = 0
    for i in range(M):
        s = S[i, :]
        ps = C[s]
        r = get_radius(ps)
        if r < a:
            Result[j, :] = s
            j += 1
    return Result[:j, :].astype(np.int32)

combos = np.array([[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]])

@jit("i4[:,:](i4[:,:], i8[:,:])", nopython=True, cache=True)
def make_simplex_set(S, combos):
    M = S.shape[0] * 4
    N = S.shape[1] - 1
    R = np.zeros((M, N), dtype=np.int32)
    for k, s in enumerate(range(0, M, 4)):
        for i in range(4):
            for j in range(3):
                R[s + i, j] = S[k, combos[i, j]]
    return R

def normal(ps):
    v1 = ps[1] - ps[0]
    v2 = ps[2] - ps[0]
    crossprod = np.cross(v1, v2)
    return crossprod / np.linalg.norm(crossprod)

class AlphaHull:

    def __init__(self, alpha=9):
        self.alpha = alpha
        self.hull = None
        self._hull = None
        self.tri = None
        self._tri = None
        self.coords = None
        self.simplices = None
        self.resindices = None

    def set_coords(self, pdb):
        type1 = isinstance(pdb, pr.atomic.selection.Selection)
        type2 = isinstance(pdb, pr.atomic.atomgroup.AtomGroup)
        if type1 or type2:
            self._set_coords(pdb)
        elif isinstance(pdb, np.ndarray):
            self.coords = pdb
        else:
            raise ValueError('*pdb* must be prody instance or numpy array')

    def _set_coords(self, pdb):
        """pdb is a prody object. pdb should have CB atoms where appropriate."""
        self.coords = pdb.select('name CB or (resname GLY and name CA)').getCoords()

    def set_tri(self):
        self.tri = Delaunay(self.coords)
        self._tri = copy.deepcopy(self.tri)

    def set_resindices(self, pdb):
        """pdb is a prody object. pdb should have CB atoms where appropriate."""
        self.resindices = pdb.select('name CB or (resname GLY and name CA)').getResindices()

    def calc_alpha_simplices(self):
        if self.tri is None:
            self.set_tri()
        self.tri.simplices.sort()
        self.tri.simplices = self.tri.simplices[self.tri.simplices[:, 0].argsort()]
        self.simplices = _calc_alpha_simplex(self.coords, self.tri.simplices, self.alpha)
        self._tri.simplices = self.simplices
        self._tri.neighbors = self.simplices

    def calc_hull(self):
        if self.simplices is None:
            self.calc_alpha_simplices()
        simpl_set = make_simplex_set(self.simplices, combos)
        un, ind, co = np.unique(simpl_set, axis=0,
                                return_counts=True, return_index=True)
        self.hull = np.array([simpl_set[i] for i in ind[co == 1]], dtype=np.int32)

    def pnts_in_hull(self, pnts):
        return self._tri.find_simplex(pnts) >= 0

    def get_pnt_distance(self, pnt):
        distances = []
        inout = self.pnts_in_hull(pnt)
        if inout:
            minmax = min
            inout = 1
        else:
            minmax = max
            inout = -1
        for i in range(len(self.hull)):
            distances.append(inout * distance(pnt, self.coords[self.hull[i]]))
        return minmax(distances)

    def get_pnts_distance(self, pnts):
        return [self.get_pnt_distance(pnt) for pnt in pnts]

def partition_res_by_burial(pdb_ala, alpha=9):
    """Returns residue indices of exposed, intermediate, and buried residues
    based on CA hull and CB hull."""
    ahull_ca = AlphaHull(alpha=alpha)
    ahull_ca.coords = pdb_ala.select('name CA').getCoords()
    ahull_ca.calc_hull()
    ahull_cb = AlphaHull(alpha=alpha)
    ahull_cb.set_coords(pdb_ala)
    ahull_cb.calc_hull()
    ahull_cb.set_resindices(pdb_ala)
    cb_in_ca_hull = ahull_ca.pnts_in_hull(ahull_cb.coords)
    resindices_cb_in_ca_hull = set(ahull_cb.resindices[cb_in_ca_hull])
    resindices_cb_hull = set(ahull_cb.resindices[np.unique(ahull_cb.hull)])
    resindices_not_cb_hull = set(ahull_cb.resindices) - resindices_cb_hull
    resindices_exposed = resindices_cb_hull - resindices_cb_in_ca_hull
    resindices_intermediate = resindices_cb_in_ca_hull - resindices_not_cb_hull
    resindices_buried = resindices_cb_in_ca_hull - resindices_intermediate
    res_ = resindices_not_cb_hull - resindices_buried
    resindices_intermediate |= res_
    return resindices_exposed, resindices_intermediate, resindices_buried

reason: import numpy,from scipy,from numba


text:
import pandas as pd
import seaborn as sns
import numpy as np

def main():

    raw_data = excel_csv_conversion()

    cleaned_data = csv_file_cleaner(raw_data)

    renamed_data = column_name_selection(cleaned_data)

    iso_data = column_extraction(renamed_data)

    northeast, div1, div2 = northeast_organizer(iso_data)

    midwest, div3, div4 = midwest_organizer(iso_data)

    south, div5, div6, div7 = south_organizer(iso_data)

    west, div8, div9 = west_organizer(iso_data)

    #Division 1-9

    #Nested in some iterative operation through regional/divisional data
    year_classes = year('some regional/divisional data')

    #Nested in some iterative operation through regional/divsional data --> iterative operation through year data
    benefit_classes = benefit('some regional/divisional data --> year data')

    #Below 2 nested in region --> year --> benfit
    percentage_classes =

    count_classes =

    #Need a variable (or multiple) that will store info yielded from the above triple nest

    #Conditional - percentage figures or total figures?

    #EDA

    #Visual display

    #Interpretation

    #Need a loop such that it iteratively calls region_organizer to obtain separated data by region

    #Idea - create functions for each region, such that respective regional data, along with the according divisional data,
    #is returned

    #Then need a loop such that it iteratively calls payment_class to obtain data specific to payment class

    #Next, pass through to EDA and obtain some interpretation of the data

    #Once we have made various interpretations (by division and by region), pass to visual_display

    #This is where we will use Seaborn to show data on a visual basis

def excel_csv_conversion():

    Yr06 = pd.read_excel('../data/raw/group36_MedicareData.xlsx', sheet_name='PUF_2006', header=1)
    Yr07 = pd.read_excel('../data/raw/group36_MedicareData.xlsx', sheet_name='PUF_2007', header=1)
    Yr08 = pd.read_excel('../data/raw/group36_MedicareData.xlsx', sheet_name='PUF_2008', header=1)
    Yr09 = pd.read_excel('../data/raw/group36_MedicareData.xlsx', sheet_name='PUF_2009', header=1)
    Yr10 = pd.read_excel('../data/raw/group36_MedicareData.xlsx', sheet_name='PUF_2010', header=1)
    Yr11 = pd.read_excel('../data/raw/group36_MedicareData.xlsx', sheet_name='PUF_2011', header=1)
    Yr12 = pd.read_excel('../data/raw/group36_MedicareData.xlsx', sheet_name='PUF_2012', header=1)

    Yrlist = [Yr06, Yr07, Yr08, Yr09, Yr10, Yr11, Yr12]

    Numlist = [2006, 2007, 2008, 2009, 2010, 2011, 2012]

    for i in range(len(Yrlist)):
        Yrlist[i]['Year'] = Numlist[i]

    Yrtable = pd.concat(Yrlist)

    Yrtable.to_csv('../data/raw/concatenated_data.csv', sep=',', index=False)

    return Yrtable

def csv_file_cleaner(raw_csv):

    cleandata = pd.read('../data/processed/concatenated_data.csv', sep=',', index=False)

    return cleandata

def column_extraction(data):

    #Goal - choose your desired columns that fit the research question(s); create new dataframes
    #fitting the scope
    #Transition - pass to region organizer, which will isolate data per state

    #population totals
    total_df = data[['total']]
    gender_df = data[['female', 'male']]
    age_df = data[['<40, %', '40-64, %', '65-84, %', '>84, %']]
    ethnicity_df = data[['non-hi wh, %', 'af am, %', 'hi, %', 'as or pi, %', 'fn, %', 'unknown, %']]

    #option 1 - substance use, mental health, related physiological ailments, poverty relation
    pr_df = data[['poverty relation, %']]
    substance_use_df = data[['tobacco, %', 'alcohol, %', 'drug use, %', 'opioid use, %']]
    mental_health_df = data[['depression, %', 'clincal depression, %', 'adhd, %', 'anxiety, %',
    'bpd, %', 'personality disorders, %', 'ptsd, %', 'schizophrenia, %', 'schizophrenia or other, %']]
    substance_complication_df = data[['opioid complicaton, %']]

    #option 2 - population, FFS population(s), services used via FFS classification, revenue generated by FFS
    #specific focus - mental health services and hospital services
    mental_health_df2 = data[['res mental health, %', 'com mental health, %', 'com mental health, ffs',
    'medicaid mental health, ffs']]
    hospital_df = data[['acute ip hosp days, m and m', 'acute ip hosp admis, m and m',
    'acute ip hosp readmis, m and m', 'readmis rate, %', 'ed visits, m and m', 'ed visits per 1k',
    'medicare ip hosp, ffs', 'medicare other ip hosp, ffs', 'medicaid ip hosp, ffs', ]]
    ffs_df = data[['total, ffs', 'female, ffs', 'male, ffs']]
    ffs_vs_partd_df = data[['ffs medicaid year-round, %', 'ffs medicare year-round, %', 'part d coverage year-round, %']]
    revenues_df = data[['medicare ffs, $', 'ip hosp admis, m and m $', 'medicare ip hosp, ffs $',
     'medicare other ip hosp, ffs $', 'medicare com mental health, ffs $', 'medicaid, ffs $', 'medicaid ip hosp, ffs $',
     'medicaid mental health, ffs $']]

    #development_df = data[['autism, %', 'intellectual, %', 'learning, %']]
    #ms_df = data[['ms, %']]

    return 'data of some kind (TBD)'

def state_organizer(data):

    #yr = int(input('Yr: '))
    #state = str(input('State: '))

    National = data.loc[(data['State'] == 'National')]
    StateAK = data.loc[(data['State'] == 'AK')]
    StateAL = data.loc[(data['State'] == 'AL')]
    StateAR = data.loc[(data['State'] == 'AR')]
    StateAZ = data.loc[(data['State'] == 'AZ')]
    StateCA = data.loc[(data['State'] == 'CA')]
    StateCO = data.loc[(data['State'] == 'CO')]
    StateCT = data.loc[(data['State'] == 'CT')]
    StateDC = data.loc[(data['State'] == 'DC')]
    StateDE = data.loc[(data['State'] == 'DE')]
    StateFL = data.loc[(data['State'] == 'FL')]
    StateGA = data.loc[(data['State'] == 'GA')]
    StateHI = data.loc[(data['State'] == 'HI')]
    StateIA = data.loc[(data['State'] == 'IA')]
    StateID = data.loc[(data['State'] == 'ID')]
    StateIL = data.loc[(data['State'] == 'IL')]
    StateIN = data.loc[(data['State'] == 'IN')]
    StateKS = data.loc[(data['State'] == 'KS')]
    StateKY = data.loc[(data['State'] == 'KY')]
    StateLA = data.loc[(data['State'] == 'LA')]
    StateMA = data.loc[(data['State'] == 'MA')]
    StateMD = data.loc[(data['State'] == 'MD')]
    StateME = data.loc[(data['State'] == 'ME')]
    StateMI = data.loc[(data['State'] == 'MI')]
    StateMN = data.loc[(data['State'] == 'MN')]
    StateMO = data.loc[(data['State'] == 'MO')]
    StateMS = data.loc[(data['State'] == 'MS')]
    StateMT = data.loc[(data['State'] == 'MT')]
    StateNC = data.loc[(data['State'] == 'NC')]
    StateND = data.loc[(data['State'] == 'ND')]
    StateNE = data.loc[(data['State'] == 'NE')]
    StateNH = data.loc[(data['State'] == 'NH')]
    StateNJ = data.loc[(data['State'] == 'NJ')]
    StateNM = data.loc[(data['State'] == 'NM')]
    StateNV = data.loc[(data['State'] == 'NV')]
    StateNY = data.loc[(data['State'] == 'NY')]
    StateOH = data.loc[(data['State'] == 'OH')]
    StateOK = data.loc[(data['State'] == 'OK')]
    StateOR = data.loc[(data['State'] == 'OR')]
    StatePA = data.loc[(data['State'] == 'PA')]
    StateRI = data.loc[(data['State'] == 'RI')]
    StateSC = data.loc[(data['State'] == 'SC')]
    StateSD = data.loc[(data['State'] == 'SD')]
    StateTN = data.loc[(data['State'] == 'TN')]
    StateTX = data.loc[(data['State'] == 'TX')]
    StateUT = data.loc[(data['State'] == 'UT')]
    StateVA = data.loc[(data['State'] == 'VA')]
    StateVT = data.loc[(data['State'] == 'VT')]
    StateWA = data.loc[(data['State'] == 'WA')]
    StateWI = data.loc[(data['State'] == 'WI')]
    StateWV = data.loc[(data['State'] == 'WV')]
    StateWY = data.loc[(data['State'] == 'WY')]

    Statelist = [National, StateAK, StateAL, StateAR, StateAZ, StateCA, StateCO, StateCT,
    StateDC, StateDE, StateFL, StateGA, StateHI, StateIA, StateID, StateIL, StateIN, StateKS, StateKY,
    StateLA, StateMA, StateMD, StateME, StateMI, StateMN, StateMO, StateMS, StateMT, StateNC, StateND,
    StateNE, StateNH, StateNJ, StateNM, StateNV, StateNY, StateOH, StateOK, StateOR, StatePA, StateRI,
    StateSC, StateSD, StateTN, StateTX, StateUT, StateVA, StateVT, StateWA, StateWI, StateWV, StateWY]

    return Statelist

def northeast_organizer(data):
    #Northeast
    div1 = [StateCT, StateME, StateMA, StateNH, StateRI, StateVT]
    df_newengland = pd.concat(div1)
    df_newengland['Division'] = 'New England'
    div2 = [StateNJ, StateNY, StatePA]
    df_midatlantic = pd.concat(div2)
    df_midatlantic['Division'] = 'Mid-Atlantic'
    northeast = [div1, div2]
    df_northeast = pd.concat(div1 + div2)
    df_northeast['Region'] = 'Northeast'

    return df_northeast, df_newengland, df_midatlantic

def midwest_organizer(data):
    #Midwest
    div3 = [StateIL, StateIN, StateMI, StateOH, StateWI]
    df_eastnorthcentral = pd.concat(div3)
    df_eastnorthcentral['Division'] = 'East North Central'
    div4 = [StateNE, StateMO, StateND, StateSD, StateIA, StateKS, StateMN]
    df_westnorthcentral = pd.concat(div4)
    df_westnorthcentral['Division'] = 'West North Central'
    midwest = [div3, div4]
    df_midwest = pd.concat(midwest)
    df_midwest['Region'] = 'Midwest'

    return df_midwest, df_eastnorthcentral, df_westnorthcentral

def south_organizer(data):
    #South
    div5 = [StateDE, StateFL, StateGA, StateMD, StateNC, StateSC, StateVA, StateWV]
    df_southatlantic = pd.concat(div5)
    df_southatlantic['Division'] = 'South Atlantic'
    div6 = [StateAL, StateKY, StateMS, StateTN]
    df_eastsouthcentral = pd.concat(div6)
    df_eastsouthcentral['Division'] = 'East South Central'
    div7 = [StateAR, StateLA, StateOK, StateTX]
    df_westsouthcentral = pd.concat(div7)
    df_westsouthcentral['Division'] = 'West South Central'
    south = [div5, div6, div7]
    df_south = pd.concat(south)
    df_south['Region'] = 'South'

    return df_south, df_southatlantic, df_eastsouthcentral, df_westsouthcentral

def west_organizer(data):
    #West
    div8 = [StateID, StateAZ, StateCO, StateMT, StateNM, StateNV, StateUT, StateWY]
    df_mountain = pd.concat(div8)
    df_mountain['Division'] = 'Mountain'
    div9 = [StateAK, StateCA, StateHI, StateOR, StateWA]
    df_pacific = pd.concat(div9)
    df_pacific['Division'] = 'Pacific'
    west = [div8, div9]
    df_west = pd.concat(west)
    df_west['Region'] = 'West'

    return df_west, df_mountain, df_pacific

def national_organizer(data):
    #National
    national = northeast + midwest + south + west
    national_df = pd.concat(national)

    return national_df

    #for comparing data across states, fixed year
    #fixed_year = data.loc[(data['Year'] == yr)]

    #for comparing data across years, fixed state
    #fixed_state = data.loc[(data['State'] == state)]

def column_name_selection(data):

    #shortening column names for ease of extraction
    abbrv_data = data.rename(columns={'Number of People': 'total', 'Number of People with FFS': 'total, ffs',
    'Number of Females with FFS': 'female, ffs', 'Number of Males with FFS': 'male, ffs',
    'Percent with all 12 months in FFS Medicaid': 'ffs medicaid year-round, %',
    'Percent with all 12 months in FFS Medicare': 'ffs medicare year-round, %',
    'Percent with all 12 months with Medicare Part D coverage': 'part d coverage year-round, %',
    'Percent under 40 Years': '<40, %', 'Percent between 40-64 Years': '40-64, %',
    'Percent between 65-84 Years': '65-84, %', 'Percent 85+ Years': '>84, %',
    'Percent under 65 Years': '<65, %', 'Percent 65+ Years': '>65, %',
    'Percent Female': 'female, %', 'Percent Male': 'male, %', 'Percent Non-Hispanic White': 'non-hi wh, %',
    'Percent African American': 'af am, %', 'Percent Hispanic': 'hi, %',
    'Percent Asian or Pacific Islander': 'as or pi, %',
    'Percent American Indian or Alaskan Native': 'fn, %',
    'Percent Other or Unknown Race': 'unknown, %',
    'Percent of Medicaid enrollees with MAS - Poverty Related': 'poverty relation, %',
    'Percent of people with depression (any instance including bipolar episodes)': 'depression, %',
    'Percent of people with Major Depressive Disorders': 'clinical depression, %',
    'Percent of people with ADHD or other conduct disorders': 'adhd, %':
    'Percent of people with anxiety': 'anxiety, %':
    'Percent of people with bipolar disorder': 'bpd, %',
    'Percent of people with personality disorders': 'personality disorders, %',
    'Percent of people with post-traumatic stress disorder': 'ptsd, %',
    'Percent of people with schizophrenia': 'schizophrenia, %',
    'Percent of people with schizophrenia or other psychotic disorders': 'schizophrenia or other, %',
    'Percent of people with tobacco use disorder': 'tobacco, %',
    'Percent of people with autism': 'autism, %',
    'Percent of people with intellectual disabilities': 'intellectual, %',
    'Percent of people with learning disabilities': 'learning, %',
    'Percent of people with multiple sclerosis and transverse myelitis': 'ms, %',
    'Percent of people with alcohol use disorder': 'alcohol, %',
    'Percent of people with drug use disorder': 'drug use, %',
    'Percent of people with opioid use disorder diagnosis': 'opioid use, %',
    'Percent of people with opioid related hospitalization or ED visit': 'opioid complication, %',
    'Total Medicare FFS payments': 'medicare ffs, $',
    'Percent with at least one Medicare or Medicaid Residential Mental Health service': 'res mental health, %',
    'Percent with at least one Medicare or Medicaid Community Mental Health Service': 'com mental health, %',
    'Number of FFS people who used Medicare Community Mental Health Clinic Services': 'com mental health, ffs',
    'Count of Acute IP Hospital Days - Medicare and Medicaid combined': 'acute ip hosp days, m and m',
    'Count of Acute IP Hospital Admissions - Medicare and Medicaid combined': 'acute ip hosp admis, m and m',
    'Total dollars (Medicare and Medicaid) associated with IP hospital admissions': 'ip hosp admis, m and m $',
    'Count of Acute IP Hospital 30-day Readmissions - Medicare and Medicaid combined': 'acute ip hosp readmis, m and m',
    'Readmission Rate ('%' of admissions that are readmissions)': 'readmis rate, %',
    'Count of ED Visits - Medicare and Medicaid combined': 'ed visits, m and m':
    'ED Visits per 1000 enrollees': 'ed visits per 1k',
    'Number of FFS people who used Medicare IP Hospital services': 'medicare ip hosp, ffs',
    'Number of FFS people who used Medicare Other IP Hospital services': 'medicare other ip hosp, ffs',
    'Total Medicare IP Hospital FFS payments': 'medicare ip hosp, ffs $',
    'Total Medicare Other IP Hospital FFS payments': 'medicare other ip hosp, ffs $',
    'Total Medicare Community Mental Health Clinic FFS Payments': 'medicare com mental health, ffs $',
    'Number of FFS people who used Medicaid Hospital services': 'medicaid ip hosp, ffs',
    'Number of FFS people who used Medicaid mental health services': 'medicaid mental health, ffs',
    'Total Medicaid FFS payments': 'medicaid, ffs $',
    'Total Medicaid IP Hospital FFS Payments': 'medicaid ip hosp, ffs $',
    'Total Medicaid mental health FFS payments': 'medicaid mental health, ffs $'})

    return abbrv_data

def new_columns():

    #Will create new columns for: regional means in a fixed year;
    #regional means across years;
    #state meaans across years;
    #national means in a fixed year;
    #national means across years;
    #ideally, a column for percent changes across each year (regional and national)

    #First, must complete EDA #1

def missing_data_filler(data):

    filled_data = data.replace('*', '0')
    filled_data = filled_data.replace('.', '0')
    filled_data = filled_data.replace(' ', '0')

    return filled_data

def year(data):

    data06 = data.loc((data['Year'] == 2006))
    data07 = data.loc((data['Year'] == 2007))
    data08 = data.loc((data['Year'] == 2008))
    data09 = data.loc((data['Year'] == 2009))
    data10 = data.loc((data['Year'] == 2010))
    data11 = data.loc((data['Year'] == 2011))
    data12 = data.loc((data['Year'] == 2012))

    year_classes = [data06, data07, data08, data09, data10, data11, data12]

    return year_classes

def benefit(data):

    #Separating data by benefit type/payment classification
    full = []
    partial = []
    medicare_only = []
    medicaid_only = []
    #Separating regionak
    for i in data_list:
        full += i.loc('Full Benefit')
        partial += i.loc('Partial Benefit')
        medicare_only += i.loc('Medicare Only')
        medicaid_only += i.loc('Medicaid Only (Disability)')

    benefit_classes = [full, partial, medicare_only, medicaid_only]

    return benefit_classes

def eda_counts(data, count_type = None):

    #first, obtain means across divsions and regions for each year;

    # gender_df = data[['female', 'male']]
    # age_df = data[['<40, %', '40-64, %', '65-84, %', '>84, %']]
    # ethnicity_df = data[['non-hi wh, %', 'af am, %', 'hi, %', 'as or pi, %', 'fn, %', 'unknown, %']]

    #Need data points for each year
    #Will graph these points over time
    #Expressed as means per division and per region
    #Need to create matrices, where rows are the years and columns are the respective states from each div.
    #data_list variable isolates years, so need to iterate through data_list
    #Assumption - either a region or division is being passed through

    #This data is pulled from ... this data covers ... amount of people nationwide,
    #... from the ... region (... from {div} and ... from {div})
    # " "
    #Question 1:

    #benefit_classes = ['Full Benefit', 'Partial Benefit', 'Medicare Only', 'Medicaid Only (Disability)'

    #Function for gathering total counts across divsions/regions, per year

    #totals_dict = {'Full Benefit': 0, 'Partial Benefit': 0, 'Medicare Only': 0, 'Medicaid Only (Disability)': 0}

    #Obtaining regional/divisional totals, per year

    array_init = data.loc(count_type)
    array_counts = np.array(array_init)
    new_total = np.sum(array_counts)

    return new_total

    #Above syntax good for counting totals across regions and divisions;
    #Can also be used for mean values of totals across years, by looping through the dictionary

def eda_total(data, percentage_type = None):

    array_init = data.loc(percentage_type)
    divisor = len(array_init)
    percentages_array = np.array(array_init)
    percentages_sum = np.sum(percentages_array)
    percentages_mean = percentages_sum/divisor

    return percentages_mean


text:
#!/usr/bin/env python
import numpy as np
import dynaphopy.interface.iofile as io
import dynaphopy
from dynaphopy.interface.phonopy_link import get_force_constants_from_file

import unittest

class TestDynaphopy(unittest.TestCase):

    def setUp(self):
        structure = io.read_from_file_structure_poscar('MgO_data/POSCAR')

        structure.set_primitive_matrix([[0.0, 0.5, 0.5],
                                        [0.5, 0.0, 0.5],
                                        [0.5, 0.5, 0.0]])

        structure.set_force_constants(get_force_constants_from_file(file_name='MgO_data/FORCE_CONSTANTS',
                                                                    fc_supercell=[[2, 0, 0],
                                                                                [0, 2, 0],
                                                                                [0, 0, 2]]))

        trajectory = io.generate_test_trajectory(structure, supercell=[2, 2, 2], total_time=5, silent=False)
        self.calculation = dynaphopy.Quasiparticle(trajectory)

    def test_force_constants_self_consistency(self):
        self.calculation.select_power_spectra_algorithm(2)
        renormalized_force_constants = self.calculation.get_renormalized_force_constants().get_array()
        harmonic_force_constants = self.calculation.dynamic.structure.get_force_constants().get_array()
        self.assertEqual(np.allclose(renormalized_force_constants, harmonic_force_constants, rtol=1, atol=1.e-2), True)

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

max_stars_repo_name: faradaymahe/DynaPhopy

max_issues_repo_name: jianskerh/DynaPhoPy

max_forks_repo_name: jianskerh/DynaPhoPy


text:
import argparse
import numpy as np
import rospy

from franka_robot import FrankaRobot
from collision_boxes_publisher import CollisionBoxesPublisher
from rrt import RRT
from rrt_connect import RRTConnect
from prm import PRM
from ob_prm import OBPRM

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--seed', '-s', type=int, default=0)
    parser.add_argument('--rrt', '-rrt', type=str2bool, const=True, nargs='?', default=False, help="Use RRT?")
    parser.add_argument('--rrtc', '-rrtc', type=str2bool, const=True, nargs='?', default=False, help="Use RRT-Connect?")
    parser.add_argument('--prm', '-prm', type=str2bool, const=True, nargs='?', default=False, help="Use PRM?")
    parser.add_argument('--obprm', '-obprm', type=str2bool, const=True, nargs='?', default=False, help="Use OBPRM?")
    parser.add_argument('--map2', '-map2', type=str2bool, const=True, nargs='?', default=False, help="Use map 2?")
    parser.add_argument('--map3', '-map3', type=str2bool, const=True, nargs='?', default=False, help="Use map 3?")
    parser.add_argument('--reuse_graph', '-reuse_graph', type=str2bool, const=True, nargs='?', default=False, help="Reuse the graph for PRM?")
    args = parser.parse_args()

    np.random.seed(args.seed)
    fr = FrankaRobot()

    rospy.init_node('planner')

    '''
    TODO: Replace obstacle box w/ the box specs in your workspace:
    [x, y, z, r, p, y, sx, sy, sz]
    '''
    if args.map3:
        boxes = np.array([
            # obstacle
            # [0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0.45, -0.45, 0.7, 0, 0, 0.78, 0.6, 0.6, 0.05],
            # sides
            [-0.7, 0.7, 0.75, 0, 0, 0.78, 2, 0.01, 1.6],
            [0.7, -0.7, 0.75, 0, 0, 0.78, 2, 0.01, 1.6],
            # back
            [-0.7, -0.7, 0.75, 0, 0, 0.78, 0.01, 2, 1.6],
            # front
            [0.7, 0.7, 0.75, 0, 0, 0.78, 0.01, 2, 1.6],
            # top
            [0, 0, 1.5, 0, 0, 0.78, 2, 2, 0.01],
            # bottom
            [0, 0, -0.05, 0, 0, 0.78, 2, 2, 0.01]
        ])
    elif args.map2:
        boxes = np.array([
            # obstacle
            [0.7, 0, 0.6, 0, 0, 0, 0.45, 0.3, 0.05],
            # sides
            [0.15, 0.66, 0.65, 0, 0, 0, 1.2, 0.01, 1.5],
            [0.15, -0.66, 0.65, 0, 0, 0, 1.2, 0.01, 1.5],
            # back
            [-0.41, 0, 0.65, 0, 0, 0, 0.01, 1.4, 1.5],
            # front
            [0.75, 0, 0.65, 0, 0, 0, 0.01, 1.4, 1.5],
            # top
            [0.2, 0, 1.35, 0, 0, 0, 1.2, 1.4, 0.01],
            # bottom
            [0.2, 0, -0.05, 0, 0, 0, 1.2, 1.4, 0.01]
        ])
    else:
        boxes = np.array([
            # obstacle
            # [0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0.4, 0, 0.25, 0, 0, 0, 0.3, 0.05, 0.5],
            # sides
            [0.15, 0.46, 0.5, 0, 0, 0, 1.2, 0.01, 1.1],
            [0.15, -0.46, 0.5, 0, 0, 0, 1.2, 0.01, 1.1],
            # back
            [-0.41, 0, 0.5, 0, 0, 0, 0.01, 1, 1.1],
            # front
            [0.75, 0, 0.5, 0, 0, 0, 0.01, 1, 1.1],
            # top
            [0.2, 0, 1, 0, 0, 0, 1.2, 1, 0.01],
            # bottom
            [0.2, 0, -0.05, 0, 0, 0, 1.2, 1, 0.01]
        ])

    def is_in_collision(joints):
        if fr.check_self_collision(joints):
            return True
        for box in boxes:
            if fr.check_box_collision(joints, box):
                return True
        return False

    desired_ee_rp = fr.ee(fr.home_joints)[3:5]

    def ee_upright_constraint(q):
        '''
        TODO: Implement constraint function and its gradient.

        This constraint should enforce the end-effector stays upright.
        Hint: Use the roll and pitch angle in desired_ee_rp. The end-effector is upright in its home state.

        Input:
            q - a joint configuration

        Output:
            err - a non-negative scalar that is 0 when the constraint is satisfied
            grad - a vector of length 6, where the ith element is the derivative of err w.r.t. the ith element of ee
        '''
        ee = fr.ee(q)
        err = np.sum((np.asarray(desired_ee_rp) - np.asarray(ee[3:5])) ** 2)
        grad = np.asarray([0, 0, 0, 2 * (ee[3] - desired_ee_rp[0]), 2 * (ee[4] - desired_ee_rp[1]), 0])
        return err, grad

    def get_plan_quality(plan):
        dist = 0
        for i in range(len(plan) - 1):
            dist += np.linalg.norm(np.array(plan[i+1]) - np.array(plan[i]))
        return dist

    '''
    TODO: Fill in start and target joint positions
    '''
    if args.map3:
        joints_start = np.array([0, 3*np.pi/8, 0, -np.pi / 8, 0, np.pi / 2, np.pi / 4])
        joints_start[0] = -np.deg2rad(45)
        joints_target = np.array([0, 0, 0, -np.pi / 4, 0, np.pi / 4, np.pi / 4])
        joints_target[0] = -np.deg2rad(45)
    elif args.map2:
        joints_start = np.array([0, np.pi/6, 0, -2*np.pi / 3, 0, 5*np.pi / 6, np.pi / 4])
        joints_target = np.array([0, 0, 0, -np.pi / 4, 0, np.pi / 4, np.pi / 4])
    else:
        joints_start = fr.home_joints.copy()
        joints_start[0] = -np.deg2rad(45)
        joints_target = joints_start.copy()
        joints_target[0] = np.deg2rad(45)

    if args.rrt:
        print("RRT: RRT planner is selected!")
        planner = RRT(fr, is_in_collision)
    elif args.rrtc:
        print("RRTC: RRT Connect planner is selected!")
        planner = RRTConnect(fr, is_in_collision)
    elif args.prm:
        print("PRM: PRM planner is selected!")
        planner = PRM(fr, is_in_collision)
    elif args.obprm:
        print("OB_PRM: OB_PRM planner is selected!")
        planner = OBPRM(fr, is_in_collision)

    constraint = ee_upright_constraint
    if args.prm or args.obprm:
        plan = planner.plan(joints_start, joints_target, constraint, args)
    else:
        plan = planner.plan(joints_start, joints_target, constraint)

    path_quality = get_plan_quality(plan)
    print("Path quality: {}".format(path_quality))

    collision_boxes_publisher = CollisionBoxesPublisher('collision_boxes')
    rate = rospy.Rate(10)
    i = 0
    while not rospy.is_shutdown():
        rate.sleep()
        joints = plan[i % len(plan)]
        fr.publish_joints(joints)
        fr.publish_collision_boxes(joints)
        collision_boxes_publisher.publish_boxes(boxes)
        i += 1


text:
"""
Growth Indicators
===============================================================================

>>> from techminer2 import *
>>> directory = "data/"
>>> growth_indicators(column="author_keywords", directory=directory).head()
                        before 2020  ...  average_growth_rate
author_keywords                      ...
fintech                          42  ...                 21.5
financial technologies           11  ...                  4.0
blockchain                        9  ...                  1.0
innovation                        9  ...                 -1.0
financial service                 7  ...                  1.0
<BLANKLINE>
[5 rows x 4 columns]

"""
import numpy as np
import pandas as pd

from ._read_records import read_filtered_records
from .load_stopwords import load_stopwords

def _num_documents_per_period(documents, column, time_window=2, sep="; "):
    #
    # Computes the total number of documents published in each period
    #
    if time_window < 2:
        raise ValueError("Time window must be greater than 1")
    year_limit = documents.pub_year.max() - time_window + 1
    result = documents.copy()
    result.index = result.record_no
    result = result[[column, "pub_year"]].copy()
    result[column] = result[column].str.split(sep)
    result = result.explode(column)
    result = result.assign(
        before=result.pub_year.map(lambda x: 1 if x < year_limit else 0)
    )
    result = result.assign(
        between=result.pub_year.map(lambda x: 1 if x >= year_limit else 0)
    )
    result = result.groupby(column, as_index=False).agg(
        {"before": np.sum, "between": np.sum}
    )
    result = result.sort_values("before", ascending=False)
    result = result.set_index(column)

    return result

def _average_growth_rate(documents, column, time_window, sep="; "):
    #
    #         sum_{i=Y_start}^Y_end  Num_Documents[i] - Num_Documents[i-1]
    #  AGR = --------------------------------------------------------------
    #                          Y_end - Y_start + 1
    #
    #
    if time_window < 2:
        raise ValueError("Time window must be greater than 1")

    # first and last years in the time window
    first_year = documents.pub_year.max() - time_window
    last_year = documents.pub_year.max()

    # generates a table of term and year.
    result = documents.copy()
    result.index = result.record_no
    result = result[[column, "pub_year"]].copy()
    result[column] = result[column].str.split(sep)
    result = result.explode(column)
    result = result[(result.pub_year == first_year) | (result.pub_year == last_year)]
    result = result.assign(num_documents=1)

    #
    result = result.groupby([column, "pub_year"], as_index=False).agg(
        {"num_documents": np.sum}
    )
    result = result.pivot(index=column, columns="pub_year", values="num_documents")
    result = result.fillna(0)
    result = result.assign(
        average_growth_rate=(result.iloc[:, 1] - result.iloc[:, 0]) / time_window
    )
    result = result[["average_growth_rate"]]

    return result

def _average_documents_per_year(documents, column, time_window, sep="; "):
    #
    #         sum_{i=Y_start}^Y_end  Num_Documents[i]
    #  ADY = -----------------------------------------
    #                  Y_end - Y_start + 1
    #
    result = _num_documents_per_period(
        documents=documents, column=column, time_window=time_window, sep=sep
    )
    result = result.assign(average_documents_per_year=result.between / time_window)
    result = result[["average_documents_per_year"]]
    return result

def growth_indicators(column, sep="; ", time_window=2, directory="./"):

    documents = read_filtered_records(directory)

    ndpp = _num_documents_per_period(
        documents=documents, column=column, time_window=time_window, sep=sep
    )
    adpy = _average_documents_per_year(
        documents=documents, column=column, time_window=time_window, sep=sep
    )
    agr = _average_growth_rate(
        documents=documents, column=column, time_window=time_window, sep=sep
    )
    result = pd.concat([ndpp, adpy, agr], axis="columns")

    year_limit = documents.pub_year.max() - time_window + 1
    result = result.rename(
        columns={
            "before": f"before {year_limit}",
            "between": f"between {year_limit}-{documents.pub_year.max()}",
        }
    )

    stopwords = load_stopwords(directory)
    result = result.drop(stopwords, axis=0)

    return result


text:
import cv2
import numpy as np
from keras.models import load_model
import solver

def image_processing_boundary(image):
    """Returns input image after processing for sudoku boundary detection"""
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = cv2.GaussianBlur(image, (3, 3), 0)
    """m = np.median(image)
    sigma = 0.33
    lower = int(max(0, (1.0 - sigma) * m))
    upper = int(min(255, (1.0 + sigma) * m))"""
    image = cv2.Canny(image, 35, 70)

    return image

def image_processing_digits(image):
    """Returns input image after processing for digit recognition"""
    image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    image = cv2.GaussianBlur(image, (3, 3), 0)
    image = cv2.adaptiveThreshold(image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 21, 10)
    #image = cv2.bitwise_not(image)
    return image

def maxContour(contours):
    """Returns the contour with maximum area, if any"""
    if not contours:
        return None

    areas = [cv2.contourArea(x) for x in contours]
    max_index = np.argmax(areas)
    max_cnt = contours[max_index]

    if areas[max_index] > 50000:
        return max_cnt
    else:
        return None

def getCorners(max_cnt):
    """Returns the edges of the contour if the approx polygon is a quadrilateral"""
    perimeter = cv2.arcLength(max_cnt, True)
    epsilon = 0.1 * perimeter
    approx = cv2.approxPolyDP(max_cnt, epsilon, True)
    boundary_pts = None

    if len(approx) == 4:
        x2, y2, x1, y1, x3, y3, x4, y4 = approx.item(0), approx.item(1), approx.item(2), approx.item(3), approx.item(
            4), approx.item(5), approx.item(6), approx.item(7)
        boundary_pts = np.float32([[x1, y1], [x2, y2], [x3, y3], [x4, y4]])

    return boundary_pts

def grid_to_batch(cell_grid):
    """Takes an array of each image cell and return a single batch for ML prediction. Also returns estimated empty cells"""
    cell_batch = np.empty([0, 28, 28, 1])
    empty_filter = np.ones(81)
    for i, row in enumerate(cell_grid):
        for j, cell in enumerate(row):
            cell_batch = np.append(cell_batch, cell.reshape(1, 28, 28, 1), axis=0)
            if np.sum(cell) < 15000:
                empty_filter[i * 9 + j] = 0
    return cell_batch, empty_filter

def predict_batch(cell_batch, empty_filter):
    """Predicts the digits for the given mask and returns it after replacing empty cells with zeroes"""
    predictions_vectorized = model.predict(cell_batch)
    predictions_raw = np.argmax(predictions_vectorized, axis=1)
    predictions = predictions_raw * empty_filter
    return predictions

def draw_over_predictions(image):
    print(type(image))
    l, w, _ = image.shape
    #image = np.zeros((image.shape), dtype=np.uint8)
    for i in range(9):
        for j in range(9):
            digit = str(int(pred_reshaped[i][j]))
            if digit == '0':
                continue
            io = int(l * i / 9 + 32)
            jo = int(w * j / 9 + 27)
            image = cv2.putText(image, digit, (jo, io), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    return image

def draw_over_solution(image):
    #black = np.zeros((342, 342), dtype=np.uint8)
    l, w, _ = image.shape
    for i in range(9):
        for j in range(9):
            sol_digit = sol_reshaped[i][j]
            pred_digit = str(int(pred_reshaped[i][j]))
            if pred_digit == '0':
                io = int(l * i / 9 + 28)
                jo = int(w * j / 9 + 12)
                image = cv2.putText(image, sol_digit, (jo, io), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
    return image

vid = cv2.VideoCapture(0)
model = load_model('model')

while True:
    ret, frame = vid.read()

    img_processed_boundary = image_processing_boundary(frame)
    contours, _ = cv2.findContours(img_processed_boundary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

    img_processed_digits = image_processing_digits(frame)

    max_contour = maxContour(contours)
    image_contours = None
    if max_contour is not None:
        cv2.drawContours(frame, max_contour, -1, (0, 0, 255), 2)
        boundary_pts = getCorners(max_contour)
        if boundary_pts is not None:
            width, height = 342, 342
            new_boundary = np.float32([[0, 0], [width, 0], [0, height], [width, height]])
            persp_mat = cv2.getPerspectiveTransform(boundary_pts, new_boundary)
            image_contours = cv2.warpPerspective(img_processed_digits, persp_mat, (width, height))

    cv2.imshow('img_processed_boundary', img_processed_boundary)
    cv2.imshow('frame', frame)
    if image_contours is not None:
        cv2.imshow('image_contours', image_contours)

        cell_grid = []
        for i in range(9):
            row = []
            for j in range(9):
                row.append(image_contours[5 + i * 38:33 + i * 38, 5 + j * 38:33 + j * 38])
            cell_grid.append(row)

        cv2.imshow('cell1', cell_grid[0][2])

        cell_batch, empty_filter = grid_to_batch(cell_grid)
        predictions = predict_batch(cell_batch, empty_filter)
        pred_reshaped = predictions.reshape(9, 9)
        print(pred_reshaped)

        frame_persp = cv2.warpPerspective(frame, persp_mat, (342, 342))
        image_write_over = draw_over_predictions(frame_persp)
        cv2.imshow("write_over", image_write_over)

        pred_formatted = np.array2string(predictions.astype('int'), separator='')
        pred_formatted = pred_formatted[1:-1]
        solution = solver.solve(pred_formatted)

        if solution:
            print(solution.values())
            sol_reshaped = np.asarray(list(solution.values())).reshape(9, 9)
            sol_write_over = draw_over_solution(image_write_over)
            cv2.imshow("sol_write_over", sol_write_over)
            #print(np.fromiter(solution.values(), dtype=str))

            persp_mat = cv2.getPerspectiveTransform(new_boundary, boundary_pts)
            sol_warped = cv2.warpPerspective(sol_write_over, persp_mat, (frame.shape[1], frame.shape[0]))
            sol_warp_overlayed = cv2.max(sol_warped, frame)
            cv2.imshow("sol_warp_overlayed", sol_warp_overlayed)
            #print(frame.shape)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

vid.release()

max_stars_repo_name: vineet-ks/sudoku-vision

max_issues_repo_name: vineet-ks/sudoku-vision

max_forks_repo_name: vineet-ks/sudoku-vision


text:
import utils
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC, LinearSVC
from sklearn.model_selection import GridSearchCV
from sklearn.feature_selection import RFE

def get_regularization_params(a=-1, b=1, c=3, d=1, e=5):
    reg_range = np.outer(np.logspace(a, b, c), np.array([d, e]))
    reg_range = reg_range.flatten()
    return reg_range

def grid_classifier(x_train, y_train, x_test, y_test, model, parameters,
                    make_feature_analysis=False, feature_names=None, top_features=0, plot_name="coeff"):
    grid = GridSearchCV(estimator=model, param_grid=parameters, verbose=0)
    grid.fit(x_train, y_train)
    sorted(grid.cv_results_.keys())
    classifier = grid.best_estimator_
    if make_feature_analysis:
        utils.plot_coefficients(classifier, feature_names, top_features, plot_name)
    y_hat = classifier.predict(x_test)
    utils.print_statistics(y_test, y_hat)

def linear_svm_grid(x_train, y_train, x_test, y_test, class_ratio,
               make_feature_analysis=False, feature_names=None, top_features=0, plot_name="coeff"):
    utils.print_model_title("Linear SVM")
    C_range = get_regularization_params()
    parameters = {'C': C_range}
    linear_svm = LinearSVC(C=1.0, class_weight=class_ratio, penalty='l2')
    grid_classifier(x_train, y_train, x_test, y_test, linear_svm, parameters,
                    make_feature_analysis, feature_names, top_features, plot_name)

def nonlinear_svm_grid(x_train, y_train, x_test, y_test, class_ratio,
                  make_feature_analysis=False, feature_names=None, top_features=0, plot_name="coeff"):
    utils.print_model_title("Nonlinear SVM")
    C_range = get_regularization_params(a=-1, b=0, c=2, d=1, e=5)
    gamma_range = get_regularization_params(a=-2, b=-1, c=2, d=1, e=5)
    parameters = {'kernel': ['rbf'], 'C': C_range, 'gamma': gamma_range}
    nonlinear_svm = SVC(class_weight=class_ratio)
    grid_classifier(x_train, y_train, x_test, y_test, nonlinear_svm, parameters,
                    make_feature_analysis, feature_names, top_features, plot_name)

def logistic_regression_grid(x_train, y_train, x_test, y_test, class_ratio,
                        make_feature_analysis=False, feature_names=None, top_features=0, plot_name="coeff"):
    utils.print_model_title("Logistic Regression")
    C_range = [0.001, 0.01, 0.1, 1, 10, 100]
    parameters = {'C': C_range}
    log_regr = LogisticRegression(C=1.0, class_weight=class_ratio, penalty='l2')
    grid_classifier(x_train, y_train, x_test, y_test, log_regr, parameters,
                    make_feature_analysis, feature_names, top_features, plot_name)

def linear_svm(x_train, y_train, x_test, y_test, class_ratio='balanced'):
    utils.print_model_title("Linear SVM")
    svm = LinearSVC(C=0.01, class_weight=class_ratio, penalty='l2')
    svm.fit(x_train, y_train)
    y_hat = svm.predict(x_test)
    utils.print_statistics(y_test, y_hat)

def logistic_regression(x_train, y_train, x_test, y_test, class_ratio='balanced'):
    utils.print_model_title("Logistic Regression")
    regr = LogisticRegression(C=0.01, class_weight=class_ratio, penalty='l2')
    regr.fit(x_train, y_train)
    y_hat = regr.predict(x_test)
    utils.print_statistics(y_test, y_hat)

def feature_selection(x_train, y_train, x_test, y_test):
    print("Feature selection with LinearSVC")
    model = LinearSVC(C=0.1, penalty='l2')
    rfe = RFE(model, 5)
    best_features_model = rfe.fit(x_train, y_train)
    y_hat = best_features_model.predict(x_test)
    utils.print_statistics(y_test, y_hat)