<commit_before>import json
import logging
import re

import elasticsearch_dsl
import requests
from collections import defaultdict

import settings

from django.contrib.admin.views.decorators import staff_member_required
from django.shortcuts import render
import elasticsearch

from seqr.models import Dataset
from settings import LOGIN_URL
from pprint import pprint

logger = logging.getLogger(__name__)

OPERATIONS_LOG = "index_operations_log"

@staff_member_required(login_url=LOGIN_URL)
def elasticsearch_status(request):
    client = elasticsearch.Elasticsearch(host=settings.ELASTICSEARCH_SERVICE_HOSTNAME)

    # look up seqr dataset records
    index_to_dataset = {}
    for dataset in Dataset.objects.all():
        if dataset.dataset_id:
            index_name = dataset.dataset_id
            index_to_dataset[index_name] = dataset

    # get index snapshots
    response = requests.get("http://{0}:{1}/_snapshot/{2}/_all".format(
        settings.ELASTICSEARCH_SERVICE_HOSTNAME, settings.ELASTICSEARCH_PORT, "callsets"))
    snapshots = json.loads(response.content)

    index_snapshot_states = defaultdict(list)
    for snapshot in snapshots["snapshots"]:
        for index_name in snapshot["indices"]:
            index_snapshot_states[index_name].append(snapshot["state"])

    # get indices
    indices = []
    for index in client.cat.indices(format="json", h="*"):
        index_name = index['index']

        # skip special indices
        if index_name in ['.kibana', 'index_operations_log']:
            continue

        index_json = {k.replace('.', '_'): v for k, v in index.items()}

        index_name = re.sub("_[0-9]{1,2}$", "", index_name)
        if index_name in index_to_dataset:
            index_json['project_guid'] = index_to_dataset[index_name].project.guid
            index_json['project_id'] = index_to_dataset[index_name].project.deprecated_project_id
            index_json['analysis_type'] = index_to_dataset[index_name].analysis_type
            index_json['genome_version'] = index_to_dataset[index_name].genome_version
            index_json['source_file_path'] = index_to_dataset[index_name].source_file_path

        if index_name in index_snapshot_states:
            index_json['snapshots'] = ", ".join(set(index_snapshot_states[index_name]))
        indices.append(index_json)

    # get operations log
    s = elasticsearch_dsl.Search(using=client, index=OPERATIONS_LOG)
    s = s.params(size=5000)
    operations = [doc.to_dict() for doc in s.execute().hits]

    #making a new list since dots in es client keys are confusing template
    disk_status=[]
    for disk in client.cat.allocation(format="json"):
        disk_json = {k.replace('.', '_'): v for k, v in disk.items()}
        disk_status.append(
                            {
                              'node_name':disk_json['node'],
                              'disk_available':disk_json['disk_avail'],
                              'disk_used':disk_json['disk_used'],
                              'disk_percent_used':disk_json['disk_percent'],
                                })
    return render(request, "staff/elasticsearch_status.html", {
        'indices': indices,
        'operations': operations,
        'disk_stats': disk_status,
        'elasticsearch_host': settings.ELASTICSEARCH_SERVER,
    })

<commit_msg>spacing<commit_after>import json
import logging
import re

import elasticsearch_dsl
import requests
from collections import defaultdict

import settings

from django.contrib.admin.views.decorators import staff_member_required
from django.shortcuts import render
import elasticsearch

from seqr.models import Dataset
from settings import LOGIN_URL
from pprint import pprint

logger = logging.getLogger(__name__)

OPERATIONS_LOG = "index_operations_log"

@staff_member_required(login_url=LOGIN_URL)
def elasticsearch_status(request):
    client = elasticsearch.Elasticsearch(host=settings.ELASTICSEARCH_SERVICE_HOSTNAME)

    # look up seqr dataset records
    index_to_dataset = {}
    for dataset in Dataset.objects.all():
        if dataset.dataset_id:
            index_name = dataset.dataset_id
            index_to_dataset[index_name] = dataset

    # get index snapshots
    response = requests.get("http://{0}:{1}/_snapshot/{2}/_all".format(
        settings.ELASTICSEARCH_SERVICE_HOSTNAME, settings.ELASTICSEARCH_PORT, "callsets"))
    snapshots = json.loads(response.content)

    index_snapshot_states = defaultdict(list)
    for snapshot in snapshots["snapshots"]:
        for index_name in snapshot["indices"]:
            index_snapshot_states[index_name].append(snapshot["state"])

    # get indices
    indices = []
    for index in client.cat.indices(format="json", h="*"):
        index_name = index['index']

        # skip special indices
        if index_name in ['.kibana', 'index_operations_log']:
            continue

        index_json = {k.replace('.', '_'): v for k, v in index.items()}

        index_name = re.sub("_[0-9]{1,2}$", "", index_name)
        if index_name in index_to_dataset:
            index_json['project_guid'] = index_to_dataset[index_name].project.guid
            index_json['project_id'] = index_to_dataset[index_name].project.deprecated_project_id
            index_json['analysis_type'] = index_to_dataset[index_name].analysis_type
            index_json['genome_version'] = index_to_dataset[index_name].genome_version
            index_json['source_file_path'] = index_to_dataset[index_name].source_file_path

        if index_name in index_snapshot_states:
            index_json['snapshots'] = ", ".join(set(index_snapshot_states[index_name]))
        indices.append(index_json)

    # get operations log
    s = elasticsearch_dsl.Search(using=client, index=OPERATIONS_LOG)
    s = s.params(size=5000)
    operations = [doc.to_dict() for doc in s.execute().hits]

    #making a new list since dots in es client keys are confusing template
    disk_status=[]
    for disk in client.cat.allocation(format="json"):
        disk_json = {k.replace('.', '_'): v for k, v in disk.items()}
        disk_status.append({
            'node_name': disk_json['node'],
            'disk_available': disk_json['disk_avail'],
            'disk_used': disk_json['disk_used'],
            'disk_percent_used': disk_json['disk_percent'],
        })

    return render(request, "staff/elasticsearch_status.html", {
        'indices': indices,
        'operations': operations,
        'disk_stats': disk_status,
        'elasticsearch_host': settings.ELASTICSEARCH_SERVER,
    })

<|endoftext|>


<commit_before>import numpy as np
from scipy.sparse import coo_matrix, issparse

from scvelo.core import l2_norm
from .. import logging as logg
from .. import settings
from ..preprocessing.moments import get_moments
from ..preprocessing.neighbors import (
    get_n_neighs,
    get_neighs,
    neighbors,
    pca,
    verify_neighbors,
)
from .utils import cosine_correlation, get_indices, get_iterative_indices
from .velocity import velocity

def vals_to_csr(vals, rows, cols, shape, split_negative=False):
    graph = coo_matrix((vals, (rows, cols)), shape=shape)

    if split_negative:
        graph_neg = graph.copy()

        graph.data = np.clip(graph.data, 0, 1)
        graph_neg.data = np.clip(graph_neg.data, -1, 0)

        graph.eliminate_zeros()
        graph_neg.eliminate_zeros()

        return graph.tocsr(), graph_neg.tocsr()

    else:
        return graph.tocsr()

class VelocityGraph:
    def __init__(
        self,
        adata,
        vkey="velocity",
        xkey="Ms",
        tkey=None,
        basis=None,
        n_neighbors=None,
        sqrt_transform=None,
        n_recurse_neighbors=None,
        random_neighbors_at_max=None,
        gene_subset=None,
        approx=None,
        report=False,
        compute_uncertainties=None,
        mode_neighbors="distances",
    ):

        subset = np.ones(adata.n_vars, bool)
        if gene_subset is not None:
            var_names_subset = adata.var_names.isin(gene_subset)
            subset &= var_names_subset if len(var_names_subset) > 0 else gene_subset
        elif f"{vkey}_genes" in adata.var.keys():
            subset &= np.array(adata.var[f"{vkey}_genes"].values, dtype=bool)

        xkey = xkey if xkey in adata.layers.keys() else "spliced"

        X = np.array(
            adata.layers[xkey].A[:, subset]
            if issparse(adata.layers[xkey])
            else adata.layers[xkey][:, subset]
        )
        V = np.array(
            adata.layers[vkey].A[:, subset]
            if issparse(adata.layers[vkey])
            else adata.layers[vkey][:, subset]
        )

        nans = np.isnan(np.sum(V, axis=0))
        if np.any(nans):
            X = X[:, ~nans]
            V = V[:, ~nans]

        if approx is True and X.shape[1] > 100:
            X_pca, PCs, _, _ = pca(X, n_comps=30, svd_solver="arpack", return_info=True)
            self.X = np.array(X_pca, dtype=np.float32)
            self.V = (V - V.mean(0)).dot(PCs.T)
            self.V[V.sum(1) == 0] = 0
        else:
            self.X = np.array(X, dtype=np.float32)
            self.V = np.array(V, dtype=np.float32)
        self.V_raw = np.array(self.V)

        self.sqrt_transform = sqrt_transform
        if self.sqrt_transform is None and f"{vkey}_params" in adata.uns.keys():
            self.sqrt_transform = adata.uns[f"{vkey}_params"]["mode"] == "stochastic"
        if self.sqrt_transform:
            self.V = np.sqrt(np.abs(self.V)) * np.sign(self.V)
        self.V -= np.nanmean(self.V, axis=1)[:, None]

        self.n_recurse_neighbors = n_recurse_neighbors
        if self.n_recurse_neighbors is None:
            if n_neighbors is not None or mode_neighbors == "connectivities":
                self.n_recurse_neighbors = 1
            else:
                self.n_recurse_neighbors = 2

        if "neighbors" not in adata.uns.keys():
            neighbors(adata)
        if np.min((get_neighs(adata, "distances") > 0).sum(1).A1) == 0:
            raise ValueError(
                "Your neighbor graph seems to be corrupted. "
                "Consider recomputing via pp.neighbors."
            )
        if n_neighbors is None or n_neighbors <= get_n_neighs(adata):
            self.indices = get_indices(
                dist=get_neighs(adata, "distances"),
                n_neighbors=n_neighbors,
                mode_neighbors=mode_neighbors,
            )[0]
        else:
            if basis is None:
                basis_keys = ["X_pca", "X_tsne", "X_umap"]
                basis = [key for key in basis_keys if key in adata.obsm.keys()][-1]
            elif f"X_{basis}" in adata.obsm.keys():
                basis = f"X_{basis}"

            if isinstance(approx, str) and approx in adata.obsm.keys():
                from sklearn.neighbors import NearestNeighbors

                neighs = NearestNeighbors(n_neighbors=n_neighbors + 1)
                neighs.fit(adata.obsm[approx])
                self.indices = neighs.kneighbors_graph(
                    mode="connectivity"
                ).indices.reshape((-1, n_neighbors + 1))
            else:
                from .. import Neighbors

                neighs = Neighbors(adata)
                neighs.compute_neighbors(
                    n_neighbors=n_neighbors, use_rep=basis, n_pcs=10
                )
                self.indices = get_indices(
                    dist=neighs.distances, mode_neighbors=mode_neighbors
                )[0]

        self.max_neighs = random_neighbors_at_max

        gkey, gkey_ = f"{vkey}_graph", f"{vkey}_graph_neg"
        self.graph = adata.uns[gkey] if gkey in adata.uns.keys() else []
        self.graph_neg = adata.uns[gkey_] if gkey_ in adata.uns.keys() else []

        if tkey in adata.obs.keys():
            self.t0 = adata.obs[tkey].copy()
            init = min(self.t0) if isinstance(min(self.t0), int) else 0
            self.t0.cat.categories = np.arange(init, len(self.t0.cat.categories))
            self.t1 = self.t0.copy()
            self.t1.cat.categories = self.t0.cat.categories + 1
        else:
            self.t0 = None

        self.compute_uncertainties = compute_uncertainties
        self.uncertainties = None
        self.self_prob = None
        self.report = report
        self.adata = adata

    def compute_cosines(self):
        vals, rows, cols, uncertainties, n_obs = [], [], [], [], self.X.shape[0]
        progress = logg.ProgressReporter(n_obs)

        if self.compute_uncertainties:
            m = get_moments(self.adata, np.sign(self.V_raw), second_order=True)

        for i in range(n_obs):
            if self.V[i].max() != 0 or self.V[i].min() != 0:
                neighs_idx = get_iterative_indices(
                    self.indices, i, self.n_recurse_neighbors, self.max_neighs
                )

                if self.t0 is not None:
                    t0, t1 = self.t0[i], self.t1[i]
                    if t0 >= 0 and t1 > 0:
                        t1_idx = np.where(self.t0 == t1)[0]
                        if len(t1_idx) > len(neighs_idx):
                            t1_idx = np.random.choice(
                                t1_idx, len(neighs_idx), replace=False
                            )
                        if len(t1_idx) > 0:
                            neighs_idx = np.unique(np.concatenate([neighs_idx, t1_idx]))

                dX = self.X[neighs_idx] - self.X[i, None]  # 60% of runtime
                if self.sqrt_transform:
                    dX = np.sqrt(np.abs(dX)) * np.sign(dX)
                val = cosine_correlation(dX, self.V[i])  # 40% of runtime

                if self.compute_uncertainties:
                    dX /= l2_norm(dX)[:, None]
                    uncertainties.extend(np.nansum(dX ** 2 * m[i][None, :], 1))

                vals.extend(val)
                rows.extend(np.ones(len(neighs_idx)) * i)
                cols.extend(neighs_idx)
                if self.report:
                    progress.update()
        if self.report:
            progress.finish()

        vals = np.hstack(vals)
        vals[np.isnan(vals)] = 0

        self.graph, self.graph_neg = vals_to_csr(
            vals, rows, cols, shape=(n_obs, n_obs), split_negative=True
        )
        if self.compute_uncertainties:
            uncertainties = np.hstack(uncertainties)
            uncertainties[np.isnan(uncertainties)] = 0
            self.uncertainties = vals_to_csr(
                uncertainties, rows, cols, shape=(n_obs, n_obs), split_negative=False
            )
            self.uncertainties.eliminate_zeros()

        confidence = self.graph.max(1).A.flatten()
        self.self_prob = np.clip(np.percentile(confidence, 98) - confidence, 0, 1)

