text:
import math
import numpy as np
from typing import Any, Tuple, Optional, Sequence
try:
  import torch
  from torch import Tensor
  import torch.nn as nn
  import torch.nn.functional as F
except ModuleNotFoundError:
  raise ImportError('These classes require PyTorch to be installed.')

try:
  from torch_scatter import scatter_mean
except ModuleNotFoundError:
  pass

class ScaleNorm(nn.Module):
  """Apply Scale Normalization to input.

  The ScaleNorm layer first computes the square root of the scale, then computes the matrix/vector norm of the input tensor.
  The norm value is calculated as `sqrt(scale) / matrix norm`.
  Finally, the result is returned as `input_tensor * norm value`.

  This layer can be used instead of LayerNorm when a scaled version of the norm is required.
  Instead of performing the scaling operation (`scale / norm`) in a lambda-like layer, we are defining it within this layer to make prototyping more efficient.

  References
  ----------
  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264

  Examples
  --------
  >>> from deepchem.models.torch_models.layers import ScaleNorm
  >>> scale = 0.35
  >>> layer = ScaleNorm(scale)
  >>> input_tensor = torch.tensor([[1.269, 39.36], [0.00918, -9.12]])
  >>> output_tensor = layer(input_tensor)
  """

  def __init__(self, scale: float, eps: float = 1e-5):
    """Initialize a ScaleNorm layer.

    Parameters
    ----------
    scale: float
      Scale magnitude.
    eps: float
      Epsilon value. Default = 1e-5.
    """
    super(ScaleNorm, self).__init__()
    self.scale = nn.Parameter(torch.tensor(math.sqrt(scale)))
    self.eps = eps

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    norm = self.scale / torch.norm(x, dim=-1, keepdim=True).clamp(min=self.eps)
    return x * norm

class MultiHeadedMATAttention(nn.Module):
  """First constructs an attention layer tailored to the Molecular Attention Transformer [1]_ and then converts it into Multi-Headed Attention.

  In Multi-Headed attention the attention mechanism multiple times parallely through the multiple attention heads.
  Thus, different subsequences of a given sequences can be processed differently.
  The query, key and value parameters are split multiple ways and each split is passed separately through a different attention head.
  References
  ----------
  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264
  Examples
  --------
  >>> from deepchem.models.torch_models.layers import MultiHeadedMATAttention, MATEmbedding
  >>> import deepchem as dc
  >>> import torch
  >>> input_smile = "CC"
  >>> feat = dc.feat.MATFeaturizer()
  >>> input_smile = "CC"
  >>> out = feat.featurize(input_smile)
  >>> node = torch.tensor(out[0].node_features).float().unsqueeze(0)
  >>> adj = torch.tensor(out[0].adjacency_matrix).float().unsqueeze(0)
  >>> dist = torch.tensor(out[0].distance_matrix).float().unsqueeze(0)
  >>> mask = torch.sum(torch.abs(node), dim=-1) != 0
  >>> layer = MultiHeadedMATAttention(
  ...    dist_kernel='softmax',
  ...    lambda_attention=0.33,
  ...    lambda_distance=0.33,
  ...    h=16,
  ...    hsize=1024,
  ...    dropout_p=0.0)
  >>> op = MATEmbedding()(node)
  >>> output = layer(op, op, op, mask, adj, dist)
  """

  def __init__(self,
               dist_kernel: str = 'softmax',
               lambda_attention: float = 0.33,
               lambda_distance: float = 0.33,
               h: int = 16,
               hsize: int = 1024,
               dropout_p: float = 0.0,
               output_bias: bool = True):
    """Initialize a multi-headed attention layer.
    Parameters
    ----------
    dist_kernel: str
      Kernel activation to be used. Can be either 'softmax' for softmax or 'exp' for exponential.
    lambda_attention: float
      Constant to be multiplied with the attention matrix.
    lambda_distance: float
      Constant to be multiplied with the distance matrix.
    h: int
      Number of attention heads.
    hsize: int
      Size of dense layer.
    dropout_p: float
      Dropout probability.
    output_bias: bool
      If True, dense layers will use bias vectors.
    """
    super().__init__()
    if dist_kernel == "softmax":
      self.dist_kernel = lambda x: torch.softmax(-x, dim=-1)
    elif dist_kernel == "exp":
      self.dist_kernel = lambda x: torch.exp(-x)
    self.lambda_attention = lambda_attention
    self.lambda_distance = lambda_distance
    self.lambda_adjacency = 1.0 - self.lambda_attention - self.lambda_distance
    self.d_k = hsize // h
    self.h = h
    linear_layer = nn.Linear(hsize, hsize)
    self.linear_layers = nn.ModuleList([linear_layer for _ in range(3)])
    self.dropout_p = nn.Dropout(dropout_p)
    self.output_linear = nn.Linear(hsize, hsize, output_bias)

  def _single_attention(self,
                        query: torch.Tensor,
                        key: torch.Tensor,
                        value: torch.Tensor,
                        mask: torch.Tensor,
                        adj_matrix: torch.Tensor,
                        distance_matrix: torch.Tensor,
                        dropout_p: float = 0.0,
                        eps: float = 1e-6,
                        inf: float = 1e12) -> Tuple[torch.Tensor, torch.Tensor]:
    """Defining and computing output for a single MAT attention layer.
    Parameters
    ----------
    query: torch.Tensor
      Standard query parameter for attention.
    key: torch.Tensor
      Standard key parameter for attention.
    value: torch.Tensor
      Standard value parameter for attention.
    mask: torch.Tensor
      Masks out padding values so that they are not taken into account when computing the attention score.
    adj_matrix: torch.Tensor
      Adjacency matrix of the input molecule, returned from dc.feat.MATFeaturizer()
    dist_matrix: torch.Tensor
      Distance matrix of the input molecule, returned from dc.feat.MATFeaturizer()
    dropout_p: float
      Dropout probability.
    eps: float
      Epsilon value
    inf: float
      Value of infinity to be used.
    """
    d_k = query.size(-1)

    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
      scores = scores.masked_fill(
          mask.unsqueeze(1).repeat(1, query.shape[1], query.shape[2], 1) == 0,
          -inf)
    p_attn = F.softmax(scores, dim=-1)

    adj_matrix = adj_matrix / (
        torch.sum(torch.tensor(adj_matrix), dim=-1).unsqueeze(2) + eps)

    if len(adj_matrix.shape) <= 3:
      p_adj = adj_matrix.unsqueeze(1).repeat(1, query.shape[1], 1, 1)
    else:
      p_adj = adj_matrix.repeat(1, query.shape[1], 1, 1)

    distance_matrix = torch.tensor(distance_matrix).squeeze().masked_fill(
        mask.repeat(1, mask.shape[-1], 1) == 0, np.inf)

    distance_matrix = self.dist_kernel(distance_matrix)

    p_dist = distance_matrix.unsqueeze(1).repeat(1, query.shape[1], 1, 1)

    p_weighted = self.lambda_attention * p_attn + self.lambda_distance * p_dist + self.lambda_adjacency * p_adj
    p_weighted = self.dropout_p(p_weighted)

    return torch.matmul(p_weighted.float(), value.float()), p_attn

  def forward(self,
              query: torch.Tensor,
              key: torch.Tensor,
              value: torch.Tensor,
              mask: torch.Tensor,
              adj_matrix: torch.Tensor,
              distance_matrix: torch.Tensor,
              dropout_p: float = 0.0,
              eps: float = 1e-6,
              inf: float = 1e12) -> torch.Tensor:
    """Output computation for the MultiHeadedAttention layer.
    Parameters
    ----------
    query: torch.Tensor
      Standard query parameter for attention.
    key: torch.Tensor
      Standard key parameter for attention.
    value: torch.Tensor
      Standard value parameter for attention.
    mask: torch.Tensor
      Masks out padding values so that they are not taken into account when computing the attention score.
    adj_matrix: torch.Tensor
      Adjacency matrix of the input molecule, returned from dc.feat.MATFeaturizer()
    dist_matrix: torch.Tensor
      Distance matrix of the input molecule, returned from dc.feat.MATFeaturizer()
    dropout_p: float
      Dropout probability.
    eps: float
      Epsilon value
    inf: float
      Value of infinity to be used.
    """
    if mask is not None and len(mask.shape) <= 2:
      mask = mask.unsqueeze(1)

    batch_size = query.size(0)
    query, key, value = [
        layer(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
        for layer, x in zip(self.linear_layers, (query, key, value))
    ]

    x, _ = self._single_attention(query, key, value, mask, adj_matrix,
                                  distance_matrix, dropout_p, eps, inf)
    x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)

    return self.output_linear(x)

