text:
#!/usr/bin/python
import sys
import logging
import re
import numpy as np

from parsevasp import utils
from parsevasp.base import BaseParser

class Outcar(BaseParser):
    ERROR_NO_ITERATIONS = 600
    BaseParser.ERROR_MESSAGES.update({
        ERROR_NO_ITERATIONS: 'A crash detected before the first SCF step.'
    })
    ERROR_MESSAGES = BaseParser.ERROR_MESSAGES

    def __init__(self,
                 file_path=None,
                 file_handler=None,
                 logger=None):
        """Initialize an OUTCAR object and set content as a dictionary.

        Parameters
        ----------
        file_path : string
            A string containing the file path to the file that is going to be parsed.
        file_handler : object
            A file like object that acts as a handler for the content to be parsed.
        logger : object
            A logger object if you would like to use an external logger for messages
            ejected inside this parser.

        """

        super(Outcar, self).__init__(file_path=file_path,
                                     file_handler=file_handler,
                                     logger=logger)

        # check that at least one is supplied
        if self._file_path is None and self._file_handler is None:
            self._logger.error(
                self.ERROR_MESSAGES[self.ERROR_USE_ONE_ARGUMENT])
            sys.exit(self.ERROR_USE_ONE_ARGUMENT)

        if self._file_path is None and self._file_handler is None:
            self._logger.error(
                self.ERROR_MESSAGES[self.ERROR_USE_ONE_ARGUMENT])
            sys.exit(self.ERROR_USE_ONE_ARGUMENT)

        self._data = {
            'elastic_moduli': {
                'non-symmetrized': None,
                'symmetrized': None,
                'total': None
            },
            'symmetry': {
                'num_space_group_operations': {
                    'static': [],
                    'dynamic': []
                },
                'original_cell_type': {
                    'static': [],
                    'dynamic': []
                },
                'symmetrized_cell_type': {
                    'static': [],
                    'dynamic': []
                }
            },
            'magnetization': {
                'sphere': {
                    'x': {
                        'site_moment': {},
                        'total_magnetization': {}
                    },
                    'y': {
                        'site_moment': {},
                        'total_magnetization': {}
                    },
                    'z': {
                        'site_moment': {},
                        'total_magnetization': {}
                    },
                },
                'full_cell': {},
            },
            'run_stats': {},
            'run_status': {'nelm': None,
                           'nsw': None,
                           'last_iteration_index': None,
                           'finished': False,
                           'ionic_converged': False,
                           'electronic_converged': False,
                           'consistent_nelm_breach': False,
                           'contains_nelm_breach': False
            }
        }

        # parse parse parse
        self._parse()

    def _write(self, *args, **kwargs):
        """Write not supported for OUTCAR."""
        raise NotImplementedError('Writing OUTCAR files is not supported.')

    def _parse(self):
        """Perform the actual parsing."""
        # Create dictionary from a file
        self._from_file()

    def _from_file(self):
        """
        Create a dictionary of entries from a
        file and store them in the this instance's data dictionary.
        """

        outcar = utils.read_from_file(self._file_path, self._file_handler)
        self._from_list(outcar)

    def _from_list(self, outcar):
        """
        Go through the list and extract what is not present in the XML file.

        Parameters
        ----------
        outcar : list
            A list of strings containing each line in the OUTCAR file.

        """
        config = ''
        s_orb = {0: 's', 1: 'p', 2: 'd', 3: 'f'}
        params = {'ibrion': -1}
        finished = False
        iter_counter = None
        nelec_steps = {}

        for index, line in enumerate(outcar):
            # Check the iteration counter
            match = re.search(r'Iteration *(\d+)\( *(\d+)\)', line)
            if match:
                iter_counter = [int(match.group(1)), int(match.group(2))]
                # Increment the counter
                if iter_counter[0] in nelec_steps:
                    nelec_steps[iter_counter[0]] += 1
                else:
                    nelec_steps[iter_counter[0]] = 1
                continue
            # Record the NELM / NSW requested
            utils.match_integer_param(self._data['run_status'], 'NSW', line)
            utils.match_integer_param(params, 'IBRION', line)
            utils.match_integer_param(self._data['run_status'], 'NELM', line)
            # Test if the end of execution has reached
            if 'timing and accounting informations' in line:
                self._data['run_status']['finished'] = True

            # Fetch the symmetry
            if line.strip().startswith(
                    'Analysis of symmetry for initial positions (statically)'):
                config = 'static'
            if line.strip().startswith('Analysis of symmetry for dynamics'):
                config = 'dynamic'
            if config:
                if line.strip().startswith('Subroutine PRICEL returns'):
                    text = outcar[index + 1].strip().lower()
                    if text:
                        self._data['symmetry']['original_cell_type'][
                            config].append('primitive cell')
                if ('primitive cells build up your supercell') in line:
                    text = '{} primitive cells'.format(line.strip().split())
                    self._data['symmetry']['original_cell_type'][
                        config].append(text)
                if line.strip().startswith(
                        'Routine SETGRP: Setting up the symmetry group for a'):
                    self._data['symmetry']['symmetrized_cell_type'][
                        config].append(outcar[index + 1].strip().lower())
                if line.strip().startswith('Subroutine GETGRP returns'):
                    self._data['symmetry']['num_space_group_operations'][
                        config].append(int(line.strip().split()[4]))

            # then the elastic tensors etc. in kBar
            if line.strip().startswith('ELASTIC MODULI  (kBar)'):
                tensor = []
                for idx in range(3, 9):
                    tensor.append([
                        float(item)
                        for item in outcar[index + idx].strip().split()[1:]
                    ])
                self._data['elastic_moduli']['non_symmetrized'] = np.asarray(
                    tensor)
            if line.strip().startswith('SYMMETRIZED ELASTIC MODULI'):
                tensor = []
                for idx in range(3, 9):
                    tensor.append([
                        float(item)
                        for item in outcar[index + idx].strip().split()[1:]
                    ])
                self._data['elastic_moduli']['symmetrized'] = np.asarray(
                    tensor)
            if line.strip().startswith('TOTAL ELASTIC MODULI'):
                tensor = []
                for idx in range(3, 9):
                    tensor.append([
                        float(item)
                        for item in outcar[index + idx].strip().split()[1:]
                    ])
                self._data['elastic_moduli']['total'] = np.asarray(tensor)
            for _proj in ['x', 'y', 'z']:
                if line.strip().startswith('magnetization ({})'.format(_proj)):
                    _counter = 0
                    mag_found = False
                    while not mag_found:
                        if outcar[index + 4 + _counter].strip().split():
                            if not outcar[index + 4 + _counter].strip(
                            ).startswith('-') and not outcar[
                                    index + 4 +
                                    _counter].strip().startswith('tot'):
                                mag_line = outcar[index + 4 + _counter].split()
                                self._data['magnetization']['sphere'][
                                    '{}'.format(_proj)]['site_moment'][int(
                                        mag_line[0])] = dict()
                                for _count, orb in enumerate(mag_line[1:-1]):
                                    self._data['magnetization']['sphere'][
                                        '{}'.format(_proj)]['site_moment'][int(
                                            mag_line[0])][
                                                s_orb[_count]] = float(orb)
                                self._data['magnetization']['sphere'][
                                    '{}'.format(_proj)]['site_moment'][int(
                                        mag_line[0])]['tot'] = float(
                                            mag_line[-1])
                            if outcar[index + 4 +
                                      _counter].strip().startswith('tot'):
                                mag_line = outcar[index + 4 + _counter].split()
                                self._data['magnetization']['sphere'][
                                    '{}'.format(
                                        _proj)]['total_magnetization'] = dict(
                                        )
                                for _count, orb in enumerate(mag_line[1:-1]):
                                    self._data['magnetization']['sphere'][
                                        '{}'.format(
                                            _proj)]['total_magnetization'][
                                                s_orb[_count]] = float(orb)
                                self._data['magnetization']['sphere'][
                                    '{}'.format(_proj)]['total_magnetization'][
                                        'tot'] = float(mag_line[-1])
                                mag_found = True
                        else:
                            self._data['magnetization']['sphere']['{}'.format(
                                _proj)]['total_magnetization'] = dict()
                            self._data['magnetization']['sphere']['{}'.format(_proj)]['total_magnetization'] =\
                                self._data['magnetization']['sphere']['{}'.format(_proj)]['site_moment'][list(self._data['magnetization']['sphere']['{}'.format(_proj)]['site_moment'].keys())[0]]
                            mag_found = True
                        _counter = _counter + 1
            if line.strip().startswith('number of electron'):
                # Only take the last value
                self._data['magnetization']['full_cell'] = [
                    float(_val) for _val in outcar[index].strip().split()[5:]
                ]

        # Check if SCF iterations are contained in the file
        if iter_counter is None:
            self._logger.error(self.ERROR_MESSAGES[self.ERROR_NO_ITERATIONS])
            sys.exit(self.ERROR_NO_ITERATIONS)

        # Work out if the ionic relaxation and electronic steps are to be considered converged
        run_status = self._data['run_status']
        run_status['last_iteration_index'] = iter_counter
        nsw = run_status['nsw']
        nelm = run_status['nelm']
        finished = run_status['finished']
        ibrion = params['ibrion']
        if finished is True:
            if ibrion > 0:
                # There are fewer number of ionic iterations than the set number of maximum number of
                # ionic iterations (NSW), thus the relaxation is considered converged.
                # Only relevant to check ionic relaxation convergence when IBRION is larger than zero
                if iter_counter[0] < nsw:
                    # Fewer iterations than requested - ionic relaxation is considered converged
                    # Note that this may include runs that has been interrupted
                    # by STOPCAR - this is a limitation of VASP
                    run_status['ionic_converged'] = True
                elif iter_counter[0] == nsw and nsw > 1:
                    # Reached the requested iterations - ionic relaxation is considered not converged
                    run_status['ionic_converged'] = False
                elif nsw <= 1:
                    # Sometimes we have no or only one ionic step, which makes it difficult to determine if the
                    # ionic relaxation is to be considered converged
                    self._logger.warning(
                        f'IBRION = {ibrion} but NSW is {nsw} - cannot deterimine if the relaxation structure is converged!')
                    run_status['ionic_converged'] = None
            else:
                # No ionic relaxation performed
                run_status['ionic_converged'] = None

            if iter_counter[1] < nelm:
                # There are fewer number of electronic steps in the last ionic iteration than the set maximum
                # number of electronic steps, thus the electronic self consistent cycle is considered converged
                run_status['electronic_converged'] = True

        # Check for consistent electronic convergence problems. VASP will not break when NELM is reached during
        # the relaxation, it will simply consider it converged. We need to detect this, which is done
        # by checking if there are any single run that have reached NELM in the history or if NELM
        # has been consistently reached.
        mask = [value >= nelm
                for sc_idx, value in sorted(nelec_steps.items(), key=lambda x: x[0])]
        if (finished and all(mask)) or \
           (not finished and all(mask[:-1]) and iter_counter[0] > 1):
            # We have consistently reached NELM. Excluded the last iteration,
            # as the calculation may not be finished
            run_status['consistent_nelm_breach'] = True
        if any(mask):
            # We have at least one ionic step where NELM was reached.
            run_status['contains_nelm_breach'] = True

        self._data['run_stats'] = self._parse_timings_memory(outcar[-50:])

    def get_symmetry(self):
        """
        Return the symmetry.

        Parameters
        ----------
        None

        Returns
        -------
        symmetry : dict
            A dictionary containing the symmetry information.

        """

        symmetry = self._data['symmetry']
        return symmetry

    def get_elastic_moduli(self):
        """
        Return the elastic moduli in kBar.

        Parameters
        ----------
        None

        Returns
        -------
        elastic : dict
            A dictionary containing the elastic moduli.

        """

        elastic = self._data['elastic_moduli']
        return elastic

    def get_magnetization(self):
        """
        Return the magnetization of the cell.

        Parameters
        ----------
        None

        Returns
        -------
        magnetic : dict
            A dictionary containing the magnetization of the cell.

        """

        magnetic = self._data['magnetization']
        return magnetic

    def get_run_stats(self):
        """
        Return the run time statistics information.

        The existence of timing and memory information also signals the calculation terminate
        gracefully.

        Parameters
        ----------
        None

        Returns
        -------
        stats : dict
            A dictionary containing timing and memory consumption information
            that are parsed from the end of the OUTCAR file. The key names are
            mostly preserved, except for the memory which is prefixed with `mem_usage_`.
            Units are preserved from OUTCAR and there are some differences between
            VASP 5 and 6.

        """

        stats = self._data['run_stats']
        return stats

    def get_run_status(self):
        """
        Return the status of the run.

        Contains information of the convergence of the ionic relaxation and electronics,
        in addition to information if the run has finished.

        Parameters
        ----------
        None

        Returns
        -------
        status : dict
            A dictionary containing the keys `finished`, which is True if the VASP calculation
            contain timing information in the end of the OUTCAR. The key `ionic_converged` is
            True if the number of ionic steps detected is smaller than the supplied NSW.
            The key `electronic_converged` is True if the number of electronic steps is smaller than
            NELM (defaults to 60 in VASP). It is also possible to check if all the ionic steps
            did reached NELM and thus did not converged if the key `consistent_nelm_breach` is True,
            while `contains_nelm_breach` is True if one or more ionic steps reached NELM and thus
            did not converge electronically.

        """

        status = self._data['run_status']
        return status

    @staticmethod
    def _parse_timings_memory(timing_lines):
        """
        Parse timing information.

        Parameters
        ----------
        timing_lines : list
            A list of lines containing the timing information.

        Returns
        -------
        info : dict
            A dictionary containing the timing information.
        """
        import re
        info = {}
        time_mem_pattern = re.compile(r"\((sec|kb)\)")
        mem_pattern = re.compile(r":.*kBytes$")
        for _, line in enumerate(timing_lines):
            if time_mem_pattern.search(line):
                tokens = line.strip().split(":")
                item_name = "_".join(tmp.lower() for tmp in tokens[0].strip().split()[:-1])
                # The entry can be empty (VASP6)
                try:
                    info[item_name] = float(tokens[1].strip())
                except ValueError:
                    info[item_name] = None

            elif mem_pattern.search(line):
                tokens = re.split(r'[: ]+', line.strip())
                try:
                    info['mem_usage_' + tokens[0]] = float(tokens[-2])
                except ValueError:
                    info['mem_usage_' + tokens[0]] = None

        return info