def velocity_graph(
    data,
    vkey="velocity",
    xkey="Ms",
    tkey=None,
    basis=None,
    n_neighbors=None,
    n_recurse_neighbors=None,
    random_neighbors_at_max=None,
    sqrt_transform=None,
    variance_stabilization=None,
    gene_subset=None,
    compute_uncertainties=None,
    approx=None,
    mode_neighbors="distances",
    copy=False,
):
    """Computes velocity graph based on cosine similarities.

    The cosine similarities are computed between velocities and potential cell state
    transitions, i.e. it measures how well a corresponding change in gene expression
    :math:`\\delta_{ij} = x_j - x_i` matches the predicted change according to the
    velocity vector :math:`\\nu_i`,

    .. math::
        \\pi_{ij} = \\cos\\angle(\\delta_{ij}, \\nu_i)
        = \\frac{\\delta_{ij}^T \\nu_i}{\\left\\lVert\\delta_{ij}\\right\\rVert
        \\left\\lVert \\nu_i \\right\\rVert}.

    Arguments
    ---------
    data: :class:`~anndata.AnnData`
        Annotated data matrix.
    vkey: `str` (default: `'velocity'`)
        Name of velocity estimates to be used.
    xkey: `str` (default: `'Ms'`)
        Layer key to extract count data from.
    tkey: `str` (default: `None`)
        Observation key to extract time data from.
    basis: `str` (default: `None`)
        Basis / Embedding to use.
    n_neighbors: `int` or `None` (default: None)
        Use fixed number of neighbors or do recursive neighbor search (if `None`).
    n_recurse_neighbors: `int` (default: `None`)
        Number of recursions for neighbors search. Defaults to
        2 if mode_neighbors is 'distances', and 1 if mode_neighbors is 'connectivities'.
    random_neighbors_at_max: `int` or `None` (default: `None`)
        If number of iterative neighbors for an individual cell is higher than this
        threshold, a random selection of such are chosen as reference neighbors.
    sqrt_transform: `bool` (default: `False`)
        Whether to variance-transform the cell states changes
        and velocities before computing cosine similarities.
    gene_subset: `list` of `str`, subset of adata.var_names or `None`(default: `None`)
        Subset of genes to compute velocity graph on exclusively.
    compute_uncertainties: `bool` (default: `None`)
        Whether to compute uncertainties along with cosine correlation.
    approx: `bool` or `None` (default: `None`)
        If True, first 30 pc's are used instead of the full count matrix
    mode_neighbors: 'str' (default: `'distances'`)
        Determines the type of KNN graph used. Options are 'distances' or
        'connectivities'. The latter yields a symmetric graph.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to adata.

    Returns
    -------
    Returns or updates `adata` with the attributes
    velocity_graph: `.uns`
        sparse matrix with transition probabilities
    """
    adata = data.copy() if copy else data
    verify_neighbors(adata)
    if vkey not in adata.layers.keys():
        velocity(adata, vkey=vkey)
    if sqrt_transform is None:
        sqrt_transform = variance_stabilization

    vgraph = VelocityGraph(
        adata,
        vkey=vkey,
        xkey=xkey,
        tkey=tkey,
        basis=basis,
        n_neighbors=n_neighbors,
        approx=approx,
        n_recurse_neighbors=n_recurse_neighbors,
        random_neighbors_at_max=random_neighbors_at_max,
        sqrt_transform=sqrt_transform,
        gene_subset=gene_subset,
        compute_uncertainties=compute_uncertainties,
        report=True,
        mode_neighbors=mode_neighbors,
    )

    if isinstance(basis, str):
        logg.warn(
            f"The velocity graph is computed on {basis} embedding coordinates.\n"
            f"        Consider computing the graph in an unbiased manner \n"
            f"        on full expression space by not specifying basis.\n"
        )

    logg.info("computing velocity graph", r=True)
    vgraph.compute_cosines()

    adata.uns[f"{vkey}_graph"] = vgraph.graph
    adata.uns[f"{vkey}_graph_neg"] = vgraph.graph_neg

    if vgraph.uncertainties is not None:
        adata.uns[f"{vkey}_graph_uncertainties"] = vgraph.uncertainties

    adata.obs[f"{vkey}_self_transition"] = vgraph.self_prob

    if f"{vkey}_params" in adata.uns.keys():
        if "embeddings" in adata.uns[f"{vkey}_params"]:
            del adata.uns[f"{vkey}_params"]["embeddings"]
    else:
        adata.uns[f"{vkey}_params"] = {}
    adata.uns[f"{vkey}_params"]["mode_neighbors"] = mode_neighbors
    adata.uns[f"{vkey}_params"]["n_recurse_neighbors"] = vgraph.n_recurse_neighbors

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint(
        "added \n"
        f"    '{vkey}_graph', sparse matrix with cosine correlations (adata.uns)"
    )

    return adata if copy else None