class MATEncoderLayer(nn.Module):
  """Encoder layer for use in the Molecular Attention Transformer [1]_.

  The MATEncoder layer primarily consists of a self-attention layer (MultiHeadedMATAttention) and a feed-forward layer (PositionwiseFeedForward).
  This layer can be stacked multiple times to form an encoder.

  References
  ----------
  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264

  Examples
  --------
  >>> from rdkit import Chem
  >>> import torch
  >>> import deepchem
  >>> from deepchem.models.torch_models.layers import MATEmbedding, MATEncoderLayer
  >>> input_smile = "CC"
  >>> feat = deepchem.feat.MATFeaturizer()
  >>> out = feat.featurize(input_smile)
  >>> node = torch.tensor(out[0].node_features).float().unsqueeze(0)
  >>> adj = torch.tensor(out[0].adjacency_matrix).float().unsqueeze(0)
  >>> dist = torch.tensor(out[0].distance_matrix).float().unsqueeze(0)
  >>> mask = torch.sum(torch.abs(node), dim=-1) != 0
  >>> layer = MATEncoderLayer()
  >>> op = MATEmbedding()(node)
  >>> output = layer(op, mask, adj, dist)
  """

  def __init__(self,
               dist_kernel: str = 'softmax',
               lambda_attention: float = 0.33,
               lambda_distance: float = 0.33,
               h: int = 16,
               sa_hsize: int = 1024,
               sa_dropout_p: float = 0.0,
               output_bias: bool = True,
               d_input: int = 1024,
               d_hidden: int = 1024,
               d_output: int = 1024,
               activation: str = 'leakyrelu',
               n_layers: int = 1,
               ff_dropout_p: float = 0.0,
               encoder_hsize: int = 1024,
               encoder_dropout_p: float = 0.0):
    """Initialize a MATEncoder layer.

    Parameters
    ----------
    dist_kernel: str
      Kernel activation to be used. Can be either 'softmax' for softmax or 'exp' for exponential, for the self-attention layer.
    lambda_attention: float
      Constant to be multiplied with the attention matrix in the self-attention layer.
    lambda_distance: float
      Constant to be multiplied with the distance matrix in the self-attention layer.
    h: int
      Number of attention heads for the self-attention layer.
    sa_hsize: int
      Size of dense layer in the self-attention layer.
    sa_dropout_p: float
      Dropout probability for the self-attention layer.
    output_bias: bool
      If True, dense layers will use bias vectors in the self-attention layer.
    d_input: int
      Size of input layer in the feed-forward layer.
    d_hidden: int
      Size of hidden layer in the feed-forward layer.
    d_output: int
      Size of output layer in the feed-forward layer.
    activation: str
      Activation function to be used in the feed-forward layer.
      Can choose between 'relu' for ReLU, 'leakyrelu' for LeakyReLU, 'prelu' for PReLU,
      'tanh' for TanH, 'selu' for SELU, 'elu' for ELU and 'linear' for linear activation.
    n_layers: int
      Number of layers in the feed-forward layer.
    dropout_p: float
      Dropout probability in the feeed-forward layer.
    encoder_hsize: int
      Size of Dense layer for the encoder itself.
    encoder_dropout_p: float
      Dropout probability for connections in the encoder layer.
    """
    super(MATEncoderLayer, self).__init__()
    self.self_attn = MultiHeadedMATAttention(dist_kernel, lambda_attention,
                                             lambda_distance, h, sa_hsize,
                                             sa_dropout_p, output_bias)
    self.feed_forward = PositionwiseFeedForward(d_input, d_hidden, d_output,
                                                activation, n_layers,
                                                ff_dropout_p)
    layer = SublayerConnection(size=encoder_hsize, dropout_p=encoder_dropout_p)
    self.sublayer = nn.ModuleList([layer for _ in range(2)])
    self.size = encoder_hsize

  def forward(self,
              x: torch.Tensor,
              mask: torch.Tensor,
              adj_matrix: torch.Tensor,
              distance_matrix: torch.Tensor,
              sa_dropout_p: float = 0.0) -> torch.Tensor:
    """Output computation for the MATEncoder layer.

    In the MATEncoderLayer intialization, self.sublayer is defined as an nn.ModuleList of 2 layers. We will be passing our computation through these layers sequentially.
    nn.ModuleList is subscriptable and thus we can access it as self.sublayer[0], for example.

    Parameters
    ----------
    x: torch.Tensor
      Input tensor.
    mask: torch.Tensor
      Masks out padding values so that they are not taken into account when computing the attention score.
    adj_matrix: torch.Tensor
      Adjacency matrix of a molecule.
    distance_matrix: torch.Tensor
      Distance matrix of a molecule.
    sa_dropout_p: float
      Dropout probability for the self-attention layer (MultiHeadedMATAttention).
    """
    x = self.sublayer[0](x,
                         self.self_attn(x,
                                        x,
                                        x,
                                        mask=mask,
                                        dropout_p=sa_dropout_p,
                                        adj_matrix=adj_matrix,
                                        distance_matrix=distance_matrix))
    return self.sublayer[1](x, self.feed_forward(x))