ext: py

max_stars_repo_name: atztogo/parsevasp

max_issues_repo_name: atztogo/parsevasp

max_forks_repo_name: atztogo/parsevasp

reason: import numpy


text:
#!/usr/bin/env python3

import numpy as np
from tqdm import tqdm

from torch import nn
import torch
import torch.utils.data
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.base import clone
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
from sklearn.utils.multiclass import unique_labels
from sklearn.metrics import accuracy_score

from .Models import Ensemble

class GradientBoostedNets(Ensemble):
    ''' (Stochastic) Gradient Boosting for Neural Networks

    Gradient Boosting sequentially trains a classifier on the residuals of the ensemble. In its most basic form this is usually a stagewise process in which a new classifier is trained in each stage / round [1] on the errors of the entire model. Related Boosting algorithms like AdaBoost can also be viewed in this framework [2,3]. In order to speed-up the training process, Stochastic Gradient boosting [4] has been proposed. Stochastic Gradient Boosting trains the individual classifiers in each round on a random sample of the entire dataset.

    This implementation is a variation of Stochastic Gradient Boosting and combines it with Stochastic Gradient Descent. The basic idea is to freeze the weights of individual models \( h^1,\dots, h^{i-1} \) and perform SGD on the i-th model's parameter using \( \ell(\\frac{1}{M}\sum_{j=1}^i h^j(x), y) \). This process is repeated for each model. As far as I know this specific training has not been discussed in literature so far, but there is some overlap with exsiting work. See e.g. [5] and references therein.

    __References__

    [1] Friedman, J. H. (2001). Greedy Function Approximation: A Gradient Boosting Machine. Retrieved from https://www.jstor.org/stable/pdf/2699986.pdf?seq=1

    [2] Mason, L., Baxter, J., Bartlet, P., & Frean, M. (1999). Boosting Algorithms as Gradient Descent in Function Space. Retrieved from http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.51.6893&rep=rep1&type=pdf

    [3] Schapire, R. E., & Freund, Y. (2012). Boosting: Foundations and algorithms. MIT press.

    [4] Friedman, J. H. (1999). Stochastic Gradient Boosting, 1(3), 1–10. Retrieved from http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.31.1666&rep=rep1&type=pdf

    [5] Hu, H., Sun, W., Venkatraman, A., Hebert, M., & Bagnell, J. A. (2017). Gradient Boosting on Stochastic Data Streams, 54. Retrieved from https://arxiv.org/pdf/1703.00377.pdf

    '''
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.estimators_ = nn.ModuleList([ self.base_estimator() for _ in range(self.n_estimators)])

        # Currently we only support average for this kind of training I guess.
        # Not sure how to change the overall training procedure to fit this tbh.
        assert self.combination_type == "average"

    def prepare_backward(self, data, target, weights = None):
        f_bar, base_preds = self.forward_with_base(data)

        staged_pred = None
        total_loss = None

        for pred in base_preds:
            if staged_pred is None:
                staged_pred = pred
            else:
                staged_pred = staged_pred.detach() + pred

            if weights is not None:
                loss = self.loss_function(staged_pred, target) * weights
            else:
                loss = self.loss_function(staged_pred, target)

            if total_loss is None:
                total_loss = loss
            else:
                total_loss += loss

        d = {
            "prediction" : f_bar,
            "backward" : total_loss,
            "metrics" :
            {
                "loss" : self.loss_function(f_bar, target),
                "accuracy" : 100.0*(f_bar.argmax(1) == target).type(self.get_float_type()),
            }

        }
        return d