<commit_msg>Parallelize cosine similarity computation (#392)<commit_after>import os

import numpy as np
from scipy.sparse import coo_matrix, issparse

from scvelo import logging as logg
from scvelo import settings
from scvelo.core import get_n_jobs, l2_norm, parallelize
from scvelo.preprocessing.moments import get_moments
from scvelo.preprocessing.neighbors import (
    get_n_neighs,
    get_neighs,
    neighbors,
    pca,
    verify_neighbors,
)
from .utils import cosine_correlation, get_indices, get_iterative_indices
from .velocity import velocity

def vals_to_csr(vals, rows, cols, shape, split_negative=False):
    graph = coo_matrix((vals, (rows, cols)), shape=shape)

    if split_negative:
        graph_neg = graph.copy()

        graph.data = np.clip(graph.data, 0, 1)
        graph_neg.data = np.clip(graph_neg.data, -1, 0)

        graph.eliminate_zeros()
        graph_neg.eliminate_zeros()

        return graph.tocsr(), graph_neg.tocsr()

    else:
        return graph.tocsr()

class VelocityGraph:
    def __init__(
        self,
        adata,
        vkey="velocity",
        xkey="Ms",
        tkey=None,
        basis=None,
        n_neighbors=None,
        sqrt_transform=None,
        n_recurse_neighbors=None,
        random_neighbors_at_max=None,
        gene_subset=None,
        approx=None,
        report=False,
        compute_uncertainties=None,
        mode_neighbors="distances",
    ):

        subset = np.ones(adata.n_vars, bool)
        if gene_subset is not None:
            var_names_subset = adata.var_names.isin(gene_subset)
            subset &= var_names_subset if len(var_names_subset) > 0 else gene_subset
        elif f"{vkey}_genes" in adata.var.keys():
            subset &= np.array(adata.var[f"{vkey}_genes"].values, dtype=bool)

        xkey = xkey if xkey in adata.layers.keys() else "spliced"

        X = np.array(
            adata.layers[xkey].A[:, subset]
            if issparse(adata.layers[xkey])
            else adata.layers[xkey][:, subset]
        )
        V = np.array(
            adata.layers[vkey].A[:, subset]
            if issparse(adata.layers[vkey])
            else adata.layers[vkey][:, subset]
        )

        nans = np.isnan(np.sum(V, axis=0))
        if np.any(nans):
            X = X[:, ~nans]
            V = V[:, ~nans]

        if approx is True and X.shape[1] > 100:
            X_pca, PCs, _, _ = pca(X, n_comps=30, svd_solver="arpack", return_info=True)
            self.X = np.array(X_pca, dtype=np.float32)
            self.V = (V - V.mean(0)).dot(PCs.T)
            self.V[V.sum(1) == 0] = 0
        else:
            self.X = np.array(X, dtype=np.float32)
            self.V = np.array(V, dtype=np.float32)
        self.V_raw = np.array(self.V)

        self.sqrt_transform = sqrt_transform
        if self.sqrt_transform is None and f"{vkey}_params" in adata.uns.keys():
            self.sqrt_transform = adata.uns[f"{vkey}_params"]["mode"] == "stochastic"
        if self.sqrt_transform:
            self.V = np.sqrt(np.abs(self.V)) * np.sign(self.V)
        self.V -= np.nanmean(self.V, axis=1)[:, None]

        self.n_recurse_neighbors = n_recurse_neighbors
        if self.n_recurse_neighbors is None:
            if n_neighbors is not None or mode_neighbors == "connectivities":
                self.n_recurse_neighbors = 1
            else:
                self.n_recurse_neighbors = 2

        if "neighbors" not in adata.uns.keys():
            neighbors(adata)
        if np.min((get_neighs(adata, "distances") > 0).sum(1).A1) == 0:
            raise ValueError(
                "Your neighbor graph seems to be corrupted. "
                "Consider recomputing via pp.neighbors."
            )
        if n_neighbors is None or n_neighbors <= get_n_neighs(adata):
            self.indices = get_indices(
                dist=get_neighs(adata, "distances"),
                n_neighbors=n_neighbors,
                mode_neighbors=mode_neighbors,
            )[0]
        else:
            if basis is None:
                basis_keys = ["X_pca", "X_tsne", "X_umap"]
                basis = [key for key in basis_keys if key in adata.obsm.keys()][-1]
            elif f"X_{basis}" in adata.obsm.keys():
                basis = f"X_{basis}"

            if isinstance(approx, str) and approx in adata.obsm.keys():
                from sklearn.neighbors import NearestNeighbors

                neighs = NearestNeighbors(n_neighbors=n_neighbors + 1)
                neighs.fit(adata.obsm[approx])
                self.indices = neighs.kneighbors_graph(
                    mode="connectivity"
                ).indices.reshape((-1, n_neighbors + 1))
            else:
                from .. import Neighbors

                neighs = Neighbors(adata)
                neighs.compute_neighbors(
                    n_neighbors=n_neighbors, use_rep=basis, n_pcs=10
                )
                self.indices = get_indices(
                    dist=neighs.distances, mode_neighbors=mode_neighbors
                )[0]

        self.max_neighs = random_neighbors_at_max

        gkey, gkey_ = f"{vkey}_graph", f"{vkey}_graph_neg"
        self.graph = adata.uns[gkey] if gkey in adata.uns.keys() else []
        self.graph_neg = adata.uns[gkey_] if gkey_ in adata.uns.keys() else []

        if tkey in adata.obs.keys():
            self.t0 = adata.obs[tkey].copy()
            init = min(self.t0) if isinstance(min(self.t0), int) else 0
            self.t0.cat.categories = np.arange(init, len(self.t0.cat.categories))
            self.t1 = self.t0.copy()
            self.t1.cat.categories = self.t0.cat.categories + 1
        else:
            self.t0 = None

        self.compute_uncertainties = compute_uncertainties
        self.uncertainties = None
        self.self_prob = None
        self.report = report
        self.adata = adata

    def compute_cosines(self, n_jobs=None, backend="loky"):
        n_jobs = get_n_jobs(n_jobs=n_jobs)

        n_obs = self.X.shape[0]

        # TODO: Use batches and vectorize calculation of dX in self._calculate_cosines
        res = parallelize(
            self._compute_cosines,
            range(self.X.shape[0]),
            n_jobs=n_jobs,
            unit="cells",
            backend=backend,
        )()
        uncertainties, vals, rows, cols = map(_flatten, zip(*res))

        vals = np.hstack(vals)
        vals[np.isnan(vals)] = 0

        self.graph, self.graph_neg = vals_to_csr(
            vals, rows, cols, shape=(n_obs, n_obs), split_negative=True
        )
        if self.compute_uncertainties:
            uncertainties = np.hstack(uncertainties)
            uncertainties[np.isnan(uncertainties)] = 0
            self.uncertainties = vals_to_csr(
                uncertainties, rows, cols, shape=(n_obs, n_obs), split_negative=False
            )
            self.uncertainties.eliminate_zeros()

        confidence = self.graph.max(1).A.flatten()
        self.self_prob = np.clip(np.percentile(confidence, 98) - confidence, 0, 1)

    def _compute_cosines(self, obs_idx, queue):
        vals, rows, cols, uncertainties = [], [], [], []
        if self.compute_uncertainties:
            moments = get_moments(self.adata, np.sign(self.V_raw), second_order=True)

        for obs_id in obs_idx:
            if self.V[obs_id].max() != 0 or self.V[obs_id].min() != 0:
                neighs_idx = get_iterative_indices(
                    self.indices, obs_id, self.n_recurse_neighbors, self.max_neighs
                )

                if self.t0 is not None:
                    t0, t1 = self.t0[obs_id], self.t1[obs_id]
                    if t0 >= 0 and t1 > 0:
                        t1_idx = np.where(self.t0 == t1)[0]
                        if len(t1_idx) > len(neighs_idx):
                            t1_idx = np.random.choice(
                                t1_idx, len(neighs_idx), replace=False
                            )
                        if len(t1_idx) > 0:
                            neighs_idx = np.unique(np.concatenate([neighs_idx, t1_idx]))

                dX = self.X[neighs_idx] - self.X[obs_id, None]  # 60% of runtime
                if self.sqrt_transform:
                    dX = np.sqrt(np.abs(dX)) * np.sign(dX)
                val = cosine_correlation(dX, self.V[obs_id])  # 40% of runtime

                if self.compute_uncertainties:
                    dX /= l2_norm(dX)[:, None]
                    uncertainties.extend(
                        np.nansum(dX ** 2 * moments[obs_id][None, :], 1)
                    )

                vals.extend(val)
                rows.extend(np.ones(len(neighs_idx)) * obs_id)
                cols.extend(neighs_idx)

            if queue is not None:
                queue.put(1)

        if queue is not None:
            queue.put(None)

        return uncertainties, vals, rows, cols

def _flatten(iterable):
    return [i for it in iterable for i in it]

def velocity_graph(
    data,
    vkey="velocity",
    xkey="Ms",
    tkey=None,
    basis=None,
    n_neighbors=None,
    n_recurse_neighbors=None,
    random_neighbors_at_max=None,
    sqrt_transform=None,
    variance_stabilization=None,
    gene_subset=None,
    compute_uncertainties=None,
    approx=None,
    mode_neighbors="distances",
    copy=False,
    n_jobs=None,
    backend="loky",
):
    """Computes velocity graph based on cosine similarities.

    The cosine similarities are computed between velocities and potential cell state
    transitions, i.e. it measures how well a corresponding change in gene expression
    :math:`\\delta_{ij} = x_j - x_i` matches the predicted change according to the
    velocity vector :math:`\\nu_i`,

    .. math::
        \\pi_{ij} = \\cos\\angle(\\delta_{ij}, \\nu_i)
        = \\frac{\\delta_{ij}^T \\nu_i}{\\left\\lVert\\delta_{ij}\\right\\rVert
        \\left\\lVert \\nu_i \\right\\rVert}.

    Arguments
    ---------
    data: :class:`~anndata.AnnData`
        Annotated data matrix.
    vkey: `str` (default: `'velocity'`)
        Name of velocity estimates to be used.
    xkey: `str` (default: `'Ms'`)
        Layer key to extract count data from.
    tkey: `str` (default: `None`)
        Observation key to extract time data from.
    basis: `str` (default: `None`)
        Basis / Embedding to use.
    n_neighbors: `int` or `None` (default: None)
        Use fixed number of neighbors or do recursive neighbor search (if `None`).
    n_recurse_neighbors: `int` (default: `None`)
        Number of recursions for neighbors search. Defaults to
        2 if mode_neighbors is 'distances', and 1 if mode_neighbors is 'connectivities'.
    random_neighbors_at_max: `int` or `None` (default: `None`)
        If number of iterative neighbors for an individual cell is higher than this
        threshold, a random selection of such are chosen as reference neighbors.
    sqrt_transform: `bool` (default: `False`)
        Whether to variance-transform the cell states changes
        and velocities before computing cosine similarities.
    gene_subset: `list` of `str`, subset of adata.var_names or `None`(default: `None`)
        Subset of genes to compute velocity graph on exclusively.
    compute_uncertainties: `bool` (default: `None`)
        Whether to compute uncertainties along with cosine correlation.
    approx: `bool` or `None` (default: `None`)
        If True, first 30 pc's are used instead of the full count matrix
    mode_neighbors: 'str' (default: `'distances'`)
        Determines the type of KNN graph used. Options are 'distances' or
        'connectivities'. The latter yields a symmetric graph.
    copy: `bool` (default: `False`)
        Return a copy instead of writing to adata.
    n_jobs: `int` or `None` (default: `None`)
        Number of parallel jobs.
    backend: `str` (default: "loky")
        Backend used for multiprocessing. See :class:`joblib.Parallel` for valid
        options.

    Returns
    -------
    Returns or updates `adata` with the attributes
    velocity_graph: `.uns`
        sparse matrix with transition probabilities
    """
    adata = data.copy() if copy else data
    verify_neighbors(adata)
    if vkey not in adata.layers.keys():
        velocity(adata, vkey=vkey)
    if sqrt_transform is None:
        sqrt_transform = variance_stabilization

    vgraph = VelocityGraph(
        adata,
        vkey=vkey,
        xkey=xkey,
        tkey=tkey,
        basis=basis,
        n_neighbors=n_neighbors,
        approx=approx,
        n_recurse_neighbors=n_recurse_neighbors,
        random_neighbors_at_max=random_neighbors_at_max,
        sqrt_transform=sqrt_transform,
        gene_subset=gene_subset,
        compute_uncertainties=compute_uncertainties,
        report=True,
        mode_neighbors=mode_neighbors,
    )

    if isinstance(basis, str):
        logg.warn(
            f"The velocity graph is computed on {basis} embedding coordinates.\n"
            f"        Consider computing the graph in an unbiased manner \n"
            f"        on full expression space by not specifying basis.\n"
        )

    n_jobs = get_n_jobs(n_jobs=n_jobs)
    logg.info(
        f"computing velocity graph (using {n_jobs}/{os.cpu_count()} cores)", r=True
    )
    vgraph.compute_cosines(n_jobs=n_jobs, backend=backend)

    adata.uns[f"{vkey}_graph"] = vgraph.graph
    adata.uns[f"{vkey}_graph_neg"] = vgraph.graph_neg

    if vgraph.uncertainties is not None:
        adata.uns[f"{vkey}_graph_uncertainties"] = vgraph.uncertainties

    adata.obs[f"{vkey}_self_transition"] = vgraph.self_prob

    if f"{vkey}_params" in adata.uns.keys():
        if "embeddings" in adata.uns[f"{vkey}_params"]:
            del adata.uns[f"{vkey}_params"]["embeddings"]
    else:
        adata.uns[f"{vkey}_params"] = {}
    adata.uns[f"{vkey}_params"]["mode_neighbors"] = mode_neighbors
    adata.uns[f"{vkey}_params"]["n_recurse_neighbors"] = vgraph.n_recurse_neighbors

    logg.info("    finished", time=True, end=" " if settings.verbosity > 2 else "\n")
    logg.hint(
        "added \n"
        f"    '{vkey}_graph', sparse matrix with cosine correlations (adata.uns)"
    )

    return adata if copy else None
<|endoftext|>


<commit_before>import numpy

from simphony.io.data_container_description import Data
from simphony.core.data_container import DataContainer
from simphony.core.cuba import CUBA

def create_data_container(restrict=None):
    """ Create a dummy data container while respecting the expected data types.

    This is a utility function to be used for testing and prototyping.

    Parameters
    ----------
    restrict : list
        The list of CUBA keys to restrict the value population. Default is to
        use all CUBA keys.

    Returns
    -------
    data : DataContainer

    """

    if restrict is None:
        restrict = CUBA
    data = {cuba: dummy_cuba_value(cuba) for cuba in restrict}
    return DataContainer(data)

def dummy_cuba_value(cuba):
    column = CUBA(cuba).name.lower()
    # get the column type
    try:
        column_type = Data.columns[column]
    except AttributeError:
        column_type = Data._v_colobjects[column]

    if numpy.issubdtype(column_type, str):
        value = column.upper()
    elif numpy.issubdtype(column_type, numpy.float):
        value = float(cuba + 3)
    elif numpy.issubdtype(column_type, numpy.integer):
        value = int(cuba + 3)
    else:
        shape = column_type.shape
        data = numpy.arange(numpy.prod(shape)) * cuba
        data = numpy.reshape(data, shape)
        if column_type.kind == 'float':
            value = numpy.ones(
                shape=shape, dtype=numpy.float64) * data
        elif column_type.kind == 'int':
            value = numpy.ones(
                shape=shape, dtype=numpy.int32) * data
        else:
            raise RuntimeError(
                'cannot create value for {}'.format(column_type))
    return value
<commit_msg>add helper functions for particles and bonds<commit_after>import uuid
from random import sample

import numpy
from numpy.testing import assert_equal

from simphony.io.data_container_description import Data
from simphony.core.data_container import DataContainer
from simphony.core.cuba import CUBA
from simphony.cuds.particles import Particle, Bond

def compare_bonds(bond, reference, msg=None, testcase=None):
    self = testcase
    self.assertEqual(bond.uid, reference.uid)
    self.assertEqual(bond.particles, reference.particles)
    compare_data_containers(bond.data, bond.data, testcase=self)

def compare_particles(particle, reference, msg=None, testcase=None):
    self = testcase
    self.assertEqual(particle.uid, reference.uid)
    self.assertEqual(particle.coordinates, reference.coordinates)
    compare_data_containers(particle.data, reference.data, testcase=self)

def compare_data_containers(data, reference, msg=None, testcase=None):
    self = testcase
    self.assertEqual(len(data), len(reference))
    for key in data:
        self.assertIn(key, reference)
        assert_equal(data[key], reference[key])

def create_particles(n=10, restrict=None):
    particle_list = []
    for i in xrange(n):
        data = create_data_container(restrict=restrict)
        data[CUBA.VELOCITY] = i
        particle_list.append(
            Particle([i, i*10, i*100], data=data))
    return particle_list

def create_bonds(n=5, restrict=None):
    bond_list = []
    for i in xrange(n):
        data = create_data_container(restrict=restrict)
        data[CUBA.VELOCITY] = i
        ids = [uuid.uuid4() for x in xrange(n)]
        bond_list.append(Bond(particles=ids, data=data))
    return bond_list

def create_data_container(restrict=None):
    """ Create a dummy data container while respecting the expected data types.

    This is a utility function to be used for testing and prototyping.

    Parameters
    ----------
    restrict : list
        The list of CUBA keys to restrict the value population. Default is to
        use all CUBA keys.

    Returns
    -------
    data : DataContainer

    """

    if restrict is None:
        restrict = CUBA
    data = {cuba: dummy_cuba_value(cuba) for cuba in restrict}
    return DataContainer(data)

def dummy_cuba_value(cuba):
    column = CUBA(cuba).name.lower()
    # get the column type
    try:
        column_type = Data.columns[column]
    except AttributeError:
        column_type = Data._v_colobjects[column]

    if numpy.issubdtype(column_type, str):
        value = column.upper()
    elif numpy.issubdtype(column_type, numpy.float):
        value = float(cuba + 3)
    elif numpy.issubdtype(column_type, numpy.integer):
        value = int(cuba + 3)
    else:
        shape = column_type.shape
        data = numpy.arange(numpy.prod(shape)) * cuba
        data = numpy.reshape(data, shape)
        if column_type.kind == 'float':
            value = numpy.ones(
                shape=shape, dtype=numpy.float64) * data
        elif column_type.kind == 'int':
            value = numpy.ones(
                shape=shape, dtype=numpy.int32) * data
        else:
            raise RuntimeError(
                'cannot create value for {}'.format(column_type))
    return value
<|endoftext|>


<commit_before>import yaml
import json
import urllib.request
import time
from wallet import Wallet
from database.database import Database

config = yaml.safe_load(open("config.yml"))

f = open('dataset.txt', 'r')
dataset = f.readlines()
dataset = [float(i) for i in dataset]

walletBitcoin = 0.0
cash_wallet = Wallet()
firstrun = True
transection = 0.002

print ('Starting simulator...')

def getSellPrice():
    with urllib.request.urlopen('https://api.bitfinex.com/v1/pubticker/btcusd') as response:
        html = response.read()
    return float(json.loads(html)['last_price'])

oldPrice = currentPrice = initPrice = dataset[0]

def getTransectionfee():
    return transection * currentPrice

def getTotalNetWorth():
    return cash_wallet.cash + (walletBitcoin * currentPrice)

def showTotalNetWorth():
    print ('Total bitcoin:', walletBitcoin, 'walletDollar', cash_wallet.cash)
    print ('Total networth:', getTotalNetWorth(), 'profit', getTotalNetWorth() - totalNetWorthStart )

totalNetWorthStart = getTotalNetWorth()

dataset.pop(0)
recent_prices = []
last_buy = 0
steps = 40

def sell(prices):
    if prices[0] < prices[steps-2]:
        if prices[steps-1] < prices[steps-2]:
            return True
    return False

min_price = 100000
max_price = 0

# Now iterate over all consequitive datapoints
for currentPrice in dataset:
    print ('old;',oldPrice,'currentPrice;',currentPrice,'min',min_price,'max',max_price)
    recent_prices.append(currentPrice)

    min_price = min(min_price, currentPrice)
    max_price = max(max_price, currentPrice)

    # nets 11281.58
    #
    if len(recent_prices) > steps-1:
        if sell(recent_prices) and walletBitcoin > 0:
            print ("sell")
            if last_buy + getTransectionfee() > currentPrice:
                walletBitcoin = walletBitcoin - 1
                cash_wallet.collect(currentPrice)
                cash_wallet.give(getTransectionfee())
        elif cash_wallet.cash > 0:
            print ("buy")
            walletBitcoin = walletBitcoin + 1
            cash_wallet.give(currentPrice)
            cash_wallet.give(getTransectionfee())
            last_buy = currentPrice

    # nets 134388.25
    #
    # if currentPrice == 5375.8:
    #     # buy
    #     number_we_afford = cash_wallet.cash / currentPrice
    #     walletBitcoin = walletBitcoin + number_we_afford
    #     cash_wallet.give(currentPrice * number_we_afford)
    #     cash_wallet.give(getTransectionfee())
    #     last_buy = currentPrice
    # elif currentPrice == 6310.0:
    #     # sell at max
    #     cash_wallet.collect(currentPrice * walletBitcoin)
    #     cash_wallet.give(getTransectionfee())

    showTotalNetWorth()
    if len(recent_prices) > steps:
            recent_prices.pop(0)
    oldPrice = currentPrice
<commit_msg>Team2 results updated version<commit_after>import yaml
import json
import urllib.request
import time
from wallet import Wallet
from database.database import Database

config = yaml.safe_load(open("config.yml"))

f = open('dataset.txt', 'r')
dataset = f.readlines()
dataset = [float(i) for i in dataset]

walletBitcoin = 0.0
cash_wallet = Wallet()
firstrun = True
transection = 0.002

print ('Starting simulator...')

def getSellPrice():
    with urllib.request.urlopen('https://api.bitfinex.com/v1/pubticker/btcusd') as response:
        html = response.read()
    return float(json.loads(html)['last_price'])

oldPrice = currentPrice = initPrice = dataset[0]

def getTransectionfee():
    return transection * currentPrice

def getTotalNetWorth():
    return cash_wallet.cash + (walletBitcoin * currentPrice)

def showTotalNetWorth():
    print ('Total bitcoin:', walletBitcoin, 'walletDollar', cash_wallet.cash)
    print ('Total networth:', getTotalNetWorth(), 'profit', getTotalNetWorth() - totalNetWorthStart )

totalNetWorthStart = getTotalNetWorth()

dataset.pop(0)
recent_prices = []
last_buy = 0
steps = 40

def sell(prices):
    if prices[0] < prices[steps-2]:
        if prices[steps-1] < prices[steps-2]:
            return True
    return False

min_price = 100000
max_price = 0
number_of_buys = 0
number_of_sells = 0

# Now iterate over all consequitive datapoints
for currentPrice in dataset:
    recent_prices.append(currentPrice)

    min_price = min(min_price, currentPrice)
    max_price = max(max_price, currentPrice)

    # nets 10621.42
    #
    if len(recent_prices) > steps-1:
        if sell(recent_prices) and walletBitcoin > 0:
            if last_buy + getTransectionfee() > currentPrice:
                print("sell")
                number_of_sells += 1
                walletBitcoin = walletBitcoin - 1
                cash_wallet.collect(currentPrice)
                cash_wallet.give(getTransectionfee())
        elif cash_wallet.cash > currentPrice:
            print ("buy")
            number_of_buys += 1
            walletBitcoin = walletBitcoin + 1
            cash_wallet.give(currentPrice)
            cash_wallet.give(getTransectionfee())
            last_buy = currentPrice

    # nets 134388.25
    #
    # if currentPrice == 5375.8:
    #     # buy
    #     number_we_afford = cash_wallet.cash / currentPrice
    #     walletBitcoin = walletBitcoin + number_we_afford
    #     cash_wallet.give(currentPrice * number_we_afford)
    #     cash_wallet.give(getTransectionfee())
    #     last_buy = currentPrice
    # elif currentPrice == 6310.0:
    #     # sell at max
    #     cash_wallet.collect(currentPrice * walletBitcoin)
    #     cash_wallet.give(getTransectionfee())

    if len(recent_prices) > steps:
            recent_prices.pop(0)
    oldPrice = currentPrice

print('buys:', number_of_buys)
print('sells:', number_of_sells)
showTotalNetWorth()
<|endoftext|>


<commit_before>"""
Sile object for reading/writing SIESTA binary files
"""
from __future__ import print_function

import numpy as np

try:
    import sisl.io.siesta._siesta
    found_module = True
except Exception as e:
    print(e)
    found_module = False

# Import sile objects
from .sile import SileBinSiesta
from ..sile import *

# Import the geometry object
from sisl import Geometry, Atom, SuperCell, Grid
from sisl.units.siesta import unit_convert
from sisl.physics import Hamiltonian

__all__ = ['tshsSileSiesta']
__all__ += ['rhoSileSiesta', 'vSileSiesta']

class tshsSileSiesta(SileBinSiesta):
    """ TranSIESTA file object """

    def read_supercell(self):
        """ Returns a SuperCell object from a siesta.TSHS file
        """

        n_s = _siesta.read_tshs_sizes(self.file)[3]
        arr = _siesta.read_tshs_cell(self.file, n_s)
        nsc = np.array(arr[0].T, np.int32)
        cell = np.array(arr[1].T, np.float64)
        cell.shape = (3, 3)
        isc = np.array(arr[2].T, np.int32)
        isc.shape = (-1, 3)

        SC = SuperCell(cell, nsc=nsc)
        SC.sc_off = isc
        return SC

    def read_geometry(self):
        """ Returns Geometry object from a siesta.TSHS file """

        # Read supercell
        sc = self.read_supercell()

        na = _siesta.read_tshs_sizes(self.file)[1]
        arr = _siesta.read_tshs_geom(self.file, na)
        xyz = np.array(arr[0].T, np.float64)
        xyz.shape = (-1, 3)
        lasto = np.array(arr[1], np.int32)

        # Create all different atoms...
        # The TSHS file does not contain the
        # atomic numbers, so we will just
        # create them individually
        orbs = np.diff(lasto)

        # Get unique orbitals
        uorb = np.unique(orbs)
        # Create atoms
        atoms = []
        for Z, orb in enumerate(uorb):
            atoms.append(Atom(Z+1, orbs=orb))

        def get_atom(atoms, orbs):
            for atom in atoms:
                if atom.orbs == orbs:
                    return atom

        atom = []
        for _, orb in enumerate(orbs):
            atom.append(get_atom(atoms, orb))

        # Create and return geometry object
        geom = Geometry(xyz, atom, sc=sc)

        return geom

    def read_hamiltonian(self, **kwargs):
        """ Returns the electronic structure from the siesta.TSHS file """

        # First read the geometry
        geom = self.read_geometry()

        # Now read the sizes used...
        sizes = _siesta.read_tshs_sizes(self.file)
        spin = sizes[0]
        no = sizes[2]
        nnz = sizes[4]
        ncol, col, dH, dS = _siesta.read_tshs_es(self.file, spin, no, nnz)

        # Create the Hamiltonian container
        H = Hamiltonian(geom, spin, nnzpr=1, orthogonal=False)

        # Create the new sparse matrix
        H._csr.ncol = ncol.astype(np.int32, copy=False)
        H._csr.ptr = np.insert(np.cumsum(ncol, dtype=np.int32), 0, 0)
        # Correct fortran indices
        H._csr.col = col.astype(np.int32, copy=False) - 1
        H._csr._nnz = len(col)

        H._csr._D = np.empty([nnz, spin+1], np.float64)
        for i in range(spin):
            # this is because of the F-ordering
            H._csr._D[:, i] = dH[:, i]
        H._csr._D[:, spin] = dS[:]

        return H

class GridSileSiesta(SileBinSiesta):
    """ Grid file object from a binary Siesta output file """

    grid_unit = 1.

    def read_supercell(self, *args, **kwargs):

        cell = _siesta.read_grid_cell(self.file)
        cell = np.array(cell.T, np.float64)
        cell.shape = (3, 3)

        SC = SuperCell(cell)
        return SC

    def read_grid(self, *args, **kwargs):
        """ Read grid contained in the Grid file """
        # Read the sizes
        nspin, mesh = _siesta.read_grid_sizes(self.file)
        # Read the cell and grid
        cell, grid = _siesta.read_grid(self.file, nspin, mesh[0], mesh[1], mesh[2])

        cell = np.array(cell.T, np.float64)
        cell.shape = (3, 3)

        g = Grid(mesh, sc=SuperCell(cell), dtype=np.float32)
        g.grid = np.array(grid.swapaxes(0, 2), np.float32) * self.grid_unit
        return g

class rhoSileSiesta(GridSileSiesta):
    """ .*RHO* file object from a binary Siesta output file """
    grid_unit = 1.

class vSileSiesta(GridSileSiesta):
    """ .V* file object from a binary Siesta output file """
    grid_unit = unit_convert('Ry', 'eV')

if found_module:
    add_sile('TSHS', tshsSileSiesta)
    add_sile('RHO', rhoSileSiesta)
    add_sile('RHOINIT', rhoSileSiesta)
    add_sile('DRHO', rhoSileSiesta)
    add_sile('IOCH', rhoSileSiesta)
    add_sile('TOCH', rhoSileSiesta)
    add_sile('VH', vSileSiesta)
    add_sile('VNA', vSileSiesta)
    add_sile('VT', vSileSiesta)
<commit_msg>bug: fixed _siesta module<commit_after>"""
Sile object for reading/writing SIESTA binary files
"""
from __future__ import print_function

import numpy as np

try:
    import sisl.io.siesta._siesta as _siesta
    found_module = True
except Exception as e:
    print(e)
    found_module = False

# Import sile objects
from .sile import SileBinSiesta
from ..sile import *

# Import the geometry object
from sisl import Geometry, Atom, SuperCell, Grid
from sisl.units.siesta import unit_convert
from sisl.physics import Hamiltonian

__all__ = ['tshsSileSiesta']
__all__ += ['rhoSileSiesta', 'vSileSiesta']

class tshsSileSiesta(SileBinSiesta):
    """ TranSIESTA file object """

    def read_supercell(self):
        """ Returns a SuperCell object from a siesta.TSHS file
        """

        n_s = _siesta.read_tshs_sizes(self.file)[3]
        arr = _siesta.read_tshs_cell(self.file, n_s)
        nsc = np.array(arr[0].T, np.int32)
        cell = np.array(arr[1].T, np.float64)
        cell.shape = (3, 3)
        isc = np.array(arr[2].T, np.int32)
        isc.shape = (-1, 3)

        SC = SuperCell(cell, nsc=nsc)
        SC.sc_off = isc
        return SC

    def read_geometry(self):
        """ Returns Geometry object from a siesta.TSHS file """

        # Read supercell
        sc = self.read_supercell()

        na = _siesta.read_tshs_sizes(self.file)[1]
        arr = _siesta.read_tshs_geom(self.file, na)
        xyz = np.array(arr[0].T, np.float64)
        xyz.shape = (-1, 3)
        lasto = np.array(arr[1], np.int32)

        # Create all different atoms...
        # The TSHS file does not contain the
        # atomic numbers, so we will just
        # create them individually
        orbs = np.diff(lasto)

        # Get unique orbitals
        uorb = np.unique(orbs)
        # Create atoms
        atoms = []
        for Z, orb in enumerate(uorb):
            atoms.append(Atom(Z+1, orbs=orb))

        def get_atom(atoms, orbs):
            for atom in atoms:
                if atom.orbs == orbs:
                    return atom

        atom = []
        for _, orb in enumerate(orbs):
            atom.append(get_atom(atoms, orb))

        # Create and return geometry object
        geom = Geometry(xyz, atom, sc=sc)

        return geom

    def read_hamiltonian(self, **kwargs):
        """ Returns the electronic structure from the siesta.TSHS file """

        # First read the geometry
        geom = self.read_geometry()

        # Now read the sizes used...
        sizes = _siesta.read_tshs_sizes(self.file)
        spin = sizes[0]
        no = sizes[2]
        nnz = sizes[4]
        ncol, col, dH, dS = _siesta.read_tshs_es(self.file, spin, no, nnz)

        # Create the Hamiltonian container
        H = Hamiltonian(geom, spin, nnzpr=1, orthogonal=False)

        # Create the new sparse matrix
        H._csr.ncol = ncol.astype(np.int32, copy=False)
        H._csr.ptr = np.insert(np.cumsum(ncol, dtype=np.int32), 0, 0)
        # Correct fortran indices
        H._csr.col = col.astype(np.int32, copy=False) - 1
        H._csr._nnz = len(col)

        H._csr._D = np.empty([nnz, spin+1], np.float64)
        for i in range(spin):
            # this is because of the F-ordering
            H._csr._D[:, i] = dH[:, i]
        H._csr._D[:, spin] = dS[:]

        return H

class GridSileSiesta(SileBinSiesta):
    """ Grid file object from a binary Siesta output file """

    grid_unit = 1.

    def read_supercell(self, *args, **kwargs):

        cell = _siesta.read_grid_cell(self.file)
        cell = np.array(cell.T, np.float64)
        cell.shape = (3, 3)

        SC = SuperCell(cell)
        return SC

    def read_grid(self, *args, **kwargs):
        """ Read grid contained in the Grid file """
        # Read the sizes
        nspin, mesh = _siesta.read_grid_sizes(self.file)
        # Read the cell and grid
        cell, grid = _siesta.read_grid(self.file, nspin, mesh[0], mesh[1], mesh[2])

        cell = np.array(cell.T, np.float64)
        cell.shape = (3, 3)

        g = Grid(mesh, sc=SuperCell(cell), dtype=np.float32)
        g.grid = np.array(grid.swapaxes(0, 2), np.float32) * self.grid_unit
        return g

class rhoSileSiesta(GridSileSiesta):
    """ .*RHO* file object from a binary Siesta output file """
    grid_unit = 1.

class vSileSiesta(GridSileSiesta):
    """ .V* file object from a binary Siesta output file """
    grid_unit = unit_convert('Ry', 'eV')

if found_module:
    add_sile('TSHS', tshsSileSiesta)
    add_sile('RHO', rhoSileSiesta)
    add_sile('RHOINIT', rhoSileSiesta)
    add_sile('DRHO', rhoSileSiesta)
    add_sile('IOCH', rhoSileSiesta)
    add_sile('TOCH', rhoSileSiesta)
    add_sile('VH', vSileSiesta)
    add_sile('VNA', vSileSiesta)
    add_sile('VT', vSileSiesta)
<|endoftext|>


<commit_before>#!/usr/bin/python3
#   suggestmovie
#   votemovie
#   nextmovie (Admin)
import random
import requests
print("Importing...")
import movie_schema
from movie_schema import User, Movie, Vote, Chatroom, botSession
import bot_secrets
key=bot_secrets.SECRET_TELEGRAM_KEY
from telegram.ext import *
from telegram import *
import logging
import re
import urllib
import json
from sqlalchemy.sql import func
updater = Updater(token=key)
dispatcher = updater.dispatcher

logging = logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)

