<commit_before># Copyright The PyTorch Lightning team.
#
# 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.
r"""
Finetuning Callback
^^^^^^^^^^^^^^^^^^^^
Freeze and unfreeze models for finetuning purposes
"""
import logging
from typing import Callable, Generator, Iterable, List, Optional, Union

import torch
from torch.nn import Module
from torch.nn.modules.batchnorm import _BatchNorm
from torch.nn.modules.container import Container, ModuleDict, ModuleList, Sequential
from torch.optim.optimizer import Optimizer

from pytorch_lightning.callbacks.base import Callback
from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.utilities import rank_zero_warn
from pytorch_lightning.utilities.exceptions import MisconfigurationException

log = logging.getLogger(__name__)

def multiplicative(epoch):
    return 2

class BaseFinetuning(Callback):
    r"""

    This class implements the base logic for writing your own Finetuning Callback.

    Override ``freeze_before_training`` and ``finetune_function`` methods with your own logic.

    ``freeze_before_training``: This method is called before ``configure_optimizers``
        and should be used to freeze any modules parameters.

    ``finetune_function``: This method is called on every train epoch start and should be used to
        ``unfreeze`` any parameters. Those parameters needs to be added in a new ``param_group``
        within the optimizer.

    .. note:: Make sure to filter the parameters based on ``requires_grad``.

    Example::

        class MyModel(LightningModule)

            ...

            def configure_optimizer(self):
                # Make sure to filter the parameters based on `requires_grad`
                return Adam(filter(lambda p: p.requires_grad, self.parameters))

        class FeatureExtractorFreezeUnfreeze(BaseFinetuning):

            def __init__(self, unfreeze_at_epoch=10)
                self._unfreeze_at_epoch = unfreeze_at_epoch

            def freeze_before_training(self, pl_module):
                # freeze any module you want
                # Here, we are freezing ``feature_extractor``
                self.freeze(pl_module.feature_extractor)

            def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):
                # When `current_epoch` is 10, feature_extractor will start training.
                if current_epoch == self._unfreeze_at_epoch:
                    self.unfreeze_and_add_param_group(
                        module=pl_module.feature_extractor,
                        optimizer=optimizer,
                        train_bn=True,
                    )
    """

    @staticmethod
    def flatten_modules(modules: Union[Module, Iterable[Union[Module, Iterable]]]) -> List[Module]:
        """
        This function is used to flatten a module or an iterable of modules into a list of its modules.

        Args:
            modules: A given module or an iterable of modules

        Returns:
            List of modules
        """
        if isinstance(modules, Iterable):
            _modules = []
            for m in modules:
                _modules.extend(BaseFinetuning.flatten_modules(m))

        else:
            _modules = modules.modules()

        return list(
            filter(
                lambda m: not isinstance(m, (Container, Sequential, ModuleDict, ModuleList, LightningModule)), _modules
            )
        )

    @staticmethod
    def filter_params(
        modules: Union[Module, Iterable[Union[Module, Iterable]]],
        train_bn: bool = True,
        requires_grad: bool = True
    ) -> Generator:
        """Yields the `requires_grad` parameters of a given module or list of modules.

        Args:
            modules: A given module or an iterable of modules
            train_bn: Whether to train BatchNorm module
            requires_grad: Whether to create a generator for trainable or non-trainable parameters.

        Returns:
            Generator
        """
        modules = BaseFinetuning.flatten_modules(modules)
        for mod in modules:
            if isinstance(mod, _BatchNorm) and not train_bn:
                continue
            for param in mod.parameters():
                if param.requires_grad == requires_grad:
                    yield param

    @staticmethod
    def make_trainable(modules: Union[Module, Iterable[Union[Module, Iterable]]]) -> None:
        """
        Unfreezes the parameters of the provided modules

        Args:
            modules: A given module or an iterable of modules
        """
        modules = BaseFinetuning.flatten_modules(modules)
        for module in modules:
            for param in module.parameters():
                param.requires_grad = True

    @staticmethod
    def freeze(modules: Union[Module, Iterable[Union[Module, Iterable]]], train_bn: bool = True) -> None:
        """
        Freezes the parameters of the provided modules

        Args:
            modules: A given module or an iterable of modules
            train_bn: If True, leave the BatchNorm layers in training mode

        Returns:
            None
        """
        modules = BaseFinetuning.flatten_modules(modules)
        for mod in modules:
            if isinstance(mod, _BatchNorm) and train_bn:
                BaseFinetuning.make_trainable(mod)
            else:
                for param in mod.parameters():
                    param.requires_grad = False

    @staticmethod
    def filter_on_optimizer(optimizer: Optimizer, params: Iterable) -> List:
        """
        This function is used to exclude any parameter which already exists in
        this optimizer

        Args:
            optimizer: Optimizer used for parameter exclusion
            params: Iterable of parameters used to check against the provided optimizer

        Returns:
            List of parameters not contained in this optimizer param groups
        """
        out_params = []
        removed_params = []
        for param in params:
            if not any(torch.equal(p, param) for group in optimizer.param_groups for p in group["params"]):
                out_params.append(param)
            else:
                removed_params.append(param)

        if removed_params:
            rank_zero_warn(
                "The provided params to be freezed already exist within another group of this optimizer."
                " Those parameters will be skipped.\n"
                "HINT: Did you init your optimizer in `configure_optimizer` as such:\n"
                f" {type(optimizer)}(filter(lambda p: p.requires_grad, self.parameters()), ...) ", UserWarning
            )
        return out_params

    @staticmethod
    def unfreeze_and_add_param_group(
        modules: Union[Module, Iterable[Union[Module, Iterable]]],
        optimizer: Optimizer,
        lr: Optional[float] = None,
        initial_denom_lr: float = 10.,
        train_bn: bool = True,
    ) -> None:
        """
        Unfreezes a module and adds its parameters to an optimizer.

        Args:

            modules: A module or iterable of modules to unfreeze.
                Their parameters will be added to an optimizer as a new param group.

            optimizer: The provided optimizer will receive new parameters and will add them to
                `add_param_group`

            lr: Learning rate for the new param group.

            initial_denom_lr: If no lr is provided, the learning from the first param group will be used
                and divided by initial_denom_lr.

            train_bn: Whether to train the BatchNormalization layers.

        Returns:
            None
        """
        BaseFinetuning.make_trainable(modules)
        params_lr = optimizer.param_groups[0]['lr'] if lr is None else float(lr)
        denom_lr = initial_denom_lr if lr is None else 1.
        params = BaseFinetuning.filter_params(modules, train_bn=train_bn, requires_grad=True)
        params = BaseFinetuning.filter_on_optimizer(optimizer, params)
        if params:
            optimizer.add_param_group({
                'params': params,
                'lr': params_lr / denom_lr,
            })

    def on_before_accelerator_backend_setup(self, trainer, pl_module):
        self.freeze_before_training(pl_module)

    def on_train_epoch_start(self, trainer, pl_module):
        """Called when the epoch begins."""
        for opt_idx, optimizer in trainer.train_loop.prepare_optimizers():
            self.finetune_function(pl_module, trainer.current_epoch, optimizer, opt_idx)

    def finetune_function(self, pl_module: LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
        """
        Override to add your unfreeze logic
        """
        raise NotImplementedError

    def freeze_before_training(self, pl_module: LightningModule):
        """
        Override to add your freeze logic
        """
        raise NotImplementedError

class BackboneFinetuning(BaseFinetuning):
    r"""

    Finetune a backbone model based on a learning rate user-defined scheduling.
    When the backbone learning rate reaches the current model learning rate
    and ``should_align`` is set to True, it will align with it for the rest of the training.

    Args:

        unfreeze_backbone_at_epoch: Epoch at which the backbone will be unfreezed.

        lambda_func: Scheduling function for increasing backbone learning rate.

        backbone_initial_ratio_lr:
            Used to scale down the backbone learning rate compared to rest of model

        backbone_initial_lr: Optional, Inital learning rate for the backbone.
            By default, we will use current_learning /  backbone_initial_ratio_lr

        should_align: Wheter to align with current learning rate when backbone learning
            reaches it.

        initial_denom_lr: When unfreezing the backbone, the intial learning rate will
            current_learning_rate /  initial_denom_lr.

        train_bn: Wheter to make Batch Normalization trainable.

        verbose: Display current learning rate for model and backbone

        round: Precision for displaying learning rate

    Example::

        >>> from pytorch_lightning import Trainer
        >>> from pytorch_lightning.callbacks import BackboneFinetuning
        >>> multiplicative = lambda epoch: 1.5
        >>> backbone_finetuning = BackboneFinetuning(200, multiplicative)
        >>> trainer = Trainer(callbacks=[backbone_finetuning])

    """

    def __init__(
        self,
        unfreeze_backbone_at_epoch: int = 10,
        lambda_func: Callable = multiplicative,
        backbone_initial_ratio_lr: float = 10e-2,
        backbone_initial_lr: Optional[float] = None,
        should_align: bool = True,
        initial_denom_lr: float = 10.,
        train_bn: bool = True,
        verbose: bool = False,
        round: int = 12,
    ):
        self.unfreeze_backbone_at_epoch = unfreeze_backbone_at_epoch
        self.backbone_initial_lr = backbone_initial_lr
        self.lambda_func = lambda_func
        self.backbone_initial_ratio_lr = backbone_initial_ratio_lr
        self.should_align = should_align
        self.initial_denom_lr = initial_denom_lr
        self.train_bn = train_bn
        self.round = round
        self.verbose = verbose

    def on_fit_start(self, trainer, pl_module):
        """
        Raises:
            MisconfigurationException:
                If LightningModule has no nn.Module `backbone` attribute.
        """
        if hasattr(pl_module, "backbone") and isinstance(pl_module.backbone, Module):
            return
        raise MisconfigurationException("The LightningModule should have a nn.Module `backbone` attribute")

    def freeze_before_training(self, pl_module: LightningModule):
        self.freeze(pl_module.backbone)

    def finetune_function(self, pl_module: LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
        """Called when the epoch begins."""

        if epoch == self.unfreeze_backbone_at_epoch:
            current_lr = optimizer.param_groups[0]['lr']
            initial_backbone_lr = self.backbone_initial_lr if self.backbone_initial_lr is not None \
                else current_lr * self.backbone_initial_ratio_lr
            self.previous_backbone_lr = initial_backbone_lr
            self.unfreeze_and_add_param_group(
                pl_module.backbone,
                optimizer,
                initial_backbone_lr,
                train_bn=self.train_bn,
                initial_denom_lr=self.initial_denom_lr
            )
            if self.verbose:
                log.info(
                    f"Current lr: {round(current_lr, self.round)}, "
                    f"Backbone lr: {round(initial_backbone_lr, self.round)}"
                )

        elif epoch > self.unfreeze_backbone_at_epoch:
            current_lr = optimizer.param_groups[0]['lr']
            next_current_backbone_lr = self.lambda_func(epoch + 1) * self.previous_backbone_lr
            next_current_backbone_lr = current_lr if (self.should_align and next_current_backbone_lr > current_lr) \
                else next_current_backbone_lr
            optimizer.param_groups[-1]["lr"] = next_current_backbone_lr
            self.previous_backbone_lr = next_current_backbone_lr
            if self.verbose:
                log.info(
                    f"Current lr: {round(current_lr, self.round)}, "
                    f"Backbone lr: {round(next_current_backbone_lr, self.round)}"
                )
<commit_msg>Fix `unfreeze_and_add_param_group` expects `modules` rather than `module` (#6822)<commit_after># Copyright The PyTorch Lightning team.
#
# 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.
r"""
Finetuning Callback
^^^^^^^^^^^^^^^^^^^^
Freeze and unfreeze models for finetuning purposes
"""
import logging
from typing import Callable, Generator, Iterable, List, Optional, Union

import torch
from torch.nn import Module
from torch.nn.modules.batchnorm import _BatchNorm
from torch.nn.modules.container import Container, ModuleDict, ModuleList, Sequential
from torch.optim.optimizer import Optimizer

from pytorch_lightning.callbacks.base import Callback
from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.utilities import rank_zero_warn
from pytorch_lightning.utilities.exceptions import MisconfigurationException

log = logging.getLogger(__name__)

def multiplicative(epoch):
    return 2

class BaseFinetuning(Callback):
    r"""

    This class implements the base logic for writing your own Finetuning Callback.

    Override ``freeze_before_training`` and ``finetune_function`` methods with your own logic.

    ``freeze_before_training``: This method is called before ``configure_optimizers``
        and should be used to freeze any modules parameters.

    ``finetune_function``: This method is called on every train epoch start and should be used to
        ``unfreeze`` any parameters. Those parameters needs to be added in a new ``param_group``
        within the optimizer.

    .. note:: Make sure to filter the parameters based on ``requires_grad``.

    Example::

        class MyModel(LightningModule)

            ...

            def configure_optimizer(self):
                # Make sure to filter the parameters based on `requires_grad`
                return Adam(filter(lambda p: p.requires_grad, self.parameters))

        class FeatureExtractorFreezeUnfreeze(BaseFinetuning):

            def __init__(self, unfreeze_at_epoch=10)
                self._unfreeze_at_epoch = unfreeze_at_epoch

            def freeze_before_training(self, pl_module):
                # freeze any module you want
                # Here, we are freezing ``feature_extractor``
                self.freeze(pl_module.feature_extractor)

            def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):
                # When `current_epoch` is 10, feature_extractor will start training.
                if current_epoch == self._unfreeze_at_epoch:
                    self.unfreeze_and_add_param_group(
                        modules=pl_module.feature_extractor,
                        optimizer=optimizer,
                        train_bn=True,
                    )
    """

    @staticmethod
    def flatten_modules(modules: Union[Module, Iterable[Union[Module, Iterable]]]) -> List[Module]:
        """
        This function is used to flatten a module or an iterable of modules into a list of its modules.

        Args:
            modules: A given module or an iterable of modules

        Returns:
            List of modules
        """
        if isinstance(modules, Iterable):
            _modules = []
            for m in modules:
                _modules.extend(BaseFinetuning.flatten_modules(m))

        else:
            _modules = modules.modules()

        return list(
            filter(
                lambda m: not isinstance(m, (Container, Sequential, ModuleDict, ModuleList, LightningModule)), _modules
            )
        )

    @staticmethod
    def filter_params(
        modules: Union[Module, Iterable[Union[Module, Iterable]]],
        train_bn: bool = True,
        requires_grad: bool = True
    ) -> Generator:
        """Yields the `requires_grad` parameters of a given module or list of modules.

        Args:
            modules: A given module or an iterable of modules
            train_bn: Whether to train BatchNorm module
            requires_grad: Whether to create a generator for trainable or non-trainable parameters.

        Returns:
            Generator
        """
        modules = BaseFinetuning.flatten_modules(modules)
        for mod in modules:
            if isinstance(mod, _BatchNorm) and not train_bn:
                continue
            for param in mod.parameters():
                if param.requires_grad == requires_grad:
                    yield param

    @staticmethod
    def make_trainable(modules: Union[Module, Iterable[Union[Module, Iterable]]]) -> None:
        """
        Unfreezes the parameters of the provided modules

        Args:
            modules: A given module or an iterable of modules
        """
        modules = BaseFinetuning.flatten_modules(modules)
        for module in modules:
            for param in module.parameters():
                param.requires_grad = True

    @staticmethod
    def freeze(modules: Union[Module, Iterable[Union[Module, Iterable]]], train_bn: bool = True) -> None:
        """
        Freezes the parameters of the provided modules

        Args:
            modules: A given module or an iterable of modules
            train_bn: If True, leave the BatchNorm layers in training mode

        Returns:
            None
        """
        modules = BaseFinetuning.flatten_modules(modules)
        for mod in modules:
            if isinstance(mod, _BatchNorm) and train_bn:
                BaseFinetuning.make_trainable(mod)
            else:
                for param in mod.parameters():
                    param.requires_grad = False

    @staticmethod
    def filter_on_optimizer(optimizer: Optimizer, params: Iterable) -> List:
        """
        This function is used to exclude any parameter which already exists in
        this optimizer

        Args:
            optimizer: Optimizer used for parameter exclusion
            params: Iterable of parameters used to check against the provided optimizer

        Returns:
            List of parameters not contained in this optimizer param groups
        """
        out_params = []
        removed_params = []
        for param in params:
            if not any(torch.equal(p, param) for group in optimizer.param_groups for p in group["params"]):
                out_params.append(param)
            else:
                removed_params.append(param)

        if removed_params:
            rank_zero_warn(
                "The provided params to be freezed already exist within another group of this optimizer."
                " Those parameters will be skipped.\n"
                "HINT: Did you init your optimizer in `configure_optimizer` as such:\n"
                f" {type(optimizer)}(filter(lambda p: p.requires_grad, self.parameters()), ...) ", UserWarning
            )
        return out_params

    @staticmethod
    def unfreeze_and_add_param_group(
        modules: Union[Module, Iterable[Union[Module, Iterable]]],
        optimizer: Optimizer,
        lr: Optional[float] = None,
        initial_denom_lr: float = 10.,
        train_bn: bool = True,
    ) -> None:
        """
        Unfreezes a module and adds its parameters to an optimizer.

        Args:

            modules: A module or iterable of modules to unfreeze.
                Their parameters will be added to an optimizer as a new param group.

            optimizer: The provided optimizer will receive new parameters and will add them to
                `add_param_group`

            lr: Learning rate for the new param group.

            initial_denom_lr: If no lr is provided, the learning from the first param group will be used
                and divided by initial_denom_lr.

            train_bn: Whether to train the BatchNormalization layers.

        Returns:
            None
        """
        BaseFinetuning.make_trainable(modules)
        params_lr = optimizer.param_groups[0]['lr'] if lr is None else float(lr)
        denom_lr = initial_denom_lr if lr is None else 1.
        params = BaseFinetuning.filter_params(modules, train_bn=train_bn, requires_grad=True)
        params = BaseFinetuning.filter_on_optimizer(optimizer, params)
        if params:
            optimizer.add_param_group({
                'params': params,
                'lr': params_lr / denom_lr,
            })

    def on_before_accelerator_backend_setup(self, trainer, pl_module):
        self.freeze_before_training(pl_module)

    def on_train_epoch_start(self, trainer, pl_module):
        """Called when the epoch begins."""
        for opt_idx, optimizer in trainer.train_loop.prepare_optimizers():
            self.finetune_function(pl_module, trainer.current_epoch, optimizer, opt_idx)

    def finetune_function(self, pl_module: LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
        """
        Override to add your unfreeze logic
        """
        raise NotImplementedError

    def freeze_before_training(self, pl_module: LightningModule):
        """
        Override to add your freeze logic
        """
        raise NotImplementedError

class BackboneFinetuning(BaseFinetuning):
    r"""

    Finetune a backbone model based on a learning rate user-defined scheduling.
    When the backbone learning rate reaches the current model learning rate
    and ``should_align`` is set to True, it will align with it for the rest of the training.

    Args:

        unfreeze_backbone_at_epoch: Epoch at which the backbone will be unfreezed.

        lambda_func: Scheduling function for increasing backbone learning rate.

        backbone_initial_ratio_lr:
            Used to scale down the backbone learning rate compared to rest of model

        backbone_initial_lr: Optional, Inital learning rate for the backbone.
            By default, we will use current_learning /  backbone_initial_ratio_lr

        should_align: Wheter to align with current learning rate when backbone learning
            reaches it.

        initial_denom_lr: When unfreezing the backbone, the intial learning rate will
            current_learning_rate /  initial_denom_lr.

        train_bn: Wheter to make Batch Normalization trainable.

        verbose: Display current learning rate for model and backbone

        round: Precision for displaying learning rate

    Example::

        >>> from pytorch_lightning import Trainer
        >>> from pytorch_lightning.callbacks import BackboneFinetuning
        >>> multiplicative = lambda epoch: 1.5
        >>> backbone_finetuning = BackboneFinetuning(200, multiplicative)
        >>> trainer = Trainer(callbacks=[backbone_finetuning])

    """

    def __init__(
        self,
        unfreeze_backbone_at_epoch: int = 10,
        lambda_func: Callable = multiplicative,
        backbone_initial_ratio_lr: float = 10e-2,
        backbone_initial_lr: Optional[float] = None,
        should_align: bool = True,
        initial_denom_lr: float = 10.,
        train_bn: bool = True,
        verbose: bool = False,
        round: int = 12,
    ):
        self.unfreeze_backbone_at_epoch = unfreeze_backbone_at_epoch
        self.backbone_initial_lr = backbone_initial_lr
        self.lambda_func = lambda_func
        self.backbone_initial_ratio_lr = backbone_initial_ratio_lr
        self.should_align = should_align
        self.initial_denom_lr = initial_denom_lr
        self.train_bn = train_bn
        self.round = round
        self.verbose = verbose

    def on_fit_start(self, trainer, pl_module):
        """
        Raises:
            MisconfigurationException:
                If LightningModule has no nn.Module `backbone` attribute.
        """
        if hasattr(pl_module, "backbone") and isinstance(pl_module.backbone, Module):
            return
        raise MisconfigurationException("The LightningModule should have a nn.Module `backbone` attribute")

    def freeze_before_training(self, pl_module: LightningModule):
        self.freeze(pl_module.backbone)

    def finetune_function(self, pl_module: LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
        """Called when the epoch begins."""

        if epoch == self.unfreeze_backbone_at_epoch:
            current_lr = optimizer.param_groups[0]['lr']
            initial_backbone_lr = self.backbone_initial_lr if self.backbone_initial_lr is not None \
                else current_lr * self.backbone_initial_ratio_lr
            self.previous_backbone_lr = initial_backbone_lr
            self.unfreeze_and_add_param_group(
                pl_module.backbone,
                optimizer,
                initial_backbone_lr,
                train_bn=self.train_bn,
                initial_denom_lr=self.initial_denom_lr
            )
            if self.verbose:
                log.info(
                    f"Current lr: {round(current_lr, self.round)}, "
                    f"Backbone lr: {round(initial_backbone_lr, self.round)}"
                )

        elif epoch > self.unfreeze_backbone_at_epoch:
            current_lr = optimizer.param_groups[0]['lr']
            next_current_backbone_lr = self.lambda_func(epoch + 1) * self.previous_backbone_lr
            next_current_backbone_lr = current_lr if (self.should_align and next_current_backbone_lr > current_lr) \
                else next_current_backbone_lr
            optimizer.param_groups[-1]["lr"] = next_current_backbone_lr
            self.previous_backbone_lr = next_current_backbone_lr
            if self.verbose:
                log.info(
                    f"Current lr: {round(current_lr, self.round)}, "
                    f"Backbone lr: {round(next_current_backbone_lr, self.round)}"
                )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Data profilers."""

import abc
import collections

from verta.data_types import (
    DiscreteHistogram,
    FloatHistogram,
)
from verta.external import six
from verta._internal_utils.importer import maybe_dependency

@six.add_metaclass(abc.ABCMeta)
class Profiler(object):
    """Produces summary data types for a data frame.

    A profiler's ``profile`` method accepts a data frame and produces a
    dictionary from profile entry names to summary data types according to its
    ``profile_column`` method.

    Attributes
    ----------
    columns : :obj:`list` of :obj:`str`
        The list of data frame columns which will be profiled.
    """

    def __init__(self, columns):
        self.columns = columns

    def __name__(self):
        """Returns \"Profiler\""""
        return "Profiler"

    def profile(self, df):
        """Profile a pandas data frame and return a dictionary.

        Profiles a pandas data frame and returns a dictionary, the values of
        which will be data types.

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.

        Returns
        -------
        dict
            Mapping from profile column name to profiler data type
        """
        return {column: histogram for (column, histogram) in self._profile_columns(df)}

    def predict(self, df):
        """Delegates to self.profile for internal Verta compatibility."""
        return self.profile(df)

    def _profile_columns(self, df):
        return [self.profile_column(df, column) for column in self.columns]

    @abc.abstractmethod
    def profile_column(self, df, column):
        """Profile a data frame column, returning a name and a summarization.

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.
        column : str
            The column in the data frame to profile.

        Returns
        -------
        dict
            A dictionary from summary data name to data types.
        """
        raise NotImplementedError("")

class MissingValuesProfiler(Profiler):
    """Produces discrete histograms for present and missing values.

    Counts the number of null and non-null values as "present" and "missing"
    values and returns DiscreteHistogram summary data for the specified
    columns.

    Parameters
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.

    Attributes
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.
    """

    def __init__(self, columns):
        super(MissingValuesProfiler, self).__init__(columns)

    def profile_column(self, df, column):
        """Profile a data frame column, returning a name and DiscreteHistogram

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.
        column : str
            The column in the data frame to profile.

        Returns
        -------
        dict
            A dictionary from summary data names to DiscreteHistogram instances.
        """
        total = df.shape[0]
        try:
            missing = sum(df[column].isnull())
        except KeyError:  # pandas raises this if the column doesn't exist
            missing = total

        return (column + "_missing", DiscreteHistogram(["present", "missing"], [total - missing, missing]))

    def profile_point(self, sample, reference):
        if sample is None: # Missing
            return DiscreteHistogram(reference._buckets, [0, 1])
        return DiscreteHistogram(reference._buckets, [1, 0])

# TODO: Rename to CategoricalHistogramProfiler?
class BinaryHistogramProfiler(Profiler):
    """Produces discrete histograms counting values per unique value in columns.

    Produces histogram columns for each unique value in the profiled column and
    counts the number of values for these keys. This should not be used for
    columns containing continuously measured values but should be used when the
    set of possible values in a column is small, e.g. boolean valued columns.

    Parameters
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.

    Attributes
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.
    """

    def __init__(self, columns):
        super(BinaryHistogramProfiler, self).__init__(columns)

    def profile_column(self, df, column):
        """Profile a data frame column, returning a name and DiscreteHistogram

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.
        column : str
            The column in the data frame to profile.

        Returns
        -------
        dict
            A dictionary from summary data names to DiscreteHistogram instances.
        """
        content = df[column].value_counts()
        keys = list(content.keys())
        values = [content[k] for k in keys]
        values = [v.item() for v in values]
        return (column + "_histogram", DiscreteHistogram(keys, values))

    def profile_point(self, sample, reference):
        if sample:
            return DiscreteHistogram(reference._buckets, [0, 1])
        return DiscreteHistogram(reference._buckets, [1, 0])

# TODO: Consider design/interface for different bins
class ContinuousHistogramProfiler(Profiler):
    """Produces float histograms counting values according to bin intervals.

    Attributes
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.
    bins
    """

    def __init__(self, columns, bins=10):
        super(ContinuousHistogramProfiler, self).__init__(columns)
        self._bins = bins
        self._np = maybe_dependency("numpy")
        if self._np is None:
            raise ImportError("numpy is not installed; try `pip install numpy`")

    def profile_column(self, df, column):
        if isinstance(self._bins, collections.Mapping):
            bins = self._bins[column]
        else:
            bins = self._bins
        values, limits = self._np.histogram(df[column], bins=bins)
        values = [v.item() for v in values]
        limits = [lim.item() for lim in limits]
        return (column + "_histogram", FloatHistogram(limits, values))

    # TODO: consider the case where the data is outside of the bucket list
    def profile_point(self, sample, reference):
        values, _ = self._np.histogram([sample], bins=reference._bucket_limits)
        return FloatHistogram(reference._bucket_limits, values)<commit_msg>Support discrete for BinaryHistProf.profile_point() (#2437)<commit_after># -*- coding: utf-8 -*-
"""Data profilers."""

import abc
import collections

from verta.data_types import (
    DiscreteHistogram,
    FloatHistogram,
)
from verta.external import six
from verta._internal_utils.importer import maybe_dependency

@six.add_metaclass(abc.ABCMeta)
class Profiler(object):
    """Produces summary data types for a data frame.

    A profiler's ``profile`` method accepts a data frame and produces a
    dictionary from profile entry names to summary data types according to its
    ``profile_column`` method.

    Attributes
    ----------
    columns : :obj:`list` of :obj:`str`
        The list of data frame columns which will be profiled.
    """

    def __init__(self, columns):
        self.columns = columns

    def __name__(self):
        """Returns \"Profiler\""""
        return "Profiler"

    def profile(self, df):
        """Profile a pandas data frame and return a dictionary.

        Profiles a pandas data frame and returns a dictionary, the values of
        which will be data types.

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.

        Returns
        -------
        dict
            Mapping from profile column name to profiler data type
        """
        return {column: histogram for (column, histogram) in self._profile_columns(df)}

    def predict(self, df):
        """Delegates to self.profile for internal Verta compatibility."""
        return self.profile(df)

    def _profile_columns(self, df):
        return [self.profile_column(df, column) for column in self.columns]

    @abc.abstractmethod
    def profile_column(self, df, column):
        """Profile a data frame column, returning a name and a summarization.

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.
        column : str
            The column in the data frame to profile.

        Returns
        -------
        dict
            A dictionary from summary data name to data types.
        """
        raise NotImplementedError("")

class MissingValuesProfiler(Profiler):
    """Produces discrete histograms for present and missing values.

    Counts the number of null and non-null values as "present" and "missing"
    values and returns DiscreteHistogram summary data for the specified
    columns.

    Parameters
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.

    Attributes
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.
    """

    def __init__(self, columns):
        super(MissingValuesProfiler, self).__init__(columns)

    def profile_column(self, df, column):
        """Profile a data frame column, returning a name and DiscreteHistogram

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.
        column : str
            The column in the data frame to profile.

        Returns
        -------
        dict
            A dictionary from summary data names to DiscreteHistogram instances.
        """
        total = df.shape[0]
        try:
            missing = sum(df[column].isnull())
        except KeyError:  # pandas raises this if the column doesn't exist
            missing = total

        return (column + "_missing", DiscreteHistogram(["present", "missing"], [total - missing, missing]))

    def profile_point(self, sample, reference):
        if sample is None: # Missing
            return DiscreteHistogram(reference._buckets, [0, 1])
        return DiscreteHistogram(reference._buckets, [1, 0])

# TODO: Rename to CategoricalHistogramProfiler?
class BinaryHistogramProfiler(Profiler):
    """Produces discrete histograms counting values per unique value in columns.

    Produces histogram columns for each unique value in the profiled column and
    counts the number of values for these keys. This should not be used for
    columns containing continuously measured values but should be used when the
    set of possible values in a column is small, e.g. boolean valued columns.

    Parameters
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.

    Attributes
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.
    """

    def __init__(self, columns):
        super(BinaryHistogramProfiler, self).__init__(columns)

    def profile_column(self, df, column):
        """Profile a data frame column, returning a name and DiscreteHistogram

        Parameters
        ----------
        df : pandas.core.frame.DataFrame
            A data frame to profile.
        column : str
            The column in the data frame to profile.

        Returns
        -------
        dict
            A dictionary from summary data names to DiscreteHistogram instances.
        """
        content = df[column].value_counts()
        keys = list(content.keys())
        values = [content[k] for k in keys]
        values = [v.item() for v in values]
        return (column + "_histogram", DiscreteHistogram(keys, values))

    # TODO: consider the case where the data is not in the bucket list
    def profile_point(self, sample, reference):
        buckets = reference._buckets
        data = [0]*len(buckets)

        if sample:
            data[buckets.index(sample)] = 1
        return DiscreteHistogram(buckets, data)

# TODO: Consider design/interface for different bins
class ContinuousHistogramProfiler(Profiler):
    """Produces float histograms counting values according to bin intervals.

    Attributes
    ----------
    columns : list of str
        The list of data frame columns which will be profiled.
    bins
    """

    def __init__(self, columns, bins=10):
        super(ContinuousHistogramProfiler, self).__init__(columns)
        self._bins = bins
        self._np = maybe_dependency("numpy")
        if self._np is None:
            raise ImportError("numpy is not installed; try `pip install numpy`")

    def profile_column(self, df, column):
        if isinstance(self._bins, collections.Mapping):
            bins = self._bins[column]
        else:
            bins = self._bins
        values, limits = self._np.histogram(df[column], bins=bins)
        values = [v.item() for v in values]
        limits = [lim.item() for lim in limits]
        return (column + "_histogram", FloatHistogram(limits, values))

    # TODO: consider the case where the data is outside of the bucket list
    def profile_point(self, sample, reference):
        values, _ = self._np.histogram([sample], bins=reference._bucket_limits)
        return FloatHistogram(reference._bucket_limits, values)
<|endoftext|>


<commit_before># This file is part of Indico.
# Copyright (C) 2002 - 2015 European Organization for Nuclear Research (CERN).
#
# Indico is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# Indico is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Indico; if not, see <http://www.gnu.org/licenses/>.

from __future__ import unicode_literals

from celery.schedules import crontab

from indico.core.celery import celery
from indico.core.db import db

from indico_livesync.models.agents import LiveSyncAgent
from indico_livesync.util import clean_old_entries

@celery.periodic_task(run_every=crontab(minute='*/15'))
def scheduled_update():
    from indico_livesync.plugin import LiveSyncPlugin
    with LiveSyncPlugin.instance.plugin_context():
        clean_old_entries()
        for agent in LiveSyncAgent.find_all():
            if agent.backend is None:
                LiveSyncPlugin.logger.warning('Skipping agent {}; backend not found'.format(agent.name))
                continue
            if not agent.initial_data_exported:
                LiveSyncPlugin.logger.warning('Skipping agent {}; initial export not performed yet'.format(agent.name))
                continue
            LiveSyncPlugin.logger.info('Running agent {}'.format(agent.name))
            agent.create_backend().run()
            db.session.commit()
<commit_msg>LiveSync: Get plugin context via task decorator<commit_after># This file is part of Indico.
# Copyright (C) 2002 - 2015 European Organization for Nuclear Research (CERN).
#
# Indico is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# Indico is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Indico; if not, see <http://www.gnu.org/licenses/>.

from __future__ import unicode_literals

from celery.schedules import crontab

from indico.core.celery import celery
from indico.core.db import db

from indico_livesync.models.agents import LiveSyncAgent
from indico_livesync.util import clean_old_entries

@celery.periodic_task(run_every=crontab(minute='*/15'), plugin='livesync')
def scheduled_update():
    from indico_livesync.plugin import LiveSyncPlugin
    clean_old_entries()
    for agent in LiveSyncAgent.find_all():
        if agent.backend is None:
            LiveSyncPlugin.logger.warning('Skipping agent {}; backend not found'.format(agent.name))
            continue
        if not agent.initial_data_exported:
            LiveSyncPlugin.logger.warning('Skipping agent {}; initial export not performed yet'.format(agent.name))
            continue
        LiveSyncPlugin.logger.info('Running agent {}'.format(agent.name))
        agent.create_backend().run()
        db.session.commit()
<|endoftext|>


<commit_before>from collections import Counter
from functools import reduce, lru_cache
from random import random
import operator

BILABIAL = (0, {'positive': ['labial'], 'negative': ['syllabic']})
APICAL = (1, {'positive': ['coronal', 'anterior'], 'negative': ['syllabic']})
PALATAL = (2, {'positive': ['distributed'], 'negative': ['anterior',
                                                         'syllabic']})
VELAR = (3, {'positive': ['dorsal'], 'negative': ['syllabic']})
GLOTTAL = (4, {'negative': ['labial', 'syllabic', 'coronal', 'dorsal']})

FRONT = (5, {'positive': ['front'], 'negative': ['back']})
CENTRAL = (6, {'negative': ['front', 'back']})
BACK = (7, {'positive': ['back'], 'negative': ['front']})

PHONETIC_CONSTANTS = [BILABIAL, APICAL, PALATAL, VELAR, GLOTTAL, FRONT,
                      CENTRAL, BACK]

LIQUID = {'positive': ['consonantal', 'approximant']}
RHOTIC_VOWEL = {'positive': ['syllabic', 'coronal', 'anterior', 'distributed'],
                'negative': ['strident']}
FRICATIVE_OR_AFFRICATE = {'positive': ['delayedrelease']}
VOICED = {'positive': ['voice']}

@lru_cache(maxsize=None)
def classify_segment(segment):
    '''Given a segment, return all phonetic constants which apply to it, in a
    list.'''
    return [constant[0] for constant in PHONETIC_CONSTANTS if
            segment.meets_conditions(constant[1])]

def phonetic_product(word):
    '''Given a word, compute the Phonetic Product outlined by Harold R. Bauer in
    "The ethologic model of phonetic development: I. Phonetic contrast
    estimators" (1988).

    '''
    phonetic_count = Counter()

    for segment in word.segments:
        phonetic_count.update(classify_segment(segment))

    # Multiply the counts together, adding one to each count. The initial value
    # is one in case of words with no special features.
    return reduce(operator.mul, map(lambda x: x + 1, phonetic_count.values()),
                  1)

def weighted_phonetic_product(word):
    '''Given a word, compute the Phonetic Product outlined by Harold R. Bauer in
    "The ethologic model of phonetic development: I. Phonetic contrast
    estimators" (1988), using weighted values by Carterette and Jones in
    "Informal Speech: Alphabetic and Phonetic Texts with Statistical Analyses
    and Tables" (1974).

    '''
    phonetic_count = Counter()

    for segment in word.segments:
        phonetic_count.update(classify_segment(segment))

    segment_types = [phonetic_count[0] * 0.1658 + 1,
                     phonetic_count[1] * 0.3149 + 1,
                     phonetic_count[2] * 0.01129 + 1,
                     phonetic_count[3] * 0.04945 + 1,
                     phonetic_count[4] * 0.04945 + 1,
                     phonetic_count[5] * 0.18 + 1,
                     phonetic_count[6] * 0.1431 + 1,
                     phonetic_count[7] * 0.0709 + 1]

    return reduce(operator.mul, segment_types)

def number_of_syllables(word):
    '''Given a word, compute the number of syllables it contains.'''
    CV_syllable_rule = {'before': {'negative': ['syllabic']},
                        'conditions': {'positive': ['syllabic']}}

    VC_syllable_rule = {'after': {'negative': ['syllabic']},
                        'conditions': {'positive': ['syllabic']}}

    syllables = 0

    for index in range(len(word.segments)):
        is_CV_syllable = word.index_applicable(index, CV_syllable_rule)
        is_VC_syllable = word.index_applicable(index, VC_syllable_rule)

        if is_VC_syllable or is_CV_syllable:
            syllables += 1

    return syllables

def number_of_consonant_clusters(word):
    '''Given a word, compute the number of consonant clusters it contains.'''
    cluster_rule = {'before': {'positive': ['consonantal']},
                    'conditions': {'positive': ['consonantal']},
                    'after': {'negative': ['consonantal']}}

    return len([index for index in range(len(word.segments))
                if word.index_applicable(index, cluster_rule)])

def word_complexity_measure(word):
    '''Given a word, compute the Word Complexity Measure outlined by Carol
    Stoel-Gammon in "The Word Complexity Measure: Description and application
    to developmental phonology and disorders" (2010).

    '''
    WCM_score = 0

    if number_of_syllables(word) > 2:
        WCM_score += 1

    if 'consonantal' in word.segments[-1].positive:
        WCM_score += 1

    WCM_score += number_of_consonant_clusters(word)

    for segment in word.segments:
        if segment.meets_conditions(VELAR[1]):
            WCM_score += 1
        elif segment.meets_conditions(LIQUID):
            WCM_score += 1
        elif segment.meets_conditions(RHOTIC_VOWEL):
            WCM_score += 1
        elif segment.meets_conditions(FRICATIVE_OR_AFFRICATE):
            WCM_score += 1
            if segment.meets_conditions(VOICED):
                WCM_score += 1

    return WCM_score

def random_value(word):
    '''Given a word, return a random value.'''
    return random()
<commit_msg>Optimise metric calculation<commit_after>from collections import Counter
from functools import reduce, lru_cache
from random import random
import operator

BILABIAL = (0, {'positive': ['labial'], 'negative': ['syllabic']})
APICAL = (1, {'positive': ['coronal', 'anterior'], 'negative': ['syllabic']})
PALATAL = (2, {'positive': ['distributed'], 'negative': ['anterior',
                                                         'syllabic']})
VELAR = (3, {'positive': ['dorsal'], 'negative': ['syllabic']})
GLOTTAL = (4, {'negative': ['labial', 'syllabic', 'coronal', 'dorsal']})

FRONT = (5, {'positive': ['front'], 'negative': ['back']})
CENTRAL = (6, {'negative': ['front', 'back']})
BACK = (7, {'positive': ['back'], 'negative': ['front']})

PHONETIC_CONSTANTS = [BILABIAL, APICAL, PALATAL, VELAR, GLOTTAL, FRONT,
                      CENTRAL, BACK]

LIQUID = {'positive': ['consonantal', 'approximant']}
RHOTIC_VOWEL = {'positive': ['syllabic', 'coronal', 'anterior', 'distributed'],
                'negative': ['strident']}
FRICATIVE_OR_AFFRICATE = {'positive': ['delayedrelease']}
VOICED = {'positive': ['voice']}

@lru_cache(maxsize=None)
def classify_segment(segment):
    '''Given a segment, return all phonetic constants which apply to it, in a
    list.'''
    return [constant[0] for constant in PHONETIC_CONSTANTS if
            segment.meets_conditions(constant[1])]

def phonetic_product(word):
    '''Given a word, compute the Phonetic Product outlined by Harold R. Bauer in
    "The ethologic model of phonetic development: I. Phonetic contrast
    estimators" (1988).

    '''

    # Count all applicable features in the word.
    features = []
    for segment in word.segments:
        features.extend(classify_segment(segment))

    feature_count = Counter(features)

    # Multiply the counts together, adding one to each count. The initial value
    # is one in case of words with no special features.
    return reduce(operator.mul, map(lambda x: x + 1, feature_count.values()),
                  1)

def weighted_phonetic_product(word):
    '''Given a word, compute the Phonetic Product outlined by Harold R. Bauer in
    "The ethologic model of phonetic development: I. Phonetic contrast
    estimators" (1988), using weighted values by Carterette and Jones in
    "Informal Speech: Alphabetic and Phonetic Texts with Statistical Analyses
    and Tables" (1974).

    '''

    # Count all applicable features in the word.
    features = []
    for segment in word.segments:
        features.extend(classify_segment(segment))

    feature_count = Counter(features)

    segment_types = [feature_count[0] * 0.1658 + 1,
                     feature_count[1] * 0.3149 + 1,
                     feature_count[2] * 0.01129 + 1,
                     feature_count[3] * 0.04945 + 1,
                     feature_count[4] * 0.04945 + 1,
                     feature_count[5] * 0.18 + 1,
                     feature_count[6] * 0.1431 + 1,
                     feature_count[7] * 0.0709 + 1]

    return reduce(operator.mul, segment_types)

def number_of_syllables(word):
    '''Given a word, compute the number of syllables it contains.'''
    CV_syllable_rule = {'before': {'negative': ['syllabic']},
                        'conditions': {'positive': ['syllabic']}}

    VC_syllable_rule = {'after': {'negative': ['syllabic']},
                        'conditions': {'positive': ['syllabic']}}

    syllables = 0

    for index in range(len(word.segments)):
        is_CV_syllable = word.index_applicable(index, CV_syllable_rule)
        is_VC_syllable = word.index_applicable(index, VC_syllable_rule)

        if is_VC_syllable or is_CV_syllable:
            syllables += 1

    return syllables

def number_of_consonant_clusters(word):
    '''Given a word, compute the number of consonant clusters it contains.'''
    cluster_rule = {'before': {'positive': ['consonantal']},
                    'conditions': {'positive': ['consonantal']},
                    'after': {'negative': ['consonantal']}}

    return len([index for index in range(len(word.segments))
                if word.index_applicable(index, cluster_rule)])

def word_complexity_measure(word):
    '''Given a word, compute the Word Complexity Measure outlined by Carol
    Stoel-Gammon in "The Word Complexity Measure: Description and application
    to developmental phonology and disorders" (2010).

    '''
    WCM_score = 0

    if number_of_syllables(word) > 2:
        WCM_score += 1

    if 'consonantal' in word.segments[-1].positive:
        WCM_score += 1

    WCM_score += number_of_consonant_clusters(word)

    for segment in word.segments:
        if segment.meets_conditions(VELAR[1]):
            WCM_score += 1
        elif segment.meets_conditions(LIQUID):
            WCM_score += 1
        elif segment.meets_conditions(RHOTIC_VOWEL):
            WCM_score += 1
        elif segment.meets_conditions(FRICATIVE_OR_AFFRICATE):
            WCM_score += 1
            if segment.meets_conditions(VOICED):
                WCM_score += 1

    return WCM_score

def random_value(word):
    '''Given a word, return a random value.'''
    return random()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This project is licensed under the New BSD License:

Copyright (c) 2012, Pierre-Yves Chibon

All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
* Neither the name of the Wageningen University nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ''AS IS'' AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
"""

"""
Unit-tests for the goutil library.
"""

import os
import sys
import unittest

sys.path.insert(0, os.path.abspath('../'))
from src import PyGoLib
from src.godistance import GoDistanceCounter
from src.oboio import OboIO

GOFILE = 'test.obo'

class GoDistanceCounterTests(unittest.TestCase):
    """ GoDistanceCounter tests. """

    def __init__(self, methodName='runTest'):
        """ Constructor. """
        unittest.TestCase.__init__(self, methodName)

    def test_get_go_terms(self):
        """ Test the get_go_terms function. """
        obio = OboIO()
        terms = obio.get_graph(GOFILE)

    def test_get_path(self):
        """ Test the get_path function. """
        obio = OboIO()
        terms = obio.get_graph(GOFILE)
        golib = PyGoLib(terms)
        term = terms['4']
        print term['id']
        output = ['4,2,1,0']
        self.assertEqual(output,
            golib.get_path(term, pred=term['id'], paths=[], verbose=True))
        self.assertEqual(output,
            golib.get_path(term, pred=term['id'], paths=[]))
        output = ['2,1,0']
        self.assertEqual(output,
                golib.get_path(term, paths=[]))

    def test_scores(self):
        """ Test the scores function. """
        obio = OboIO()
        terms = obio.get_graph(GOFILE)
        gdc = GoDistanceCounter(terms)
        self.assertEqual(5.5, gdc.scores('11', '0'))
        self.assertEqual(4.4, gdc.scores('11', '1'))
        self.assertEqual(6.0, gdc.scores('9', '5'))

suite = unittest.TestLoader().loadTestsFromTestCase(GoDistanceCounterTests)
unittest.TextTestRunner(verbosity=2).run(suite)
<commit_msg>Get the obo file for test within the tests directory<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This project is licensed under the New BSD License:

Copyright (c) 2012, Pierre-Yves Chibon

All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
* Redistributions of source code must retain the above copyright notice,
this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
* Neither the name of the Wageningen University nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ''AS IS'' AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
"""

"""
Unit-tests for the goutil library.
"""

import os
import sys
import unittest

sys.path.insert(0, os.path.abspath('../'))
from src import PyGoLib
from src.godistance import GoDistanceCounter
from src.oboio import OboIO

GOFILE = '%s/test.obo' % os.path.dirname(__file__)

class GoDistanceCounterTests(unittest.TestCase):
    """ GoDistanceCounter tests. """

    def __init__(self, methodName='runTest'):
        """ Constructor. """
        unittest.TestCase.__init__(self, methodName)

    def test_get_go_terms(self):
        """ Test the get_go_terms function. """
        obio = OboIO()
        terms = obio.get_graph(GOFILE)

    def test_get_path(self):
        """ Test the get_path function. """
        obio = OboIO()
        terms = obio.get_graph(GOFILE)
        golib = PyGoLib(terms)
        term = terms['4']
        print term['id']
        output = ['4,2,1,0']
        self.assertEqual(output,
            golib.get_path(term, pred=term['id'], paths=[], verbose=True))
        self.assertEqual(output,
            golib.get_path(term, pred=term['id'], paths=[]))
        output = ['2,1,0']
        self.assertEqual(output,
                golib.get_path(term, paths=[]))

    def test_scores(self):
        """ Test the scores function. """
        obio = OboIO()
        terms = obio.get_graph(GOFILE)
        gdc = GoDistanceCounter(terms)
        self.assertEqual(5.5, gdc.scores('11', '0'))
        self.assertEqual(4.4, gdc.scores('11', '1'))
        self.assertEqual(6.0, gdc.scores('9', '5'))

suite = unittest.TestLoader().loadTestsFromTestCase(GoDistanceCounterTests)
unittest.TextTestRunner(verbosity=2).run(suite)
<|endoftext|>


<commit_before>#-*- coding: utf-8 -*-

import viennagrid_wrapper as _wrapper

_SUPPORTED_NORMS = (1, 2, 'inf')

def inner_prod(point1, point2):
	# Try to get method 'inner_prod' from 'point1'. If it doesn't have the method,
	# it means it's not a cartesian point. Thus, convert to cartesian coordinates
	# and get the method. If it still doesn't have the method, raise an exception.
	try:
		inner_prod_fn = point1.__getattribute__('inner_prod')
	except AttributeError:
		casted_pnt1 = point1.to_cartesian()
		try:
			inner_prod_fn = casted_pnt1.__getattribute__('inner_prod')
		except AttributeError:
			raise TypeError('point1 has no method named inner_prod')
	else:
		casted_pnt1 = point1

	# If point types are equal, simply calculate the inner product. If they're not
	# equal, try to convert 'point2' to the type of 'point1'. If types are still
	# different, it means that both points are of incompatible types
	# (i.e. incompatible dimensions).
	if casted_pnt1.__class__ is point2.__class__:
		return inner_prod_fun(point2)
	else:
		casted_pnt2 = point2.to_cartesian()
		if casted_pnt1.__class__ is casted_pnt2.__class__:
			return inner_prod_fun(casted_pnt2)
		else:
			raise TypeError('incompatible point types')

def cross_prod(point1, point2):
	# Try to get method 'cross_prod' from 'point1'. If it doesn't have the method,
	# it means it's not a cartesian point. Thus, convert to cartesian coordinates
	# and get the method. If it still doesn't have the method, raise an exception.
	try:
		cross_prod_fn = point1.__getattribute__('cross_prod')
	except AttributeError:
		casted_pnt1 = point1.to_cartesian()
		try:
			cross_prod_fn = casted_pnt1.__getattribute__('cross_prod')
		except AttributeError:
			raise TypeError('point1 has no method named cross_prod')
	else:
		casted_pnt1 = point1

	# If point types are equal, simply calculate the cross product. If they're not
	# equal, try to convert 'point2' to the type of 'point1'. If types are still
	# different, it means that both points are of incompatible types
	# (i.e. incompatible dimensions).
	if casted_pnt1.__class__ is point2.__class__:
		return cross_prod_fn(point2)
	else:
		casted_pnt2 = point2.to_cartesian()
		if casted_pnt1.__class__ is casted_pnt2.__class__:
			return cross_prod_fn(casted_pnt2)
		else:
			raise TypeError('incompatible point types')

def norm(point, norm_type=2):
	if norm_type in _SUPPORTED_NORMS:
		norm_type = str(norm_type)
		norm_fn = _wrapper.__getattribute__('norm_%(norm_type)s' % locals())
		return norm_fn(point)
	else:
		raise ValueError('unsupported norm type: %(norm_type)s')

def apply_voronoi(dom):
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	_wrapper.apply_voronoi(dom)

def centroid(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	return _wrapper.centroid(cell)

def cell_refine(dom, seg, predicate):
	"""Returns a tuple containing the output domain and segmentation after the refination."""
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	if isinstance(seg, viennagrid.Segmentation):
		seg = seg._segmentation
	return _wrapper.cell_refine(dom, seg, predicate)

def circumcenter(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	return _wrapper.circumcenter(cell)

def is_boundary(domseg, interface_elem):
	if isinstance(domseg, viennagrid.Domain):
		domseg = domseg._domain
	elif isinstance(domseg, viennagrid.Segment):
		domseg = domseg._segment

	if isinstance(interface_elem, viennagrid.Facet):
		interface_elem = interface_elem._facet
	elif isinstance(interface_elem, viennagrid.Edge):
		interface_elem = interface_elem._edge
	elif isinstance(interface_elem, viennagrid.Vertex):
		interface_elem = interface_elem._vertex

	return _wrapper.is_boundary(domseg, interface_elem)

def is_interface(seg0, seg1, boundary_elem):
	if isinstance(seg0, viennagrid.Segment):
		seg0 = seg0._segment

	if isinstance(seg1, viennagrid.Segment):
		seg1 = seg1._segment

	if isinstance(boundary_elem, viennagrid.Facet):
		boundary_elem = boundary_elem._facet
	elif isinstance(boundary_elem, viennagrid.Edge):
		boundary_elem = boundary_elem._edge
	elif isinstance(boundary_elem, viennagrid.Vertex):
		boundary_elem = boundary_elem._vertex

	return _wrapper.is_interface(seg0, seg1, boundary_elem)

def refine(dom, seg, predicate):
	"""Returns a tuple containing the output domain and segmentation after the refination."""
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	if isinstance(seg, viennagrid.Segmentation):
		seg = seg._segmentation
	return _wrapper.refine(dom, seg, predicate)

def refine_uniformly(dom, seg):
	"""Returns a tuple containing the output domain and segmentation after the refination."""
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	if isinstance(seg, viennagrid.Segmentation):
		seg = seg._segmentation
	return _wrapper.refine_uniformly(dom, seg)

def surface(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	elif isinstance(cell, viennagrid.Domain):
		cell = cell._domain
	elif isinstance(cell, viennagrid.Segment):
		cell = cell._segment
	return _wrapper.surface(cell)

def volume(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	elif isinstance(cell, viennagrid.Domain):
		cell = cell._domain
	elif isinstance(cell, viennagrid.Segment):
		cell = cell._segment
	return _wrapper.volume(cell)
<commit_msg>Write a high-level wrapper function for 'spanned_volume'.<commit_after>#-*- coding: utf-8 -*-

import viennagrid_wrapper as _wrapper

_SUPPORTED_NORMS = (1, 2, 'inf')

def inner_prod(point1, point2):
	# Try to get method 'inner_prod' from 'point1'. If it doesn't have the method,
	# it means it's not a cartesian point. Thus, convert to cartesian coordinates
	# and get the method. If it still doesn't have the method, raise an exception.
	try:
		inner_prod_fn = point1.__getattribute__('inner_prod')
	except AttributeError:
		casted_pnt1 = point1.to_cartesian()
		try:
			inner_prod_fn = casted_pnt1.__getattribute__('inner_prod')
		except AttributeError:
			raise TypeError('point1 has no method named inner_prod')
	else:
		casted_pnt1 = point1

	# If point types are equal, simply calculate the inner product. If they're not
	# equal, try to convert 'point2' to the type of 'point1'. If types are still
	# different, it means that both points are of incompatible types
	# (i.e. incompatible dimensions).
	if casted_pnt1.__class__ is point2.__class__:
		return inner_prod_fun(point2)
	else:
		casted_pnt2 = point2.to_cartesian()
		if casted_pnt1.__class__ is casted_pnt2.__class__:
			return inner_prod_fun(casted_pnt2)
		else:
			raise TypeError('incompatible point types')

def cross_prod(point1, point2):
	# Try to get method 'cross_prod' from 'point1'. If it doesn't have the method,
	# it means it's not a cartesian point. Thus, convert to cartesian coordinates
	# and get the method. If it still doesn't have the method, raise an exception.
	try:
		cross_prod_fn = point1.__getattribute__('cross_prod')
	except AttributeError:
		casted_pnt1 = point1.to_cartesian()
		try:
			cross_prod_fn = casted_pnt1.__getattribute__('cross_prod')
		except AttributeError:
			raise TypeError('point1 has no method named cross_prod')
	else:
		casted_pnt1 = point1

	# If point types are equal, simply calculate the cross product. If they're not
	# equal, try to convert 'point2' to the type of 'point1'. If types are still
	# different, it means that both points are of incompatible types
	# (i.e. incompatible dimensions).
	if casted_pnt1.__class__ is point2.__class__:
		return cross_prod_fn(point2)
	else:
		casted_pnt2 = point2.to_cartesian()
		if casted_pnt1.__class__ is casted_pnt2.__class__:
			return cross_prod_fn(casted_pnt2)
		else:
			raise TypeError('incompatible point types')

def norm(point, norm_type=2):
	if norm_type in _SUPPORTED_NORMS:
		norm_type = str(norm_type)
		norm_fn = _wrapper.__getattribute__('norm_%(norm_type)s' % locals())
		return norm_fn(point)
	else:
		raise ValueError('unsupported norm type: %(norm_type)s')

def apply_voronoi(dom):
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	_wrapper.apply_voronoi(dom)

def centroid(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	return _wrapper.centroid(cell)

def cell_refine(dom, seg, predicate):
	"""Returns a tuple containing the output domain and segmentation after the refination."""
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	if isinstance(seg, viennagrid.Segmentation):
		seg = seg._segmentation
	return _wrapper.cell_refine(dom, seg, predicate)

def circumcenter(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	return _wrapper.circumcenter(cell)

def is_boundary(domseg, interface_elem):
	if isinstance(domseg, viennagrid.Domain):
		domseg = domseg._domain
	elif isinstance(domseg, viennagrid.Segment):
		domseg = domseg._segment

	if isinstance(interface_elem, viennagrid.Facet):
		interface_elem = interface_elem._facet
	elif isinstance(interface_elem, viennagrid.Edge):
		interface_elem = interface_elem._edge
	elif isinstance(interface_elem, viennagrid.Vertex):
		interface_elem = interface_elem._vertex

	return _wrapper.is_boundary(domseg, interface_elem)

def is_interface(seg0, seg1, boundary_elem):
	if isinstance(seg0, viennagrid.Segment):
		seg0 = seg0._segment

	if isinstance(seg1, viennagrid.Segment):
		seg1 = seg1._segment

	if isinstance(boundary_elem, viennagrid.Facet):
		boundary_elem = boundary_elem._facet
	elif isinstance(boundary_elem, viennagrid.Edge):
		boundary_elem = boundary_elem._edge
	elif isinstance(boundary_elem, viennagrid.Vertex):
		boundary_elem = boundary_elem._vertex

	return _wrapper.is_interface(seg0, seg1, boundary_elem)

def refine(dom, seg, predicate):
	"""Returns a tuple containing the output domain and segmentation after the refination."""
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	if isinstance(seg, viennagrid.Segmentation):
		seg = seg._segmentation
	return _wrapper.refine(dom, seg, predicate)

def refine_uniformly(dom, seg):
	"""Returns a tuple containing the output domain and segmentation after the refination."""
	if isinstance(dom, viennagrid.Domain):
		dom = dom._domain
	if isinstance(seg, viennagrid.Segmentation):
		seg = seg._segmentation
	return _wrapper.refine_uniformly(dom, seg)

def surface(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	elif isinstance(cell, viennagrid.Domain):
		cell = cell._domain
	elif isinstance(cell, viennagrid.Segment):
		cell = cell._segment
	return _wrapper.surface(cell)

def volume(cell):
	if isinstance(cell, viennagrid.Cell):
		cell = cell._cell
	elif isinstance(cell, viennagrid.Domain):
		cell = cell._domain
	elif isinstance(cell, viennagrid.Segment):
		cell = cell._segment
	return _wrapper.volume(cell)

def spanned_volume(*args):
	point_list = []

	for point in args:
		# If the point is a high-level point, get the low-level point
		if isinstance(point, viennagrid.Point):
			point = point._point

		# If point is not cartesian, convert it
		if point.coord_system != 'cartesian':
			point = point.to_cartesian()

		# Append the point to the list
		point_list.append(point)

	return _wrapper.spanned_volume(point_list)
<|endoftext|>


<commit_before>"""celery.backends.base"""
import time

from celery import states
from celery.exceptions import TimeoutError, TaskRevokedError
from celery.serialization import pickle, get_pickled_exception
from celery.serialization import get_pickleable_exception
from celery.datastructures import LocalCache

class BaseBackend(object):
    """The base backend class. All backends should inherit from this."""

    READY_STATES = states.READY_STATES
    UNREADY_STATES = states.UNREADY_STATES
    EXCEPTION_STATES = states.EXCEPTION_STATES

    TimeoutError = TimeoutError

    def __init__(self, *args, **kwargs):
        from celery.app import app_or_default
        self.app = app_or_default(kwargs.get("app"))

    def encode_result(self, result, status):
        if status in self.EXCEPTION_STATES:
            return self.prepare_exception(result)
        else:
            return self.prepare_value(result)

    def store_result(self, task_id, result, status):
        """Store the result and status of a task."""
        raise NotImplementedError(
                "store_result is not supported by this backend.")

    def mark_as_started(self, task_id, **meta):
        """Mark a task as started"""
        return self.store_result(task_id, meta, status=states.STARTED)

    def mark_as_done(self, task_id, result):
        """Mark task as successfully executed."""
        return self.store_result(task_id, result, status=states.SUCCESS)

    def mark_as_failure(self, task_id, exc, traceback=None):
        """Mark task as executed with failure. Stores the execption."""
        return self.store_result(task_id, exc, status=states.FAILURE,
                                 traceback=traceback)

    def mark_as_retry(self, task_id, exc, traceback=None):
        """Mark task as being retries. Stores the current
        exception (if any)."""
        return self.store_result(task_id, exc, status=states.RETRY,
                                 traceback=traceback)

    def mark_as_revoked(self, task_id):
        return self.store_result(task_id, TaskRevokedError(),
                                 status=states.REVOKED, traceback=None)

    def prepare_exception(self, exc):
        """Prepare exception for serialization."""
        return get_pickleable_exception(exc)

    def exception_to_python(self, exc):
        """Convert serialized exception to Python exception."""
        return get_pickled_exception(exc)

    def prepare_value(self, result):
        """Prepare value for storage."""
        return result

    def forget(self, task_id):
        raise NotImplementedError("%s does not implement forget." % (
                    self.__class__))

    def wait_for(self, task_id, timeout=None):
        """Wait for task and return its result.

        If the task raises an exception, this exception
        will be re-raised by :func:`wait_for`.

        If ``timeout`` is not ``None``, this raises the
        :class:`celery.exceptions.TimeoutError` exception if the operation
        takes longer than ``timeout`` seconds.

        """

        sleep_inbetween = 0.5
        time_elapsed = 0.0

        while True:
            status = self.get_status(task_id)
            if status == states.SUCCESS:
                return self.get_result(task_id)
            elif status in states.PROPAGATE_STATES:
                raise self.get_result(task_id)
            # avoid hammering the CPU checking status.
            time.sleep(sleep_inbetween)
            time_elapsed += sleep_inbetween
            if timeout and time_elapsed >= timeout:
                raise TimeoutError("The operation timed out.")

    def cleanup(self):
        """Backend cleanup. Is run by
        :class:`celery.task.DeleteExpiredTaskMetaTask`."""
        pass

    def process_cleanup(self):
        """Cleanup actions to do at the end of a task worker process."""
        pass

    def get_status(self, task_id):
        """Get the status of a task."""
        raise NotImplementedError(
                "get_status is not supported by this backend.")

    def get_result(self, task_id):
        """Get the result of a task."""
        raise NotImplementedError(
                "get_result is not supported by this backend.")

    def get_traceback(self, task_id):
        """Get the traceback for a failed task."""
        raise NotImplementedError(
                "get_traceback is not supported by this backend.")

    def save_taskset(self, taskset_id, result):
        """Store the result and status of a task."""
        raise NotImplementedError(
                "save_taskset is not supported by this backend.")

    def restore_taskset(self, taskset_id, cache=True):
        """Get the result of a taskset."""
        raise NotImplementedError(
                "restore_taskset is not supported by this backend.")

    def reload_task_result(self, task_id):
        """Reload task result, even if it has been previously fetched."""
        raise NotImplementedError(
                "reload_task_result is not supported by this backend.")

    def reload_taskset_result(self, task_id):
        """Reload taskset result, even if it has been previously fetched."""
        raise NotImplementedError(
                "reload_taskset_result is not supported by this backend.")

class BaseDictBackend(BaseBackend):

    def __init__(self, *args, **kwargs):
        super(BaseDictBackend, self).__init__(*args, **kwargs)
        self._cache = LocalCache(limit=kwargs.get("max_cached_results") or
                                 self.app.conf.CELERY_MAX_CACHED_RESULTS)

    def store_result(self, task_id, result, status, traceback=None):
        """Store task result and status."""
        result = self.encode_result(result, status)
        return self._store_result(task_id, result, status, traceback)

    def forget(self, task_id):
        self._cache.pop(task_id, None)
        self._forget(task_id)

    def get_status(self, task_id):
        """Get the status of a task."""
        return self.get_task_meta(task_id)["status"]

    def get_traceback(self, task_id):
        """Get the traceback for a failed task."""
        return self.get_task_meta(task_id).get("traceback")

    def get_result(self, task_id):
        """Get the result of a task."""
        meta = self.get_task_meta(task_id)
        if meta["status"] in self.EXCEPTION_STATES:
            return self.exception_to_python(meta["result"])
        else:
            return meta["result"]

    def get_task_meta(self, task_id, cache=True):
        if cache and task_id in self._cache:
            return self._cache[task_id]

        meta = self._get_task_meta_for(task_id)
        if cache and meta.get("status") == states.SUCCESS:
            self._cache[task_id] = meta
        return meta

    def reload_task_result(self, task_id):
        self._cache[task_id] = self.get_task_meta(task_id, cache=False)

    def reload_taskset_result(self, taskset_id):
        self._cache[taskset_id] = self.get_taskset_meta(taskset_id,
                                                        cache=False)

    def get_taskset_meta(self, taskset_id, cache=True):
        if cache and taskset_id in self._cache:
            return self._cache[taskset_id]

        meta = self._restore_taskset(taskset_id)
        if cache and meta is not None:
            self._cache[taskset_id] = meta
        return meta

    def restore_taskset(self, taskset_id, cache=True):
        """Get the result for a taskset."""
        meta = self.get_taskset_meta(taskset_id, cache=cache)
        if meta:
            return meta["result"]

    def save_taskset(self, taskset_id, result):
        """Store the result of an executed taskset."""
        return self._save_taskset(taskset_id, result)

class KeyValueStoreBackend(BaseDictBackend):

    def get(self, key):
        raise NotImplementedError("Must implement the get method.")

    def set(self, key, value):
        raise NotImplementedError("Must implement the set method.")

    def delete(self, key):
        raise NotImplementedError("Must implement the delete method")

    def get_key_for_task(self, task_id):
        """Get the cache key for a task by id."""
        return "celery-task-meta-%s" % task_id

    def get_key_for_taskset(self, task_id):
        """Get the cache key for a task by id."""
        return "celery-taskset-meta-%s" % task_id

    def _forget(self, task_id):
        self.delete(task_id)

    def _store_result(self, task_id, result, status, traceback=None):
        meta = {"status": status, "result": result, "traceback": traceback}
        self.set(self.get_key_for_task(task_id), pickle.dumps(meta))
        return result

    def _save_taskset(self, taskset_id, result):
        meta = {"result": result}
        self.set(self.get_key_for_taskset(taskset_id), pickle.dumps(meta))
        return result

    def _get_task_meta_for(self, task_id):
        """Get task metadata for a task by id."""
        meta = self.get(self.get_key_for_task(task_id))
        if not meta:
            return {"status": states.PENDING, "result": None}
        return pickle.loads(str(meta))

    def _restore_taskset(self, taskset_id):
        """Get task metadata for a task by id."""
        meta = self.get(self.get_key_for_taskset(taskset_id))
        if meta:
            meta = pickle.loads(str(meta))
            return meta
<commit_msg>KeyValueStoreBackend.forget: delete(key_for_task(id)) not delete(id)<commit_after>"""celery.backends.base"""
import time

from celery import states
from celery.exceptions import TimeoutError, TaskRevokedError
from celery.serialization import pickle, get_pickled_exception
from celery.serialization import get_pickleable_exception
from celery.datastructures import LocalCache

class BaseBackend(object):
    """The base backend class. All backends should inherit from this."""

    READY_STATES = states.READY_STATES
    UNREADY_STATES = states.UNREADY_STATES
    EXCEPTION_STATES = states.EXCEPTION_STATES

    TimeoutError = TimeoutError

    def __init__(self, *args, **kwargs):
        from celery.app import app_or_default
        self.app = app_or_default(kwargs.get("app"))

    def encode_result(self, result, status):
        if status in self.EXCEPTION_STATES:
            return self.prepare_exception(result)
        else:
            return self.prepare_value(result)

    def store_result(self, task_id, result, status):
        """Store the result and status of a task."""
        raise NotImplementedError(
                "store_result is not supported by this backend.")

    def mark_as_started(self, task_id, **meta):
        """Mark a task as started"""
        return self.store_result(task_id, meta, status=states.STARTED)

    def mark_as_done(self, task_id, result):
        """Mark task as successfully executed."""
        return self.store_result(task_id, result, status=states.SUCCESS)

    def mark_as_failure(self, task_id, exc, traceback=None):
        """Mark task as executed with failure. Stores the execption."""
        return self.store_result(task_id, exc, status=states.FAILURE,
                                 traceback=traceback)

    def mark_as_retry(self, task_id, exc, traceback=None):
        """Mark task as being retries. Stores the current
        exception (if any)."""
        return self.store_result(task_id, exc, status=states.RETRY,
                                 traceback=traceback)

    def mark_as_revoked(self, task_id):
        return self.store_result(task_id, TaskRevokedError(),
                                 status=states.REVOKED, traceback=None)

    def prepare_exception(self, exc):
        """Prepare exception for serialization."""
        return get_pickleable_exception(exc)

    def exception_to_python(self, exc):
        """Convert serialized exception to Python exception."""
        return get_pickled_exception(exc)

    def prepare_value(self, result):
        """Prepare value for storage."""
        return result

    def forget(self, task_id):
        raise NotImplementedError("%s does not implement forget." % (
                    self.__class__))

    def wait_for(self, task_id, timeout=None):
        """Wait for task and return its result.

        If the task raises an exception, this exception
        will be re-raised by :func:`wait_for`.

        If ``timeout`` is not ``None``, this raises the
        :class:`celery.exceptions.TimeoutError` exception if the operation
        takes longer than ``timeout`` seconds.

        """

        sleep_inbetween = 0.5
        time_elapsed = 0.0

        while True:
            status = self.get_status(task_id)
            if status == states.SUCCESS:
                return self.get_result(task_id)
            elif status in states.PROPAGATE_STATES:
                raise self.get_result(task_id)
            # avoid hammering the CPU checking status.
            time.sleep(sleep_inbetween)
            time_elapsed += sleep_inbetween
            if timeout and time_elapsed >= timeout:
                raise TimeoutError("The operation timed out.")

    def cleanup(self):
        """Backend cleanup. Is run by
        :class:`celery.task.DeleteExpiredTaskMetaTask`."""
        pass

    def process_cleanup(self):
        """Cleanup actions to do at the end of a task worker process."""
        pass

    def get_status(self, task_id):
        """Get the status of a task."""
        raise NotImplementedError(
                "get_status is not supported by this backend.")

    def get_result(self, task_id):
        """Get the result of a task."""
        raise NotImplementedError(
                "get_result is not supported by this backend.")

    def get_traceback(self, task_id):
        """Get the traceback for a failed task."""
        raise NotImplementedError(
                "get_traceback is not supported by this backend.")

    def save_taskset(self, taskset_id, result):
        """Store the result and status of a task."""
        raise NotImplementedError(
                "save_taskset is not supported by this backend.")

    def restore_taskset(self, taskset_id, cache=True):
        """Get the result of a taskset."""
        raise NotImplementedError(
                "restore_taskset is not supported by this backend.")

    def reload_task_result(self, task_id):
        """Reload task result, even if it has been previously fetched."""
        raise NotImplementedError(
                "reload_task_result is not supported by this backend.")

    def reload_taskset_result(self, task_id):
        """Reload taskset result, even if it has been previously fetched."""
        raise NotImplementedError(
                "reload_taskset_result is not supported by this backend.")

class BaseDictBackend(BaseBackend):

    def __init__(self, *args, **kwargs):
        super(BaseDictBackend, self).__init__(*args, **kwargs)
        self._cache = LocalCache(limit=kwargs.get("max_cached_results") or
                                 self.app.conf.CELERY_MAX_CACHED_RESULTS)

    def store_result(self, task_id, result, status, traceback=None):
        """Store task result and status."""
        result = self.encode_result(result, status)
        return self._store_result(task_id, result, status, traceback)

    def forget(self, task_id):
        self._cache.pop(task_id, None)
        self._forget(task_id)

    def get_status(self, task_id):
        """Get the status of a task."""
        return self.get_task_meta(task_id)["status"]

    def get_traceback(self, task_id):
        """Get the traceback for a failed task."""
        return self.get_task_meta(task_id).get("traceback")

    def get_result(self, task_id):
        """Get the result of a task."""
        meta = self.get_task_meta(task_id)
        if meta["status"] in self.EXCEPTION_STATES:
            return self.exception_to_python(meta["result"])
        else:
            return meta["result"]

    def get_task_meta(self, task_id, cache=True):
        if cache and task_id in self._cache:
            return self._cache[task_id]

        meta = self._get_task_meta_for(task_id)
        if cache and meta.get("status") == states.SUCCESS:
            self._cache[task_id] = meta
        return meta

    def reload_task_result(self, task_id):
        self._cache[task_id] = self.get_task_meta(task_id, cache=False)

    def reload_taskset_result(self, taskset_id):
        self._cache[taskset_id] = self.get_taskset_meta(taskset_id,
                                                        cache=False)

    def get_taskset_meta(self, taskset_id, cache=True):
        if cache and taskset_id in self._cache:
            return self._cache[taskset_id]

        meta = self._restore_taskset(taskset_id)
        if cache and meta is not None:
            self._cache[taskset_id] = meta
        return meta

    def restore_taskset(self, taskset_id, cache=True):
        """Get the result for a taskset."""
        meta = self.get_taskset_meta(taskset_id, cache=cache)
        if meta:
            return meta["result"]

    def save_taskset(self, taskset_id, result):
        """Store the result of an executed taskset."""
        return self._save_taskset(taskset_id, result)

class KeyValueStoreBackend(BaseDictBackend):

    def get(self, key):
        raise NotImplementedError("Must implement the get method.")

    def set(self, key, value):
        raise NotImplementedError("Must implement the set method.")

    def delete(self, key):
        raise NotImplementedError("Must implement the delete method")

    def get_key_for_task(self, task_id):
        """Get the cache key for a task by id."""
        return "celery-task-meta-%s" % task_id

    def get_key_for_taskset(self, task_id):
        """Get the cache key for a task by id."""
        return "celery-taskset-meta-%s" % task_id

    def _forget(self, task_id):
        self.delete(self.get_key_for_task(task_id))

    def _store_result(self, task_id, result, status, traceback=None):
        meta = {"status": status, "result": result, "traceback": traceback}
        self.set(self.get_key_for_task(task_id), pickle.dumps(meta))
        return result

    def _save_taskset(self, taskset_id, result):
        meta = {"result": result}
        self.set(self.get_key_for_taskset(taskset_id), pickle.dumps(meta))
        return result

    def _get_task_meta_for(self, task_id):
        """Get task metadata for a task by id."""
        meta = self.get(self.get_key_for_task(task_id))
        if not meta:
            return {"status": states.PENDING, "result": None}
        return pickle.loads(str(meta))

    def _restore_taskset(self, taskset_id):
        """Get task metadata for a task by id."""
        meta = self.get(self.get_key_for_taskset(taskset_id))
        if meta:
            meta = pickle.loads(str(meta))
            return meta
<|endoftext|>