max_stars_repo_name: sbuschjaeger/Pysembles

max_issues_repo_name: sbuschjaeger/Pysembles

max_forks_repo_name: sbuschjaeger/Pysembles


text:
import numpy as np
import pandas as pd

from typing import Tuple, List, Dict
from datetime import datetime
from abc import ABCMeta, abstractmethod

from lib.data.providers.dates import ProviderDateFormat

class BaseDataProvider(object, metaclass=ABCMeta):
    columns = ['Date', 'Open', 'High', 'Low', 'Close', 'Volume']
    in_columns = ['Date', 'Open', 'High', 'Low', 'Close', 'Volume']
    custom_datetime_format = None

    @abstractmethod
    def __init__(self, date_format: ProviderDateFormat, **kwargs):
        self.date_format = date_format

        self.custom_datetime_format: str = kwargs.get('custom_datetime_format', None)

        data_columns: Dict[str, str] = kwargs.get('data_columns', None)

        if data_columns is not None:
            self.data_columns = data_columns
            self.columns = list(data_columns.keys())
            self.in_columns = list(data_columns.values())
        else:
            self.data_columns = dict(zip(self.columns, self.in_columns))

    @abstractmethod
    def split_data_train_test(self, train_split_percentage: float = 0.8) -> Tuple:
        raise NotImplementedError

    @abstractmethod
    def historical_ohlcv(self) -> pd.DataFrame:
        raise NotImplementedError

    @abstractmethod
    def reset_ohlcv_index(self) -> int:
        raise NotImplementedError

    @abstractmethod
    def has_next_ohlcv(self) -> bool:
        raise NotImplementedError

    @abstractmethod
    def next_ohlcv(self) -> pd.DataFrame:
        raise NotImplementedError

    def prepare_data(self, data_frame: pd.DataFrame, inplace: bool = True) -> pd.DataFrame:
        column_map = dict(zip(self.in_columns, self.columns))

        formatted = data_frame[self.in_columns]
        formatted = formatted.rename(index=str, columns=column_map)

        formatted = self._format_date_column(formatted, inplace=inplace)
        formatted = self._sort_by_date(formatted, inplace=inplace)

        return formatted

    def _sort_by_date(self, data_frame: pd.DataFrame, inplace: bool = True) -> pd.DataFrame:
        if inplace is True:
            formatted = data_frame
        else:
            formatted = data_frame.copy()

        formatted = formatted.sort_values(self.data_columns['Date'])

        return formatted

    def _format_date_column(self, data_frame: pd.DataFrame, inplace: bool = True) -> pd.DataFrame:
        if inplace is True:
            formatted = data_frame
        else:
            formatted = data_frame.copy()

        date_col = self.data_columns['Date']
        date_frame = formatted.loc[:, date_col]

        if self.date_format is ProviderDateFormat.TIMESTAMP_UTC:
            formatted[date_col] = date_frame.apply(
                lambda x: datetime.utcfromtimestamp(x).strftime('%Y-%m-%d %H:%M'))
            formatted[date_col] = pd.to_datetime(date_frame, format='%Y-%m-%d %H:%M')
        elif self.date_format is ProviderDateFormat.TIMESTAMP_MS:
            formatted[date_col] = pd.to_datetime(date_frame, unit='ms')
        elif self.date_format is ProviderDateFormat.DATETIME_HOUR_12:
            formatted[date_col] = pd.to_datetime(date_frame, format='%Y-%m-%d %I-%p')
        elif self.date_format is ProviderDateFormat.DATETIME_HOUR_24:
            formatted[date_col] = pd.to_datetime(date_frame, format='%Y-%m-%d %H')
        elif self.date_format is ProviderDateFormat.DATETIME_MINUTE_12:
            formatted[date_col] = pd.to_datetime(date_frame, format='%Y-%m-%d %I:%M-%p')
        elif self.date_format is ProviderDateFormat.DATETIME_MINUTE_24:
            formatted[date_col] = pd.to_datetime(date_frame, format='%Y-%m-%d %H:%M')
        elif self.date_format is ProviderDateFormat.DATE:
            formatted[date_col] = pd.to_datetime(date_frame, format='%Y-%m-%d')
        elif self.date_format is ProviderDateFormat.CUSTOM_DATIME:
            formatted[date_col] = pd.to_datetime(
                date_frame, format=self.custom_datetime_format, infer_datetime_format=True)
        else:
            raise NotImplementedError

        formatted[date_col] = formatted[date_col].values.astype(np.int64) // 10 ** 9

        return formatted