class SublayerConnection(nn.Module):
  """SublayerConnection layer which establishes a residual connection, as used in the Molecular Attention Transformer [1]_.

  The SublayerConnection layer is a residual layer which is then passed through Layer Normalization.
  The residual connection is established by computing the dropout-adjusted layer output of a normalized tensor and adding this to the original input tensor.

  References
  ----------
  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264

  Examples
  --------
  >>> from deepchem.models.torch_models.layers import SublayerConnection
  >>> scale = 0.35
  >>> layer = SublayerConnection(2, 0.)
  >>> input_ar = torch.tensor([[1., 2.], [5., 6.]])
  >>> output = layer(input_ar, input_ar)
  """

  def __init__(self, size: int, dropout_p: float = 0.0):
    """Initialize a SublayerConnection Layer.

    Parameters
    ----------
    size: int
      Size of layer.
    dropout_p: float
      Dropout probability.
    """
    super(SublayerConnection, self).__init__()
    self.norm = nn.LayerNorm(size)
    self.dropout_p = nn.Dropout(dropout_p)

  def forward(self, x: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
    """Output computation for the SublayerConnection layer.

    Takes an input tensor x, then adds the dropout-adjusted sublayer output for normalized x to it.
    This is done to add a residual connection followed by LayerNorm.

    Parameters
    ----------
    x: torch.Tensor
      Input tensor.
    output: torch.Tensor
      Layer whose normalized output will be added to x.
    """
    if x is None:
      return self.dropout(self.norm(output))
    return x + self.dropout_p(self.norm(output))

class PositionwiseFeedForward(nn.Module):
  """PositionwiseFeedForward is a layer used to define the position-wise feed-forward (FFN) algorithm for the Molecular Attention Transformer [1]_

  Each layer in the MAT encoder contains a fully connected feed-forward network which applies two linear transformations and the given activation function.
  This is done in addition to the SublayerConnection module.

  References
  ----------
  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264

  Examples
  --------
  >>> from deepchem.models.torch_models.layers import PositionwiseFeedForward
  >>> feed_fwd_layer = PositionwiseFeedForward(d_input = 2, d_hidden = 2, d_output = 2, activation = 'relu', n_layers = 1, dropout_p = 0.1)
  >>> input_tensor = torch.tensor([[1., 2.], [5., 6.]])
  >>> output_tensor = feed_fwd_layer(input_tensor)
  """

  def __init__(self,
               d_input: int = 1024,
               d_hidden: int = 1024,
               d_output: int = 1024,
               activation: str = 'leakyrelu',
               n_layers: int = 1,
               dropout_p: float = 0.0):
    """Initialize a PositionwiseFeedForward layer.

    Parameters
    ----------
    d_input: int
      Size of input layer.
    d_hidden: int (same as d_input if d_output = 0)
      Size of hidden layer.
    d_output: int (same as d_input if d_output = 0)
      Size of output layer.
    activation: str
      Activation function to be used. Can choose between 'relu' for ReLU, 'leakyrelu' for LeakyReLU, 'prelu' for PReLU,
      'tanh' for TanH, 'selu' for SELU, 'elu' for ELU and 'linear' for linear activation.
    n_layers: int
      Number of layers.
    dropout_p: float
      Dropout probability.
    """
    super(PositionwiseFeedForward, self).__init__()

    if activation == 'relu':
      self.activation: Any = nn.ReLU()

    elif activation == 'leakyrelu':
      self.activation = nn.LeakyReLU(0.1)

    elif activation == 'prelu':
      self.activation = nn.PReLU()

    elif activation == 'tanh':
      self.activation = nn.Tanh()

    elif activation == 'selu':
      self.activation = nn.SELU()

    elif activation == 'elu':
      self.activation = nn.ELU()

    elif activation == "linear":
      self.activation = lambda x: x

    self.n_layers: int = n_layers
    d_output = d_output if d_output != 0 else d_input
    d_hidden = d_hidden if d_hidden != 0 else d_input

    if n_layers == 1:
      self.linears: Any = [nn.Linear(d_input, d_output)]

    else:
      self.linears = [nn.Linear(d_input, d_hidden)] + \
                      [nn.Linear(d_hidden, d_hidden) for _ in range(n_layers - 2)] + \
                      [nn.Linear(d_hidden, d_output)]

    self.linears = nn.ModuleList(self.linears)
    dropout_layer = nn.Dropout(dropout_p)
    self.dropout_p = nn.ModuleList([dropout_layer for _ in range(n_layers)])

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Output Computation for the PositionwiseFeedForward layer.

    Parameters
    ----------
    x: torch.Tensor
      Input tensor.
    """
    if not self.n_layers:
      return x

    if self.n_layers == 1:
      return self.dropout_p[0](self.activation(self.linears[0](x)))

    else:
      for i in range(self.n_layers - 1):
        x = self.dropout_p[i](self.activation(self.linears[i](x)))
      return self.linears[-1](x)

class MATEmbedding(nn.Module):
  """Embedding layer to create embedding for inputs.

  In an embedding layer, input is taken and converted to a vector representation for each input.
  In the MATEmbedding layer, an input tensor is processed through a dropout-adjusted linear layer and the resultant vector is returned.

  References
  ----------
  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264

  Examples
  --------
  >>> from deepchem.models.torch_models.layers import MATEmbedding
  >>> layer = MATEmbedding(d_input = 3, d_output = 3, dropout_p = 0.2)
  >>> input_tensor = torch.tensor([1., 2., 3.])
  >>> output = layer(input_tensor)
  """

  def __init__(self,
               d_input: int = 36,
               d_output: int = 1024,
               dropout_p: float = 0.0):
    """Initialize a MATEmbedding layer.

    Parameters
    ----------
    d_input: int
      Size of input layer.
    d_output: int
      Size of output layer.
    dropout_p: float
      Dropout probability for layer.
    """
    super(MATEmbedding, self).__init__()
    self.linear_unit = nn.Linear(d_input, d_output)
    self.dropout = nn.Dropout(dropout_p)

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Computation for the MATEmbedding layer.

    Parameters
    ----------
    x: torch.Tensor
      Input tensor to be converted into a vector.
    """
    return self.dropout(self.linear_unit(x))

class MATGenerator(nn.Module):
  """MATGenerator defines the linear and softmax generator step for the Molecular Attention Transformer [1]_.

  In the MATGenerator, a Generator is defined which performs the Linear + Softmax generation step.
  Depending on the type of aggregation selected, the attention output layer performs different operations.

  References
  ----------
  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer" Graph Representation Learning workshop and Machine Learning and the Physical Sciences workshop at NeurIPS 2019. 2020. https://arxiv.org/abs/2002.08264

  Examples
  --------
  >>> from deepchem.models.torch_models.layers import MATGenerator
  >>> layer = MATGenerator(hsize = 3, aggregation_type = 'mean', d_output = 1, n_layers = 1, dropout_p = 0.3, attn_hidden = 128, attn_out = 4)
  >>> input_tensor = torch.tensor([1., 2., 3.])
  >>> mask = torch.tensor([1., 1., 1.])
  >>> output = layer(input_tensor, mask)
  """

  def __init__(self,
               hsize: int = 1024,
               aggregation_type: str = 'mean',
               d_output: int = 1,
               n_layers: int = 1,
               dropout_p: float = 0.0,
               attn_hidden: int = 128,
               attn_out: int = 4):
    """Initialize a MATGenerator.

    Parameters
    ----------
    hsize: int
      Size of input layer.
    aggregation_type: str
      Type of aggregation to be used. Can be 'grover', 'mean' or 'contextual'.
    d_output: int
      Size of output layer.
    n_layers: int
      Number of layers in MATGenerator.
    dropout_p: float
      Dropout probability for layer.
    attn_hidden: int
      Size of hidden attention layer.
    attn_out: int
      Size of output attention layer.
    """
    super(MATGenerator, self).__init__()

    if aggregation_type == 'grover':
      self.att_net = nn.Sequential(
          nn.Linear(hsize, attn_hidden, bias=False),
          nn.Tanh(),
          nn.Linear(attn_hidden, attn_out, bias=False),
      )
      hsize *= attn_out

    if n_layers == 1:
      self.proj: Any = nn.Linear(hsize, d_output)

    else:
      self.proj = []

      for i in range(n_layers - 1):
        self.proj.append(nn.Linear(hsize, attn_hidden))
        self.proj.append(nn.LeakyReLU(negative_slope=0.1))
        self.proj.append(nn.LayerNorm(attn_hidden))
        self.proj.append(nn.Dropout(dropout_p))
      self.proj.append(nn.Linear(attn_hidden, d_output))
      self.proj = torch.nn.Sequential(*self.proj)
    self.aggregation_type = aggregation_type

  def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Computation for the MATGenerator layer.

    Parameters
    ----------
    x: torch.Tensor
      Input tensor.
    mask: torch.Tensor
      Mask for padding so that padded values do not get included in attention score calculation.
    """

    mask = mask.unsqueeze(-1).float()
    out_masked = x * mask

    if self.aggregation_type == 'mean':
      out_sum = out_masked.sum(dim=1)
      mask_sum = mask.sum(dim=(1))
      out_avg_pooling = out_sum / mask_sum

    elif self.aggregation_type == 'grover':
      out_attn = self.att_net(out_masked)
      out_attn = out_attn.masked_fill(mask == 0, -1e9)
      out_attn = F.softmax(out_attn, dim=1)
      out_avg_pooling = torch.matmul(torch.transpose(out_attn, -1, -2),
                                     out_masked)
      out_avg_pooling = out_avg_pooling.view(out_avg_pooling.size(0), -1)

    elif self.aggregation_type == 'contextual':
      out_avg_pooling = x

    return self.proj(out_avg_pooling)

class GraphNetwork(torch.nn.Module):
  """Graph Networks

  A Graph Network [1]_ takes a graph as input and returns an updated graph
  as output. The output graph has same structure as input graph but it
  has updated node features, edge features and global state features.

  Parameters
  ----------
  n_node_features: int
    Number of features in a node
  n_edge_features: int
    Number of features in a edge
  n_global_features: int
    Number of global features
  is_undirected: bool, optional (default True)
    Directed or undirected graph
  residual_connection: bool, optional (default True)
    If True, the layer uses a residual connection during training

  Example
  -------
  >>> import torch
  >>> from deepchem.models.torch_models.layers import GraphNetwork as GN
  >>> n_nodes, n_node_features = 5, 10
  >>> n_edges, n_edge_features = 5, 2
  >>> n_global_features = 4
  >>> node_features = torch.randn(n_nodes, n_node_features)
  >>> edge_features = torch.randn(n_edges, n_edge_features)
  >>> edge_index = torch.tensor([[0, 1, 2, 3, 4], [1, 2, 3, 4, 0]]).long()
  >>> global_features = torch.randn(1, n_global_features)
  >>> gn = GN(n_node_features=n_node_features, n_edge_features=n_edge_features, n_global_features=n_global_features)
  >>> node_features, edge_features, global_features = gn(node_features, edge_index, edge_features, global_features)

  References
  ----------
  .. [1] Battaglia et al, Relational inductive biases, deep learning, and graph networks. https://arxiv.org/abs/1806.01261 (2018)
  """

  def __init__(self,
               n_node_features: int = 32,
               n_edge_features: int = 32,
               n_global_features: int = 32,
               is_undirected: bool = True,
               residual_connection: bool = True):
    super().__init__()
    self.n_node_features = n_node_features
    self.n_edge_features = n_edge_features
    self.n_global_features = n_global_features
    self.is_undirected = is_undirected
    self.residual_connection = residual_connection

    self.edge_models, self.node_models, self.global_models = torch.nn.ModuleList(
    ), torch.nn.ModuleList(), torch.nn.ModuleList()
    self.edge_models.append(
        nn.Linear(in_features=n_node_features * 2 + n_edge_features +
                  n_global_features,
                  out_features=32))
    self.node_models.append(
        nn.Linear(in_features=n_node_features + n_edge_features +
                  n_global_features,
                  out_features=32))
    self.global_models.append(
        nn.Linear(in_features=n_node_features + n_edge_features +
                  n_global_features,
                  out_features=32))

    # Used for converting edges back to their original shape
    self.edge_dense = nn.Linear(in_features=32, out_features=n_edge_features)
    self.node_dense = nn.Linear(in_features=32, out_features=n_node_features)
    self.global_dense = nn.Linear(in_features=32,
                                  out_features=n_global_features)

  def reset_parameters(self) -> None:
    self.edge_dense.reset_parameters()
    self.node_dense.reset_parameters()
    self.global_dense.reset_parameters()
    for i in range(0, len(self.edge_models)):
      self.edge_models[i].reset_parameters()
    for i in range(0, len(self.node_models)):
      self.node_models[i].reset_parameters()
    for i in range(0, len(self.global_models)):
      self.global_models[i].reset_parameters()

  def _update_edge_features(self, node_features, edge_index, edge_features,
                            global_features, batch):
    src_index, dst_index = edge_index
    out = torch.cat((node_features[src_index], node_features[dst_index],
                     edge_features, global_features[batch]),
                    dim=1)
    assert out.shape[
        1] == self.n_node_features * 2 + self.n_edge_features + self.n_global_features
    for model in self.edge_models:
      out = model(out)
    return self.edge_dense(out)

  def _update_node_features(self, node_features, edge_index, edge_features,
                            global_features, batch):

    src_index, dst_index = edge_index
    # Compute mean edge features for each node by dst_index (each node
    # receives information from edges which have that node as its destination,
    # hence the computation uses dst_index to aggregate information)
    edge_features_mean_by_node = scatter_mean(edge_features, dst_index, dim=0)
    out = torch.cat(
        (node_features, edge_features_mean_by_node, global_features[batch]),
        dim=1)
    for model in self.node_models:
      out = model(out)
    return self.node_dense(out)

  def _update_global_features(self, node_features, edge_features,
                              global_features, node_batch_map, edge_batch_map):
    edge_features_mean = scatter_mean(edge_features, edge_batch_map, dim=0)
    node_features_mean = scatter_mean(node_features, node_batch_map, dim=0)
    out = torch.cat((edge_features_mean, node_features_mean, global_features),
                    dim=1)
    for model in self.global_models:
      out = model(out)
    return self.global_dense(out)

  def forward(self,
              node_features: Tensor,
              edge_index: Tensor,
              edge_features: Tensor,
              global_features: Tensor,
              batch: Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]:
    """Output computation for a GraphNetwork

    Parameters
    ----------
    node_features: torch.Tensor
      Input node features of shape :math:`(|\mathcal{V}|, F_n)`
    edge_index: torch.Tensor
      Edge indexes of shape :math:`(2, |\mathcal{E}|)`
    edge_features: torch.Tensor
      Edge features of the graph, shape: :math:`(|\mathcal{E}|, F_e)`
    global_features: torch.Tensor
      Global features of the graph, shape: :math:`(F_g, 1)` where, :math:`|\mathcal{V}|` and :math:`|\mathcal{E}|` denotes the number of nodes and edges in the graph, :math:`F_n`, :math:`F_e`, :math:`F_g` denotes the number of node features, edge features and global state features respectively.
    batch: torch.LongTensor (optional, default: None)
      A vector that maps each node to its respective graph identifier. The attribute is used only when more than one graph are batched together during a single forward pass.
    """
    if batch is None:
      batch = node_features.new_zeros(node_features.size(0), dtype=torch.int64)

    node_features_copy, edge_features_copy, global_features_copy = node_features, edge_features, global_features
    if self.is_undirected is True:
      # holding bi-directional edges in case of undirected graphs
      edge_index = torch.cat((edge_index, edge_index.flip([0])), dim=1)
      edge_features_len = edge_features.shape[0]
      edge_features = torch.cat((edge_features, edge_features), dim=0)
    edge_batch_map = batch[edge_index[0]]
    edge_features = self._update_edge_features(node_features, edge_index,
                                               edge_features, global_features,
                                               edge_batch_map)
    node_features = self._update_node_features(node_features, edge_index,
                                               edge_features, global_features,
                                               batch)
    global_features = self._update_global_features(node_features, edge_features,
                                                   global_features, batch,
                                                   edge_batch_map)

    if self.is_undirected is True:
      # coonverting edge features to its original shape
      split = torch.split(edge_features, (edge_features_len, edge_features_len))
      edge_features = (split[0] + split[1]) / 2

    if self.residual_connection:
      edge_features += edge_features_copy
      node_features += node_features_copy
      global_features += global_features_copy

    return node_features, edge_features, global_features

  def __repr__(self) -> str:
    return (
        f'{self.__class__.__name__}(n_node_features={self.n_node_features}, n_edge_features={self.n_edge_features}, n_global_features={self.n_global_features}, is_undirected={self.is_undirected}, residual_connection={self.residual_connection})'
    )

class Affine(nn.Module):
  """Class which performs the Affine transformation.

  This transformation is based on the affinity of the base distribution with
  the target distribution. A geometric transformation is applied where
  the parameters performs changes on the scale and shift of a function
  (inputs).

  Normalizing Flow transformations must be bijective in order to compute
  the logarithm of jacobian's determinant. For this reason, transformations
  must perform a forward and inverse pass.

  Example
  --------
  >>> import deepchem as dc
  >>> from deepchem.models.torch_models.layers import Affine
  >>> import torch
  >>> from torch.distributions import MultivariateNormal
  >>> # initialize the transformation layer's parameters
  >>> dim = 2
  >>> samples = 96
  >>> transforms = Affine(dim)
  >>> # forward pass based on a given distribution
  >>> distribution = MultivariateNormal(torch.zeros(dim), torch.eye(dim))
  >>> tensor = distribution.sample(torch.Size((samples, dim)))
  >>> len(transforms.forward(tensor))
  2
  >>> # inverse pass based on a distribution
  >>> len(transforms.inverse(tensor))
  2

  """

  def __init__(self, dim: int) -> None:
    """Create a Affine transform layer.

    Parameters
    ----------
    dim: int
      Value of the Nth dimension of the dataset.

    """

    super().__init__()
    self.dim = dim
    self.scale = nn.Parameter(torch.zeros(self.dim))
    self.shift = nn.Parameter(torch.zeros(self.dim))

  def forward(self, x: Sequence) -> Tuple[torch.Tensor, torch.Tensor]:
    """Performs a transformation between two different distributions. This
    particular transformation represents the following function:
    y = x * exp(a) + b, where a is scale parameter and b performs a shift.
    This class also returns the logarithm of the jacobians determinant
    which is useful when invert a transformation and compute the
    probability of the transformation.
    input shape: (samples, dim)
    output shape: (samples, dim)
    """

    y = torch.exp(self.scale) * x + self.shift
    det_jacobian = torch.exp(self.scale.sum())
    log_det_jacobian = torch.ones(y.shape[0]) * torch.log(det_jacobian)

    return y, log_det_jacobian

  def inverse(self, y: Sequence) -> Tuple[torch.Tensor, torch.Tensor]:
    """Performs a transformation between two different distributions.
    This transformation represents the bacward pass of the function
    mention before. Its mathematical representation is x = (y - b) / exp(a)
    , where "a" is scale parameter and "b" performs a shift. This class
    also returns the logarithm of the jacobians determinant which is
    useful when invert a transformation and compute the probability of
    the transformation.
    input shape: (samples, dim)
    output shape: (samples, dim)
    """

    x = (y - self.shift) / torch.exp(self.scale)
    det_jacobian = 1 / torch.exp(self.scale.sum())
    inverse_log_det_jacobian = torch.ones(x.shape[0]) * torch.log(det_jacobian)

    return x, inverse_log_det_jacobian

max_stars_repo_name: JoseAntonioSiguenza/deepchem

max_issues_repo_name: JoseAntonioSiguenza/deepchem

max_forks_repo_name: JoseAntonioSiguenza/deepchem


text:
import numpy as np

from multiprocessing import Pool

class BaseModel:
    """
    Base class for a data generating model.
    """
    def __init__(self, dim_param=1, sample_size=10, n_workers=1, seed=None):
        """
        :param dim_param: number of parameters in the model
        :param sample_size: number of samples drawn in a single generation step
        :param n_workers: number of cores for multiprocessing
        :param seed:
        """
        self.dim_param = dim_param
        self.sample_size = sample_size
        self.n_workers = n_workers
        self.run_parallel = self.n_workers > 1

        self.seed = seed
        self.rng = np.random.RandomState(seed=seed)

    def gen(self, params_list):
        # generate data given parameter list

        params_list = np.array(params_list)
        n_params = params_list.shape[0]

        if self.run_parallel:
            # for multiprocessing the seeds for the individual processes have to
            # be prellocated.
            process_seeds = self.rng.randint(low=10000000, size=n_params)
            # add them to the parameter vector of every sample
            params_list_with_seeds = np.hstack((params_list.reshape(params_list.shape[0], -1),
                                                process_seeds.reshape(process_seeds.shape[0], -1))).tolist()
            # run in parallel
            p = Pool(processes=self.n_workers)
            data_list = list(p.map(self.gen_single, params_list_with_seeds))
        else:
            # if sequential, run as for loop over parameters
            data_list = []
            for param in params_list:
                data_list.append(self.gen_single(param))

        return np.array(data_list)

    def gen_single(self, params):
        pass

ext: py

max_stars_repo_name: janfb/mcabc

max_issues_repo_name: janfb/mcabc

max_forks_repo_name: janfb/mcabc

reason: import numpy


text:
# Copyright (c) 2016-present, Facebook, 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.
##############################################################################

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import numpy as np
from caffe2.python import core, workspace
from caffe2.python.test_util import TestCase, rand_array

class TestPartitionOps(TestCase):

    def test_configs(self):
        # (main dims, partitions,  main type, [list of (extra dims, type)])
        configs = [
            ((10, ), 3),
            ((4, ), 10),
            ((10, 10), 4),
            ((100, ), 2),
            ((5, ), 1),
            ((1, ), 1),
            ((2, 10), 2),
        ]
        suffixes = [
            [],
            [((2, 2), np.float32)],
            [((3, ), np.int64), ((2, ), np.float32)],
        ]
        return [
            (main_dims, parts, main_type, extra, pack)
            for main_dims, parts in configs
            for main_type in [np.int32, np.int64] for extra in suffixes
            for pack in [False, True]
        ]

    def testPartition(self):
        for main_dims, parts, main_type, extra_ins, pack in self.test_configs():
            ins = ['in' + str(i) for i in range(1 + len(extra_ins))]
            outs = [
                'in{}_p{}'.format(j, i)
                for i in range(parts) for j in range(1 + len(extra_ins))
            ]
            op = core.CreateOperator(
                'Partition', ins, outs, pack_first_input=(1 if pack else 0))
            x = []
            for i, (dims, t) in enumerate([((), main_type)] + extra_ins):
                if t in [np.float32, np.float64]:
                    d = rand_array(*(main_dims + dims))
                else:
                    d = np.random.randint(-100, 100, (main_dims + dims))
                d = d.astype(t)
                workspace.FeedBlob(ins[i], d)
                x.append(d)

            def sharding(x):
                # numpy has proper modulo op that yields non-negative results
                shards = (x[0] % parts).reshape([-1])
                out = []
                for i in range(parts):
                    for ind, v in enumerate(x):
                        suffix_shape = v.shape[len(x[0].shape):]
                        accum = []
                        data = v.reshape((-1, ) + suffix_shape)

                        if pack and ind == 0:
                            data = data // parts

                        for j, s in enumerate(shards):
                            if s == i:
                                accum.append(data[j])

                        def join(a):
                            if not a:
                                return np.empty(shape=(0, ) + suffix_shape)
                            return np.stack(a)

                        out.append(join(accum))
                return out

            workspace.RunOperatorOnce(op)
            ref = sharding(x)
            print(x)
            print(ref)
            for name, expected in zip(outs, ref):
                np.testing.assert_array_equal(
                    expected, workspace.FetchBlob(name)
                )

            # test inverse operation (GatherByKey)
            if len(main_dims) == 1:
                # currently only 1D key tensor supported
                for i in range(len(extra_ins)):
                    expected_out = ins[i + 1]
                    gather_ins = [ins[0]] + [
                        outs[len(ins) * p + i + 1] for p in range(parts)]
                    actual_out = expected_out + '_actual'
                    op = core.CreateOperator(
                        'GatherByKey', gather_ins, actual_out)
                    workspace.RunOperatorOnce(op)
                    expected = workspace.FetchBlob(expected_out)
                    actual = workspace.FetchBlob(actual_out)
                    np.testing.assert_array_equal(expected, actual)

    def testLengthsPartition(self):
        for main_dims, parts, main_type, extra_ins, pack in self.test_configs():
            # For LengthsSharding only 1-D tensors supported as a first input
            if len(main_dims) > 1:
                continue
            ins = ['in' + str(i) for i in range(2 + len(extra_ins))]
            outs = [
                'in{}_p{}'.format(j, i)
                for i in range(parts) for j in range(2 + len(extra_ins))
            ]
            op = core.CreateOperator(
                'LengthsPartition', ins, outs,
                pack_first_input=(1 if pack else 0)
            )
            x = []
            for i, (dims, t) in enumerate([((), main_type)] + extra_ins):
                if t in [np.float32, np.float64]:
                    d = rand_array(*(main_dims + dims))
                else:
                    d = np.random.randint(-100, 100, (main_dims + dims))
                d = d.astype(t)
                workspace.FeedBlob(ins[i + 1], d)
                x.append(d)

            # Randomly generate length tensor as well
            elements = np.random.randint(2, 10)
            lengths = []
            total_length = 0
            for _ in range(elements - 1):
                lengths.append(np.random.randint(main_dims[0] - total_length))
                total_length += lengths[-1]
            lengths.append(main_dims[0] - total_length)
            workspace.FeedBlob(ins[0], np.array(lengths, dtype=np.int32))

            def sharding(x):
                # numpy has proper modulo op that yields non-negative results
                shards = (x[0] % parts).reshape([-1])
                out = []
                for i in range(parts):
                    idx = 0
                    sharded_lengths = np.zeros(elements)
                    for ind, length in enumerate(lengths):
                        for _ in range(length):
                            if shards[idx] == i:
                                sharded_lengths[ind] += 1
                            idx += 1
                    out.append(sharded_lengths)

                    for ind, v in enumerate(x):
                        suffix_shape = v.shape[len(x[0].shape):]
                        accum = []
                        data = v.reshape((-1, ) + suffix_shape)

                        if pack and ind == 0:
                            data = data // parts

                        for j, s in enumerate(shards):
                            if s == i:
                                accum.append(data[j])

                        def join(a):
                            if not a:
                                return np.empty(shape=(0, ) + suffix_shape)
                            return np.stack(a)

                        out.append(join(accum))
                return out

            workspace.RunOperatorOnce(op)
            ref = sharding(x)
            for name, expected in zip(outs, ref):
                np.testing.assert_array_equal(
                    expected, workspace.FetchBlob(name)
                )

if __name__ == "__main__":
    import unittest
    unittest.main()


text:
"""
Module for useful general functions.

ECE 4250 Final Project.
Brian Richard (bcr53), Gautam Mekkat (gm484), Sameer Lal (sjl328).
April 21th, 2018.
"""

import numpy as np
import SimpleITK as sitk

def get_image_array(im):
    """
    Convert Image object to image array.

    :param im: SimpleITK Image.
    :return: Image array as numpy array.
    """
    return sitk.GetArrayFromImage(im)

def standardize(img_arr):
    """
    Standardize a numpy array so that the mean is 0 and std. dev. is 1.

    :param img_arr: The numpy array to be standardized.
    :return: The standardized image as a numpy array.
    """
    return (img_arr - np.mean(img_arr)) / np.std(img_arr)

def rescale(img_arr, min=0, max=1):
    """
    Rescale a numpy array so that all values are between min and max (default: 0 and 1).

    :param img_arr: The numpy array to be normalized.
    :param min: The new minimum value in the image. Default: 0.
    :param max: The new maximum value in the image. Default: 1.
    :return: The normalized image as a numpy array with values between min and max.
    """
    normalized = (img_arr - np.min(img_arr)) / (np.max(img_arr) - np.min(img_arr))
    return normalized * (max - min) + min

def standardize_and_remove_bg(img_arr):
    """
    Standardize lung image and remove black background.

    :param img_arr: N-dimensional image as numpy array.
    :return: Image with standardized values and background removed.
    """
    standardized = standardize(img_arr)
    standardized[standardized < 0] = 0
    return standardized

def distance(point1, point2):
    """Return the distance between two (x, y, z) points."""
    xd = point1[0] - point2[0]
    yd = point1[1] - point2[1]
    zd = point1[2] - point2[2]
    return (xd ** 2 + yd ** 2 + zd ** 2) ** 0.5

def average_intensity(img, coord, diameter):
    """
    Return the average intensity of a nodule.

    :param img: N-dimensional image as numpy array.
    :param coord: The coordinates of the nodules, (x, y, z).
    :param radius: The approximate radius of the nodule
    :return: Image with standardized values and background removed.
    """
    radius = int(diameter / 2)
    img = np.asarray(img)
    coord = np.asarray(coord, dtype=int)
    z1 = max(0, coord[2] - radius)
    z2 = min(img.shape[0], coord[2] + radius)
    y1 = max(0, coord[1] - radius)
    y2 = min(img.shape[1], coord[1] + radius)
    x1 = max(0, coord[0] - radius)
    x2 = min(img.shape[2], coord[0] + radius)
    return np.mean(img[z1:z2, y1:y2, x1:x2])

max_stars_repo_name: sameerlal/lung-cancer-detection

max_issues_repo_name: sameerlal/lung-cancer-detection

max_forks_repo_name: sameerlal/lung-cancer-detection


text:
# coding=utf-8
# Copyright 2020 The Google AI Perception Team Authors.
#
# 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.
"""Estimate AIST++ camera parameters."""
import json
import math
import os
import random

from absl import app
from absl import flags
from aist_plusplus.loader import AISTDataset
import aniposelib
import numpy as np

FLAGS = flags.FLAGS
flags.DEFINE_string(
    'anno_dir',
    '/usr/local/google/home/ruilongli/data/public/aist_plusplus_final/',
    'input local dictionary for AIST++ annotations.')
flags.DEFINE_string(
    'save_dir',
    '/usr/local/google/home/ruilongli/data/public/aist_plusplus_final/cameras/',
    'output local dictionary that stores AIST++ camera parameters.')
random.seed(0)
np.random.seed(0)

def init_env_cameras():
  """Trys to estimate the environment manually."""
  cams = []
  for i, view in enumerate(AISTDataset.VIEWS):
    f = 1600
    cx = 1920 // 2
    cy = 1080 // 2
    if view == 'c09':
      rvec = [0, -math.radians(0), 0]
      tvec = [0, -170, -500]
    else:
      rvec = [0, -math.radians(360 // 8 * i), 0]
      tvec = [0, -180, -500]

    matrix = np.array([
        [f, 0, cx],
        [0, f, cy],
        [0, 0, 1],
    ], dtype=np.float32)
    cams.append(
        aniposelib.cameras.Camera(
            matrix=matrix, rvec=rvec, tvec=tvec, name=view, size=(1920, 1080)))
  cgroup = aniposelib.cameras.CameraGroup(cams)
  return cgroup

def main(_):
  aist_dataset = AISTDataset(anno_dir=FLAGS.anno_dir)

  for env_name, seq_names in aist_dataset.mapping_env2seq.items():
    # Init camera parameters
    cgroup = init_env_cameras()

    # Select a set of sequences for optimizing camera parameters.
    seq_names = random.choices(seq_names, k=20)

    # Load 2D keypoints
    keypoints2d_all = []
    for seq_name in seq_names:
      keypoints2d_raw, _, _ = AISTDataset.load_keypoint2d(
          aist_dataset.keypoint2d_dir, seq_name=seq_name)
      # Special cases
      if seq_name == 'gBR_sBM_cAll_d04_mBR0_ch01':
        keypoints2d_raw[4] = np.nan  # not synced view
      if seq_name == 'gJB_sBM_cAll_d07_mJB3_ch05':
        keypoints2d_raw[6] = np.nan  # size 640x480
      keypoints2d_all.append(keypoints2d_raw)
    keypoints2d_all = np.concatenate(keypoints2d_all, axis=1)

    # Filter keypoints to select those best points
    kpt_thre = 0.5
    ignore_idxs = np.where(keypoints2d_all[:, :, :, 2] < kpt_thre)
    keypoints2d_all[ignore_idxs[0], ignore_idxs[1], ignore_idxs[2], :] = np.nan
    keypoints2d_all = keypoints2d_all[..., 0:2]

    # Apply bundle adjustment and dump the camera parameters
    nviews = keypoints2d_all.shape[0]
    cgroup.bundle_adjust_iter(
        keypoints2d_all.reshape(nviews, -1, 2),
        n_iters=20,
        n_samp_iter=500,
        n_samp_full=5000,
        verbose=True)
    os.makedirs(FLAGS.save_dir, exist_ok=True)
    camera_file = os.path.join(FLAGS.save_dir, f'{env_name}.json')
    with open(camera_file, 'w') as f:
      json.dump([camera.get_dict() for camera in cgroup.cameras], f)

if __name__ == '__main__':
  app.run(main)


text:
from eagles.Supervised import config
from eagles.Supervised.utils.feature_selection import EaglesFeatureSelection
from sklearn.feature_selection import SelectFromModel, SelectKBest
from sklearn.ensemble import (
    RandomForestClassifier,
    ExtraTreesClassifier,
    GradientBoostingClassifier,
    AdaBoostClassifier,
    RandomForestRegressor,
    ExtraTreesRegressor,
    GradientBoostingRegressor,
    AdaBoostRegressor,
    VotingClassifier,
    VotingRegressor,
)
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.linear_model import (
    LogisticRegression,
    LinearRegression,
    Lasso,
    ElasticNet,
    Ridge,
    PoissonRegressor,
)
from sklearn.svm import SVC, SVR
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor

from sklearn.impute import (
    SimpleImputer,
    MissingIndicator,
    KNNImputer,
)  # , IterativeImputer
from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler
from sklearn.pipeline import Pipeline
from sklearn.model_selection import KFold, TimeSeriesSplit, StratifiedKFold

from xgboost import XGBRegressor, XGBClassifier

import numpy as np

import logging
import warnings

logger = logging.getLogger(__name__)

def init_cv_splitter(cv_method=None, num_cv: int = 5, random_seed: int = None):
    if type(cv_method) == str:
        if cv_method == "kfold":
            splitter = KFold(n_splits=num_cv, shuffle=True, random_state=random_seed)
        elif cv_method == "time":
            splitter = TimeSeriesSplit(n_splits=num_cv)
        elif cv_method == "stratified_k":
            splitter = StratifiedKFold(n_splits=num_cv, random_state=random_seed)
    else:
        splitter = cv_method
    return splitter

def init_model(model=None, params={}, random_seed=None, tune_test=False):

    if model is None:
        logger.warning("NO MODEL PASSED IN")
        return

    random_state_flag = [True if "random_state" in pr else False for pr in params]
    random_state_flag = any(random_state_flag)

    if (
        model
        not in [
            "linear",
            "svr",
            "vc_clf",
            "vc_regress",
            "knn_clf",
            "knn_regress",
            "poisson",
        ]
        and ("random_state" not in params.keys() and random_state_flag is False)
    ):
        if tune_test:
            params["random_state"] = [random_seed]
        else:
            params["random_state"] = random_seed

    if model == "rf_clf":
        mod = RandomForestClassifier(**params)
    elif model == "et_clf":
        mod = ExtraTreesClassifier(**params)
    elif model == "gb_clf":
        mod = GradientBoostingClassifier(**params)
    elif model == "dt_clf":
        mod = DecisionTreeClassifier(**params)
    elif model == "logistic":
        mod = LogisticRegression(**params)
    elif model == "svc":
        mod = SVC(**params)
    elif model == "knn_clf":
        mod = KNeighborsClassifier(**params)
    elif model == "ada_clf":
        mod = AdaBoostClassifier(**params)
    elif model == "vc_clf":
        if "estimators" not in params.keys():
            params["estimators"] = [
                ("rf", RandomForestClassifier()),
                ("lr", LogisticRegression()),
            ]
        mod = VotingClassifier(**params)
    elif model == "xgb_clf":
        mod = XGBClassifier(**params)
    elif model == "rf_regress":
        mod = RandomForestRegressor(**params)
    elif model == "et_regress":
        mod = ExtraTreesRegressor(**params)
    elif model == "gb_regress":
        mod = GradientBoostingRegressor(**params)
    elif model == "dt_regress":
        mod = DecisionTreeRegressor(**params)
    elif model == "linear":
        mod = LinearRegression(**params)
    elif model == "lasso":
        mod = Lasso(**params)
    elif model == "ridge":
        mod = Ridge(**params)
    elif model == "elastic":
        mod = ElasticNet(**params)
    elif model == "poisson":
        mod = PoissonRegressor(**params)
    elif model == "svr":
        mod = SVR(**params)
    elif model == "knn_regress":
        mod = KNeighborsRegressor(**params)
    elif model == "ada_regress":
        mod = AdaBoostRegressor(**params)
    elif model == "vc_regress":
        if "estimators" not in params.keys():
            params["estimators"] = [
                ("rf", RandomForestRegressor()),
                ("linear", LinearRegression()),
            ]
        mod = VotingRegressor(**params)
    elif model == "xgb_regress":
        mod = XGBRegressor(**params)
    else:
        mod = model

    return mod

def build_pipes(
    mod=None,
    pipe=None,
    params: dict = None,
    imputer: str or object = None,
    scale: str or object = None,
    select_features: str or object = None,
    mod_type: str = "clf",
    num_features: int = None,
):

    # If pipeline passed in then add on the classifier
    # else init the pipeline with the model
    if pipe:
        pipe.steps.append((mod_type, mod))
        mod = pipe
    else:
        pipe = Pipeline(steps=[(mod_type, mod)])
        mod = pipe

    # inserts imputation method into the first position in the pipe
    if imputer:
        if type(imputer) == str:
            if "simple" in imputer:
                mod.steps.insert(
                    0, ("impute", SimpleImputer(strategy=imputer.split("_")[1]))
                )
            elif imputer == "missing_indicator":
                mod.steps.insert(0, ("impute", MissingIndicator()))
            elif imputer == "knn":
                mod.steps.insert(0, ("impute", KNNImputer()))
        else:
            mod.steps.insert(0, ("impute", imputer))

    # If scaling wanted adds the scaling
    if scale:
        if imputer:
            insert_pos = 1
        else:
            insert_pos = 0
        if type(scale) == str:
            if scale == "standard":
                mod.steps.insert(insert_pos, ("scale", StandardScaler()))
            elif scale == "minmax":
                mod.steps.insert(insert_pos, ("scale", MinMaxScaler()))
            elif scale == "robust":
                mod.steps.insert(insert_pos, ("scale", RobustScaler()))
        else:
            mod.steps.insert(insert_pos, ("scale", scale))

    # Appends the feature selection wanted
    # if wanted scaling then feature selection is second step (i.e. position 1) else first step (i.e. position 0)

    if select_features:

        if scale and imputer:
            insert_position = 2
        elif scale or imputer:
            insert_position = 1
        else:
            insert_position = 0

        if type(select_features) == str:
            if select_features not in ["eagles", "select_from_model", "selectkbest"]:
                warnings.warn(
                    "select_features not supported expects eagles or select_from_model got: "
                    + str(select_features)
                )

            if select_features == "eagles":
                mod.steps.insert(
                    insert_position,
                    (
                        "feature_selection",
                        EaglesFeatureSelection(
                            methods=["correlation", "regress"], problem_type=mod_type
                        ),
                    ),
                )
            elif select_features == "select_from_model":
                if mod_type == "clf":
                    mod.steps.insert(
                        insert_position,
                        (
                            "feature_selection",
                            SelectFromModel(
                                estimator=LogisticRegression(
                                    solver="liblinear", penalty="l1"
                                )
                            ),
                        ),
                    )
                elif mod_type == "rgr":
                    mod.steps.insert(
                        insert_position,
                        (
                            "feature_selection",
                            SelectFromModel(estimator=Lasso()),
                        ),
                    )
            elif select_features == "selectkbest":
                if num_features < 10:
                    k = np.ceil(num_features / 2).astype(int)
                else:
                    k = 10

                mod.steps.insert(
                    insert_position,
                    (
                        "feature_selection",
                        SelectKBest(k=k),
                    ),
                )
        else:
            mod.steps.insert(
                insert_position,
                select_features,
            )

    pipe_steps = ""
    for k in mod.named_steps.keys():
        pipe_steps = pipe_steps + type(mod.named_steps[k]).__name__ + ", "
    print("Final pipeline: " + pipe_steps)

    # Adjust the params for the model to make sure have appropriate prefix
    if len(params) > 0:
        param_prefix = mod_type + "__"

        params = {
            k if param_prefix in k else param_prefix + k: v for k, v in params.items()
        }
        return [mod, params]
    else:
        return [mod, params]

max_stars_repo_name: JFLandrigan/eagles

max_issues_repo_name: JFLandrigan/eagles

max_forks_repo_name: JFLandrigan/eagles


text:
import numba as nb
from numba import jit, prange, int_, int64, float64
import numpy as np

@jit('int_(int_, int_)', nopython=True, nogil=True, fastmath=True, cache=True)
def _get_rindex(col_dims, index):
    return int_(np.floor(float64(index) / float64(col_dims)))

@jit('int_(int_, int_, int_)', nopython=True, nogil=True, fastmath=True, cache=True)
def _get_cindex(col_dims, index, row_index):
    return int_(index - float64(col_dims) * row_index)

@jit('float64(float64, float64, float64)', nopython=True, nogil=True, fastmath=True, cache=True)
def edist(xloc, yloc, hw):
    return ((xloc - hw)**2 + (yloc - hw)**2)**0.5

@jit('float64(float64, float64, float64)', nopython=True, nogil=True, fastmath=True, cache=True)
def logistic_func(x, x0, r):
    denom = 1.0 + np.exp(-r * (x-x0))
    return 1.0 / denom if denom != 0 else 0.0

@jit('float64(int_, '
     'float64[:, :], '
     'float64, '
     'float64, '
     'float64)', nopython=True, nogil=True, boundscheck=False, cache=True)
def _apply(f__,
           results_,
           center_avg,
           nodata,
           dev_thresh):

    """
    Applies the least squares solution to estimate the missing value
    """

    y_mean = results_[f__, 0]
    beta1 = results_[f__, 0]
    beta2 = results_[f__, 0]
    alpha = results_[f__, 0]

    # y = b0*x + b1*x^2 + a
    estimate = beta1*center_avg + beta2*center_avg**2 + alpha

    if estimate > 1:
        estimate = nodata
    elif estimate < 0:
        estimate = nodata
    elif estimate > (y_mean + dev_thresh*10.0):
        estimate = y_mean
    elif estimate < (y_mean - dev_thresh*10.0):
        estimate = y_mean

    return estimate

@jit('void(int_, '
     'float64[:, :], '
     'float64[:, :], '
     'float64, '
     'float64, '
     'float64, '
     'int_, '
     'float64[:, :], '
     'float64[:, :])', nopython=True, nogil=True, boundscheck=False, cache=True)
def _calc_wlsq(f_,
               xdata,
               ydata,
               xavg1,
               xavg2,
               yavg,
               count,
               weights,
               lin_coeffs_):

    """
    Calculates the least squares slope and intercept
    """

    xvar1 = 0.0
    xvar2 = 0.0
    yvar = 0.0
    xx_cov = 0.0
    x1y_cov = 0.0
    x2y_cov = 0.0

    for s in range(0, count):

        xdev1 = weights[f_, s] * (xdata[f_, s] - xavg1)
        xdev2 = weights[f_, s] * (xdata[f_, s]**2 - xavg2)
        ydev = weights[f_, s] * (ydata[f_, s] - yavg)

        # Variance
        xvar1 += xdev1**2
        xvar2 += xdev2**2
        yvar += ydev**2

        x1y_cov += (xdev1 * ydev)
        x2y_cov += (xdev2 * ydev)

        xx_cov += (xdev1 * xdev2)

    denom = (xvar1 * xvar2) - xx_cov**2

    # Polynomial least squares
    beta1 = ((xvar2 * x1y_cov) - (xx_cov * x2y_cov)) / denom if denom != 0 else 0.0
    beta2 = ((xvar1 * x2y_cov) - (xx_cov * x1y_cov)) / denom if denom != 0 else 0.0
    alpha = yavg - (beta1 * xavg1) - (beta2 * xavg2)

    lin_coeffs_[f_, 0] = yavg
    lin_coeffs_[f_, 1] = beta1
    lin_coeffs_[f_, 2] = beta2
    lin_coeffs_[f_, 3] = alpha

@jit('void(float64[:, :, :, :], '
     'int64[:, :], '
     'int64[:], '
     'float64[:, :], '
     'int_, int_, int_, int_, int_, int_, int_, '
     'float64,'
     'float64[:, :])', nopython=True, nogil=True, boundscheck=False, cache=True)
def _get_center_mean(indata,
                     clusters,
                     lc_classes,
                     x_cluster_,
                     b,
                     fff,
                     iii,
                     jjj,
                     dims,
                     ci,
                     hw,
                     nodata,
                     count_vector):

    """
    Gets the center average of the reference data
    """

    offset = hw - int64(ci / 2.0)

    # sum
    count_vector[fff, 0] = 0.0

    # average
    count_vector[fff, 1] = 0.0

    for d in range(0, dims):

        for m in range(0, ci):
            for n in range(0, ci):

                # Reference X value
                ref_xvalue = indata[d, b, iii+m+offset, jjj+n+offset]

                # Reference X class
                ref_xclass = clusters[iii+m+offset, jjj+n+offset]

                # Update the class count
                # [X mean, X class, <class counts>]
                x_cluster_[fff, 2+ref_xclass] += 1.0

                if ref_xvalue != nodata:

                    count_vector[fff, 0] += ref_xvalue
                    count_vector[fff, 1] += 1

        # Break if there is data
        if count_vector[fff, 0] > 0:
            break

    if count_vector[fff, 0] > 0:

        ref_xclass = lc_classes[0]
        ref_xclass_mode = x_cluster_[fff, 2]

        # Get the class with the highest count
        for cidx in range(3, lc_classes.shape[0]):

            if x_cluster_[fff, cidx] > ref_xclass_mode:
                ref_xclass = lc_classes[cidx-2]
                ref_xclass_mode = x_cluster_[fff, cidx]

        x_cluster_[fff, 0] = count_vector[fff, 0] / count_vector[fff, 1]
        x_cluster_[fff, 1] = float64(ref_xclass)

        # Get the updated average for the target class

        count_vector[fff, 0] = 0.0
        count_vector[fff, 1] = 0.0

        for d in range(0, dims):
            for m in range(0, ci):
                for n in range(0, ci):

                    # Reference X value
                    ref_xvalue = indata[d, b, iii+m+offset, jjj+n+offset]

                    # Reference X class
                    ref_xclass_spass = clusters[iii+m+offset, jjj+n+offset]

                    if (ref_xvalue != nodata) and (ref_xclass_spass == ref_xclass):

                        count_vector[fff, 0] += ref_xvalue
                        count_vector[fff, 1] += 1

            # Break if there is data
            if count_vector[fff, 0] > 0:
                break

        if count_vector[fff, 1] > 0:
            x_cluster_[fff, 0] = count_vector[fff, 0] / count_vector[fff, 1]

@jit('void(float64[:, :, :, :], '
     'int64[:, :], '
     'float64[:, :, :], '
     'int_, int_, int_, int_, int_, int_, int_, int_, int_, '
     'float64, '
     'float64, '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :])', nopython=True, nogil=True, boundscheck=False, cache=True)
def _estimate_gap(indata,
                  clusters,
                  output__,
                  b,
                  ffff,
                  iiii,
                  jjjj,
                  dims,
                  wmin,
                  wmax,
                  hw,
                  min_thresh,
                  nodata,
                  dev_thresh,
                  x_cluster__,
                  xdata,
                  ydata,
                  weights,
                  spatial_weights,
                  lin_coeffs):

    count = 0
    yavg = 0.0

    max_vct_size = int64(wmax*wmax*dims)

    tar_xclass = int64(x_cluster__[ffff, 1])

    for rzi in range(0, max_vct_size):
        xdata[ffff, rzi] = 0.0
        ydata[ffff, rzi] = 0.0

    xavg1 = 0.0
    xavg2 = 0.0

    # Search for data over increasing windows
    for wi in range(wmin, wmax+2, 2):

        # Adjustment for different sized windows
        offset = hw - int64(wi / 2.0)

        # Iterate over each reference file to fill the window
        for d in range(1, dims):

            for m in range(0, wi):
                for n in range(0, wi):

                    # Skip the center that has already been checked
                    if wi > wmin:

                        if (m > 0) and (m < wi-1) and (n > 0) and (n < wi-1):
                            continue

                    # Get the reference cluster value
                    ref_xclass = clusters[iiii+m+offset, jjjj+n+offset]

                    # Only take samples with the same cluster as the target
                    if ref_xclass == tar_xclass:

                        # Target value
                        yvalue = indata[0, b, iiii+m+offset, jjjj+n+offset]

                        # Reference value
                        xvalue = indata[d, b, iiii+m+offset, jjjj+n+offset]

                        # X and y must have data
                        if (xvalue != nodata) and (yvalue != nodata):

                            # Squared difference
                            sq_diff = (xvalue - yvalue)**2

                            # X and y must have a low deviation
                            if sq_diff <= dev_thresh:

                                xdata[ffff, count] = xvalue
                                ydata[ffff, count] = yvalue
                                weights[ffff, count] = spatial_weights[m, n]

                                xavg1 += xvalue
                                xavg2 += xvalue**2
                                yavg += yvalue

                                count += 1

            if count >= min_thresh:
                break

        if count >= min_thresh:
            break

    if count >= min_thresh:

        countf = float64(count)

        # Window average
        xavg1 /= countf
        xavg2 /= countf
        yavg /= countf

        # Std. dev. of [x, y], slope, intercept
        _calc_wlsq(ffff,
                   xdata,
                   ydata,
                   xavg1,
                   xavg2,
                   yavg,
                   count,
                   weights,
                   lin_coeffs)

        # Calculate the least squares solution
        estimate = _apply(ffff,
                          lin_coeffs,
                          x_cluster__[ffff, 0],
                          nodata,
                          dev_thresh)

        output__[b, iiii+hw, jjjj+hw] = estimate

@jit('void(int_, '
     'float64[:, :, :, :], '
     'int64[:, :], '
     'int64[:], '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :, :], '
     'int_, int_, int_, int_, int_, int_, int_, '
     'float64, '
     'float64, '
     'float64[:, :], '
     'float64[:, :], '
     'float64[:, :])', nopython=True, nogil=True, boundscheck=False, cache=True)
def _fill_gap(ff,
              indata,
              clusters,
              lc_classes,
              x_cluster,
              xdata,
              ydata,
              weights,
              output_,
              b,
              col_dims,
              wmin,
              wmax,
              hw,
              dims,
              min_count,
              nodata,
              dev_thresh,
              spatial_weights,
              count_vector,
              lin_coeffs):

    ii = _get_rindex(col_dims, ff)
    jj = _get_cindex(col_dims, ff, ii)

    # Center target sample
    if indata[0, b, ii+hw, jj+hw] == nodata:

        for xxi in range(0, 2+lc_classes.shape[0]):
            x_cluster[ff, xxi] = 0.0

        # Get an average of the center value
        for ci in range(3, 9, 2):

            # center_avg = X
            _get_center_mean(indata, clusters, lc_classes, x_cluster, b, ff, ii, jj, dims, ci, hw, nodata, count_vector)

            if x_cluster[ff, 0] > 0:
                break

        if x_cluster[ff, 0] > 0:

            _estimate_gap(indata,
                          clusters,
                          output_,
                          b,
                          ff,
                          ii,
                          jj,
                          dims,
                          wmin,
                          wmax,
                          hw,
                          min_count,
                          nodata,
                          dev_thresh,
                          x_cluster,
                          xdata,
                          ydata,
                          weights,
                          spatial_weights,
                          lin_coeffs)

@jit('float64[:, :, :](float64[:, :, :, :], '
     'int64[:, :], '
     'int_, '
     'int_, '
     'float64, '
     'int_, '
     'float64, '
     'int_)', parallel=True, cache=True)
def _fill_gaps(indata,
               clusters,
               wmax,
               wmin,
               nodata,
               min_count,
               dev_thresh,
               n_jobs):

    eps = 1e-08

    dims = indata.shape[0]
    bands = indata.shape[1]
    rows = indata.shape[2]
    cols = indata.shape[3]
    hw = int64(wmax / 2.0)
    row_dims = rows - int64(hw*2.0)
    col_dims = cols - int64(hw*2.0)
    nsamples = int64(row_dims * col_dims)

    num_threads = nb.set_num_threads(n_jobs)

    # Get the maximum distance for the spatial weights
    max_edist = edist(0.0, 0.0, hw) + eps

    max_vct_size = int64(wmax*wmax*dims)

    # Initialize arrays
    output = indata[0].copy()
    lc_classes = np.int64(np.unique(clusters))
    x_cluster = np.zeros((nsamples, 2+lc_classes.shape[0]), dtype='float64')
    xdata = np.zeros((nsamples, max_vct_size), dtype='float64')
    ydata = np.zeros((nsamples, max_vct_size), dtype='float64')
    weights = np.zeros((nsamples, max_vct_size), dtype='float64')
    lin_coeffs = np.zeros((nsamples, 4), dtype='float64')
    spatial_weights = np.zeros((wmax, wmax), dtype='float64')
    count_vector = np.zeros((nsamples, 2), dtype='float64')

    # Fill the spatial weights
    for mm in range(0, wmax):
        for nn in range(0, wmax):
            spatial_weights[mm, nn] = logistic_func(1.0 - ((edist(float64(nn), float64(mm), hw)+eps) / max_edist), 0.5, 10.0)

    for f in prange(0, nsamples):

        for b in range(0, bands):

            _fill_gap(f,
                      indata,
                      clusters,
                      lc_classes,
                      x_cluster,
                      xdata,
                      ydata,
                      weights,
                      output,
                      b,
                      col_dims,
                      wmin,
                      wmax,
                      hw,
                      dims,
                      min_count,
                      nodata,
                      dev_thresh,
                      spatial_weights,
                      count_vector,
                      lin_coeffs)

    return output

@jit('float64[:, :, :](float64[:, :, :, :], '
     'int64[:, :], '
     'float64[:], '
     'int_, '
     'int_, '
     'float64, '
     'int_, '
     'float64, '
     'int_, '
     'int_)', cache=True)
def fill_gaps(gap_data,
              cluster_data,
              days,
              wmax=25,
              wmin=9,
              nodata=0,
              min_count=20,
              dev_thresh=0.02,
              n_jobs=1,
              chunksize=1):

    """
    Fills data gaps using spatial-temporal weighted least squares linear regression

    Args:
        gap_data (4d array): Layers x bands x rows x columns. The first layer is the target and the remaining layers
            are the references. The reference layers should be sorted from the date closest to the target to the
            date furthest from the target date.
        cluster_data (2d array): Land cover clusters.
        wmax (Optional[int]): The maximum window size.
        wmin (Optional[int]): The minimum window size.
        nodata (Optional[int]): The 'no data' value to fill.
        min_count (Optional[int]): The minimum required number of samples to fit a window.
        dev_thresh (Optional[float]): The squared deviation threshold for similar values.
        n_jobs (Optional[int]): The number of bands to process in parallel.
        chunksize (Optional[int]): The parallel thread chunksize.

    Returns:
        3d ``numpy.ndarray`` (bands x rows x columns)
    """

    return np.float64(_fill_gaps(gap_data,
                                 cluster_data,
                                 wmax,
                                 wmin,
                                 float64(nodata),
                                 min_count,
                                 dev_thresh,
                                 n_jobs))

max_stars_repo_name: jgrss/satsmooth

max_issues_repo_name: jgrss/satsmooth

max_forks_repo_name: jgrss/satsmooth

reason: import numpy,import numba,from numba


text:
from collections import Counter
from collections import defaultdict
from functools import lru_cache
from functools import reduce

import numpy as np

from ..convo import Conversation
from .harmonic import mixing
from .harmonic import novelty
from .params import CACHE_SIZE
from .post import type_frequency_distribution as post_freq
from .post_in_conv import PostInConvoFeatures as PICF
from .post_in_conv import agg_post_stats
from .post_in_conv import avg_token_entropy_conv

class UserInConvoFeatures:
    """
    Container for feature extraction on users situated within a single conversation.
    """

    @staticmethod
    def bools(user, convo):
        return {
            'is_source_author': is_source_author(user, convo),
        }

    @staticmethod
    def categoricals(user, convo):
        return {

        }

    @staticmethod
    def counter(user, convo):
        return {
            'type_freq': type_frequency_distribution(user, convo),
        }

    @staticmethod
    def floats(user, convo):
        out = mixing_features(user, convo)
        out['avg_user_token_entropy'] = avg_user_token_entropy(user, convo)

        for k, v in agg_post_stats(convo, filter_by=lambda p: p.author == user).items():
            out[k] = v

        return out

    @staticmethod
    def ints(user, convo):
        out = {
            'message_count': messages_by_user(user, convo),
            'types':         len(type_frequency_distribution(user, convo)),
        }

        for k, v in sum_post_bools_by_user(user, convo).items():
            out[k] = v

        for k, v in sum_post_ints_by_user(user, convo).items():
            out[k] = v

        return out

    @staticmethod
    def strs(user, convo):
        return {}

def iter_over_users(convo):
    """
    Creates an iterator over the unique users within a conversation

    Parameters
    ----------
    convo : Conversation

    Yields
    -------
    str
        A user
    """
    users = set()
    for pid in convo.posts:
        user = convo.posts[pid].author

        if user in users:
            continue

        yield user

        users.add(user)

@lru_cache(maxsize=CACHE_SIZE)
def is_source_author(user, convo):
    """
    Returns if this user created a source message
    for the conversation

    Parameters
    ----------
    user : str
    convo : Conversation

    Returns
    -------
    bool
    """
    auths = {convo.posts[pid].author for pid in convo.get_sources()}
    return user in auths

@lru_cache(maxsize=CACHE_SIZE)
def messages_per_user(conv):
    """
    Returns the user-distribution of posts written per user

    Parameters
    ----------
    conv : Conversation
        A collection of posts

    Returns
    -------
    Counter
        The counts of messages written per user (keyed by author name)
    """
    return Counter([p.author for p in conv.posts.values()])

@lru_cache(maxsize=CACHE_SIZE)
def messages_by_user(user, conv):
    """
    Returns number of messages created by this user
    Parameters
    ----------
    user : str
    conv : Conversation

    Returns
    -------
    int
    """
    return messages_per_user(conv)[user]

@lru_cache(maxsize=CACHE_SIZE)
def get_user_posts(user, conv):
    """
    Filters to just this users messages within the conversation

    Parameters
    ----------
    user : str
    conv : Conversation

    Returns
    -------
    Conversation
    """
    if user is None:
        return Conversation()
    return Conversation(posts={pid: conv.posts[pid] for pid in conv.filter(by_author=user)},
                        convo_id=f'{conv.convo_id}-{user}')

@lru_cache(maxsize=CACHE_SIZE)
def type_frequency_distribution(user, convo):
    """
    Returns the type frequency (unigram) distribution for the user's posts in the convo.

    Parameters
    ----------
    user : str
    convo : Conversation

    Returns
    -------
    collections.Counter
    """
    return reduce(lambda x, y: x + y, map(post_freq, get_user_posts(user, convo).posts.values()), Counter())

@lru_cache(maxsize=CACHE_SIZE)
def mixing_features(user, convo):
    """
    Returns the measured parameters using the harmonic mixing law.

    Parameters
    ----------
    user : str
    convo : Conversation

    Returns
    -------
    dict(str, float)
    """
    freq = type_frequency_distribution(user, convo)
    if len(freq) == 0:
        return {
            'k1':      float(0),
            'theta':   float(0),
            'entropy': float(0),
            'N_avg':   float(0),
            'M_avg':   float(0),
        }

    return mixing(freq)

@lru_cache(maxsize=CACHE_SIZE)
def novelty_vector(user, convo):
    """
    Returns the novelty vector measured from the convo text, filtered to just the user.

    Parameters
    ----------
    user : str
    convo : Conversation

    Returns
    -------
    np.array
    """
    freq = type_frequency_distribution(user, convo)
    if len(freq) == 0:
        return []

    return novelty(freq)

@lru_cache(maxsize=CACHE_SIZE)
def avg_user_token_entropy(user, convo):
    """
    Returns the average token entropy when comparing a user
    to a conversation (containing said user's posts)

    Parameters
    ----------
    user : str
    convo : Conversation

    Returns
    -------
    float
    """
    return float(avg_token_entropy_conv(get_user_posts(user, convo), convo))

def agg_user_stats(convo):
    """
    Computes a set of aggregate user statistical measures.
    This is only computed for the integer and float subsets.
    Specifically, the following stats are measured:
    min, max, mean, median, standard deviation

    Parameters
    ----------
    convo : Conversation

    Returns
    -------
    dict(str, dict(str, float))
    """
    agg = defaultdict(list)
    fs = [UserInConvoFeatures.floats, UserInConvoFeatures.ints]
    for user in iter_over_users(convo):
        for f in fs:
            for k, v in f(user, convo).items():
                agg[k].append(v)

    out = {}
    for k, vs in agg.items():
        out[f'user_min_{k}'] = float(np.nanmin(vs))
        out[f'user_max_{k}'] = float(np.nanmax(vs))
        out[f'user_mean_{k}'] = float(np.nanmean(vs))
        out[f'user_median_{k}'] = float(np.median(vs))
        out[f'user_std_{k}'] = float(np.nanstd(vs) if len(vs) > 1 else 1)

    return out

def sum_post_bools_by_user(user, convo):
    cnt = Counter()

    for p in convo.posts.values():
        if p.author != user:
            continue

        for k, v in PICF.bools(p, convo).items():
            kx = k.replace('is_', '') + '_count'
            cnt[kx] += 1 if v else 0

    return dict(cnt)

def sum_post_ints_by_user(user, convo):
    skipset = {
        'type_count',  # must be aggregated in set theoretic way
        'depth', 'width',  # nonsensical accumulation stats
    }
    cnt = Counter()
    for p in convo.posts.values():
        if p.author != user:
            continue

        for k, v in PICF.ints(p, convo).items():
            if k in skipset:
                continue

            cnt[k] += v

    return dict(cnt)

max_stars_repo_name: hunter-heidenreich/pyconversations

max_issues_repo_name: hunter-heidenreich/pyconversations

max_forks_repo_name: hunter-heidenreich/pyconversations