g_chat = ""
g_moviecache = {}

def assign_chat(chat_id = -173831978):
	sesh = botSession();
	global g_chat
	g_chat = sesh.query(Chatroom).first()

	if not g_chat:
		print("Using default chatroom")
		g_chat = Chatroom(id=chat_id)
		sesh.add(g_chat)
		sesh.commit()
	botSession.remove()

assign_chat()
print("chatting in room: {}".format(g_chat.id))

it_me = 97858058

def vote_here(bot, update):
	if(update.message.from_user.id != it_me):
		bot.sendMessage(update.message.chat.id, "Can't let you do that.", reply_to_message_id=update.message.message_id)
		return
	assign_chat(update.message.chat.id)
	print("Chaning vote room to {} ".format(g_chat.id))
	bot.sendMessage(update.message.chat.id, "Movie party happening here!", reply_to_message_id=update.message.message_id)

def list_movies(searchterm):
	if len(searchterm) <=2:
		return []
	escapeterm = urllib.parse.quote(searchterm)
	req = requests.get("http://www.omdbapi.com/?s={}".format(escapeterm))
	movies = req.json()
	if "Search" in movies:
		return movies["Search"]
	else:
		return []

def suggest_movie(bot, update):
	global g_moviecache
	query = update.inline_query.query
	results = list()
	movielist = list_movies(query)
	print(query)
	for movie in movielist[:5]:
		full_title = "http://www.imdb.com/title/{imdbID} {Title} ({Year})".format(**movie)
		poster = movie["Poster"]
		if poster == "N/A":
			poster = "http://66.media.tumblr.com/0fd6af995e183fc99e4f44128930c632/tumblr_oaskifwYyi1t95h1uo1_500.jpg"

		g_moviecache[movie["imdbID"]] = movie;
		results.append( InlineQueryResultPhoto(id=movie["imdbID"], photo_url=poster, thumb_url=poster, caption=full_title))
	bot.answerInlineQuery(update.inline_query.id, results=results)