text:
"""
MIT License

Copyright (c) 2020 Hyeonki Hong <hhk7734@gmail.com>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
import numpy as np

def DIoU_NMS(candidates, threshold):
    """
    Distance Intersection over Union(DIoU)
    Non-Maximum Suppression(NMS)

    @param `candidates`:
            [[center_x, center_y, w, h, class_id, propability], ...]
    @param `threshold`: If DIoU is above the threshold, bboxes are considered
            the same. All but the bbox with the highest probability are removed.
    """
    bboxes = []
    for class_id in set(candidates[:, 4]):
        class_bboxes = candidates[candidates[:, 4] == class_id]
        if class_bboxes.shape[0] == 1:
            # One candidate
            bboxes.append(class_bboxes)
            continue

        while True:
            half_wh = class_bboxes[:, 2:4] * 0.5
            max_index = np.argmax(class_bboxes[:, 5])
            max_bbox = class_bboxes[max_index, :]
            max_half_wh = half_wh[max_index, :]
            # Max probability
            bboxes.append(max_bbox[np.newaxis, :])

            enclose_left = np.minimum(
                class_bboxes[:, 0] - half_wh[:, 0],
                max_bbox[0] - max_half_wh[0],
            )
            enclose_right = np.maximum(
                class_bboxes[:, 0] + half_wh[:, 0],
                max_bbox[0] + max_half_wh[0],
            )
            enclose_top = np.minimum(
                class_bboxes[:, 1] - half_wh[:, 1],
                max_bbox[1] - max_half_wh[1],
            )
            enclose_bottom = np.maximum(
                class_bboxes[:, 1] + half_wh[:, 1],
                max_bbox[1] + max_half_wh[1],
            )

            enclose_width = enclose_right - enclose_left
            enclose_height = enclose_bottom - enclose_top

            width_mask = enclose_width >= class_bboxes[:, 2] + max_bbox[2]
            height_mask = enclose_height >= class_bboxes[:, 3] + max_bbox[3]
            # bboxes with no overlap with max_bbox
            no_overlap_mask = np.logical_or(width_mask, height_mask)
            no_overlap_bboxes = class_bboxes[no_overlap_mask]

            overlap_mask = np.logical_not(no_overlap_mask)
            class_bboxes = class_bboxes[overlap_mask]
            if class_bboxes.shape[0] == 1:
                if no_overlap_bboxes.shape[0] == 1:
                    bboxes.append(no_overlap_bboxes)
                    break

                class_bboxes = no_overlap_bboxes
                continue

            half_wh = half_wh[overlap_mask]
            enclose_left = enclose_left[overlap_mask]
            enclose_right = enclose_right[overlap_mask]
            enclose_top = enclose_top[overlap_mask]
            enclose_bottom = enclose_bottom[overlap_mask]

            inter_left = np.maximum(
                class_bboxes[:, 0] - half_wh[:, 0],
                max_bbox[0] - max_half_wh[0],
            )
            inter_right = np.minimum(
                class_bboxes[:, 0] + half_wh[:, 0],
                max_bbox[0] + max_half_wh[0],
            )
            inter_top = np.maximum(
                class_bboxes[:, 1] - half_wh[:, 1],
                max_bbox[1] - max_half_wh[1],
            )
            inter_bottom = np.minimum(
                class_bboxes[:, 1] + half_wh[:, 1],
                max_bbox[1] + max_half_wh[1],
            )

            class_bboxes_area = class_bboxes[:, 2] * class_bboxes[:, 3]
            max_bbox_area = max_bbox[2] * max_bbox[3]
            inter_area = (inter_right - inter_left) * (inter_bottom - inter_top)
            iou = inter_area / (class_bboxes_area + max_bbox_area - inter_area)

            c_squared = (enclose_right - enclose_left) * (
                enclose_right - enclose_left
            ) + (enclose_bottom - enclose_top) * (enclose_bottom - enclose_top)
            d_squared = (class_bboxes[:, 0] - max_bbox[0]) * (
                class_bboxes[:, 0] - max_bbox[0]
            ) + (class_bboxes[:, 1] - max_bbox[1]) * (
                class_bboxes[:, 1] - max_bbox[1]
            )

            # DIoU = IoU - d^2 / c^2
            little_overlap_mask = iou - d_squared / c_squared < threshold
            little_overlap_bboxes = class_bboxes[little_overlap_mask]
            if (
                no_overlap_bboxes.shape[0] > 0
                and little_overlap_bboxes.shape[0] > 0
            ):
                class_bboxes = np.concatenate(
                    [no_overlap_bboxes, little_overlap_bboxes], axis=0
                )
                continue

            if no_overlap_bboxes.shape[0] > 0:
                if no_overlap_bboxes.shape[0] == 1:
                    bboxes.append(no_overlap_bboxes)
                    break

                class_bboxes = no_overlap_bboxes
                continue

            if little_overlap_bboxes.shape[0] > 0:
                if little_overlap_bboxes.shape[0] == 1:
                    bboxes.append(little_overlap_bboxes)
                    break

                class_bboxes = little_overlap_bboxes
                continue

            break

    if len(bboxes) == 0:
        return np.zeros(shape=(1, 6))

    return np.concatenate(bboxes, axis=0)

def candidates_to_pred_bboxes(
    candidates,
    input_size,
    iou_threshold: float = 0.3,
    score_threshold: float = 0.25,
):
    """
    @param candidates: Dim(-1, (x, y, w, h, obj_score, probabilities))

    @return Dim(-1, (x, y, w, h, class_id, class_probability))
    """
    # Remove low socre candidates
    # This step should be the first !!
    class_ids = np.argmax(candidates[:, 5:], axis=-1)
    # class_prob = obj_score * max_probability
    class_prob = (
        candidates[:, 4] * candidates[np.arange(len(candidates)), class_ids + 5]
    )
    candidates = candidates[class_prob > score_threshold, :]

    # Remove out of range candidates
    half_wh = candidates[:, 2:4] * 0.5
    mask = candidates[:, 0] - half_wh[:, 0] >= 0
    candidates = candidates[mask, :]
    half_wh = half_wh[mask, :]
    mask = candidates[:, 0] + half_wh[:, 0] <= 1
    candidates = candidates[mask, :]
    half_wh = half_wh[mask, :]
    mask = candidates[:, 1] - half_wh[:, 1] >= 0
    candidates = candidates[mask, :]
    half_wh = half_wh[mask, :]
    mask = candidates[:, 1] + half_wh[:, 1] <= 1
    candidates = candidates[mask, :]

    # Remove small candidates
    candidates = candidates[
        np.logical_and(
            candidates[:, 2] > (4 / input_size[0]),  # width
            candidates[:, 3] > (4 / input_size[1]),  # height
        ),
        :,
    ]

    class_ids = np.argmax(candidates[:, 5:], axis=-1)
    class_prob = (
        candidates[:, 4] * candidates[np.arange(len(candidates)), class_ids + 5]
    )

    # x, y, w, h, class_id, class_probability
    candidates = np.concatenate(
        [
            candidates[:, :4],
            class_ids[:, np.newaxis],
            class_prob[:, np.newaxis],
        ],
        axis=-1,
    )

    return DIoU_NMS(candidates, iou_threshold)

def fit_pred_bboxes_to_original(bboxes, input_size, original_shape):
    """
    @param `bboxes`: Dim(-1, (x, y, w, h, class_id, probability))
    @param `input_size`: (width, height)
    @param `original_shape`: (height, width, channels)
    """

    height, width, _ = original_shape
    bboxes = np.copy(bboxes)

    w_h = width / height
    iw_ih = input_size[0] / input_size[1]

    if w_h > iw_ih:
        scale = w_h / iw_ih
        bboxes[:, 1] = scale * (bboxes[:, 1] - 0.5) + 0.5
        bboxes[:, 3] = scale * bboxes[:, 3]
    elif w_h < iw_ih:
        scale = iw_ih / w_h
        bboxes[:, 0] = scale * (bboxes[:, 0] - 0.5) + 0.5
        bboxes[:, 2] = scale * bboxes[:, 2]

    return bboxes