<commit_before>"""
A ``TokenEmbedder`` which uses one of the BERT models
(https://github.com/google-research/bert)
to produce embeddings.

At its core it uses Hugging Face's PyTorch implementation
(https://github.com/huggingface/pytorch-pretrained-BERT),
so thanks to them!
"""
from typing import Dict
import logging

import torch
import torch.nn.functional as F

from pytorch_pretrained_bert.modeling import BertModel

from allennlp.modules.scalar_mix import ScalarMix
from allennlp.modules.token_embedders.token_embedder import TokenEmbedder
from allennlp.nn import util

logger = logging.getLogger(__name__)

class PretrainedBertModel:
    """
    In some instances you may want to load the same BERT model twice
    (e.g. to use as a token embedder and also as a pooling layer).
    This factory provides a cache so that you don't actually have to load the model twice.
    """
    _cache: Dict[str, BertModel] = {}

    @classmethod
    def load(cls, model_name: str, cache_model: bool = True) -> BertModel:
        if model_name in cls._cache:
            return PretrainedBertModel._cache[model_name]

        model = BertModel.from_pretrained(model_name)
        if cache_model:
            cls._cache[model_name] = model

        return model

class BertEmbedder(TokenEmbedder):
    """
    A ``TokenEmbedder`` that produces BERT embeddings for your tokens.
    Should be paired with a ``BertIndexer``, which produces wordpiece ids.

    Most likely you probably want to use ``PretrainedBertEmbedder``
    for one of the named pretrained models, not this base class.

    Parameters
    ----------
    bert_model: ``BertModel``
        The BERT model being wrapped.
    top_layer_only: ``bool``, optional (default = ``False``)
        If ``True``, then only return the top layer instead of apply the scalar mix.
    max_pieces : int, optional (default: 512)
        The BERT embedder uses positional embeddings and so has a corresponding
        maximum length for its input ids. Assuming the inputs are windowed
        and padded appropriately by this length, the embedder will split them into a
        large batch, feed them into BERT, and recombine the output as if it was a
        longer sequence.
    num_start_tokens : int, optional (default: 1)
        The number of starting special tokens input to BERT (usually 1, i.e., [CLS])
    num_end_tokens : int, optional (default: 1)
        The number of ending tokens input to BERT (usually 1, i.e., [SEP])
    """
    def __init__(self,
                 bert_model: BertModel,
                 top_layer_only: bool = False,
                 max_pieces: int = 512,
                 num_start_tokens: int = 1,
                 num_end_tokens: int = 1) -> None:
        super().__init__()
        self.bert_model = bert_model
        self.output_dim = bert_model.config.hidden_size
        self.max_pieces = max_pieces
        self.num_start_tokens = num_start_tokens
        self.num_end_tokens = num_end_tokens

        if not top_layer_only:
            self._scalar_mix = ScalarMix(bert_model.config.num_hidden_layers,
                                         do_layer_norm=False)
        else:
            self._scalar_mix = None

    def get_output_dim(self) -> int:
        return self.output_dim

    def forward(self,
                input_ids: torch.LongTensor,
                offsets: torch.LongTensor = None,
                token_type_ids: torch.LongTensor = None) -> torch.Tensor:
        """
        Parameters
        ----------
        input_ids : ``torch.LongTensor``
            The (batch_size, ..., max_sequence_length) tensor of wordpiece ids.
        offsets : ``torch.LongTensor``, optional
            The BERT embeddings are one per wordpiece. However it's possible/likely
            you might want one per original token. In that case, ``offsets``
            represents the indices of the desired wordpiece for each original token.
            Depending on how your token indexer is configured, this could be the
            position of the last wordpiece for each token, or it could be the position
            of the first wordpiece for each token.

            For example, if you had the sentence "Definitely not", and if the corresponding
            wordpieces were ["Def", "##in", "##ite", "##ly", "not"], then the input_ids
            would be 5 wordpiece ids, and the "last wordpiece" offsets would be [3, 4].
            If offsets are provided, the returned tensor will contain only the wordpiece
            embeddings at those positions, and (in particular) will contain one embedding
            per token. If offsets are not provided, the entire tensor of wordpiece embeddings
            will be returned.
        token_type_ids : ``torch.LongTensor``, optional
            If an input consists of two sentences (as in the BERT paper),
            tokens from the first sentence should have type 0 and tokens from
            the second sentence should have type 1.  If you don't provide this
            (the default BertIndexer doesn't) then it's assumed to be all 0s.
        """
        # pylint: disable=arguments-differ
        batch_size, full_seq_len = input_ids.size(0), input_ids.size(-1)
        initial_dims = list(input_ids.shape[:-1])

        # The embedder may receive an input tensor that has a sequence length longer than can
        # be fit. In that case, we should expect the wordpiece indexer to create padded windows
        # of length `self.max_pieces` for us, and have them concatenated into one long sequence.
        # E.g., "[CLS] I went to the [SEP] [CLS] to the store to [SEP] ..."
        # We can then split the sequence into sub-sequences of that length, and concatenate them
        # along the batch dimension so we effectively have one huge batch of partial sentences.
        # This can then be fed into BERT without any sentence length issues. Keep in mind
        # that the memory consumption can dramatically increase for large batches with extremely
        # long sentences.
        needs_split = full_seq_len > self.max_pieces
        last_window_size = 0
        if needs_split:
            # Split the flattened list by the window size, `max_pieces`
            split_input_ids = list(input_ids.split(self.max_pieces, dim=-1))

            # We want all sequences to be the same length, so pad the last sequence
            last_window_size = split_input_ids[-1].size(-1)
            padding_amount = self.max_pieces - last_window_size
            split_input_ids[-1] = F.pad(split_input_ids[-1], pad=[0, padding_amount], value=0)

            # Now combine the sequences along the batch dimension
            input_ids = torch.cat(split_input_ids, dim=0)

            if token_type_ids is not None:
                # Same for token_type_ids
                split_token_type_ids = list(token_type_ids.split(self.max_pieces, dim=-1))

                last_window_size = split_token_type_ids[-1].size(-1)
                padding_amount = self.max_pieces - last_window_size
                split_token_type_ids[-1] = F.pad(split_token_type_ids[-1], pad=[0, padding_amount], value=0)

                token_type_ids = torch.cat(split_token_type_ids, dim=0)

        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)

        input_mask = (input_ids != 0).long()

        # input_ids may have extra dimensions, so we reshape down to 2-d
        # before calling the BERT model and then reshape back at the end.
        all_encoder_layers, _ = self.bert_model(input_ids=util.combine_initial_dims(input_ids),
                                                token_type_ids=util.combine_initial_dims(token_type_ids),
                                                attention_mask=util.combine_initial_dims(input_mask))
        all_encoder_layers = torch.stack(all_encoder_layers)

        if needs_split:
            # First, unpack the output embeddings into one long sequence again
            unpacked_embeddings = torch.split(all_encoder_layers, batch_size, dim=1)
            unpacked_embeddings = torch.cat(unpacked_embeddings, dim=2)

            # Next, select indices of the sequence such that it will result in embeddings representing the original
            # sentence. To capture maximal context, the indices will be the middle part of each embedded window
            # sub-sequence (plus any leftover start and final edge windows), e.g.,
            #  0     1 2    3  4   5    6    7     8     9   10   11   12    13 14  15
            # "[CLS] I went to the very fine [SEP] [CLS] the very fine store to eat [SEP]"
            # with max_pieces = 8 should produce max context indices [2, 3, 4, 10, 11, 12] with additional start
            # and final windows with indices [0, 1] and [14, 15] respectively.

            # Find the stride as half the max pieces, ignoring the special start and end tokens
            # Calculate an offset to extract the centermost embeddings of each window
            stride = (self.max_pieces - self.num_start_tokens - self.num_end_tokens) // 2
            stride_offset = stride // 2 + self.num_start_tokens

            first_window = list(range(stride_offset))

            max_context_windows = [i for i in range(full_seq_len)
                                   if stride_offset - 1 < i % self.max_pieces < stride_offset + stride]

            # Lookback what's left, unless it's the whole self.max_pieces window
            if full_seq_len % self.max_pieces == 0:
                lookback = self.max_pieces
            else:
                lookback = full_seq_len % self.max_pieces

            final_window_start = full_seq_len - lookback + stride_offset + stride
            final_window = list(range(final_window_start, full_seq_len))

            select_indices = first_window + max_context_windows + final_window

            initial_dims.append(len(select_indices))

            recombined_embeddings = unpacked_embeddings[:, :, select_indices]
        else:
            recombined_embeddings = all_encoder_layers

        # Recombine the outputs of all layers
        # (layers, batch_size * d1 * ... * dn, sequence_length, embedding_dim)
        # recombined = torch.cat(combined, dim=2)
        input_mask = (recombined_embeddings != 0).long()

        if self._scalar_mix is not None:
            mix = self._scalar_mix(recombined_embeddings, input_mask)
        else:
            mix = recombined_embeddings[-1]

        # At this point, mix is (batch_size * d1 * ... * dn, sequence_length, embedding_dim)

        if offsets is None:
            # Resize to (batch_size, d1, ..., dn, sequence_length, embedding_dim)
            dims = initial_dims if needs_split else input_ids.size()
            return util.uncombine_initial_dims(mix, dims)
        else:
            # offsets is (batch_size, d1, ..., dn, orig_sequence_length)
            offsets2d = util.combine_initial_dims(offsets)
            # now offsets is (batch_size * d1 * ... * dn, orig_sequence_length)
            range_vector = util.get_range_vector(offsets2d.size(0),
                                                 device=util.get_device_of(mix)).unsqueeze(1)
            # selected embeddings is also (batch_size * d1 * ... * dn, orig_sequence_length)
            selected_embeddings = mix[range_vector, offsets2d]

            return util.uncombine_initial_dims(selected_embeddings, offsets.size())