def create_votes(imdbID):
	keyboard = [[InlineKeyboardButton("Yes", callback_data=imdbID +'|Y')],
				[InlineKeyboardButton("No", callback_data=imdbID  +'|N')] ]
	return keyboard

def movie_selected(bot,update):

	print(update.to_dict())
	inline_message_id = update.chosen_inline_result["result_id"]

	sesh = botSession();
	selectedMovie = sesh.query(Movie).filter(Movie.imdbID == inline_message_id).first()
	# add movie if not exists
	if not selectedMovie:  # TODO: race condition. do the transaction correctly.
		print("Adding movie: " + inline_message_id)
		selectedMovie = Movie( g_moviecache[inline_message_id] )
		sesh.add(selectedMovie)
		sesh.commit()
	botSession.remove()

	#movie = update.chosen_inline_result.id
	keyboard = create_votes(inline_message_id);
	reply_markup = InlineKeyboardMarkup(keyboard)
	bot.sendMessage(g_chat.id, text="Like this movie?: ", reply_markup=reply_markup)

def made_vote(bot, update):
	callback = update.to_dict()["callback_query"];
	movieID, vote = callback["data"].split("|")
	userid = callback["from"]["id"]
	if vote == "Y":
		vote = 1
	else:
		vote = -.3

	sesh = botSession()
	usr = sesh.query(movie_schema.User).filter(movie_schema.User.id == userid).first()
	if not usr:
		usr = movie_schema.User(userid,callback["from"]["username"])
		sesh.add(usr)

	curvote = sesh.query(Vote).filter(Vote.user == usr
		).filter(Vote.movie_id == movieID).first()
	if not curvote:
		print("Adding vote")
		curvote = Vote()
		curvote.movie_id = movieID;
		curvote.user_id = userid
		sesh.add(curvote)
	else:
		print("Changing vote")
	curvote.voterank = vote
	sesh.commit()

	botSession.remove()
	print("vote made {}({}) {}".format(userid,callback["from"]["username"], callback["data"] ))

def get_scores(bot, update):
	reply = "Movies suggested: \n"
	sesh = botSession()
	sum_votes = func.sum(Vote.voterank).label("Score")
	qry = sesh.query(Movie.Title, sum_votes
		).join(Movie.votes
		).group_by(Movie.Title
		).order_by(sum_votes.desc())

	i=0
	for m in qry.all():
		i+=1
		reply += "{}: {} ({})\n".format(i, m.Title, m.Score)
	botSession.remove()

	bot.sendMessage(update.message.chat.id, reply, reply_to_message_id=update.message.message_id)

def error(bot, update, error):
	print('>>>Update "%s" caused error "%s"' % (update, error))

suggestion_handler = InlineQueryHandler(suggest_movie)
dispatcher.add_handler(suggestion_handler)

dispatcher.add_handler(ChosenInlineResultHandler(movie_selected))
dispatcher.add_error_handler(error)

dispatcher.add_handler(CallbackQueryHandler(made_vote))

dispatcher.add_handler(CommandHandler('votehere', vote_here))
dispatcher.add_handler(CommandHandler('scores', get_scores))

print("Waiting...")
updater.start_polling()
updater.idle()

def this_doesnt_work():
	#Get most popular N movies
	sum_votes = func.sum(Vote.voterank).label("Score")
	qry = session.query(Movie.Title, sum_votes
		).join(Movie.votes
		).group_by(Movie.Title
		).order_by(sum_votes.desc())
	#find vote by movie and user
	qry = session.query(Vote).filter(Vote.user==userDB[0]
		).filter(Vote.movie == movieDB[1])

<commit_msg>changing voting metrics<commit_after>#!/usr/bin/python3
#   suggestmovie
#   votemovie
#   nextmovie (Admin)
import random
import requests
print("Importing...")
import movie_schema
from movie_schema import User, Movie, Vote, Chatroom, botSession
import bot_secrets
key=bot_secrets.SECRET_TELEGRAM_KEY
from telegram.ext import *
from telegram import *
import logging
import re
import urllib
import json
from sqlalchemy.sql import func
updater = Updater(token=key)
dispatcher = updater.dispatcher

logging = logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)

g_chat = ""
g_moviecache = {}

def assign_chat(chat_id = -173831978):
	sesh = botSession();
	global g_chat
	g_chat = sesh.query(Chatroom).first()

	if not g_chat:
		print("Using default chatroom")
		g_chat = Chatroom(id=chat_id)
		sesh.add(g_chat)
		sesh.commit()
	botSession.remove()

assign_chat()
print("chatting in room: {}".format(g_chat.id))

it_me = 97858058

def make_title(movie):
	return "http://www.imdb.com/title/{imdbID} \n {Title} ({Year})".format(**movie)

def vote_here(bot, update):
	if(update.message.from_user.id != it_me):
		bot.sendMessage(update.message.chat.id, "Can't let you do that.", reply_to_message_id=update.message.message_id)
		return
	assign_chat(update.message.chat.id)
	print("Chaning vote room to {} ".format(g_chat.id))
	bot.sendMessage(update.message.chat.id, "Movie party happening here!", reply_to_message_id=update.message.message_id)

def list_movies(searchterm):
	if len(searchterm) <=2:
		return []
	escapeterm = urllib.parse.quote(searchterm)
	req = requests.get("http://www.omdbapi.com/?s={}".format(escapeterm))
	movies = req.json()
	if "Search" in movies:
		return movies["Search"]
	else:
		return []

def suggest_movie(bot, update):
	global g_moviecache
	query = update.inline_query.query
	results = list()
	movielist = list_movies(query)
	print(query)
	for movie in movielist[:5]:
		full_title =  make_title(movie)
		poster = movie["Poster"]
		if poster == "N/A":
			poster = "http://66.media.tumblr.com/0fd6af995e183fc99e4f44128930c632/tumblr_oaskifwYyi1t95h1uo1_500.jpg"

		g_moviecache[movie["imdbID"]] = movie;
		results.append( InlineQueryResultPhoto(id=movie["imdbID"], photo_url=poster, thumb_url=poster, caption=full_title))
	bot.answerInlineQuery(update.inline_query.id, results=results)

def create_votes(imdbID):
	keyboard = [[InlineKeyboardButton("Yes!", callback_data=imdbID +'|2'),
				InlineKeyboardButton("yeah", callback_data=imdbID +'|1'),
				InlineKeyboardButton("meh", callback_data=imdbID  +'|-.2'),
				InlineKeyboardButton("No!", callback_data=imdbID  +'|-.5')] ]
	return keyboard

def movie_selected(bot,update):

	print(update.to_dict())
	inline_message_id = update.chosen_inline_result["result_id"]

	sesh = botSession();
	selectedMovie = sesh.query(Movie).filter(Movie.imdbID == inline_message_id).first()
	# add movie if not exists
	if not selectedMovie:  # TODO: race condition. do the transaction correctly.
		print("Adding movie: " + inline_message_id)
		selectedMovie = Movie( g_moviecache[inline_message_id] )
		sesh.add(selectedMovie)
		sesh.commit()
	botSession.remove()

	#movie = update.chosen_inline_result.id
	keyboard = create_votes(inline_message_id);
	reply_markup = InlineKeyboardMarkup(keyboard)
	votename = make_title(g_moviecache[inline_message_id])

	bot.sendMessage(g_chat.id, text= votename+"?" , reply_markup=reply_markup)

def made_vote(bot, update):
	callback = update.to_dict()["callback_query"];
	movieID, vote = callback["data"].split("|")
	userid = callback["from"]["id"]

	vote = float(vote)

	sesh = botSession()
	usr = sesh.query(movie_schema.User).filter(movie_schema.User.id == userid).first()
	if not usr:
		usr = movie_schema.User(userid,callback["from"]["username"])
		sesh.add(usr)

	curvote = sesh.query(Vote).filter(Vote.user == usr
		).filter(Vote.movie_id == movieID).first()
	if not curvote:
		print("Adding vote")
		curvote = Vote()
		curvote.movie_id = movieID;
		curvote.user_id = userid
		sesh.add(curvote)
	else:
		print("Changing vote")
	curvote.voterank = vote
	sesh.commit()

	botSession.remove()
	print("vote made {}({}) {}".format(userid,callback["from"]["username"], callback["data"] ))

	#bot.editMessageText(text)

def get_scores(bot, update):
	reply = "Movies suggested: \n"
	sesh = botSession()
	sum_votes = func.sum(Vote.voterank).label("Score")
	qry = sesh.query(Movie.Title, sum_votes
		).join(Movie.votes
		).group_by(Movie.Title
		).order_by(sum_votes.desc())

	i=0
	for m in qry.all():
		i+=1
		reply += "{}: {} ({})\n".format(i, m.Title, m.Score)
	botSession.remove()

	bot.sendMessage(update.message.chat.id, reply, reply_to_message_id=update.message.message_id)

def error(bot, update, error):
	print('>>>Update "%s" caused error "%s"' % (update, error))

suggestion_handler = InlineQueryHandler(suggest_movie)
dispatcher.add_handler(suggestion_handler)

dispatcher.add_handler(ChosenInlineResultHandler(movie_selected))
dispatcher.add_error_handler(error)

dispatcher.add_handler(CallbackQueryHandler(made_vote))

dispatcher.add_handler(CommandHandler('votehere', vote_here))
dispatcher.add_handler(CommandHandler('scores', get_scores))

print("Waiting...")
updater.start_polling()
updater.idle()