@TokenEmbedder.register("bert-pretrained")
class PretrainedBertEmbedder(BertEmbedder):
    # pylint: disable=line-too-long
    """
    Parameters
    ----------
    pretrained_model: ``str``
        Either the name of the pretrained model to use (e.g. 'bert-base-uncased'),
        or the path to the .tar.gz file with the model weights.

        If the name is a key in the list of pretrained models at
        https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/pytorch_pretrained_bert/modeling.py#L41
        the corresponding path will be used; otherwise it will be interpreted as a path or URL.
    requires_grad : ``bool``, optional (default = False)
        If True, compute gradient of BERT parameters for fine tuning.
    top_layer_only: ``bool``, optional (default = ``False``)
        If ``True``, then only return the top layer instead of apply the scalar mix.
    """
    def __init__(self, pretrained_model: str, requires_grad: bool = False, top_layer_only: bool = False) -> None:
        model = PretrainedBertModel.load(pretrained_model)

        for param in model.parameters():
            param.requires_grad = requires_grad

        super().__init__(bert_model=model, top_layer_only=top_layer_only)
<commit_msg>code for mixed bert embedding layers (#3199)<commit_after>"""
A ``TokenEmbedder`` which uses one of the BERT models
(https://github.com/google-research/bert)
to produce embeddings.

At its core it uses Hugging Face's PyTorch implementation
(https://github.com/huggingface/pytorch-pretrained-BERT),
so thanks to them!
"""
from typing import Dict, List
import logging

import torch
import torch.nn.functional as F

from pytorch_pretrained_bert.modeling import BertModel

from allennlp.modules.scalar_mix import ScalarMix
from allennlp.modules.token_embedders.token_embedder import TokenEmbedder
from allennlp.nn import util

logger = logging.getLogger(__name__)

class PretrainedBertModel:
    """
    In some instances you may want to load the same BERT model twice
    (e.g. to use as a token embedder and also as a pooling layer).
    This factory provides a cache so that you don't actually have to load the model twice.
    """
    _cache: Dict[str, BertModel] = {}

    @classmethod
    def load(cls, model_name: str, cache_model: bool = True) -> BertModel:
        if model_name in cls._cache:
            return PretrainedBertModel._cache[model_name]

        model = BertModel.from_pretrained(model_name)
        if cache_model:
            cls._cache[model_name] = model

        return model

class BertEmbedder(TokenEmbedder):
    """
    A ``TokenEmbedder`` that produces BERT embeddings for your tokens.
    Should be paired with a ``BertIndexer``, which produces wordpiece ids.

    Most likely you probably want to use ``PretrainedBertEmbedder``
    for one of the named pretrained models, not this base class.

    Parameters
    ----------
    bert_model: ``BertModel``
        The BERT model being wrapped.
    top_layer_only: ``bool``, optional (default = ``False``)
        If ``True``, then only return the top layer instead of apply the scalar mix.
    max_pieces : int, optional (default: 512)
        The BERT embedder uses positional embeddings and so has a corresponding
        maximum length for its input ids. Assuming the inputs are windowed
        and padded appropriately by this length, the embedder will split them into a
        large batch, feed them into BERT, and recombine the output as if it was a
        longer sequence.
    num_start_tokens : int, optional (default: 1)
        The number of starting special tokens input to BERT (usually 1, i.e., [CLS])
    num_end_tokens : int, optional (default: 1)
        The number of ending tokens input to BERT (usually 1, i.e., [SEP])
    scalar_mix_parameters: ``List[float]``, optional, (default = None)
        If not ``None``, use these scalar mix parameters to weight the representations
        produced by different layers. These mixing weights are not updated during
        training.
    """
    def __init__(self,
                 bert_model: BertModel,
                 top_layer_only: bool = False,
                 max_pieces: int = 512,
                 num_start_tokens: int = 1,
                 num_end_tokens: int = 1,
                 scalar_mix_parameters: List[float] = None) -> None:
        super().__init__()
        self.bert_model = bert_model
        self.output_dim = bert_model.config.hidden_size
        self.max_pieces = max_pieces
        self.num_start_tokens = num_start_tokens
        self.num_end_tokens = num_end_tokens

        if not top_layer_only:
            self._scalar_mix = ScalarMix(bert_model.config.num_hidden_layers,
                                         do_layer_norm=False,
                                         initial_scalar_parameters=scalar_mix_parameters,
                                         trainable=scalar_mix_parameters is None)
        else:
            self._scalar_mix = None

    def get_output_dim(self) -> int:
        return self.output_dim

    def forward(self,
                input_ids: torch.LongTensor,
                offsets: torch.LongTensor = None,
                token_type_ids: torch.LongTensor = None) -> torch.Tensor:
        """
        Parameters
        ----------
        input_ids : ``torch.LongTensor``
            The (batch_size, ..., max_sequence_length) tensor of wordpiece ids.
        offsets : ``torch.LongTensor``, optional
            The BERT embeddings are one per wordpiece. However it's possible/likely
            you might want one per original token. In that case, ``offsets``
            represents the indices of the desired wordpiece for each original token.
            Depending on how your token indexer is configured, this could be the
            position of the last wordpiece for each token, or it could be the position
            of the first wordpiece for each token.

            For example, if you had the sentence "Definitely not", and if the corresponding
            wordpieces were ["Def", "##in", "##ite", "##ly", "not"], then the input_ids
            would be 5 wordpiece ids, and the "last wordpiece" offsets would be [3, 4].
            If offsets are provided, the returned tensor will contain only the wordpiece
            embeddings at those positions, and (in particular) will contain one embedding
            per token. If offsets are not provided, the entire tensor of wordpiece embeddings
            will be returned.
        token_type_ids : ``torch.LongTensor``, optional
            If an input consists of two sentences (as in the BERT paper),
            tokens from the first sentence should have type 0 and tokens from
            the second sentence should have type 1.  If you don't provide this
            (the default BertIndexer doesn't) then it's assumed to be all 0s.
        """
        # pylint: disable=arguments-differ
        batch_size, full_seq_len = input_ids.size(0), input_ids.size(-1)
        initial_dims = list(input_ids.shape[:-1])

        # The embedder may receive an input tensor that has a sequence length longer than can
        # be fit. In that case, we should expect the wordpiece indexer to create padded windows
        # of length `self.max_pieces` for us, and have them concatenated into one long sequence.
        # E.g., "[CLS] I went to the [SEP] [CLS] to the store to [SEP] ..."
        # We can then split the sequence into sub-sequences of that length, and concatenate them
        # along the batch dimension so we effectively have one huge batch of partial sentences.
        # This can then be fed into BERT without any sentence length issues. Keep in mind
        # that the memory consumption can dramatically increase for large batches with extremely
        # long sentences.
        needs_split = full_seq_len > self.max_pieces
        last_window_size = 0
        if needs_split:
            # Split the flattened list by the window size, `max_pieces`
            split_input_ids = list(input_ids.split(self.max_pieces, dim=-1))

            # We want all sequences to be the same length, so pad the last sequence
            last_window_size = split_input_ids[-1].size(-1)
            padding_amount = self.max_pieces - last_window_size
            split_input_ids[-1] = F.pad(split_input_ids[-1], pad=[0, padding_amount], value=0)

            # Now combine the sequences along the batch dimension
            input_ids = torch.cat(split_input_ids, dim=0)

            if token_type_ids is not None:
                # Same for token_type_ids
                split_token_type_ids = list(token_type_ids.split(self.max_pieces, dim=-1))

                last_window_size = split_token_type_ids[-1].size(-1)
                padding_amount = self.max_pieces - last_window_size
                split_token_type_ids[-1] = F.pad(split_token_type_ids[-1], pad=[0, padding_amount], value=0)

                token_type_ids = torch.cat(split_token_type_ids, dim=0)

        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)

        input_mask = (input_ids != 0).long()

        # input_ids may have extra dimensions, so we reshape down to 2-d
        # before calling the BERT model and then reshape back at the end.
        all_encoder_layers, _ = self.bert_model(input_ids=util.combine_initial_dims(input_ids),
                                                token_type_ids=util.combine_initial_dims(token_type_ids),
                                                attention_mask=util.combine_initial_dims(input_mask))
        all_encoder_layers = torch.stack(all_encoder_layers)

        if needs_split:
            # First, unpack the output embeddings into one long sequence again
            unpacked_embeddings = torch.split(all_encoder_layers, batch_size, dim=1)
            unpacked_embeddings = torch.cat(unpacked_embeddings, dim=2)

            # Next, select indices of the sequence such that it will result in embeddings representing the original
            # sentence. To capture maximal context, the indices will be the middle part of each embedded window
            # sub-sequence (plus any leftover start and final edge windows), e.g.,
            #  0     1 2    3  4   5    6    7     8     9   10   11   12    13 14  15
            # "[CLS] I went to the very fine [SEP] [CLS] the very fine store to eat [SEP]"
            # with max_pieces = 8 should produce max context indices [2, 3, 4, 10, 11, 12] with additional start
            # and final windows with indices [0, 1] and [14, 15] respectively.

            # Find the stride as half the max pieces, ignoring the special start and end tokens
            # Calculate an offset to extract the centermost embeddings of each window
            stride = (self.max_pieces - self.num_start_tokens - self.num_end_tokens) // 2
            stride_offset = stride // 2 + self.num_start_tokens

            first_window = list(range(stride_offset))

            max_context_windows = [i for i in range(full_seq_len)
                                   if stride_offset - 1 < i % self.max_pieces < stride_offset + stride]

            # Lookback what's left, unless it's the whole self.max_pieces window
            if full_seq_len % self.max_pieces == 0:
                lookback = self.max_pieces
            else:
                lookback = full_seq_len % self.max_pieces

            final_window_start = full_seq_len - lookback + stride_offset + stride
            final_window = list(range(final_window_start, full_seq_len))

            select_indices = first_window + max_context_windows + final_window

            initial_dims.append(len(select_indices))

            recombined_embeddings = unpacked_embeddings[:, :, select_indices]
        else:
            recombined_embeddings = all_encoder_layers

        # Recombine the outputs of all layers
        # (layers, batch_size * d1 * ... * dn, sequence_length, embedding_dim)
        # recombined = torch.cat(combined, dim=2)
        input_mask = (recombined_embeddings != 0).long()

        if self._scalar_mix is not None:
            mix = self._scalar_mix(recombined_embeddings, input_mask)
        else:
            mix = recombined_embeddings[-1]

        # At this point, mix is (batch_size * d1 * ... * dn, sequence_length, embedding_dim)

        if offsets is None:
            # Resize to (batch_size, d1, ..., dn, sequence_length, embedding_dim)
            dims = initial_dims if needs_split else input_ids.size()
            return util.uncombine_initial_dims(mix, dims)
        else:
            # offsets is (batch_size, d1, ..., dn, orig_sequence_length)
            offsets2d = util.combine_initial_dims(offsets)
            # now offsets is (batch_size * d1 * ... * dn, orig_sequence_length)
            range_vector = util.get_range_vector(offsets2d.size(0),
                                                 device=util.get_device_of(mix)).unsqueeze(1)
            # selected embeddings is also (batch_size * d1 * ... * dn, orig_sequence_length)
            selected_embeddings = mix[range_vector, offsets2d]

            return util.uncombine_initial_dims(selected_embeddings, offsets.size())

@TokenEmbedder.register("bert-pretrained")
class PretrainedBertEmbedder(BertEmbedder):
    # pylint: disable=line-too-long
    """
    Parameters
    ----------
    pretrained_model: ``str``
        Either the name of the pretrained model to use (e.g. 'bert-base-uncased'),
        or the path to the .tar.gz file with the model weights.

        If the name is a key in the list of pretrained models at
        https://github.com/huggingface/pytorch-pretrained-BERT/blob/master/pytorch_pretrained_bert/modeling.py#L41
        the corresponding path will be used; otherwise it will be interpreted as a path or URL.
    requires_grad : ``bool``, optional (default = False)
        If True, compute gradient of BERT parameters for fine tuning.
    top_layer_only: ``bool``, optional (default = ``False``)
        If ``True``, then only return the top layer instead of apply the scalar mix.
    scalar_mix_parameters: ``List[float]``, optional, (default = None)
        If not ``None``, use these scalar mix parameters to weight the representations
        produced by different layers. These mixing weights are not updated during
        training.
    """
    def __init__(self, pretrained_model: str, requires_grad: bool = False, top_layer_only: bool = False,
                 scalar_mix_parameters: List[float] = None) -> None:
        model = PretrainedBertModel.load(pretrained_model)

        for param in model.parameters():
            param.requires_grad = requires_grad

        super().__init__(bert_model=model, top_layer_only=top_layer_only, scalar_mix_parameters=scalar_mix_parameters)
<|endoftext|>


<commit_before>from netCDF4 import Dataset
from os.path import isdir, exists
import numpy as np
from glob import glob

class WRFModelGrid(object):
    """
    Load data from raw WRF model output files.
    """
    def __init__(self, forecast_date, variable, domain, path):
        self.forecast_date = forecast_date
        self.variable = variable
        self.domain = domain
        self.path = path
        self.wrf_filename = "wrfout_{0}_{1}".format(self.domain, self.forecast_date.strftime("%Y-%m-%d_%H:%M:%S"))
        self.patch_files = False
        if exists(self.path + self.wrf_filename):
            if isdir(self.path + self.wrf_filename):
                self.patch_files = True
        else:
            raise IOError(self.path + self.wrf_filename + " not found.")
        return

    def get_global_attributes(self):
        attributes = {}
        if self.patch_files:
            patch_files = sorted(glob(self.path + "{0}/{0}_".format(self.wrf_filename) + "[0-9]" * 5))
            wrf_data = Dataset(patch_files[0])
        else:
            wrf_data = Dataset(self.wrf_filename)
        for attr in wrf_data.ncattrs():
            attributes[attr] = getattr(wrf_data, attr)
        wrf_data.close()
        return attributes

    def load_time_var(self, time_var="XTIME"):
        var_attrs = dict()
        var_val = 0
        if self.patch_files:
            patch_file_list = sorted(glob(self.path + "{0}/{0}_".format(self.wrf_filename) + "[0-9]" * 5))
            patch_zero = Dataset(patch_file_list[0])
            var_list = patch_zero.variables.keys()
            if time_var in var_list:
                for attr in patch_zero.variables[time_var].ncattrs():
                    var_attrs[attr] = getattr(patch_zero.variables[time_var], attr)
                var_val = patch_zero.variables[time_var][:]
            patch_zero.close()
        else:
            wrf_data = Dataset(self.path + self.wrf_filename)
            if time_var in var_list:
                for attr in wrf_data.variables[time_var].ncattrs():
                    var_attrs[attr] = getattr(wrf_data.variables[time_var], attr)
                var_val = wrf_data.variables[time_var][:]
            wrf_data.close()
        return var_val, var_attrs

    def load_full_grid(self):
        var_data = None
        if self.patch_files:
            patch_file_list = sorted(glob(self.path + "{0}/{0}_".format(self.wrf_filename) + "[0-9]" * 5))
            patch_zero = Dataset(patch_file_list[0])
            var_list = patch_zero.variables.keys()
            var_attrs = dict()
            if self.variable in var_list:
                dimension_names = patch_zero.variables[self.variable].dimensions
                is_stag = np.array(["stag" in x for x in patch_zero.variables[self.variable].dimensions])
                for attr in patch_zero.variables[self.variable].ncattrs():
                    if attr == "coordinates":
                        var_attrs[attr] = getattr(patch_zero.variables[self.variable], attr)
                        if "_U" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_U", "")
                        elif "_V" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_V", "")
                    else:
                        var_attrs[attr] = getattr(patch_zero.variables[self.variable], attr)
                if len(dimension_names) == 4:
                    grid_dim = (1,
                                getattr(patch_zero, "BOTTOM-TOP_GRID_DIMENSION") - 1,
                                getattr(patch_zero, "SOUTH-NORTH_GRID_DIMENSION") - 1,
                                getattr(patch_zero, "WEST-EAST_GRID_DIMENSION") - 1)
                else:
                    grid_dim = (1,
                                getattr(patch_zero, "SOUTH-NORTH_GRID_DIMENSION") - 1,
                                getattr(patch_zero, "WEST-EAST_GRID_DIMENSION") - 1)
                patch_zero.close()
                if np.any(is_stag):
                    grid_dim_arr = np.array(grid_dim)
                    grid_dim_arr[is_stag] += 1
                    var_data = np.zeros(grid_dim_arr, dtype=np.float32)
                    stag_dim = np.where(is_stag)[0][0]
                else:
                    var_data = np.zeros(grid_dim, dtype=np.float32)
                    stag_dim = None
                for p, patch_file in enumerate(patch_file_list):
                    wrf_patch_file = Dataset(patch_file)
                    sub_data = wrf_patch_file.variables[self.variable][:]
                    if stag_dim is not None:
                        if len(sub_data.shape) - stag_dim == 1:
                            we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_STAG") - 1,
                                             getattr(wrf_patch_file, "WEST-EAST_PATCH_END_STAG"))
                            sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_UNSTAG"))
                        elif len(sub_data.shape) - stag_dim == 2:
                            we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "WEST-EAST_PATCH_END_UNSTAG"))
                            sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_STAG") - 1,
                                             getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_STAG"))
                        else:
                            we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "WEST-EAST_PATCH_END_UNSTAG"))
                            sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_UNSTAG"))
                    else:
                        we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_UNSTAG") - 1,
                                         getattr(wrf_patch_file, "WEST-EAST_PATCH_END_UNSTAG"))
                        sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_UNSTAG") - 1,
                                         getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_UNSTAG"))
                    if len(var_data.shape) == 3:
                        var_data[:, sn_slice, we_slice] = sub_data
                    else:
                        var_data[:, :, sn_slice, we_slice] = sub_data
                    wrf_patch_file.close()
                if stag_dim is not None:
                    if stag_dim == 1:
                        var_data = 0.5 * (var_data[:, :-1] + var_data[:, 1:])
                    elif stag_dim == 2:
                        var_data = 0.5 * (var_data[:, :, :-1] + var_data[:, :, 1:])
                    else:
                        var_data = 0.5 * (var_data[:, :, :, :-1] + var_data[:, :, :, 1:])
            else:
                patch_zero.close()

        else:
            wrf_data = Dataset(self.path + self.wrf_filename)
            is_stag = np.array(["stag" in x for x in wrf_data.variables[self.variable].dimensions])
            if self.variable in wrf_data.variables.keys():
                for attr in wrf_data.variables[self.variable].ncattrs():
                    if attr == "coordinates":
                        var_attrs[attr] = getattr(wrf_data.variables[self.variable], attr)
                        if "_U" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_U", "")
                        elif "_V" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_V", "")
                    else:
                        var_attrs[attr] = getattr(wrf_data.variables[self.variable], attr)

                var_data = wrf_data.variables[self.variable][:]
                if np.any(is_stag):
                    stag_dim = np.where(is_stag)[0][0]
                    if stag_dim == 1:
                        var_data = 0.5 * (var_data[:, :-1] + var_data[:, 1:])
                    elif stag_dim == 2:
                        var_data = 0.5 * (var_data[:, :, :-1] + var_data[:, :, 1:])
                    else:
                        var_data = 0.5 * (var_data[:, :, :, :-1] + var_data[:, :, :, 1:])
            wrf_data.close()

        return var_data, var_attrs
<commit_msg>Fixed minor errors in WRFModelGrid<commit_after>from netCDF4 import Dataset
from os.path import isdir, exists
import numpy as np
from glob import glob

class WRFModelGrid(object):
    """
    Load data from raw WRF model output files.
    """
    def __init__(self, forecast_date, variable, domain, path):
        self.forecast_date = forecast_date
        self.variable = variable
        self.domain = domain
        self.path = path
        self.wrf_filename = "wrfout_{0}_{1}".format(self.domain, self.forecast_date.strftime("%Y-%m-%d_%H:%M:%S"))
        self.patch_files = False
        if exists(self.path + self.wrf_filename):
            if isdir(self.path + self.wrf_filename):
                self.patch_files = True
        else:
            raise IOError(self.path + self.wrf_filename + " not found.")
        return

    def get_global_attributes(self):
        attributes = {}
        if self.patch_files:
            patch_files = sorted(glob(self.path + "{0}/{0}_".format(self.wrf_filename) + "[0-9]" * 5))
            wrf_data = Dataset(patch_files[0])
        else:
            wrf_data = Dataset(self.wrf_filename)
        for attr in wrf_data.ncattrs():
            attributes[attr] = getattr(wrf_data, attr)
        wrf_data.close()
        return attributes

    def load_time_var(self, time_var="XTIME"):
        var_attrs = dict()
        var_val = 0
        if self.patch_files:
            patch_file_list = sorted(glob(self.path + "{0}/{0}_".format(self.wrf_filename) + "[0-9]" * 5))
            patch_zero = Dataset(patch_file_list[0])
            var_list = patch_zero.variables.keys()
            if time_var in var_list:
                for attr in patch_zero.variables[time_var].ncattrs():
                    var_attrs[attr] = getattr(patch_zero.variables[time_var], attr)
                var_val = patch_zero.variables[time_var][:]
            patch_zero.close()
        else:
            wrf_data = Dataset(self.path + self.wrf_filename)
            var_list = wrf_data.variables.keys()
            if time_var in var_list:
                for attr in wrf_data.variables[time_var].ncattrs():
                    var_attrs[attr] = getattr(wrf_data.variables[time_var], attr)
                var_val = wrf_data.variables[time_var][:]
            wrf_data.close()
        return var_val, var_attrs

    def load_full_grid(self):
        var_data = None
        var_attrs = dict()
        if self.patch_files:
            patch_file_list = sorted(glob(self.path + "{0}/{0}_".format(self.wrf_filename) + "[0-9]" * 5))
            patch_zero = Dataset(patch_file_list[0])
            var_list = patch_zero.variables.keys()
            if self.variable in var_list:
                dimension_names = patch_zero.variables[self.variable].dimensions
                is_stag = np.array(["stag" in x for x in patch_zero.variables[self.variable].dimensions])
                for attr in patch_zero.variables[self.variable].ncattrs():
                    if attr == "coordinates":
                        var_attrs[attr] = getattr(patch_zero.variables[self.variable], attr)
                        if "_U" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_U", "")
                        elif "_V" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_V", "")
                    else:
                        var_attrs[attr] = getattr(patch_zero.variables[self.variable], attr)
                if len(dimension_names) == 4:
                    grid_dim = (1,
                                getattr(patch_zero, "BOTTOM-TOP_GRID_DIMENSION") - 1,
                                getattr(patch_zero, "SOUTH-NORTH_GRID_DIMENSION") - 1,
                                getattr(patch_zero, "WEST-EAST_GRID_DIMENSION") - 1)
                else:
                    grid_dim = (1,
                                getattr(patch_zero, "SOUTH-NORTH_GRID_DIMENSION") - 1,
                                getattr(patch_zero, "WEST-EAST_GRID_DIMENSION") - 1)
                patch_zero.close()
                if np.any(is_stag):
                    grid_dim_arr = np.array(grid_dim)
                    grid_dim_arr[is_stag] += 1
                    var_data = np.zeros(grid_dim_arr, dtype=np.float32)
                    stag_dim = np.where(is_stag)[0][0]
                else:
                    var_data = np.zeros(grid_dim, dtype=np.float32)
                    stag_dim = None
                for p, patch_file in enumerate(patch_file_list):
                    wrf_patch_file = Dataset(patch_file)
                    sub_data = wrf_patch_file.variables[self.variable][:]
                    if stag_dim is not None:
                        if len(sub_data.shape) - stag_dim == 1:
                            we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_STAG") - 1,
                                             getattr(wrf_patch_file, "WEST-EAST_PATCH_END_STAG"))
                            sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_UNSTAG"))
                        elif len(sub_data.shape) - stag_dim == 2:
                            we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "WEST-EAST_PATCH_END_UNSTAG"))
                            sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_STAG") - 1,
                                             getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_STAG"))
                        else:
                            we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "WEST-EAST_PATCH_END_UNSTAG"))
                            sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_UNSTAG") - 1,
                                             getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_UNSTAG"))
                    else:
                        we_slice = slice(getattr(wrf_patch_file, "WEST-EAST_PATCH_START_UNSTAG") - 1,
                                         getattr(wrf_patch_file, "WEST-EAST_PATCH_END_UNSTAG"))
                        sn_slice = slice(getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_START_UNSTAG") - 1,
                                         getattr(wrf_patch_file, "SOUTH-NORTH_PATCH_END_UNSTAG"))
                    if len(var_data.shape) == 3:
                        var_data[:, sn_slice, we_slice] = sub_data
                    else:
                        var_data[:, :, sn_slice, we_slice] = sub_data
                    wrf_patch_file.close()
                if stag_dim is not None:
                    if stag_dim == 1:
                        var_data = 0.5 * (var_data[:, :-1] + var_data[:, 1:])
                    elif stag_dim == 2:
                        var_data = 0.5 * (var_data[:, :, :-1] + var_data[:, :, 1:])
                    else:
                        var_data = 0.5 * (var_data[:, :, :, :-1] + var_data[:, :, :, 1:])
            else:
                patch_zero.close()

        else:
            wrf_data = Dataset(self.path + self.wrf_filename)
            is_stag = np.array(["stag" in x for x in wrf_data.variables[self.variable].dimensions])
            if self.variable in wrf_data.variables.keys():
                for attr in wrf_data.variables[self.variable].ncattrs():
                    if attr == "coordinates":
                        var_attrs[attr] = getattr(wrf_data.variables[self.variable], attr)
                        if "_U" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_U", "")
                        elif "_V" in var_attrs[attr]:
                            var_attrs[attr] = var_attrs[attr].replace("_V", "")
                    else:
                        var_attrs[attr] = getattr(wrf_data.variables[self.variable], attr)

                var_data = wrf_data.variables[self.variable][:]
                if np.any(is_stag):
                    stag_dim = np.where(is_stag)[0][0]
                    if stag_dim == 1:
                        var_data = 0.5 * (var_data[:, :-1] + var_data[:, 1:])
                    elif stag_dim == 2:
                        var_data = 0.5 * (var_data[:, :, :-1] + var_data[:, :, 1:])
                    else:
                        var_data = 0.5 * (var_data[:, :, :, :-1] + var_data[:, :, :, 1:])
            wrf_data.close()

        return var_data, var_attrs
<|endoftext|>


<commit_before>##########################################################################
# Copyright 2013-2022 Aerospike, Inc.
#
# 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.
##########################################################################
'''
records.py defines objects for use with aerospike client batch APIs. Currently batch_write, batch_operate,
batch_remove, and batch_apply make use of objects in this file. Typically BatchReacords and underlying
BatchRecord objects are used as input and output for the aformentioned client methods.

    .. note:: APIs that utitlize these objects require server >= 6.0.0.

Example::

    import aerospike
    from aerospike import exception as ex
    from aerospike_helpers.batch import records as br
    import aerospike_helpers.expressions as exp
    from aerospike_helpers.operations import operations as op
    import sys

    # Configure the client.
    config = {"hosts": [("127.0.0.1", 3000)]}

    # Create a client and connect it to the cluster.
    try:
        client = aerospike.client(config).connect()
    except ex.ClientError as e:
        print("Error: {0} [{1}]".format(e.msg, e.code))
        sys.exit(1)

    # setup records
    namespace = "test"
    set = "demo"

    keys = [(namespace, set, i) for i in range(1000)]
    records = [{"id": i, "balance": i * 10} for i in range(1000)]
    for key, rec in zip(keys, records):
        client.put(key, rec)

    print("===== BATCH_OPERATE EXAMPLE =====")
    # Batch add 10 to balance and read it if it's over
    # 1000 NOTE: batch_operate ops must include a write op
    # get_batch_ops or get_many can be used for all read ops.
    expr = exp.GT(exp.IntBin("balance"), 1000).compile()
    ops = [
        op.increment("balance", 10),
        op.read("balance")
    ]
    policy_batch = {"expressions": expr}
    res = client.batch_operate(keys, ops, policy_batch)

    # res is an instance of BatchRecords
    # the field, batch_records, contains a BatchRecord instance
    # for each key used by the batch_operate call.
    # the field, results, is 0 if all batch subtransactions completed succesfully
    # or the only failures are FILTERED_OUT or RECORD_NOT_FOUND.
    # Otherwise its value corresponds to an as_status and signifies that
    # one or more of the batch subtransactions failed. Each BatchRecord instance
    # also has a results field that signifies the status of that batch subtransaction.

    if res.result == 0:
        # BatchRecord 100 should have a result code of 27 meaning it was filtered out by an expression.
        print("BatchRecord 100 result: {result}".format(result=res.batch_records[100].result))
        # BatchRecord 100,record should be None.
        print("BatchRecord 100 record: {record}".format(record=res.batch_records[100].record))
        # BatchRecord 101 should have a result code of 0 meaning it succeeded.
        print("BatchRecord 101 result: {result}".format(result=res.batch_records[101].result))
        # BatchRecord 101, record should be populated.
        print("BatchRecord 101 record: {record}".format(record=res.batch_records[101].record))
    else:
        # Some batch sub transaction failed.
        print("res result: {result}".format(result=res.result))

    print("===== BATCH_WRITE EXAMPLE =====")
    # Apply different operations to different keys
    # using batch_write.
    batch_writes = br.BatchRecords(
        [
            br.Remove(
                key=(namespace, set, 1),
                policy={}
            ),
            br.Write(
                key=(namespace, set, 100),
                ops=[
                    op.write("id", 100),
                    op.write("balance", 100),
                    op.read("id"),
                    op.read("balance"),
                ],
                policy={"expressions": exp.GT(exp.IntBin("balance"), 2000).compile()}
            ),
            br.Read(
                key=(namespace, set, 333),
                ops=[
                    op.read("id")
                ],
                policy=None
            ),
        ]
    )

    # batch_write modifies its BatchRecords argument.
    # Results for each requested key will be set in
    # their coresponding BatchRecord result,
    # record, and in_doubt fields.
    client.batch_write(batch_writes)
    print("batch_writes result: {result}".format(result=batch_writes.result))

    # should have bins {'id': 333}.
    print("batch_writes batch Write record: {result}".format(result=batch_writes.batch_records[2].record))

    print("===== BATCH_APPLY EXAMPLE =====")
    # Apply a user defined function (UDF) to a batch
    # of records using batch_apply.
    module = "test_record_udf"
    path_to_module = "/path/to/test_record_udf.lua"
    function = "bin_udf_operation_integer"
    args = ["balance", 10, 5]

    client.udf_put(path_to_module)

    # This should add 15 to each balance bin.
    res = client.batch_apply(keys, module, function, args)
    # NOTE res.result should be -16 (one or more batch sub transactions failed)
    # because the UDF failed on record 1 which was previously removed.
    print("res result: {result}".format(result=res.result))

    res_rec = res.batch_records[90].record
    bins = res_rec[2]
    # Should be 915.
    print("res BatchRecord 90 bins: {result}".format(result=bins))

    print("===== BATCH_REMOVE EXAMPLE =====")
    # Delete the records using batch_remove.
    res = client.batch_remove(keys)
    # Should be 0 signifying success.
    print("BatchRecords result: {result}".format(result=res.result))

'''

import typing as ty

TypeOps = ty.List[ty.Dict]
TypeBatchPolicyWrite = ty.Union[ty.Dict, None]
TypeBatchPolicyRemove = ty.Union[ty.Dict, None]
TypeBatchPolicyApply = ty.Union[ty.Dict, None]
TypeBatchPolicyRead = ty.Union[ty.Dict, None]
TypeRecord = ty.Union[ty.Tuple, None]
TypeUDFArgs = ty.List[ty.Any]

class _Types():
    READ = 0
    WRITE = 1
    APPLY = 2
    REMOVE = 3

#### BatchRecord ####
class BatchRecord:
    """ BatchRecord provides the base fields for BatchRecord objects.

        BatchRecord should usually be read from as a result and not created by the user. Its subclasses can be used as input to batch_write.
        Client methods :meth:`~Client.batch_apply`, :meth:`~Client.batch_operate`, :meth:`~Client.batch_remove`
        with batch_records field as a list of these BatchRecord objects containing the batch request results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (TypeRecord): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (int): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
    """
    def __init__(self, key: tuple) -> None:
        self.key = key
        self.record = None
        self.result = 0
        self.in_doubt = False

class Write(BatchRecord):
    """ Write is used for executing Batch write operations with batch_write and retrieving batch write results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (:obj:`tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (int): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (TypeOps): A list of aerospike operation dictionaries to perform on the record at key.
            policy (TypeBatchPolicyWrite, optional): An optional dictionary of batch write policy flags.
    """

    def __init__(self, key: tuple, ops: 'TypeOps', policy: 'TypeBatchPolicyWrite' = None) -> None:
        """
            Example::

                # Create a batch Write to increment bin "a" by 10 and read the result from the record.
                import aerospike_helpers.operations as op

                bin_name = "a"

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                ops = [
                    op.increment(bin_name, 10),
                    op.read(bin_name)
                ]

                bw = Write(key, ops)
        """
        super().__init__(key)
        self.ops = ops
        self._type = _Types.WRITE
        self._has_write = True
        self.policy = policy

class Read(BatchRecord):
    """ Read is used for executing Batch read operations with batch_write and retrieving results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (:obj:`tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (int): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (TypeOps): list of aerospike operation dictionaries to perform on the record at key.
            read_all_bins (bool, optional): An optional bool, if True, read all bins in the record.
            policy (TypeBatchPolicyRead, optional): An optional dictionary of batch read policy flags.
    """

    def __init__(self, key: tuple, ops: ty.Union[TypeOps, None], read_all_bins: bool = False, policy: 'TypeBatchPolicyRead' = None) -> None:
        """
            Example::

                # Create a batch Read to read bin "a" from the record.
                import aerospike_helpers.operations as op

                bin_name = "a"

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                ops = [
                    op.read(bin_name)
                ]

                br = Write(key, ops)
        """
        super().__init__(key)
        self.ops = ops
        self.read_all_bins = read_all_bins
        self._type = _Types.READ
        self._has_write = False
        self.policy = policy

class Apply(BatchRecord):
    """ BatchApply is used for executing Batch UDF (user defined function) apply operations with batch_write and retrieving results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            module (str): Name of the lua module previously registerd with the server.
            function (str): Name of the UDF to invoke.
            args (TypeUDFArgs): List of arguments to pass to the UDF.
            record (:obj:`tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (int): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (TypeOps): A list of aerospike operation dictionaries to perform on the record at key.
            policy (TypeBatchPolicyApply, optional): An optional dictionary of batch apply policy flags.
    """

    def __init__(self, key: tuple, module: str, function: str, args: 'TypeUDFArgs', policy: 'TypeBatchPolicyApply' = None) -> None:
        """
            Example::

                # Create a batch Apply to apply UDF "test_func" to bin "a" from the record.
                # Assume that "test_func" takes a bin name string as an argument.
                # Assume the appropriate UDF module has already been registerd.
                import aerospike_helpers.operations as op

                module = "my_lua"
                function = "test_func"

                bin_name = "a"
                args = [
                    bin_name
                ]

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                ba = Apply(key, module, function, args)
        """
        super().__init__(key)
        self._type = _Types.APPLY
        self._has_write = True
        self.module = module
        self.function = function
        self.args = args
        self.policy = policy

class Remove(BatchRecord):
    """ Remove is used for executing Batch remove operations with batch_write and retrieving results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (:obj:`tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (int): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (TypeOps): A list of aerospike operation dictionaries to perform on the record at key.
            policy (TypeBatchPolicyRemove, optional): An optional dictionary of batch remove policy flags.
    """

    def __init__(self, key: tuple, policy: 'TypeBatchPolicyRemove' = None) -> None:
        """
            Example::

                # Create a batch Remove to remove the record.
                import aerospike_helpers.operations as op

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                br = Remove(key, ops)
        """
        super().__init__(key)
        self._type = _Types.REMOVE
        self._has_write = True
        self.policy = policy

#### BatchRecords ####
TypeBatchRecordList = ty.List[BatchRecord]

class BatchRecords:
    """ BatchRecords is used as input and output for multiple batch APIs.

        Attributes:
            batch_records (TypeBatchRecordList): A list of BatchRecord subtype objects used to \
            define batch operations and hold results. BatchRecord Types can be Remove, Write, \
            Read, and Apply.
            result (int): The status code of the last batch call that used this BatchRecords.
            0 if all batch subtransactions succeeded (or if the only failures were FILTERED_OUT or RECORD_NOT_FOUND)
            non 0 if an error occured. The most common error being -16 (One or more batch sub transactions failed).
    """

    def __init__(self, batch_records: TypeBatchRecordList = None) -> None:
        """
            Example::

                # Create a BatchRecords to remove a record, write a bin, and read a bin.
                # Assume client is an instantiated and connected aerospike cleint.
                import aerospike_helpers.operations as op

                namespace = "test"
                set = "demo"
                bin_name = "id"
                keys = [
                    (namespace, set, 1),
                    (namespace, set, 2),
                    (namespace, set, 3)
                ]

                brs = BatchRecords(
                    [
                        Remove(
                            key=(namespace, set, 1),
                        ),
                        Write(
                            key=(namespace, set, 100),
                            ops=[
                                op.write(bin_name, 100),
                                op.read(bin_name),
                            ]
                        ),
                        BatchRead(
                            key=(namespace, set, 333),
                            ops=[
                                op.read(bin_name)
                            ]
                        )
                    ]
                )

                # Note this call will mutate brs and set results in it.
                client.batch_write(brs)
        """

        if batch_records is None:
            batch_records = []

        self.batch_records = batch_records
        self.result = 0<commit_msg>Remove example already in client.rst and fix types<commit_after>##########################################################################
# Copyright 2013-2022 Aerospike, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##########################################################################

import typing as ty

TypeOps = ty.List[ty.Dict]
TypeBatchPolicyWrite = ty.Union[ty.Dict, None]
TypeBatchPolicyRemove = ty.Union[ty.Dict, None]
TypeBatchPolicyApply = ty.Union[ty.Dict, None]
TypeBatchPolicyRead = ty.Union[ty.Dict, None]
TypeRecord = ty.Union[ty.Tuple, None]
TypeUDFArgs = ty.List[ty.Any]

class _Types():
    READ = 0
    WRITE = 1
    APPLY = 2
    REMOVE = 3

#### BatchRecord ####
class BatchRecord:
    """ BatchRecord provides the base fields for BatchRecord objects.

        BatchRecord should usually be read from as a result and not created by the user. Its subclasses can be used as input to batch_write.
        Client methods :meth:`~Client.batch_apply`, :meth:`~Client.batch_operate`, :meth:`~Client.batch_remove`
        with batch_records field as a list of these BatchRecord objects containing the batch request results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (:ref:`aerospike_record_tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (bool): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
    """
    def __init__(self, key: tuple) -> None:
        self.key = key
        self.record = None
        self.result = 0
        self.in_doubt = False

class Write(BatchRecord):
    """ Write is used for executing Batch write operations with batch_write and retrieving batch write results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (:obj:`tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (bool): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (:ref:`aerospike_operation_helpers.operations`): A list of aerospike operation dictionaries to perform on the record at key.
            policy (:ref:`aerospike_batch_write_policies`, optional): An optional dictionary of batch write policy flags.
    """

    def __init__(self, key: tuple, ops: 'TypeOps', policy: 'TypeBatchPolicyWrite' = None) -> None:
        """
            Example::

                # Create a batch Write to increment bin "a" by 10 and read the result from the record.
                import aerospike_helpers.operations as op

                bin_name = "a"

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                ops = [
                    op.increment(bin_name, 10),
                    op.read(bin_name)
                ]

                bw = Write(key, ops)
        """
        super().__init__(key)
        self.ops = ops
        self._type = _Types.WRITE
        self._has_write = True
        self.policy = policy

class Read(BatchRecord):
    """ Read is used for executing Batch read operations with batch_write and retrieving results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (:obj:`tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (bool): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (:ref:`aerospike_operation_helpers.operations`): list of aerospike operation dictionaries to perform on the record at key.
            read_all_bins (bool, optional): An optional bool, if True, read all bins in the record.
            policy (:ref:`aerospike_batch_read_policies`, optional): An optional dictionary of batch read policy flags.
    """

    def __init__(self, key: tuple, ops: ty.Union[TypeOps, None], read_all_bins: bool = False, policy: 'TypeBatchPolicyRead' = None) -> None:
        """
            Example::

                # Create a batch Read to read bin "a" from the record.
                import aerospike_helpers.operations as op

                bin_name = "a"

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                ops = [
                    op.read(bin_name)
                ]

                br = Read(key, ops)
        """
        super().__init__(key)
        self.ops = ops
        self.read_all_bins = read_all_bins
        self._type = _Types.READ
        self._has_write = False
        self.policy = policy

class Apply(BatchRecord):
    """ BatchApply is used for executing Batch UDF (user defined function) apply operations with batch_write and retrieving results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            module (str): Name of the lua module previously registerd with the server.
            function (str): Name of the UDF to invoke.
            args (:obj:`list`): List of arguments to pass to the UDF.
            record (:ref:`aerospike_record_tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (bool): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (:ref:`aerospike_operation_helpers.operations`): A list of aerospike operation dictionaries to perform on the record at key.
            policy (:ref:`aerospike_batch_apply_policies`, optional): An optional dictionary of batch apply policy flags.
    """

    def __init__(self, key: tuple, module: str, function: str, args: 'TypeUDFArgs', policy: 'TypeBatchPolicyApply' = None) -> None:
        """
            Example::

                # Create a batch Apply to apply UDF "test_func" to bin "a" from the record.
                # Assume that "test_func" takes a bin name string as an argument.
                # Assume the appropriate UDF module has already been registerd.
                import aerospike_helpers.operations as op

                module = "my_lua"
                function = "test_func"

                bin_name = "a"
                args = [
                    bin_name
                ]

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                ba = Apply(key, module, function, args)
        """
        super().__init__(key)
        self._type = _Types.APPLY
        self._has_write = True
        self.module = module
        self.function = function
        self.args = args
        self.policy = policy

class Remove(BatchRecord):
    """ Remove is used for executing Batch remove operations with batch_write and retrieving results.

        Attributes:
            key (:obj:`tuple`): The aerospike key to operate on.
            record (:ref:`aerospike_record_tuple`): The record corresponding to the requested key.
            result (int): The status code of the operation.
            in_doubt (bool): Is it possible that the write transaction completed even though an error was generated. \
            This may be the case when a client error occurs (like timeout) after the command was sent \
            to the server.
            ops (:ref:`aerospike_operation_helpers.operations`): A list of aerospike operation dictionaries to perform on the record at key.
            policy (:ref:`aerospike_batch_remove_policies`, optional): An optional dictionary of batch remove policy flags.
    """

    def __init__(self, key: tuple, policy: 'TypeBatchPolicyRemove' = None) -> None:
        """
            Example::

                # Create a batch Remove to remove the record.
                import aerospike_helpers.operations as op

                namespace = "test"
                set = "demo"
                user_key = 1
                key = (namespace, set, user_key)

                br = Remove(key, ops)
        """
        super().__init__(key)
        self._type = _Types.REMOVE
        self._has_write = True
        self.policy = policy

#### BatchRecords ####
TypeBatchRecordList = ty.List[BatchRecord]

class BatchRecords:
    """ BatchRecords is used as input and output for multiple batch APIs.

        Attributes:
            batch_records (list): A list of BatchRecord subtype objects used to \
            define batch operations and hold results. BatchRecord Types can be Remove, Write, \
            Read, and Apply.
            result (int): The status code of the last batch call that used this BatchRecords.
            0 if all batch subtransactions succeeded (or if the only failures were FILTERED_OUT or RECORD_NOT_FOUND)
            non 0 if an error occured. The most common error being -16 (One or more batch sub transactions failed).
    """

    def __init__(self, batch_records: TypeBatchRecordList = None) -> None:
        """
            Example::

                # Create a BatchRecords to remove a record, write a bin, and read a bin.
                # Assume client is an instantiated and connected aerospike cleint.
                import aerospike_helpers.operations as op

                namespace = "test"
                set = "demo"
                bin_name = "id"
                keys = [
                    (namespace, set, 1),
                    (namespace, set, 2),
                    (namespace, set, 3)
                ]

                brs = BatchRecords(
                    [
                        Remove(
                            key=(namespace, set, 1),
                        ),
                        Write(
                            key=(namespace, set, 100),
                            ops=[
                                op.write(bin_name, 100),
                                op.read(bin_name),
                            ]
                        ),
                        BatchRead(
                            key=(namespace, set, 333),
                            ops=[
                                op.read(bin_name)
                            ]
                        )
                    ]
                )

                # Note this call will mutate brs and set results in it.
                client.batch_write(brs)
        """

        if batch_records is None:
            batch_records = []

        self.batch_records = batch_records
        self.result = 0<|endoftext|>


<commit_before>
from random import choice, randint

from django.contrib.auth.models import User
from django import forms
from gevent import sleep, spawn

from core.forms import GameForm
from core.settings import BETTING_PERIOD

game_registry = {}

class GameBase(type):
    """
    Whenever Game is subclassed, add a new game instance to the registry.
    """
    def __new__(cls, name, bases, attrs):
        attrs.setdefault("Form", GameForm)
        new = super(GameBase, cls).__new__(cls, name, bases, attrs)
        name = name.lower()
        if name != "game":
            if name != "dummy":
                del game_registry["dummy"]
            game_registry[name] = new(name)
        return new

class Game(object):
    """
    Base class for games. Subclasses should override the bet()
    method to perform validation of game specific betting args,
    and the outcome() method, which receives each player's betting
    args and returns True if the args perform a win.
    """

    __metaclass__ = GameBase

    def __init__(self, name):
        self.form = self.__class__.Form()
        self.name = name
        self.reset()

    def bet(self, namespace, amount, bet_args):
        """
        Called each time a player bets. When there are no players,
        a greenlet is spawned that will trigger the play of a
        turn of the game.

        Player's socketio namespace is passed through so that we
        have a way to broadcast back results once the turn is complete.

        The bet_args argument will be specific for each type of game.

        Subsequent calls to bet() for the same player prior to the turn
        being played, will update their betting arguments, and increase
        their bet.
        """
        if not self.players:
            spawn(self.turn)
        user_id = namespace.user["id"]
        self.players.setdefault(user_id, {"amount": 0})
        self.players[user_id]["namespace"] = namespace
        self.players[user_id]["amount"] += amount
        self.players[user_id]["bet_args"] = bet_args
        return True

    def turn(self):
        """
        Takes an actual turn of the game - called within a separate
        greenlet on the first call to bet(). We iterate through each of
        the players passing their bet_args to outcome() which multiplies
        the amount bet, and build a results dict we can broadcast back
        to all sockets. We piggyback the first player we find that's
        still connected, to broadcast the results back.
        """
        sleep(BETTING_PERIOD)
        broadcaster = None
        results = {}
        for user_id, player in self.players.items():
            result = self.outcome(*player["bet_args"])
            results[user_id] = player["amount"] * result
            if results[user_id] > 0:
                # Put a positive bet amount back into the user's account.
                user = User.objects.get(id=user_id)
                user.account.balance += (results[user_id] * 2)
                user.account.save()
            if broadcaster is None and player["namespace"].socket.connected:
                broadcaster = player["namespace"]
        if broadcaster is not None:
            broadcaster.broadcast_event("game_end", self.name, results)
        self.reset()

    def reset(self):
        self.players = {}

class Dummy(Game):
    """
    Default game with no input - just randomly win lose or draw.
    If any other games are registered, this one is removed.
    """
    def outcome(self):
        # Give back 0 (lose), 1 (even) or 2 (win) times the bet.
        return randint(0, 2)

ROULETTE_CHOICES = range(0, 37)

class Roulette(Game):

    class Form(GameForm):
        number = forms.ChoiceField(choices=[(c, c) for c in ROULETTE_CHOICES])

    def turn(self):
        self.landed_on = str(choice(ROULETTE_CHOICES))
        super(Roulette, self).turn()

    def outcome(self, choice):
        return (choice == self.landed_on) * len(ROULETTE_CHOICES)
<commit_msg>Refactor game broadcast out into its own method.<commit_after>
from random import choice, randint

from django.contrib.auth.models import User
from django import forms
from gevent import sleep, spawn

from core.forms import GameForm
from core.settings import BETTING_PERIOD

game_registry = {}

class GameBase(type):
    """
    Whenever Game is subclassed, add a new game instance to the registry.
    """
    def __new__(cls, name, bases, attrs):
        attrs.setdefault("Form", GameForm)
        new = super(GameBase, cls).__new__(cls, name, bases, attrs)
        name = name.lower()
        if name != "game":
            if name != "dummy":
                del game_registry["dummy"]
            game_registry[name] = new(name)
        return new

class Game(object):
    """
    Base class for games. Subclasses should override the bet()
    method to perform validation of game specific betting args,
    and the outcome() method, which receives each player's betting
    args and returns True if the args perform a win.
    """

    __metaclass__ = GameBase

    def __init__(self, name):
        self.form = self.__class__.Form()
        self.name = name
        self.reset()

    def bet(self, namespace, amount, bet_args):
        """
        Called each time a player bets. When there are no players,
        a greenlet is spawned that will trigger the play of a
        turn of the game.

        Player's socketio namespace is passed through so that we
        have a way to broadcast back results once the turn is complete.

        The bet_args argument will be specific for each type of game.

        Subsequent calls to bet() for the same player prior to the turn
        being played, will update their betting arguments, and increase
        their bet.
        """
        if not self.players:
            spawn(self.turn)
        user_id = namespace.user["id"]
        self.players.setdefault(user_id, {"amount": 0})
        self.players[user_id]["namespace"] = namespace
        self.players[user_id]["amount"] += amount
        self.players[user_id]["bet_args"] = bet_args
        return True

    def broadcast(self, *args):
        """
        Piggyback the first connected player's namespace to
        broadcast a message.
        """
        for player in self.players.values():
            if player["namespace"].socket.connected
                player["namespace"].broadcast_event(*args)
                break

    def turn(self):
        """
        Takes an actual turn of the game - called within a separate
        greenlet on the first call to bet(). We iterate through each of
        the players passing their bet_args to outcome() which multiplies
        the amount bet, and build a results dict we can broadcast back
        to all sockets.
        """
        sleep(BETTING_PERIOD)
        results = {}
        for user_id, player in self.players.items():
            result = self.outcome(*player["bet_args"])
            results[user_id] = player["amount"] * result
            if results[user_id] > 0:
                # Put a positive bet amount back into the user's account.
                user = User.objects.get(id=user_id)
                user.account.balance += (results[user_id] * 2)
                user.account.save()
        self.broadcast("game_end", self.name, results)
        self.reset()

    def reset(self):
        self.players = {}

class Dummy(Game):
    """
    Default game with no input - just randomly win lose or draw.
    If any other games are registered, this one is removed.
    """
    def outcome(self):
        # Give back 0 (lose), 1 (even) or 2 (win) times the bet.
        return randint(0, 2)

ROULETTE_CHOICES = range(0, 37)

class Roulette(Game):

    class Form(GameForm):
        number = forms.ChoiceField(choices=[(c, c) for c in ROULETTE_CHOICES])

    def turn(self):
        self.landed_on = str(choice(ROULETTE_CHOICES))
        super(Roulette, self).turn()

    def outcome(self, choice):
        return (choice == self.landed_on) * len(ROULETTE_CHOICES)
<|endoftext|>