def this_doesnt_work():
	#Get most popular N movies
	sum_votes = func.sum(Vote.voterank).label("Score")
	qry = session.query(Movie.Title, sum_votes
		).join(Movie.votes
		).group_by(Movie.Title
		).order_by(sum_votes.desc())
	#find vote by movie and user
	qry = session.query(Vote).filter(Vote.user==userDB[0]
		).filter(Vote.movie == movieDB[1])

<|endoftext|>


<commit_before>#!/usr/bin/env python

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# Copyright (c) 2014 Mozilla Corporation
# Author: gdestuynder@mozilla.com
# Author: ameihm@mozilla.com

import os
import sys
import copy
from datetime import datetime
import pytz
import json
import socket
import syslog
try:
    from requests_futures.sessions import FuturesSession as Session
    futures_loaded = True
except ImportError:
    from requests import Session
    futures_loaded = False
import unittest

class MozDefError(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return repr(self.msg)

class MozDefMessage(object):
    # Supported message types
    MSGTYPE_NONE = 0
    MSGTYPE_EVENT = 1
    MSGTYPE_COMPLIANCE = 2
    MSGTYPE_VULNERABILITY = 3

    def __init__(self, url):
        self._msgtype = self.MSGTYPE_NONE

        self.log = {}
        self._sendlog = {}

        self._httpsession = Session()
        self._httpsession.trust_env = False
        self._url = url

        # Set some default options
        self._send_to_syslog = False
        self._syslog_only = False
        self._fire_and_forget = False
        self._verify_certificate = False

    def validate(self):
        return True

    def validate_log(self):
        return True

    def set_fire_and_forget(self, f):
        self._fire_and_forget = f

    def set_send_to_syslog(self, f, only_syslog=False):
        self._send_to_syslog = f
        self._syslog_only = only_syslog

    def syslog_convert(self):
        raise MozDefError('message type does not support syslog conversion')

    def construct(self):
        raise MozDefError('subclass of MozDefMessage must override construct()')

    def _httpsession_cb(self, session, response):
        if response.result().status_code != 200:
            if not self._fire_and_forget:
                raise MozDefError('POST failed with code %r' % \
                    response.result().status_code)

    def send_syslog(self):
        raise MozDefError('message type does not support syslog submission')

    def send(self):
        if not self.validate():
            raise MozDefError('message failed validation')
        self.construct()
        if not self.validate_log():
            raise MozDefError('message failed post construct validation')

        if self._send_to_syslog:
            self.send_syslog()
            if self._syslog_only:
                return

        buf = json.dumps(self._sendlog, sort_keys=True, indent=4)
        if futures_loaded:
            self._httpsession.post(self._url, buf,
                verify=self._verify_certificate,
                background_callback=self._httpsession_cb)
        else:
            self._httpsession.post(self._url, buf,
                verify=self._verify_certificate)

class MozDefCompliance(MozDefMessage):
    def validate_log(self):
        for k in ['target', 'policy', 'check', 'compliance', 'link',
            'utctimestamp']:
            if k not in self._sendlog.keys():
                return False
        for k in ['level', 'name', 'url']:
            if k not in self._sendlog['policy'].keys():
                return False
        for k in ['description', 'location', 'name', 'test']:
            if k not in self._sendlog['check'].keys():
                return False
        for k in ['type', 'value']:
            if k not in self._sendlog['check']['test'].keys():
                return False
        return True

    def construct(self):
        self._sendlog = self.log

    def __init__(self, url):
        MozDefMessage.__init__(self, url)
        self._msgtype = self.MSGTYPE_COMPLIANCE

class MozDefVulnerability(MozDefMessage):
    def validate_log(self):
        for k in ['utctimestamp', 'description', 'vuln', 'asset',
            'sourcename']:
            if k not in self._sendlog.keys():
                return False
        for k in ['assetid', 'ipv4address', 'hostname', 'macaddress']:
            if k not in self._sendlog['asset'].keys():
                return False
        for k in ['status', 'vulnid', 'title', 'discovery_time', 'age_days',
            'known_malware', 'known_exploits', 'cvss', 'cves']:
            if k not in self._sendlog['vuln'].keys():
                return False
        return True

    def construct(self):
        self._sendlog = self.log

    def __init__(self, url):
        MozDefMessage.__init__(self, url)
        self._msgtype = self.MSGTYPE_VULNERABILITY

class MozDefEvent(MozDefMessage):
    # Flags associated with legacy syslog conversion which is currently not
    # implemented
    SEVERITY_INFO = 0
    SEVERITY_WARNING = 1
    SEVERITY_CRITICAL = 2
    SEVERITY_ERROR = 3
    SEVERITY_DEBUG = 4

    _sevmap = {
        SEVERITY_INFO: ['INFO', syslog.LOG_INFO],
        SEVERITY_WARNING: ['WARNING', syslog.LOG_WARNING],
        SEVERITY_CRITICAL: ['CRIT', syslog.LOG_CRIT],
        SEVERITY_ERROR: ['ERR', syslog.LOG_ERR],
        SEVERITY_DEBUG: ['DEBUG', syslog.LOG_DEBUG],
    }

    def validate(self):
        if self.summary == None or self.summary == '':
            return False
        return True

    def set_severity(self, x):
        self._severity = x

    def syslog_convert(self):
        s = json.dumps(self._sendlog)
        return s

    def send_syslog(self):
        syspri = syslog.LOG_INFO
        for i in self._sevmap:
            if i == self._severity:
                syspri = self._sevmap[i][1]
        syslog.syslog(self.syslog_convert())

    def construct(self):
        self._sendlog = {}
        self._sendlog['timestamp'] = \
            pytz.timezone('UTC').localize(datetime.utcnow()).isoformat()
        self._sendlog['category'] = self._category
        self._sendlog['details'] = self.details
        self._sendlog['summary'] = self.summary
        self._sendlog['tags'] = self.tags

        for i in self._sevmap:
            if i == self._severity:
                self._sendlog['severity'] = self._sevmap[i][0]

    def __init__(self, url):
        MozDefMessage.__init__(self, url)
        self._msgtype = self.MSGTYPE_EVENT
        self._category = 'event'
        self._process_name = sys.argv[0]
        self._process_id = os.getpid()
        self._hostname = socket.getfqdn()
        self._severity = self.SEVERITY_INFO

        self.summary = None
        self.tags = []
        self.details = {}

class MozDefTests(unittest.TestCase):
    def create_valid_event(self):
        self.emsg_summary = 'a test event'
        self.emsg_tags = ['generic', 'test']
        self.emsg_details = {'one': 1, 'two': 'two'}

    def create_valid_vuln(self):
        self.vulnmsg = {}
        self.vulnmsg['description'] = 'system vulnerability management automation'
        self.vulnmsg['utctimestamp'] = '2015-01-21T15:33:51.136378+00:00'
        self.vulnmsg['sourcename'] = 'development'
        self.vulnmsg['asset'] = {}
        self.vulnmsg['asset']['assetid'] = 23
        self.vulnmsg['asset']['ipv4address'] = '1.2.3.4'
        self.vulnmsg['asset']['macaddress'] = ''
        self.vulnmsg['asset']['hostname'] = 'git.mozilla.com'
        self.vulnmsg['vuln'] = {}
        self.vulnmsg['vuln']['status'] = 'new'
        self.vulnmsg['vuln']['vulnid'] = 'nexpose:43883'
        self.vulnmsg['vuln']['title'] = \
            'RHSA-2013:1475: postgresql and postgresql84 security update'
        self.vulnmsg['vuln']['discovery_time'] = 1421845863
        self.vulnmsg['vuln']['age_days'] = 32.7
        self.vulnmsg['vuln']['known_malware'] = False
        self.vulnmsg['vuln']['known_exploits'] = False
        self.vulnmsg['vuln']['cvss'] = 8.5
        self.vulnmsg['vuln']['cves'] = ['CVE-2013-022', 'CVE-2013-1900']

    def create_valid_comp(self):
        self.compmsg = {}
        self.compmsg['target'] = 'www.mozilla.com'
        self.compmsg['utctimestamp'] = '2015-03-04T18:25:52.849272+00:00'
        self.compmsg['tags'] = {
            'operator': 'it',
            'autogroup': 'opsec'
        }
        self.compmsg['compliance'] = True
        self.compmsg['link'] = 'http://a.url'
        self.compmsg['policy'] = {
            'url': 'http://another.url',
            'name': 'system',
            'level': 'medium'
        }
        self.compmsg['check'] = {
            'test': {
                'type': 'nexpose',
                'name': 'assess',
                'value': 'nexpose'
            },
            'location': 'endpoint',
            'ref': 'sysmediumupdates1',
            'name': 'vulnerability scanner check',
            'description': 'validate system patch level'
        }

    def setUp(self):
        self.create_valid_vuln()
        self.create_valid_comp()
        self.create_valid_event()

    def testFailMessageSend(self):
        m = MozDefMessage('http://127.0.0.1')
        with self.assertRaises(MozDefError):
            m.send()

    def testFailEventSend(self):
        m = MozDefEvent('http://127.0.0.1:1/nonexistent')
        with self.assertRaises(Exception):
            m.send()

    def testMozdefEvent(self):
        m = MozDefEvent('http://127.0.0.1')
        self.assertIsNotNone(m)
        self.assertEqual(m._msgtype, MozDefMessage.MSGTYPE_EVENT)

    def testMozdefEventValidate(self):
        m = MozDefEvent('http://127.0.0.1')
        self.assertFalse(m.validate())
        m.summary = 'test event'
        self.assertTrue(m.validate())

    def testMozdefEventConstruct(self):
        m = MozDefEvent('http://127.0.0.1')
        m.summary = 'test event'
        m.construct()
        self.assertEqual(m._sendlog['category'], 'event')
        self.assertEqual(m._sendlog['summary'], 'test event')

    def testMozdefVulnValidate(self):
        m = MozDefVulnerability('http://127.0.0.1')
        self.assertTrue(m.validate())
        m.construct()
        self.assertFalse(m.validate_log())
        m.log = self.vulnmsg
        m.construct()
        self.assertTrue(m.validate_log())

    def testMozdefComplianceValidate(self):
        m = MozDefCompliance('http://127.0.0.1')
        self.assertTrue(m.validate())
        m.construct()
        self.assertFalse(m.validate_log())
        m.log = self.compmsg
        m.construct()
        self.assertTrue(m.validate_log())

    def testMozdefEventSyslog(self):
        m = MozDefEvent('http://127.0.0.1')
        m.summary = self.emsg_summary
        m.tags = self.emsg_tags
        m.details = self.emsg_details
        m.set_severity(MozDefEvent.SEVERITY_CRITICAL)
        m.construct()
        s = m.syslog_convert()
        self.assertIsNotNone(s)
        m.set_send_to_syslog(True, only_syslog=True)
        m.send()

    def testMozdefCompSyslog(self):
        m = MozDefCompliance('http://127.0.0.1')
        m.log = self.compmsg
        with self.assertRaises(MozDefError):
            m.syslog_convert()

    def testMozdefCompSyslogSend(self):
        m = MozDefCompliance('http://127.0.0.1')
        m.log = self.compmsg
        m.set_send_to_syslog(True, only_syslog=True)
        with self.assertRaises(MozDefError):
            m.send()

if __name__ == "__main__":
    unittest.main(verbosity=2)
<commit_msg>remove comment that is no longer applicable<commit_after>#!/usr/bin/env python

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# Copyright (c) 2014 Mozilla Corporation
# Author: gdestuynder@mozilla.com
# Author: ameihm@mozilla.com

import os
import sys
import copy
from datetime import datetime
import pytz
import json
import socket
import syslog
try:
    from requests_futures.sessions import FuturesSession as Session
    futures_loaded = True
except ImportError:
    from requests import Session
    futures_loaded = False
import unittest

class MozDefError(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return repr(self.msg)

class MozDefMessage(object):
    # Supported message types
    MSGTYPE_NONE = 0
    MSGTYPE_EVENT = 1
    MSGTYPE_COMPLIANCE = 2
    MSGTYPE_VULNERABILITY = 3

    def __init__(self, url):
        self._msgtype = self.MSGTYPE_NONE

        self.log = {}
        self._sendlog = {}

        self._httpsession = Session()
        self._httpsession.trust_env = False
        self._url = url

        # Set some default options
        self._send_to_syslog = False
        self._syslog_only = False
        self._fire_and_forget = False
        self._verify_certificate = False

    def validate(self):
        return True

    def validate_log(self):
        return True

    def set_fire_and_forget(self, f):
        self._fire_and_forget = f

    def set_send_to_syslog(self, f, only_syslog=False):
        self._send_to_syslog = f
        self._syslog_only = only_syslog

    def syslog_convert(self):
        raise MozDefError('message type does not support syslog conversion')

    def construct(self):
        raise MozDefError('subclass of MozDefMessage must override construct()')

    def _httpsession_cb(self, session, response):
        if response.result().status_code != 200:
            if not self._fire_and_forget:
                raise MozDefError('POST failed with code %r' % \
                    response.result().status_code)

    def send_syslog(self):
        raise MozDefError('message type does not support syslog submission')

    def send(self):
        if not self.validate():
            raise MozDefError('message failed validation')
        self.construct()
        if not self.validate_log():
            raise MozDefError('message failed post construct validation')

        if self._send_to_syslog:
            self.send_syslog()
            if self._syslog_only:
                return

        buf = json.dumps(self._sendlog, sort_keys=True, indent=4)
        if futures_loaded:
            self._httpsession.post(self._url, buf,
                verify=self._verify_certificate,
                background_callback=self._httpsession_cb)
        else:
            self._httpsession.post(self._url, buf,
                verify=self._verify_certificate)

class MozDefCompliance(MozDefMessage):
    def validate_log(self):
        for k in ['target', 'policy', 'check', 'compliance', 'link',
            'utctimestamp']:
            if k not in self._sendlog.keys():
                return False
        for k in ['level', 'name', 'url']:
            if k not in self._sendlog['policy'].keys():
                return False
        for k in ['description', 'location', 'name', 'test']:
            if k not in self._sendlog['check'].keys():
                return False
        for k in ['type', 'value']:
            if k not in self._sendlog['check']['test'].keys():
                return False
        return True

    def construct(self):
        self._sendlog = self.log

    def __init__(self, url):
        MozDefMessage.__init__(self, url)
        self._msgtype = self.MSGTYPE_COMPLIANCE

class MozDefVulnerability(MozDefMessage):
    def validate_log(self):
        for k in ['utctimestamp', 'description', 'vuln', 'asset',
            'sourcename']:
            if k not in self._sendlog.keys():
                return False
        for k in ['assetid', 'ipv4address', 'hostname', 'macaddress']:
            if k not in self._sendlog['asset'].keys():
                return False
        for k in ['status', 'vulnid', 'title', 'discovery_time', 'age_days',
            'known_malware', 'known_exploits', 'cvss', 'cves']:
            if k not in self._sendlog['vuln'].keys():
                return False
        return True

    def construct(self):
        self._sendlog = self.log

    def __init__(self, url):
        MozDefMessage.__init__(self, url)
        self._msgtype = self.MSGTYPE_VULNERABILITY

class MozDefEvent(MozDefMessage):
    SEVERITY_INFO = 0
    SEVERITY_WARNING = 1
    SEVERITY_CRITICAL = 2
    SEVERITY_ERROR = 3
    SEVERITY_DEBUG = 4

    _sevmap = {
        SEVERITY_INFO: ['INFO', syslog.LOG_INFO],
        SEVERITY_WARNING: ['WARNING', syslog.LOG_WARNING],
        SEVERITY_CRITICAL: ['CRIT', syslog.LOG_CRIT],
        SEVERITY_ERROR: ['ERR', syslog.LOG_ERR],
        SEVERITY_DEBUG: ['DEBUG', syslog.LOG_DEBUG],
    }

    def validate(self):
        if self.summary == None or self.summary == '':
            return False
        return True

    def set_severity(self, x):
        self._severity = x

    def syslog_convert(self):
        s = json.dumps(self._sendlog)
        return s

    def send_syslog(self):
        syspri = syslog.LOG_INFO
        for i in self._sevmap:
            if i == self._severity:
                syspri = self._sevmap[i][1]
        syslog.syslog(self.syslog_convert())

    def construct(self):
        self._sendlog = {}
        self._sendlog['timestamp'] = \
            pytz.timezone('UTC').localize(datetime.utcnow()).isoformat()
        self._sendlog['category'] = self._category
        self._sendlog['details'] = self.details
        self._sendlog['summary'] = self.summary
        self._sendlog['tags'] = self.tags

        for i in self._sevmap:
            if i == self._severity:
                self._sendlog['severity'] = self._sevmap[i][0]

    def __init__(self, url):
        MozDefMessage.__init__(self, url)
        self._msgtype = self.MSGTYPE_EVENT
        self._category = 'event'
        self._process_name = sys.argv[0]
        self._process_id = os.getpid()
        self._hostname = socket.getfqdn()
        self._severity = self.SEVERITY_INFO

        self.summary = None
        self.tags = []
        self.details = {}

class MozDefTests(unittest.TestCase):
    def create_valid_event(self):
        self.emsg_summary = 'a test event'
        self.emsg_tags = ['generic', 'test']
        self.emsg_details = {'one': 1, 'two': 'two'}

    def create_valid_vuln(self):
        self.vulnmsg = {}
        self.vulnmsg['description'] = 'system vulnerability management automation'
        self.vulnmsg['utctimestamp'] = '2015-01-21T15:33:51.136378+00:00'
        self.vulnmsg['sourcename'] = 'development'
        self.vulnmsg['asset'] = {}
        self.vulnmsg['asset']['assetid'] = 23
        self.vulnmsg['asset']['ipv4address'] = '1.2.3.4'
        self.vulnmsg['asset']['macaddress'] = ''
        self.vulnmsg['asset']['hostname'] = 'git.mozilla.com'
        self.vulnmsg['vuln'] = {}
        self.vulnmsg['vuln']['status'] = 'new'
        self.vulnmsg['vuln']['vulnid'] = 'nexpose:43883'
        self.vulnmsg['vuln']['title'] = \
            'RHSA-2013:1475: postgresql and postgresql84 security update'
        self.vulnmsg['vuln']['discovery_time'] = 1421845863
        self.vulnmsg['vuln']['age_days'] = 32.7
        self.vulnmsg['vuln']['known_malware'] = False
        self.vulnmsg['vuln']['known_exploits'] = False
        self.vulnmsg['vuln']['cvss'] = 8.5
        self.vulnmsg['vuln']['cves'] = ['CVE-2013-022', 'CVE-2013-1900']

    def create_valid_comp(self):
        self.compmsg = {}
        self.compmsg['target'] = 'www.mozilla.com'
        self.compmsg['utctimestamp'] = '2015-03-04T18:25:52.849272+00:00'
        self.compmsg['tags'] = {
            'operator': 'it',
            'autogroup': 'opsec'
        }
        self.compmsg['compliance'] = True
        self.compmsg['link'] = 'http://a.url'
        self.compmsg['policy'] = {
            'url': 'http://another.url',
            'name': 'system',
            'level': 'medium'
        }
        self.compmsg['check'] = {
            'test': {
                'type': 'nexpose',
                'name': 'assess',
                'value': 'nexpose'
            },
            'location': 'endpoint',
            'ref': 'sysmediumupdates1',
            'name': 'vulnerability scanner check',
            'description': 'validate system patch level'
        }

    def setUp(self):
        self.create_valid_vuln()
        self.create_valid_comp()
        self.create_valid_event()

    def testFailMessageSend(self):
        m = MozDefMessage('http://127.0.0.1')
        with self.assertRaises(MozDefError):
            m.send()

    def testFailEventSend(self):
        m = MozDefEvent('http://127.0.0.1:1/nonexistent')
        with self.assertRaises(Exception):
            m.send()

    def testMozdefEvent(self):
        m = MozDefEvent('http://127.0.0.1')
        self.assertIsNotNone(m)
        self.assertEqual(m._msgtype, MozDefMessage.MSGTYPE_EVENT)

    def testMozdefEventValidate(self):
        m = MozDefEvent('http://127.0.0.1')
        self.assertFalse(m.validate())
        m.summary = 'test event'
        self.assertTrue(m.validate())

    def testMozdefEventConstruct(self):
        m = MozDefEvent('http://127.0.0.1')
        m.summary = 'test event'
        m.construct()
        self.assertEqual(m._sendlog['category'], 'event')
        self.assertEqual(m._sendlog['summary'], 'test event')

    def testMozdefVulnValidate(self):
        m = MozDefVulnerability('http://127.0.0.1')
        self.assertTrue(m.validate())
        m.construct()
        self.assertFalse(m.validate_log())
        m.log = self.vulnmsg
        m.construct()
        self.assertTrue(m.validate_log())

    def testMozdefComplianceValidate(self):
        m = MozDefCompliance('http://127.0.0.1')
        self.assertTrue(m.validate())
        m.construct()
        self.assertFalse(m.validate_log())
        m.log = self.compmsg
        m.construct()
        self.assertTrue(m.validate_log())

    def testMozdefEventSyslog(self):
        m = MozDefEvent('http://127.0.0.1')
        m.summary = self.emsg_summary
        m.tags = self.emsg_tags
        m.details = self.emsg_details
        m.set_severity(MozDefEvent.SEVERITY_CRITICAL)
        m.construct()
        s = m.syslog_convert()
        self.assertIsNotNone(s)
        m.set_send_to_syslog(True, only_syslog=True)
        m.send()

    def testMozdefCompSyslog(self):
        m = MozDefCompliance('http://127.0.0.1')
        m.log = self.compmsg
        with self.assertRaises(MozDefError):
            m.syslog_convert()

    def testMozdefCompSyslogSend(self):
        m = MozDefCompliance('http://127.0.0.1')
        m.log = self.compmsg
        m.set_send_to_syslog(True, only_syslog=True)
        with self.assertRaises(MozDefError):
            m.send()

if __name__ == "__main__":
    unittest.main(verbosity=2)
<|endoftext|>