<commit_before># This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright (c) 2020 ScyllaDB

# pylint: disable=too-many-arguments

import os
import time
import queue
import logging
import threading
from typing import Optional
from functools import cached_property

import kubernetes as k8s
from urllib3.util.retry import Retry

from sdcm.remote import LOCALRUNNER
from sdcm.sct_config import sct_abs_path
from sdcm.utils.docker_utils import ContainerManager, DockerException, Container

KUBECTL_BIN = "kubectl"
HELM_IMAGE = "alpine/helm:3.3.4"

KUBECTL_TIMEOUT = 300  # seconds

K8S_CONFIGS = sct_abs_path("sdcm/k8s_configs")

JSON_PATCH_TYPE = "application/json-patch+json"

LOGGER = logging.getLogger(__name__)

logging.getLogger("kubernetes.client.rest").setLevel(logging.INFO)

class ApiLimiterClient(k8s.client.ApiClient):
    _api_rate_limiter: 'ApiCallRateLimiter' = None

    def call_api(self, *args, **kwargs):
        if self._api_rate_limiter:
            self._api_rate_limiter.wait()
        return super().call_api(*args, **kwargs)

    def bind_api_limiter(self, instance: 'ApiCallRateLimiter'):
        self._api_rate_limiter = instance

class ApiLimiterRetry(Retry):
    _api_rate_limiter: 'ApiCallRateLimiter' = None

    def sleep(self, *args, **kwargs):
        super().sleep(*args, **kwargs)
        if self._api_rate_limiter:
            self._api_rate_limiter.wait()

    def new(self, *args, **kwargs):
        result = super().new(*args, **kwargs)
        if self._api_rate_limiter:
            ApiLimiterRetry.bind_api_limiter(result, self._api_rate_limiter)
        return result

    def bind_api_limiter(self, instance: 'ApiCallRateLimiter'):
        self._api_rate_limiter = instance

class ApiCallRateLimiter(threading.Thread):
    """Simple and not very accurate rate limiter.

    Allow 1 call each `1 / rate_limit' seconds interval.
    If some call not able to start after `queue_size / rate_limit' seconds then raise `queue.Full' for caller.
    """

    def __init__(self, rate_limit: float, queue_size: int, urllib_retry: int):
        super().__init__(name=type(self).__name__, daemon=True)
        self._lock = threading.Lock()
        self.rate_limit = rate_limit  # ops/s
        self.queue_size = queue_size
        self.urllib_retry = urllib_retry
        self.running = threading.Event()

    def wait(self):
        if not self._lock.acquire(timeout=self.queue_size / self.rate_limit):  # deepcode ignore E1123: deepcode error
            LOGGER.error("k8s API call rate limiter queue size limit has been reached")
            raise queue.Full

    def stop(self):
        self.running.clear()
        self.join()

    def run(self) -> None:
        LOGGER.info("k8s API call rate limiter started: rate_limit=%s, queue_size=%s",
                    self.rate_limit, self.queue_size)
        self.running.set()
        while self.running.is_set():
            if self._lock.locked():
                self._lock.release()
            time.sleep(1 / self.rate_limit)

    def get_k8s_configuration(self, kluster) -> k8s.client.Configuration:
        output = KubernetesOps.create_k8s_configuration(kluster)
        output.retries = ApiLimiterRetry(self.urllib_retry)
        output.retries.bind_api_limiter(self)
        return output

    def get_api_client(self, k8s_configuration: k8s.client.Configuration) -> ApiLimiterClient:
        output = ApiLimiterClient(k8s_configuration)
        output.bind_api_limiter(self)
        return output

class KubernetesOps:

    @staticmethod
    def create_k8s_configuration(kluster) -> k8s.client.Configuration:
        k8s_configuration = k8s.client.Configuration()
        if kluster.k8s_server_url:
            k8s_configuration.host = kluster.k8s_server_url
        else:
            k8s.config.load_kube_config(
                config_file=os.environ.get('KUBECONFIG', '~/.kube/config'),
                client_configuration=k8s_configuration)
        return k8s_configuration

    @classmethod
    def api_client(cls, k8s_configuration: k8s.client.Configuration) -> k8s.client.ApiClient:
        return k8s.client.ApiClient(k8s_configuration)

    @classmethod
    def dynamic_client(cls, api_client: k8s.client.ApiClient) -> k8s.dynamic.DynamicClient:
        return k8s.dynamic.DynamicClient(api_client)

    @classmethod
    def dynamic_api(cls, dynamic_client: k8s.dynamic.DynamicClient, api_version, kind):
        return dynamic_client.resources.get(api_version=api_version, kind=kind)

    @classmethod
    def apps_v1_api(cls, api_client: k8s.client.ApiClient) -> k8s.client.AppsV1Api:
        return k8s.client.AppsV1Api(api_client)

    @classmethod
    def core_v1_api(cls, api_client: k8s.client.ApiClient) -> k8s.client.CoreV1Api:
        return k8s.client.CoreV1Api(api_client)

    @classmethod
    def list_statefulsets(cls, kluster, namespace=None, **kwargs):
        if namespace is None:
            return kluster.k8s_apps_v1_api.list_stateful_set_for_all_namespaces(watch=False, **kwargs).items
        return kluster.k8s_apps_v1_api.list_namespaced_stateful_set(namespace=namespace, watch=False, **kwargs).items

    @classmethod
    def list_pods(cls, kluster, namespace=None, **kwargs):
        if namespace is None:
            return kluster.k8s_core_v1_api.list_pod_for_all_namespaces(watch=False, **kwargs).items
        return kluster.k8s_core_v1_api.list_namespaced_pod(namespace=namespace, watch=False, **kwargs).items

    @classmethod
    def list_services(cls, kluster, namespace=None, **kwargs):
        if namespace is None:
            return kluster.k8s_core_v1_api.list_service_all_namespaces(watch=False, **kwargs).items
        return kluster.k8s_core_v1_api.list_namespaced_service(namespace=namespace, watch=False, **kwargs).items

    @staticmethod
    def kubectl_cmd(kluster, *command, namespace=None, ignore_k8s_server_url=False):
        cmd = [KUBECTL_BIN, ]
        if not ignore_k8s_server_url and kluster.k8s_server_url is not None:
            cmd.append(f"--server={kluster.k8s_server_url}")
        if namespace:
            cmd.append(f"--namespace={namespace}")
        cmd.extend(command)
        return " ".join(cmd)

    @classmethod
    def kubectl(cls, kluster, *command, namespace: Optional[str] = None, timeout: int = KUBECTL_TIMEOUT,
                remoter: Optional['KubernetesCmdRunner'] = None, ignore_status: bool = False, verbose: bool = True):
        cmd = cls.kubectl_cmd(kluster, *command, namespace=namespace, ignore_k8s_server_url=bool(remoter))
        if remoter is None:
            remoter = LOCALRUNNER
        return remoter.run(cmd, timeout=timeout, ignore_status=ignore_status, verbose=verbose)

    @classmethod
    def kubectl_multi_cmd(cls, kluster, *command, namespace: Optional[str] = None, timeout: int = KUBECTL_TIMEOUT,
                          remoter: Optional['KubernetesCmdRunner'] = None, ignore_status: bool = False,
                          verbose: bool = True):
        total_command = ' '.join(command)
        final_command = []
        for cmd in total_command.split(' '):
            if cmd == 'kubectl':
                final_command.append(
                    cls.kubectl_cmd(kluster, namespace=namespace, ignore_k8s_server_url=bool(remoter)))
            else:
                final_command.append(cmd)
        if remoter is None:
            remoter = LOCALRUNNER
        final_command = ' '.join(final_command)
        return remoter.run(final_command, timeout=timeout, ignore_status=ignore_status, verbose=verbose)

    @classmethod
    def apply_file(cls, kluster, config_path, namespace=None, timeout=KUBECTL_TIMEOUT, envsubst=True):
        if envsubst:
            config_path = f"<(envsubst<{config_path})"
        cls.kubectl(kluster, "apply", "-f",  config_path, namespace=namespace, timeout=timeout)

    @classmethod
    def copy_file(cls, kluster, src, dst, container=None, timeout=KUBECTL_TIMEOUT):
        command = ["cp", src, dst]
        if container:
            command.extend(("-c", container))
        cls.kubectl(kluster, *command, timeout=timeout)

    @classmethod
    def expose_pod_ports(cls, kluster, pod_name, ports, labels=None, selector=None, namespace=None, timeout=KUBECTL_TIMEOUT):
        command = ["expose pod", pod_name, "--type=LoadBalancer",
                   "--port", ",".join(map(str, ports)),
                   f"--name={pod_name}-loadbalancer", ]
        if labels:
            command.extend(("--labels", labels))
        if selector:
            command.extend(("--selector", selector))
        cls.kubectl(kluster, *command, namespace=namespace, timeout=timeout)

    @classmethod
    def unexpose_pod_ports(cls, kluster, pod_name, namespace=None, timeout=KUBECTL_TIMEOUT):
        cls.kubectl(kluster, f"delete service {pod_name}-loadbalancer", namespace=namespace, timeout=timeout)

class HelmContainerMixin:
    def helm_container_run_args(self) -> dict:
        kube_config_path = os.path.expanduser(os.environ.get('KUBECONFIG', '~/.kube/config'))
        helm_config_path = os.path.expanduser(os.environ.get('HELM_CONFIG_HOME', '~/.helm'))
        volumes = {
            os.path.dirname(kube_config_path): {"bind": "/root/.kube", "mode": "rw"},
            helm_config_path: {"bind": "/root/.helm", "mode": "rw"},
            K8S_CONFIGS: {"bind": "/apps", "mode": "rw"}, }
        return dict(image=HELM_IMAGE,
                    entrypoint="/bin/cat",
                    tty=True,
                    name=f"{self.name}-helm",
                    network_mode="host",
                    volumes=volumes)

    @cached_property
    def _helm_container(self) -> Container:
        return ContainerManager.run_container(self, "helm")

    def helm(self, kluster, *command: str, namespace: Optional[str] = None) -> str:
        cmd = ["helm", ]
        if kluster.k8s_server_url:
            cmd.extend(("--kube-apiserver", kluster.k8s_server_url, ))
        if namespace:
            cmd.extend(("--namespace", namespace, ))
        cmd.extend(command)
        cmd = " ".join(cmd)
        LOGGER.debug("Execute `%s'", cmd)
        res = self._helm_container.exec_run(["sh", "-c", cmd])
        if res.exit_code:
            raise DockerException(f"{self._helm_container}: {res.output.decode('utf-8')}")
        return res.output.decode("utf-8")
<commit_msg>fix(k8s.py): retry on api calls<commit_after># This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright (c) 2020 ScyllaDB

# pylint: disable=too-many-arguments

import os
import time
import queue
import logging
import threading
from typing import Optional
from functools import cached_property

import kubernetes as k8s
from urllib3.util.retry import Retry

from sdcm.remote import LOCALRUNNER
from sdcm.sct_config import sct_abs_path
from sdcm.utils.decorators import timeout as timeout_decor
from sdcm.utils.docker_utils import ContainerManager, DockerException, Container

KUBECTL_BIN = "kubectl"
HELM_IMAGE = "alpine/helm:3.3.4"

KUBECTL_TIMEOUT = 300  # seconds

K8S_CONFIGS = sct_abs_path("sdcm/k8s_configs")

JSON_PATCH_TYPE = "application/json-patch+json"

LOGGER = logging.getLogger(__name__)

logging.getLogger("kubernetes.client.rest").setLevel(logging.INFO)

class ApiLimiterClient(k8s.client.ApiClient):
    _api_rate_limiter: 'ApiCallRateLimiter' = None

    def call_api(self, *args, **kwargs):
        if self._api_rate_limiter:
            self._api_rate_limiter.wait()
        return super().call_api(*args, **kwargs)

    def bind_api_limiter(self, instance: 'ApiCallRateLimiter'):
        self._api_rate_limiter = instance

class ApiLimiterRetry(Retry):
    _api_rate_limiter: 'ApiCallRateLimiter' = None

    def sleep(self, *args, **kwargs):
        super().sleep(*args, **kwargs)
        if self._api_rate_limiter:
            self._api_rate_limiter.wait()

    def new(self, *args, **kwargs):
        result = super().new(*args, **kwargs)
        if self._api_rate_limiter:
            ApiLimiterRetry.bind_api_limiter(result, self._api_rate_limiter)
        return result

    def bind_api_limiter(self, instance: 'ApiCallRateLimiter'):
        self._api_rate_limiter = instance

class ApiCallRateLimiter(threading.Thread):
    """Simple and not very accurate rate limiter.

    Allow 1 call each `1 / rate_limit' seconds interval.
    If some call not able to start after `queue_size / rate_limit' seconds then raise `queue.Full' for caller.
    """

    def __init__(self, rate_limit: float, queue_size: int, urllib_retry: int):
        super().__init__(name=type(self).__name__, daemon=True)
        self._lock = threading.Lock()
        self.rate_limit = rate_limit  # ops/s
        self.queue_size = queue_size
        self.urllib_retry = urllib_retry
        self.running = threading.Event()

    def wait(self):
        if not self._lock.acquire(timeout=self.queue_size / self.rate_limit):  # deepcode ignore E1123: deepcode error
            LOGGER.error("k8s API call rate limiter queue size limit has been reached")
            raise queue.Full

    def stop(self):
        self.running.clear()
        self.join()

    def run(self) -> None:
        LOGGER.info("k8s API call rate limiter started: rate_limit=%s, queue_size=%s",
                    self.rate_limit, self.queue_size)
        self.running.set()
        while self.running.is_set():
            if self._lock.locked():
                self._lock.release()
            time.sleep(1 / self.rate_limit)

    def get_k8s_configuration(self, kluster) -> k8s.client.Configuration:
        output = KubernetesOps.create_k8s_configuration(kluster)
        output.retries = ApiLimiterRetry(self.urllib_retry)
        output.retries.bind_api_limiter(self)
        return output

    def get_api_client(self, k8s_configuration: k8s.client.Configuration) -> ApiLimiterClient:
        output = ApiLimiterClient(k8s_configuration)
        output.bind_api_limiter(self)
        return output

class KubernetesOps:

    @staticmethod
    def create_k8s_configuration(kluster) -> k8s.client.Configuration:
        k8s_configuration = k8s.client.Configuration()
        if kluster.k8s_server_url:
            k8s_configuration.host = kluster.k8s_server_url
        else:
            k8s.config.load_kube_config(
                config_file=os.environ.get('KUBECONFIG', '~/.kube/config'),
                client_configuration=k8s_configuration)
        return k8s_configuration

    @classmethod
    def api_client(cls, k8s_configuration: k8s.client.Configuration) -> k8s.client.ApiClient:
        return k8s.client.ApiClient(k8s_configuration)

    @classmethod
    def dynamic_client(cls, api_client: k8s.client.ApiClient) -> k8s.dynamic.DynamicClient:
        return k8s.dynamic.DynamicClient(api_client)

    @classmethod
    def dynamic_api(cls, dynamic_client: k8s.dynamic.DynamicClient, api_version, kind):
        return dynamic_client.resources.get(api_version=api_version, kind=kind)

    @classmethod
    def apps_v1_api(cls, api_client: k8s.client.ApiClient) -> k8s.client.AppsV1Api:
        return k8s.client.AppsV1Api(api_client)

    @classmethod
    def core_v1_api(cls, api_client: k8s.client.ApiClient) -> k8s.client.CoreV1Api:
        return k8s.client.CoreV1Api(api_client)

    @classmethod
    @timeout_decor(timeout=600)
    def list_statefulsets(cls, kluster, namespace=None, **kwargs):
        if namespace is None:
            return kluster.k8s_apps_v1_api.list_stateful_set_for_all_namespaces(watch=False, **kwargs).items
        return kluster.k8s_apps_v1_api.list_namespaced_stateful_set(namespace=namespace, watch=False, **kwargs).items

    @classmethod
    @timeout_decor(timeout=600)
    def list_pods(cls, kluster, namespace=None, **kwargs):
        if namespace is None:
            return kluster.k8s_core_v1_api.list_pod_for_all_namespaces(watch=False, **kwargs).items
        return kluster.k8s_core_v1_api.list_namespaced_pod(namespace=namespace, watch=False, **kwargs).items

    @classmethod
    @timeout_decor(timeout=600)
    def list_services(cls, kluster, namespace=None, **kwargs):
        if namespace is None:
            return kluster.k8s_core_v1_api.list_service_all_namespaces(watch=False, **kwargs).items
        return kluster.k8s_core_v1_api.list_namespaced_service(namespace=namespace, watch=False, **kwargs).items

    @staticmethod
    def kubectl_cmd(kluster, *command, namespace=None, ignore_k8s_server_url=False):
        cmd = [KUBECTL_BIN, ]
        if not ignore_k8s_server_url and kluster.k8s_server_url is not None:
            cmd.append(f"--server={kluster.k8s_server_url}")
        if namespace:
            cmd.append(f"--namespace={namespace}")
        cmd.extend(command)
        return " ".join(cmd)

    @classmethod
    def kubectl(cls, kluster, *command, namespace: Optional[str] = None, timeout: int = KUBECTL_TIMEOUT,
                remoter: Optional['KubernetesCmdRunner'] = None, ignore_status: bool = False, verbose: bool = True):
        cmd = cls.kubectl_cmd(kluster, *command, namespace=namespace, ignore_k8s_server_url=bool(remoter))
        if remoter is None:
            remoter = LOCALRUNNER
        return remoter.run(cmd, timeout=timeout, ignore_status=ignore_status, verbose=verbose)

    @classmethod
    def kubectl_multi_cmd(cls, kluster, *command, namespace: Optional[str] = None, timeout: int = KUBECTL_TIMEOUT,
                          remoter: Optional['KubernetesCmdRunner'] = None, ignore_status: bool = False,
                          verbose: bool = True):
        total_command = ' '.join(command)
        final_command = []
        for cmd in total_command.split(' '):
            if cmd == 'kubectl':
                final_command.append(
                    cls.kubectl_cmd(kluster, namespace=namespace, ignore_k8s_server_url=bool(remoter)))
            else:
                final_command.append(cmd)
        if remoter is None:
            remoter = LOCALRUNNER
        final_command = ' '.join(final_command)
        return remoter.run(final_command, timeout=timeout, ignore_status=ignore_status, verbose=verbose)

    @classmethod
    def apply_file(cls, kluster, config_path, namespace=None, timeout=KUBECTL_TIMEOUT, envsubst=True):
        if envsubst:
            config_path = f"<(envsubst<{config_path})"
        cls.kubectl(kluster, "apply", "-f",  config_path, namespace=namespace, timeout=timeout)

    @classmethod
    def copy_file(cls, kluster, src, dst, container=None, timeout=KUBECTL_TIMEOUT):
        command = ["cp", src, dst]
        if container:
            command.extend(("-c", container))
        cls.kubectl(kluster, *command, timeout=timeout)

    @classmethod
    def expose_pod_ports(cls, kluster, pod_name, ports, labels=None, selector=None, namespace=None, timeout=KUBECTL_TIMEOUT):
        command = ["expose pod", pod_name, "--type=LoadBalancer",
                   "--port", ",".join(map(str, ports)),
                   f"--name={pod_name}-loadbalancer", ]
        if labels:
            command.extend(("--labels", labels))
        if selector:
            command.extend(("--selector", selector))
        cls.kubectl(kluster, *command, namespace=namespace, timeout=timeout)

    @classmethod
    def unexpose_pod_ports(cls, kluster, pod_name, namespace=None, timeout=KUBECTL_TIMEOUT):
        cls.kubectl(kluster, f"delete service {pod_name}-loadbalancer", namespace=namespace, timeout=timeout)

class HelmContainerMixin:
    def helm_container_run_args(self) -> dict:
        kube_config_path = os.path.expanduser(os.environ.get('KUBECONFIG', '~/.kube/config'))
        helm_config_path = os.path.expanduser(os.environ.get('HELM_CONFIG_HOME', '~/.helm'))
        volumes = {
            os.path.dirname(kube_config_path): {"bind": "/root/.kube", "mode": "rw"},
            helm_config_path: {"bind": "/root/.helm", "mode": "rw"},
            K8S_CONFIGS: {"bind": "/apps", "mode": "rw"}, }
        return dict(image=HELM_IMAGE,
                    entrypoint="/bin/cat",
                    tty=True,
                    name=f"{self.name}-helm",
                    network_mode="host",
                    volumes=volumes)

    @cached_property
    def _helm_container(self) -> Container:
        return ContainerManager.run_container(self, "helm")

    def helm(self, kluster, *command: str, namespace: Optional[str] = None) -> str:
        cmd = ["helm", ]
        if kluster.k8s_server_url:
            cmd.extend(("--kube-apiserver", kluster.k8s_server_url, ))
        if namespace:
            cmd.extend(("--namespace", namespace, ))
        cmd.extend(command)
        cmd = " ".join(cmd)
        LOGGER.debug("Execute `%s'", cmd)
        res = self._helm_container.exec_run(["sh", "-c", cmd])
        if res.exit_code:
            raise DockerException(f"{self._helm_container}: {res.output.decode('utf-8')}")
        return res.output.decode("utf-8")
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-,
from __future__ import absolute_import
from __future__ import print_function

import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True

import csv
import numpy as np

def init_palette():
  from rootpy.plotting.style import set_style, get_style
  atlas = get_style('ATLAS')
  atlas.SetPalette(51)
  set_style(atlas)
  return True

def nbins(start, stop, step):
    return abs(int( (stop - start) / step ))

def init_canvas(x, y, name="c", topMargin=0.07, rightMargin=0.16):
    c = ROOT.TCanvas("c", "", 0, 0, x, y)
    c.SetTopMargin(topMargin)
    c.SetRightMargin(rightMargin)
    return c

def axis_labels(x_label="m_{#tilde{g}} [GeV]", y_label="m_{#tilde{#chi}^{0}_{1}} [GeV]", z_label="", title=""):
    return ';'.join([x_label, y_label, z_label, title])

def init_hist(label, x_min, x_max, y_min, y_max, bin_size, name="grid"):
    return ROOT.TH2F(name,
                axis_labels(z_label=label),
                nbins(x_min, x_max, bin_size),
                x_min,
                x_max,
                nbins(y_min, y_max, bin_size),
                y_min,
                y_max)

def fill_hist(hist,plot_array,label,skipNegativeSig=True):

  for i in range(len(plot_array[label])):
      g = int(plot_array['mgluino'][i])
      l = int(plot_array['mlsp'][i])
      z = plot_array[label][i]
      sig = plot_array['sig'][i]
      b = hist.FindFixBin(g,l)
      if(sig>0) or not(skipNegativeSig):
        xx=ROOT.Long(0)
        yy=ROOT.Long(0)
        zz=ROOT.Long(0)
        hist.GetBinXYZ(b,xx,yy,zz)
        z_old =  hist.GetBinContent(xx,yy)
        newz = max(z_old,z) #for significances this makes sense. For the other quantities not so much. Oh well.
        hist.SetBinContent(b,newz)
      else:
        hist.SetBinContent(b,0.01)

def draw_hist(hist, nSigs=1, markercolor=0, drawOpts="TEXT45 COLZ"):
    hist.SetMarkerSize(800)
    hist.SetMarkerColor(markercolor)
    #gStyle.SetPalette(51)
    ROOT.gStyle.SetPaintTextFormat("1.{0:d}f".format(nSigs));
    hist.Draw(drawOpts)

def draw_labels(lumi):
    txt = ROOT.TLatex()
    txt.SetNDC()
    txt.DrawText(0.32,0.87,"Internal")
    txt.DrawText(0.2,0.82,"Simulation")
    #txt.SetTextSize(0.030)
    txt.SetTextSize(18)
    txt.DrawLatex(0.16,0.95,"#tilde{g}-#tilde{g} production, #tilde{g} #rightarrow t #bar{t} + #tilde{#chi}^{0}_{1}")
    txt.DrawLatex(0.62,0.95,"L_{int} = %d fb^{-1}, #sqrt{s} = 13 TeV"% lumi)
    txt.SetTextFont(72)
    txt.SetTextSize(0.05)
    txt.DrawText(0.2,0.87,"ATLAS")

def draw_text(path):

    if path is None:
        return

    txt = TLatex()
    txt.SetNDC()
    txt.SetTextSize(0.030)

    with open(path,'r') as f:
        reader = csv.reader(f,delimiter=",")
        for row in reader:
            txt.DrawLatex(float(row[0]), float(row[1]), row[2])

def draw_line(x_min, y_min, x_max, y_max, topmass=173.34):
  l=ROOT.TLine(1000,1000,2000,2000)
  l.SetLineStyle(2)
  if x_min - 2*topmass > y_min:
    line_min_x = x_min
    line_min_y = x_min-2*topmass
  else:
    line_min_x = y_min + 2*topmass
    line_min_y = y_min

  if x_max - 2*topmass > y_max:
    line_max_x = y_max + 2*topmass
    line_max_y = y_max
  else:
    line_max_x = x_max
    line_max_y = x_max - 2*topmass

  l.DrawLine(line_min_x, line_min_y, line_max_x, line_max_y)
  # slope should be one as it's: LSP < Gluino - 2*topmass
  slope = float(line_max_y - line_min_y)/(line_max_x - line_min_x)

  # Draw Kinematically Forbidden as well
  txt = ROOT.TLatex()
  #txt.SetNDC()
  txt.SetTextFont(12)
  txt.SetTextAngle(np.degrees(np.arctan(slope)))
  txt.SetTextSize(0.02)
  txt.DrawText((line_max_x+line_min_x)/2., (line_max_x+line_min_x)/2. - 2*topmass + 125, "Kinematically Forbidden")

from array import *
def get_run1(filename,linestyle,linewidth,linecolor):
  x = array('f')
  y = array('f')
  n = 0
  with open(filename,'r') as csvfile:
    reader = csv.reader(csvfile, delimiter = ' ')
    for row in reader:
      n += 1
      x.append(float(row[0]))
      y.append(float(row[1]))

  gr = ROOT.TGraph(n,x,y)
  gr.SetLineColor(linecolor)
  gr.SetLineWidth(linewidth)
  gr.SetLineStyle(linestyle)
  return gr

def draw_run1_text(color):
    txt = ROOT.TLatex()
    txt.SetNDC()
    txt.SetTextFont(22)
    txt.SetTextSize(0.04)
    txt.SetTextColor(color)
    txt.DrawText(0.2,0.2,"Run 1 Limit")

def exclusion():
  x = array('d',[1400,1600,1600,1400])
  y = array('d',[600,600,800,600])
  p=ROOT.TPolyLine(4,x,y)
  p.SetFillColor(1)
  p.SetFillStyle(3001)
  #p.DrawPolyLine(4,x,y)
  return p
<commit_msg>titles not set right<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-,
from __future__ import absolute_import
from __future__ import print_function

import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True

import csv
import numpy as np

def init_palette():
  from rootpy.plotting.style import set_style, get_style
  atlas = get_style('ATLAS')
  atlas.SetPalette(51)
  set_style(atlas)
  return True

def nbins(start, stop, step):
    return abs(int( (stop - start) / step ))

def init_canvas(x, y, name="c", topMargin=0.07, rightMargin=0.16):
    c = ROOT.TCanvas("c", "", 0, 0, x, y)
    c.SetTopMargin(topMargin)
    c.SetRightMargin(rightMargin)
    return c

def axis_labels(x_label="m_{#tilde{g}} [GeV]", y_label="m_{#tilde{#chi}^{0}_{1}} [GeV]", z_label="", title=""):
    return ';'.join([title, x_label, y_label, z_label])

def init_hist(label, x_min, x_max, y_min, y_max, bin_size, name="grid"):
    return ROOT.TH2F(name,
                axis_labels(z_label=label),
                nbins(x_min, x_max, bin_size),
                x_min,
                x_max,
                nbins(y_min, y_max, bin_size),
                y_min,
                y_max)

def fill_hist(hist,plot_array,label,skipNegativeSig=True):

  for i in range(len(plot_array[label])):
      g = int(plot_array['mgluino'][i])
      l = int(plot_array['mlsp'][i])
      z = plot_array[label][i]
      sig = plot_array['sig'][i]
      b = hist.FindFixBin(g,l)
      if(sig>0) or not(skipNegativeSig):
        xx=ROOT.Long(0)
        yy=ROOT.Long(0)
        zz=ROOT.Long(0)
        hist.GetBinXYZ(b,xx,yy,zz)
        z_old =  hist.GetBinContent(xx,yy)
        newz = max(z_old,z) #for significances this makes sense. For the other quantities not so much. Oh well.
        hist.SetBinContent(b,newz)
      else:
        hist.SetBinContent(b,0.01)

def draw_hist(hist, nSigs=1, markercolor=0, drawOpts="TEXT45 COLZ"):
    hist.SetMarkerSize(800)
    hist.SetMarkerColor(markercolor)
    #gStyle.SetPalette(51)
    ROOT.gStyle.SetPaintTextFormat("1.{0:d}f".format(nSigs));
    hist.Draw(drawOpts)

def draw_labels(lumi):
    txt = ROOT.TLatex()
    txt.SetNDC()
    txt.DrawText(0.32,0.87,"Internal")
    txt.DrawText(0.2,0.82,"Simulation")
    #txt.SetTextSize(0.030)
    txt.SetTextSize(18)
    txt.DrawLatex(0.16,0.95,"#tilde{g}-#tilde{g} production, #tilde{g} #rightarrow t #bar{t} + #tilde{#chi}^{0}_{1}")
    txt.DrawLatex(0.62,0.95,"L_{int} = %d fb^{-1}, #sqrt{s} = 13 TeV"% lumi)
    txt.SetTextFont(72)
    txt.SetTextSize(0.05)
    txt.DrawText(0.2,0.87,"ATLAS")

def draw_text(path):

    if path is None:
        return

    txt = TLatex()
    txt.SetNDC()
    txt.SetTextSize(0.030)

    with open(path,'r') as f:
        reader = csv.reader(f,delimiter=",")
        for row in reader:
            txt.DrawLatex(float(row[0]), float(row[1]), row[2])

def draw_line(x_min, y_min, x_max, y_max, topmass=173.34):
  l=ROOT.TLine(1000,1000,2000,2000)
  l.SetLineStyle(2)
  if x_min - 2*topmass > y_min:
    line_min_x = x_min
    line_min_y = x_min-2*topmass
  else:
    line_min_x = y_min + 2*topmass
    line_min_y = y_min

  if x_max - 2*topmass > y_max:
    line_max_x = y_max + 2*topmass
    line_max_y = y_max
  else:
    line_max_x = x_max
    line_max_y = x_max - 2*topmass

  l.DrawLine(line_min_x, line_min_y, line_max_x, line_max_y)
  # slope should be one as it's: LSP < Gluino - 2*topmass
  slope = float(line_max_y - line_min_y)/(line_max_x - line_min_x)

  # Draw Kinematically Forbidden as well
  txt = ROOT.TLatex()
  #txt.SetNDC()
  txt.SetTextFont(12)
  txt.SetTextAngle(np.degrees(np.arctan(slope)))
  txt.SetTextSize(0.02)
  txt.DrawText((line_max_x+line_min_x)/2., (line_max_x+line_min_x)/2. - 2*topmass + 125, "Kinematically Forbidden")

from array import *
def get_run1(filename,linestyle,linewidth,linecolor):
  x = array('f')
  y = array('f')
  n = 0
  with open(filename,'r') as csvfile:
    reader = csv.reader(csvfile, delimiter = ' ')
    for row in reader:
      n += 1
      x.append(float(row[0]))
      y.append(float(row[1]))

  gr = ROOT.TGraph(n,x,y)
  gr.SetLineColor(linecolor)
  gr.SetLineWidth(linewidth)
  gr.SetLineStyle(linestyle)
  return gr

def draw_run1_text(color):
    txt = ROOT.TLatex()
    txt.SetNDC()
    txt.SetTextFont(22)
    txt.SetTextSize(0.04)
    txt.SetTextColor(color)
    txt.DrawText(0.2,0.2,"Run 1 Limit")

def exclusion():
  x = array('d',[1400,1600,1600,1400])
  y = array('d',[600,600,800,600])
  p=ROOT.TPolyLine(4,x,y)
  p.SetFillColor(1)
  p.SetFillStyle(3001)
  #p.DrawPolyLine(4,x,y)
  return p
<|endoftext|>


<commit_before>from django.conf.urls import patterns, url
from seadssite import views as v

urlpatterns = patterns('seadssite',
                       url(r'^$', v.IndexView.as_view()),
                       url(r'^dashboard/$', v.DashboardView),
                       url(r'^dashboard/[0-9]+/$', 'views.graph', name='graph'),
                       url(r'^dashboard/[0-9]+/appliances/$', v.DevicesView),
                       url(r'^help/$', v.help),)
<commit_msg>add back the register url<commit_after>from django.conf.urls import patterns, url
from seadssite import views as v

urlpatterns = patterns('seadssite',
                       url(r'^$', v.IndexView.as_view()),
                       url(r'^dashboard/$', v.DashboardView),
                       url(r'^dashboard/[0-9]+/$', 'views.graph', name='graph'),
                       url(r'^dashboard/[0-9]+/appliances/$', v.DevicesView),
                       url(r'^register/', v.register),
                       url(r'^help/$', v.help),)
<|endoftext|>


<commit_before>#!/usr/bin/env python
import roslib
roslib.load_manifest('roswww')
import rospy

import socket
import subprocess

import tornado.ioloop
import tornado.web

def run(*args):
    '''run the provided command and return its stdout'''
    args = sum([(arg if type(arg) == list else [arg]) for arg in args], [])
    return subprocess.Popen(args, stdout=subprocess.PIPE).communicate()[0].strip()

def split_words(text):
    '''return a list of lines where each line is a list of words'''
    return [line.strip().split() for line in text.split('\n')]

def get_packages():
    ''' Find the names and locations of all packages '''
    lines = split_words(run('rospack', 'list'))
    packages = [{ 'name': name, 'path': path } for name, path in lines]
    return packages

class MainHandler(tornado.web.RequestHandler):

    def initialize(self, packages):
        self.packages = packages

    def get(self):
        self.write("
ROS web server successfully started.

Package List
")
        for package in self.packages:
            self.write("
"+package['name']+"
")

def create_webserver(packages):
    handlers = [(r"/", MainHandler, {"packages": packages})]

    for package in packages:
        handler = ("/"+package['name']+"/(.*)", tornado.web.StaticFileHandler, {"path": package['path']+"/www", "default_filename": "index.html" })
        handlers.append(handler)

    rospy.loginfo("Webserver initialized for %d packages", len(packages))
    application = tornado.web.Application(handlers)

    return application

def bind_webserver(application):
    """ See if there's a default port, use 80 if not """
    default_port, start_port, end_port = get_webserver_params()

    """ First, we try the default http port 80 """
    bound = bind_to_port(application, default_port)

    if not bound:
        """ Otherwise bind any available port within the specified range """
        bound = bind_in_range(application, start_port, end_port)

    return bound

def get_webserver_params():
    try:
        default_port = rospy.get_param("http/default", 80)
        start_port = rospy.get_param("http/range_start", 8000)
        end_port = rospy.get_param("http/range_end", 9000)
        return (default_port, start_port, end_port)
    except socket.error as err:
        if err.errno == 111:
            # Roscore isn't started or cannot be contacted
            rospy.logwarn("Could not contact ROS master. Is a roscore running? Error: %s", err.strerror)
            return 80, 8000, 9000
        else:
            raise

def start_webserver(application):
    try:
        tornado.ioloop.IOLoop.instance().start()
    except KeyboardInterrupt:
        rospy.loginfo("Webserver shutting down")

def bind_to_port(application, portno):
    rospy.loginfo("Attempting to start webserver on port %d", portno)
    try:
        application.listen(portno)
        rospy.loginfo("Webserver successfully started on port %d", portno)
    except socket.error as err:
        # Socket exceptions get handled, all other exceptions propagated
        if err.errno==13:
            rospy.logwarn("Insufficient priveliges to run webserver on port %d. Error: %s", portno, err.strerror)
            rospy.loginfo("-- Try re-running as super-user: sudo su; source ~/.bashrc)")
        elif err.errno==98:
            rospy.logwarn("There is already a webserver running on port %d. Error: %s", portno, err.strerror)
            rospy.loginfo("-- Try stopping your web server. For example, to stop apache: sudo /etc/init.d/apache2 stop")
        else:
            rospy.logerr("An error occurred attempting to listen on port %d: %s", portno, err.strerror)
        return False
    return True

def bind_in_range(application, start_port, end_port):
    if (end_port > start_port):
        for i in range(start_port, end_port):
            if bind_to_port(application, i):
                return True
    return False

def run_webserver():
    try:
        packages = get_packages()
        server = create_webserver(packages)
        bound = bind_webserver(server)
        if (bound):
            start_webserver(server)
        else:
            raise Exception()
    except Exception as exc:
        rospy.logerr("Unable to bind webserver.  Exiting.  %s" % exc)

if __name__=='__main__':
    run_webserver()<commit_msg>add handler of /ROS_PACKAGE/?<commit_after>#!/usr/bin/env python
import roslib
roslib.load_manifest('roswww')
import rospy

import socket
import subprocess

import tornado.ioloop
import tornado.web

def run(*args):
    '''run the provided command and return its stdout'''
    args = sum([(arg if type(arg) == list else [arg]) for arg in args], [])
    return subprocess.Popen(args, stdout=subprocess.PIPE).communicate()[0].strip()

def split_words(text):
    '''return a list of lines where each line is a list of words'''
    return [line.strip().split() for line in text.split('\n')]

def get_packages():
    ''' Find the names and locations of all packages '''
    lines = split_words(run('rospack', 'list'))
    packages = [{ 'name': name, 'path': path } for name, path in lines]
    return packages

class MainHandler(tornado.web.RequestHandler):

    def initialize(self, packages):
        self.packages = packages

    def get(self):
        self.write("
ROS web server successfully started.

Package List
")
        for package in self.packages:
            self.write("
"+package['name']+"
")

def create_webserver(packages):
    handlers = [(r"/", MainHandler, {"packages": packages})]

    for package in packages:
        handler_root = ("/"+package['name']+"/?()", tornado.web.StaticFileHandler, {"path": package['path']+"/www/index.html" })
        handlers.append(handler_root)
        handler = ("/"+package['name']+"/(.*)", tornado.web.StaticFileHandler, {"path": package['path']+"/www", "default_filename": "index.html" })
        handlers.append(handler)

    rospy.loginfo("Webserver initialized for %d packages", len(packages))
    application = tornado.web.Application(handlers)

    return application

def bind_webserver(application):
    """ See if there's a default port, use 80 if not """
    default_port, start_port, end_port = get_webserver_params()

    """ First, we try the default http port 80 """
    bound = bind_to_port(application, default_port)

    if not bound:
        """ Otherwise bind any available port within the specified range """
        bound = bind_in_range(application, start_port, end_port)

    return bound

def get_webserver_params():
    try:
        default_port = rospy.get_param("http/default", 80)
        start_port = rospy.get_param("http/range_start", 8000)
        end_port = rospy.get_param("http/range_end", 9000)
        return (default_port, start_port, end_port)
    except socket.error as err:
        if err.errno == 111:
            # Roscore isn't started or cannot be contacted
            rospy.logwarn("Could not contact ROS master. Is a roscore running? Error: %s", err.strerror)
            return 80, 8000, 9000
        else:
            raise

def start_webserver(application):
    try:
        tornado.ioloop.IOLoop.instance().start()
    except KeyboardInterrupt:
        rospy.loginfo("Webserver shutting down")

def bind_to_port(application, portno):
    rospy.loginfo("Attempting to start webserver on port %d", portno)
    try:
        application.listen(portno)
        rospy.loginfo("Webserver successfully started on port %d", portno)
    except socket.error as err:
        # Socket exceptions get handled, all other exceptions propagated
        if err.errno==13:
            rospy.logwarn("Insufficient priveliges to run webserver on port %d. Error: %s", portno, err.strerror)
            rospy.loginfo("-- Try re-running as super-user: sudo su; source ~/.bashrc)")
        elif err.errno==98:
            rospy.logwarn("There is already a webserver running on port %d. Error: %s", portno, err.strerror)
            rospy.loginfo("-- Try stopping your web server. For example, to stop apache: sudo /etc/init.d/apache2 stop")
        else:
            rospy.logerr("An error occurred attempting to listen on port %d: %s", portno, err.strerror)
        return False
    return True

def bind_in_range(application, start_port, end_port):
    if (end_port > start_port):
        for i in range(start_port, end_port):
            if bind_to_port(application, i):
                return True
    return False

def run_webserver():
    try:
        packages = get_packages()
        server = create_webserver(packages)
        bound = bind_webserver(server)
        if (bound):
            start_webserver(server)
        else:
            raise Exception()
    except Exception as exc:
        rospy.logerr("Unable to bind webserver.  Exiting.  %s" % exc)

if __name__=='__main__':
    run_webserver()
<|endoftext|>


<commit_before>import os
import sys

from sea.utils import import_string
from sea.local import Proxy

__version__ = '0.11.0'
_app = None

def create_app(root_path=None):
    global _app
    if _app is not None:
        return _app

    if root_path is None:
        root_path = os.getcwd()
    sys.path.append(root_path)
    sys.path.append(os.path.join(root_path, 'protos'))

    env = os.environ.get('SEA_ENV', 'development')
    config = import_string('configs.{}'.format(env))

    app_class = import_string('app:App')
    _app = app_class(root_path, env=env)
    _app.config.from_object(config)

    _app.load_middlewares()
    _app.load_extensions_in_module(import_string('app.extensions'))
    _app.load_servicers_in_module(import_string('app.servicers'))

    _app.ready()

    return _app

current_app = Proxy(lambda: _app)
<commit_msg>upgrade<commit_after>import os
import sys

from sea.utils import import_string
from sea.local import Proxy

__version__ = '0.12.0'
_app = None

def create_app(root_path=None):
    global _app
    if _app is not None:
        return _app

    if root_path is None:
        root_path = os.getcwd()
    sys.path.append(root_path)
    sys.path.append(os.path.join(root_path, 'protos'))

    env = os.environ.get('SEA_ENV', 'development')
    config = import_string('configs.{}'.format(env))

    app_class = import_string('app:App')
    _app = app_class(root_path, env=env)
    _app.config.from_object(config)

    _app.load_middlewares()
    _app.load_extensions_in_module(import_string('app.extensions'))
    _app.load_servicers_in_module(import_string('app.servicers'))

    _app.ready()

    return _app

current_app = Proxy(lambda: _app)
<|endoftext|>


<commit_before>import argparse
import re
import sys
from subprocess import PIPE, STDOUT, Popen
import threading
import Queue

from login import _print, webwxsendmsgtome, dic, init
from sendmessage import webwxsendmsg

mode = {'p': '([0-9]+%)|Done'}
powershell = 'C:\Windows.old\WINDOWS\SysWOW64\WindowsPowerShell\v1.0\powershell.exe'

class run_subprocess(threading.Thread):
    def __init__(self, cmd):
        threading.Thread.__init__(self)
        self.cmd = cmd

    def run(self):
        global q
        launch_str = self.cmd
        p = Popen(launch_str, stdout=PIPE, stderr=STDOUT, shell=True)
        _print('Running ' + launch_str)
        while True:
            line = p.stdout.readline()
            if not line:
                break
            if not q.full():
                q.put(line)
            else:
                q.get()
                q.put(line)

q = Queue.Queue(100)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--user', help='username or nickname to send(default send to filehelper)', type=str)
    parser.add_argument('--file', help='if cmd will get a file name')
    parser.add_argument('--cmd', '-c', help='cmd or filename', type=str)
    parser.add_argument('-m',
                        '--mode', help='re mode, choose this --re will not work', type=str)
    parser.add_argument('-r', '--re', help='regular express', type=str)
    args = parser.parse_args()
    init()
    if args.re != None:
        patten = re.compile(args.re)
    if args.mode != None:
        try:
            patten = re.compile(mode[args.mode])
        except:
            _print('No mode named {}'.format(args.mode))
            sys.exit()
    if args.file:
        # todo
        pass
    else:
        thread_subprocess = run_subprocess(args.cmd)
        thread_subprocess.start()
    flag = True
    _print('start to send')
    while True:
        if q.empty():
            time.sleep(1)
            continue
        line = q.get()
        try:
            if args.re != None or args.mode != None:
                try:
                    line = patten.findall(line)[0]
                except:
                    _print('Ignore:\t' + line.split('\n')[0])
                    continue
            if args.user != None:
                if not webwxsendmsg(friend, line):
                    _print('Fail')
                else:
                    _print('Send\t' + line.split('\n')[0])
            else:
                if not webwxsendmsgtome(line):
                    _print('Fail')
                else:
                    _print('Send\t' + line.split('\n')[0])
        except:
            _print('Send failed')

if __name__ == '__main__':
    main()
<commit_msg>missing import time<commit_after>import argparse
import re
import sys
from subprocess import PIPE, STDOUT, Popen
import threading
import Queue
import time

from login import _print, webwxsendmsgtome, dic, init
from sendmessage import webwxsendmsg

mode = {'p': '([0-9]+%)|Done'}
powershell = 'C:\Windows.old\WINDOWS\SysWOW64\WindowsPowerShell\v1.0\powershell.exe'

class run_subprocess(threading.Thread):
    def __init__(self, cmd):
        threading.Thread.__init__(self)
        self.cmd = cmd

    def run(self):
        global q
        launch_str = self.cmd
        p = Popen(launch_str, stdout=PIPE, stderr=STDOUT, shell=True)
        _print('Running ' + launch_str)
        while True:
            line = p.stdout.readline()
            if not line:
                break
            if not q.full():
                q.put(line)
            else:
                q.get()
                q.put(line)

q = Queue.Queue(100)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--user', help='username or nickname to send(default send to filehelper)', type=str)
    parser.add_argument('--file', help='if cmd will get a file name')
    parser.add_argument('--cmd', '-c', help='cmd or filename', type=str)
    parser.add_argument('-m',
                        '--mode', help='re mode, choose this --re will not work', type=str)
    parser.add_argument('-r', '--re', help='regular express', type=str)
    args = parser.parse_args()
    init()
    if args.re != None:
        patten = re.compile(args.re)
    if args.mode != None:
        try:
            patten = re.compile(mode[args.mode])
        except:
            _print('No mode named {}'.format(args.mode))
            sys.exit()
    if args.file:
        # todo
        pass
    else:
        thread_subprocess = run_subprocess(args.cmd)
        thread_subprocess.start()
    flag = True
    _print('start to send')
    while True:
        if q.empty():
            time.sleep(1)
            continue
        line = q.get()
        try:
            if args.re != None or args.mode != None:
                try:
                    line = patten.findall(line)[0]
                except:
                    _print('Ignore:\t' + line.split('\n')[0])
                    continue
            if args.user != None:
                if not webwxsendmsg(friend, line):
                    _print('Fail')
                else:
                    _print('Send\t' + line.split('\n')[0])
            else:
                if not webwxsendmsgtome(line):
                    _print('Fail')
                else:
                    _print('Send\t' + line.split('\n')[0])
        except:
            _print('Send failed')

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>import pysam

class PacBioInputBAM(object):

    def __init__(self, filename):
        self.filename = filename
        self.data = pysam.AlignmentFile(filename, check_sq=False)

    def reset(self):
        self.data.close()
        self.data = pysam.AlignmentFile(self.filename, check_sq=False)

    def stride(self, output_filename, stride=10):
        with pysam.AlignmentFile(output_filename,  "wb", template=self.data) as fh:

            for i, read in enumerate(self.data):
                if i % stride == 0:
                    fh.write(read)

<commit_msg>add doc<commit_after>import pysam

class PacBioInputBAM(object):
    """PacBio utilities

    Downsample PacBio base-call BAM file

    """
    def __init__(self, filename):
        self.filename = filename
        self.data = pysam.AlignmentFile(filename, check_sq=False)

    def reset(self):
        self.data.close()
        self.data = pysam.AlignmentFile(self.filename, check_sq=False)

    def stride(self, output_filename, stride=10):
        with pysam.AlignmentFile(output_filename,  "wb", template=self.data) as fh:

            for i, read in enumerate(self.data):
                if i % stride == 0:
                    fh.write(read)

<|endoftext|>


<commit_before>"""
Utility classes and functions related to computing fairness evaluations

:author: Anastassia Loukina (aloukina@ets.org)
:author: Jeremy Biggs (jbiggs@ets.org)
:author: Nitin Madnani (nmadnani@ets.org)

:date: 07/29/2019
:organization: ETS
"""

import pandas as pd
import pickle
import numpy as np

from os.path import join

import statsmodels.formula.api as smf
from statsmodels.stats.anova import anova_lm

from rsmtool.writer import DataWriter
from rsmtool.container import DataContainer

def convert_to_ordered_category(group_values):
    """
    Convert supplied series to an ordered category
    with levels ordered by category size.
    If multiple categories have the same size,
    the order is determined alphabetically.

    Parameters
    ----------
    group_values: pandas Series

    Returns
    -------
    group_ordered_category: pandas Series
    """

    # get ordered list by size

    # We convert the  value counts to data frame to allow for multilevel sorting.
    # This makes sure that the order is consistent and reproducible across runs
    # when there are more than
    # one group with the maximum number of occurrences.
    df_groups_by_size = pd.DataFrame(group_values.value_counts()).reset_index()
    df_groups_by_size.columns = ['group_name', 'count']
    df_groups_by_size_sorted = df_groups_by_size.sort_values(['count', 'group_name'],
                                                            ascending=[False, True])
    groups_by_size = df_groups_by_size_sorted['group_name'].values

    #convert to category and reorder
    group_category = group_values.astype("category")
    group_ordered_category = group_category.cat.reorder_categories(groups_by_size,
                                                                   ordered=True)
    return group_ordered_category

def get_coefficients(fit,
                     base_category):

    """
    Extract estimates, significance and confidence intervals
    for members of the group of interest. The names of the
    predictors are processed to remove the prefix added by Statmodels.
    The name of the base category is added in parenthesis to the Intercept.

    Parameters:
    -----------
    fit: fit
        Linear regression object fitted using statsmodels

    base_category: str
        Name of the group used as reference category when fitting the model

    Returns:
    -------
    df_results: pandas DataFrame
        A dataframe with rows for each category and the following columns:
        - estimate
        - P>[t]
        - Lower and upper boundaries for 95% confidence intervals
    """

    # extract the data we need
    df_results = pd.concat([fit.params, fit.pvalues, fit.conf_int()], axis=1)

    df_results.columns = ['estimate', 'P>[t]', '[0.025', '0.975]']

    # identify the rows we are interested in
    groups = ['Intercept'] + [v for v in df_results.index if 'group' in v]

    df_results = df_results.loc[groups]

    # rename the rows
    df_results.index = [v.split('.')[1].strip(']')
                            if not v == 'Intercept'
                            else 'Intercept ({})'.format(base_category)
                            for v in df_results.index]

    return df_results

def write_fairness_results(fit_dictionary,
                           fairness_container,
                           group,
                           output_dir,
                           experiment_id,
                           file_format):
    """
    Save the results of fairness analysis to disk.

    Parameters
    ----------
    fit_dictionary: dict
        A dictionary of fitted models generated by `get_fairness_analysis()`
    fairness_container: container.DataContainer
        A data container with the results of fairness analysis generated by get_fairness_analysis
    group: str
        The subgroup considered in this analysis
    output_dir: str
        The directory where the results will be saved
    experiment_id: str
        Experiment id
    file_format: str
        File format for data files
    """

    # let's first save model files and summaries
    for model in fit_dictionary:
        fit = fit_dictionary[model]

        ols_file = join(output_dir, '{}_{}_by_{}.ols'.format(experiment_id,
                                                             model,
                                                             group))
        summary_file = join(output_dir, '{}_{}_by_{}_ols_summary.txt'.format(experiment_id,
                                                                             model,
                                                                             group))
        with open(ols_file, 'wb') as olsf, open(summary_file, 'w') as summf:
            pickle.dump(fit, olsf)
            summf.write(str(fit.summary()))

    # Now let's write out the content of the data container
    writer = DataWriter(experiment_id)
    writer.write_experiment_output(output_dir,
                                   fairness_container,
                                   file_format=file_format,
                                   index=True)

def get_fairness_analysis(df,
                          group,
                          system_score_column,
                          human_score_column='sc1'):
    """
    Compute fairness analyses described in Loukina et al. (2019) (https://aclweb.org/anthology/papers/W/W19/W19-4401/)
    The functions computes how muach variance group memebership explains in
    overall score accuracy (osa), overall score difference (osd) and conditional score difference (csd).
    See paper for further detail.

    Parameters
    ----------
    df: pandas DataFrame
        A dataframe containing columns with numeric human scores, numeric system scores
        and a column showing group membership
    group: str
        Name of the column containing group membership
    system_score_column: str
        Name of the column containing system scores
    human_score_column: str
        Name of the column containing human scores

    Returns
    -------
    model_dict: dictionary
        A dictionary with different proposed metrics as keys
        and fitted models as values
    fairness_container: container.DataContainer
        A datacontainer with the following datasets:
            - `estimates_METRICS_by_GROUP` where `GROUP` corresponds to group
            and `METRICS` can be osa, osd and csd
            Estimates for each group computed by the respective models
            - `fairness_metrics_by_GROUP` - a summary of model fits (R2 and p values)
    """

    # compute error and squared error

    df['error'] = df[system_score_column]-df[human_score_column]
    df['SE'] = df['error']**2

    # convert group values to category and reorder them using
    # the largest category as reference

    df['group'] = convert_to_ordered_category(df[group])
    base_group = df['group'].cat.categories[0]
    df['sc1_cat'] = convert_to_ordered_category(df[human_score_column])

    # Overall score accuracy (OSA)
    # Variance in squared error explained by L1

    # fit the model
    osa_model = smf.ols(formula='SE ~ group', data=df)
    osa_fit = osa_model.fit()

    # collect the results
    osa_dict = {'R2': osa_fit.rsquared_adj,
               'sig': osa_fit.f_pvalue}

    osa_results = pd.Series(osa_dict, name='Overall score accuracy')

    df_coefficients_osa = get_coefficients(osa_fit, base_group)

    # Overall score difference (OSD)
    # variance in signed residuals (raw error) explained by L1

    # fit the model
    osd_model = smf.ols(formula='error ~ group', data=df)
    osd_fit = osd_model.fit()

    # collect the results
    osd_dict = {'R2': osd_fit.rsquared_adj,
               'sig': osd_fit.f_pvalue}

    osd_results = pd.Series(osd_dict, name='Overall score difference')

    df_coefficients_osd = get_coefficients(osd_fit, base_group)

    # conditional score difference CSD
    # Variance in score difference conditioned on Native language

    # fit "null" model with human score only
    csd_null_mod = smf.ols(formula='error ~ sc1_cat', data=df)
    csd_null_fit = csd_null_mod.fit()

    # fit model with both human score and group
    csd_mod = smf.ols(formula='error ~ group + sc1_cat', data=df)
    csd_fit = csd_mod.fit()

    # compare the two models using anova_lm
    # we filter warnings for this function because we get
    # runtime warning due to NaNs in the data.
    # these seem to be by design: https://groups.google.com/forum/#!topic/pystatsmodels/-flY0cNnb3k
    np.warnings.filterwarnings('ignore')
    anova_results = anova_lm(csd_null_fit, csd_fit)
    # we reset warnings
    np.warnings.resetwarnings()

    # collect the results. Note that R2 in this case is a difference
    # in R2 between the two models and significance is obtained from anova
    csd_dict = {'R2': csd_fit.rsquared_adj - csd_null_fit.rsquared_adj,
               'sig': anova_results.values[1][-1]}
    csd_results = pd.Series(csd_dict, name='Conditional score difference')

    df_coefficients_csd = get_coefficients(osd_fit, base_group)

    # create a summary table

    df_r2_all = pd.concat([osa_results, osd_results, csd_results], axis=1, sort=True)
    df_r2_all['base_category'] = base_group

    # assemble all datasets into a DataContainer

    datasets = [{'name': 'estimates_osa_by_{}'.format(group),
                 'frame': df_coefficients_osa},
                {'name': 'estimates_osd_by_{}'.format(group),
                 'frame': df_coefficients_osd},
                {'name': 'estimates_csd_by_{}'.format(group),
                 'frame': df_coefficients_csd},
                {'name': 'fairness_metrics_by_{}'.format(group),
                 'frame': df_r2_all}]

    # assemble all models into a dictionary
    model_dict = {'osa': osa_fit,
                  'osd': osd_fit,
                  'csd': csd_fit}

    return model_dict, DataContainer(datasets=datasets)

<commit_msg>Update rsmtool/fairness_utils.py<commit_after>"""
Utility classes and functions related to computing fairness evaluations

:author: Anastassia Loukina (aloukina@ets.org)
:author: Jeremy Biggs (jbiggs@ets.org)
:author: Nitin Madnani (nmadnani@ets.org)

:date: 07/29/2019
:organization: ETS
"""

import pandas as pd
import pickle
import numpy as np

from os.path import join

import statsmodels.formula.api as smf
from statsmodels.stats.anova import anova_lm

from rsmtool.writer import DataWriter
from rsmtool.container import DataContainer

def convert_to_ordered_category(group_values):
    """
    Convert supplied series to an ordered category
    with levels ordered by category size.
    If multiple categories have the same size,
    the order is determined alphabetically.

    Parameters
    ----------
    group_values: pandas Series

    Returns
    -------
    group_ordered_category: pandas Series
    """

    # get ordered list by size

    # We convert the  value counts to data frame to allow for multilevel sorting.
    # This makes sure that the order is consistent and reproducible across runs
    # when there are more than
    # one group with the maximum number of occurrences.
    df_groups_by_size = pd.DataFrame(group_values.value_counts()).reset_index()
    df_groups_by_size.columns = ['group_name', 'count']
    df_groups_by_size_sorted = df_groups_by_size.sort_values(['count', 'group_name'],
                                                            ascending=[False, True])
    groups_by_size = df_groups_by_size_sorted['group_name'].values

    #convert to category and reorder
    group_category = group_values.astype("category")
    group_ordered_category = group_category.cat.reorder_categories(groups_by_size,
                                                                   ordered=True)
    return group_ordered_category

def get_coefficients(fit,
                     base_category):

    """
    Extract estimates, significance and confidence intervals
    for members of the group of interest. The names of the
    predictors are processed to remove the prefix added by Statmodels.
    The name of the base category is added in parenthesis to the Intercept.

    Parameters:
    -----------
    fit: fit
        Linear regression object fitted using statsmodels

    base_category: str
        Name of the group used as reference category when fitting the model

    Returns:
    -------
    df_results: pandas DataFrame
        A dataframe with rows for each category and the following columns:
        - estimate
        - P>[t]
        - Lower and upper boundaries for 95% confidence intervals
    """

    # extract the data we need
    df_results = pd.concat([fit.params, fit.pvalues, fit.conf_int()], axis=1)

    df_results.columns = ['estimate', 'P>[t]', '[0.025', '0.975]']

    # identify the rows we are interested in
    groups = ['Intercept'] + [v for v in df_results.index if 'group' in v]

    df_results = df_results.loc[groups]

    # rename the rows
    df_results.index = [v.split('.')[1].strip(']')
                            if not v == 'Intercept'
                            else 'Intercept ({})'.format(base_category)
                            for v in df_results.index]

    return df_results

def write_fairness_results(fit_dictionary,
                           fairness_container,
                           group,
                           output_dir,
                           experiment_id,
                           file_format):
    """
    Save the results of fairness analysis to disk.

    Parameters
    ----------
    fit_dictionary: dict
        A dictionary of fitted models generated by `get_fairness_analysis()`
    fairness_container: container.DataContainer
        A data container with the results of fairness analysis generated by `get_fairness_analysis()`
    group: str
        The subgroup considered in this analysis
    output_dir: str
        The directory where the results will be saved
    experiment_id: str
        Experiment id
    file_format: str
        File format for data files
    """

    # let's first save model files and summaries
    for model in fit_dictionary:
        fit = fit_dictionary[model]

        ols_file = join(output_dir, '{}_{}_by_{}.ols'.format(experiment_id,
                                                             model,
                                                             group))
        summary_file = join(output_dir, '{}_{}_by_{}_ols_summary.txt'.format(experiment_id,
                                                                             model,
                                                                             group))
        with open(ols_file, 'wb') as olsf, open(summary_file, 'w') as summf:
            pickle.dump(fit, olsf)
            summf.write(str(fit.summary()))

    # Now let's write out the content of the data container
    writer = DataWriter(experiment_id)
    writer.write_experiment_output(output_dir,
                                   fairness_container,
                                   file_format=file_format,
                                   index=True)

def get_fairness_analysis(df,
                          group,
                          system_score_column,
                          human_score_column='sc1'):
    """
    Compute fairness analyses described in Loukina et al. (2019) (https://aclweb.org/anthology/papers/W/W19/W19-4401/)
    The functions computes how muach variance group memebership explains in
    overall score accuracy (osa), overall score difference (osd) and conditional score difference (csd).
    See paper for further detail.

    Parameters
    ----------
    df: pandas DataFrame
        A dataframe containing columns with numeric human scores, numeric system scores
        and a column showing group membership
    group: str
        Name of the column containing group membership
    system_score_column: str
        Name of the column containing system scores
    human_score_column: str
        Name of the column containing human scores

    Returns
    -------
    model_dict: dictionary
        A dictionary with different proposed metrics as keys
        and fitted models as values
    fairness_container: container.DataContainer
        A datacontainer with the following datasets:
            - `estimates_METRICS_by_GROUP` where `GROUP` corresponds to group
            and `METRICS` can be osa, osd and csd
            Estimates for each group computed by the respective models
            - `fairness_metrics_by_GROUP` - a summary of model fits (R2 and p values)
    """

    # compute error and squared error

    df['error'] = df[system_score_column]-df[human_score_column]
    df['SE'] = df['error']**2

    # convert group values to category and reorder them using
    # the largest category as reference

    df['group'] = convert_to_ordered_category(df[group])
    base_group = df['group'].cat.categories[0]
    df['sc1_cat'] = convert_to_ordered_category(df[human_score_column])

    # Overall score accuracy (OSA)
    # Variance in squared error explained by L1

    # fit the model
    osa_model = smf.ols(formula='SE ~ group', data=df)
    osa_fit = osa_model.fit()

    # collect the results
    osa_dict = {'R2': osa_fit.rsquared_adj,
               'sig': osa_fit.f_pvalue}

    osa_results = pd.Series(osa_dict, name='Overall score accuracy')

    df_coefficients_osa = get_coefficients(osa_fit, base_group)

    # Overall score difference (OSD)
    # variance in signed residuals (raw error) explained by L1

    # fit the model
    osd_model = smf.ols(formula='error ~ group', data=df)
    osd_fit = osd_model.fit()

    # collect the results
    osd_dict = {'R2': osd_fit.rsquared_adj,
               'sig': osd_fit.f_pvalue}

    osd_results = pd.Series(osd_dict, name='Overall score difference')

    df_coefficients_osd = get_coefficients(osd_fit, base_group)

    # conditional score difference CSD
    # Variance in score difference conditioned on Native language

    # fit "null" model with human score only
    csd_null_mod = smf.ols(formula='error ~ sc1_cat', data=df)
    csd_null_fit = csd_null_mod.fit()

    # fit model with both human score and group
    csd_mod = smf.ols(formula='error ~ group + sc1_cat', data=df)
    csd_fit = csd_mod.fit()

    # compare the two models using anova_lm
    # we filter warnings for this function because we get
    # runtime warning due to NaNs in the data.
    # these seem to be by design: https://groups.google.com/forum/#!topic/pystatsmodels/-flY0cNnb3k
    np.warnings.filterwarnings('ignore')
    anova_results = anova_lm(csd_null_fit, csd_fit)
    # we reset warnings
    np.warnings.resetwarnings()

    # collect the results. Note that R2 in this case is a difference
    # in R2 between the two models and significance is obtained from anova
    csd_dict = {'R2': csd_fit.rsquared_adj - csd_null_fit.rsquared_adj,
               'sig': anova_results.values[1][-1]}
    csd_results = pd.Series(csd_dict, name='Conditional score difference')

    df_coefficients_csd = get_coefficients(osd_fit, base_group)

    # create a summary table

    df_r2_all = pd.concat([osa_results, osd_results, csd_results], axis=1, sort=True)
    df_r2_all['base_category'] = base_group

    # assemble all datasets into a DataContainer

    datasets = [{'name': 'estimates_osa_by_{}'.format(group),
                 'frame': df_coefficients_osa},
                {'name': 'estimates_osd_by_{}'.format(group),
                 'frame': df_coefficients_osd},
                {'name': 'estimates_csd_by_{}'.format(group),
                 'frame': df_coefficients_csd},
                {'name': 'fairness_metrics_by_{}'.format(group),
                 'frame': df_r2_all}]

    # assemble all models into a dictionary
    model_dict = {'osa': osa_fit,
                  'osd': osd_fit,
                  'csd': csd_fit}

    return model_dict, DataContainer(datasets=datasets)

<|endoftext|>


<commit_before>"""
Utility classes and functions related to computing fairness evaluations

:author: Anastassia Loukina (aloukina@ets.org)
:author: Jeremy Biggs (jbiggs@ets.org)
:author: Nitin Madnani (nmadnani@ets.org)

:date: 07/29/2019
:organization: ETS
"""

import pandas as pd
import pickle
import numpy as np

from os.path import join

import statsmodels.formula.api as smf
from statsmodels.stats.anova import anova_lm

from rsmtool.writer import DataWriter
from rsmtool.container import DataContainer

def convert_to_ordered_category(group_values):
    """
    Convert supplied series to an ordered category
    with levels ordered by category size.
    If multiple categories have the same size,
    the order is determined alphabetically.

    Parameters
    ----------
    group_values: pandas Series

    Returns
    -------
    group_ordered_category: pandas Series
    """

    # get ordered list by size

    # We convert the  value counts to data frame to allow for multilevel sorting.
    # This makes sure that the order is consistent and reproducible across runs
    # when there are more than
    # one group with the maximum number of occurrences.
    df_groups_by_size = pd.DataFrame(group_values.value_counts()).reset_index()
    df_groups_by_size.columns = ['group_name', 'count']
    df_groups_by_size_sorted = df_groups_by_size.sort_values(['count', 'group_name'],
                                                             ascending=[False, True])
    groups_by_size = df_groups_by_size_sorted['group_name'].values

    #convert to category and reorder
    group_category = group_values.astype("category")
    group_ordered_category = group_category.cat.reorder_categories(groups_by_size,
                                                                   ordered=True)
    return group_ordered_category

def get_coefficients(fit, base_category):

    """
    Extract estimates, significance, and confidence intervals
    for members of the group of interest. The names of the
    predictors are processed to remove the prefix added by Statmodels.
    The name of the base category is added in parenthesis to the Intercept.

    Parameters:
    -----------
    fit: fit
        Linear regression object fitted using statsmodels

    base_category: str
        Name of the group used as reference category when fitting the model

    Returns:
    -------
    df_results: pandas DataFrame
        A dataframe with rows for each category and the following columns:
        - estimate
        - P>[t]
        - Lower and upper boundaries for 95% confidence intervals
    """

    # extract the data we need
    df_results = pd.concat([fit.params, fit.pvalues, fit.conf_int()], axis=1)

    df_results.columns = ['estimate', 'P>[t]', '[0.025', '0.975]']

    # identify the rows we are interested in
    groups = ['Intercept'] + [v for v in df_results.index if 'group' in v]

    df_results = df_results.loc[groups]

    # rename the rows
    df_results.index = [v.split('.')[1].strip(']')
                            if not v == 'Intercept'
                            else 'Intercept ({})'.format(base_category)
                            for v in df_results.index]

    return df_results

def write_fairness_results(fit_dictionary,
                           fairness_container,
                           group,
                           output_dir,
                           experiment_id,
                           file_format):
    """
    Save the results of fairness analysis to disk.

    Parameters
    ----------
    fit_dictionary: dict
        A dictionary of fitted models generated by `get_fairness_analysis()`
    fairness_container: container.DataContainer
        A data container with the results of fairness analysis generated by `get_fairness_analysis()`
    group: str
        The subgroup considered in this analysis
    output_dir: str
        The directory where the results will be saved
    experiment_id: str
        Experiment id
    file_format: str
        File format for data files
    """

    # let's first save model files and summaries
    for model in fit_dictionary:
        fit = fit_dictionary[model]

        ols_file = join(output_dir, '{}_{}_by_{}.ols'.format(experiment_id,
                                                             model,
                                                             group))
        summary_file = join(output_dir, '{}_{}_by_{}_ols_summary.txt'.format(experiment_id,
                                                                             model,
                                                                             group))
        with open(ols_file, 'wb') as olsf, open(summary_file, 'w') as summf:
            pickle.dump(fit, olsf)
            summf.write(str(fit.summary()))

    # Now let's write out the content of the data container
    writer = DataWriter(experiment_id)
    writer.write_experiment_output(output_dir,
                                   fairness_container,
                                   file_format=file_format,
                                   index=True)

def get_fairness_analyses(df,
                          group,
                          system_score_column,
                          human_score_column='sc1'):
    """
    Compute fairness analyses described in Loukina et al. (2019) https://aclweb.org/anthology/papers/W/W19/W19-4401/
    The functions computes how muach variance group memebership explains in
    overall score accuracy (osa), overall score difference (osd) and conditional score difference (csd).
    See paper for further detail.

    Parameters
    ----------
    df: pandas DataFrame
        A dataframe containing columns with numeric human scores, numeric system scores
        and a column showing group membership
    group: str
        Name of the column containing group membership
    system_score_column: str
        Name of the column containing system scores
    human_score_column: str
        Name of the column containing human scores

    Returns
    -------
    model_dict: dictionary
        A dictionary with different proposed metrics as keys
        and fitted models as values
    fairness_container: container.DataContainer
        A datacontainer with the following datasets:
            - `estimates_METRICS_by_GROUP` where `GROUP` corresponds to group
            and `METRICS` can be osa, osd and csd
            Estimates for each group computed by the respective models
            - `fairness_metrics_by_GROUP` - a summary of model fits (R2 and p values)
    """

    # compute error and squared error

    df['error'] = df[system_score_column]-df[human_score_column]
    df['SE'] = df['error']**2

    # convert group values to category and reorder them using
    # the largest category as reference

    df['group'] = convert_to_ordered_category(df[group])
    base_group = df['group'].cat.categories[0]
    df['sc1_cat'] = convert_to_ordered_category(df[human_score_column])

    # Overall score accuracy (OSA)
    # Variance in squared error explained by L1

    # fit the model
    osa_model = smf.ols(formula='SE ~ group', data=df)
    osa_fit = osa_model.fit()

    # collect the results
    osa_dict = {'R2': osa_fit.rsquared_adj,
               'sig': osa_fit.f_pvalue}

    osa_results = pd.Series(osa_dict, name='Overall score accuracy')

    df_coefficients_osa = get_coefficients(osa_fit, base_group)

    # Overall score difference (OSD)
    # variance in signed residuals (raw error) explained by L1

    # fit the model
    osd_model = smf.ols(formula='error ~ group', data=df)
    osd_fit = osd_model.fit()

    # collect the results
    osd_dict = {'R2': osd_fit.rsquared_adj,
                'sig': osd_fit.f_pvalue}

    osd_results = pd.Series(osd_dict, name='Overall score difference')

    df_coefficients_osd = get_coefficients(osd_fit, base_group)

    # conditional score difference CSD
    # Variance in score difference conditioned on Native language

    # fit "null" model with human score only
    csd_null_mod = smf.ols(formula='error ~ sc1_cat', data=df)
    csd_null_fit = csd_null_mod.fit()

    # fit model with both human score and group
    csd_mod = smf.ols(formula='error ~ group + sc1_cat', data=df)
    csd_fit = csd_mod.fit()

    # compare the two models using anova_lm
    # we filter warnings for this function because we get
    # runtime warning due to NaNs in the data.
    # these seem to be by design: https://groups.google.com/forum/#!topic/pystatsmodels/-flY0cNnb3k
    np.warnings.filterwarnings('ignore')
    anova_results = anova_lm(csd_null_fit, csd_fit)
    # we reset warnings
    np.warnings.resetwarnings()

    # collect the results. Note that R2 in this case is a difference
    # in R2 between the two models and significance is obtained from anova
    csd_dict = {'R2': csd_fit.rsquared_adj - csd_null_fit.rsquared_adj,
               'sig': anova_results.values[1][-1]}
    csd_results = pd.Series(csd_dict, name='Conditional score difference')

    df_coefficients_csd = get_coefficients(osd_fit, base_group)

    # create a summary table

    df_r2_all = pd.concat([osa_results, osd_results, csd_results], axis=1, sort=True)
    df_r2_all['base_category'] = base_group

    # assemble all datasets into a DataContainer

    datasets = [{'name': 'estimates_osa_by_{}'.format(group),
                 'frame': df_coefficients_osa},
                {'name': 'estimates_osd_by_{}'.format(group),
                 'frame': df_coefficients_osd},
                {'name': 'estimates_csd_by_{}'.format(group),
                 'frame': df_coefficients_csd},
                {'name': 'fairness_metrics_by_{}'.format(group),
                 'frame': df_r2_all}]

    # assemble all models into a dictionary
    model_dict = {'osa': osa_fit,
                  'osd': osd_fit,
                  'csd': csd_fit}

    return model_dict, DataContainer(datasets=datasets)

<commit_msg>Some more alignment fixes<commit_after>"""
Utility classes and functions related to computing fairness evaluations

:author: Anastassia Loukina (aloukina@ets.org)
:author: Jeremy Biggs (jbiggs@ets.org)
:author: Nitin Madnani (nmadnani@ets.org)

:date: 07/29/2019
:organization: ETS
"""

import pandas as pd
import pickle
import numpy as np

from os.path import join

import statsmodels.formula.api as smf
from statsmodels.stats.anova import anova_lm

from rsmtool.writer import DataWriter
from rsmtool.container import DataContainer

def convert_to_ordered_category(group_values):
    """
    Convert supplied series to an ordered category
    with levels ordered by category size.
    If multiple categories have the same size,
    the order is determined alphabetically.

    Parameters
    ----------
    group_values: pandas Series

    Returns
    -------
    group_ordered_category: pandas Series
    """

    # get ordered list by size

    # We convert the  value counts to data frame to allow for multilevel sorting.
    # This makes sure that the order is consistent and reproducible across runs
    # when there are more than
    # one group with the maximum number of occurrences.
    df_groups_by_size = pd.DataFrame(group_values.value_counts()).reset_index()
    df_groups_by_size.columns = ['group_name', 'count']
    df_groups_by_size_sorted = df_groups_by_size.sort_values(['count', 'group_name'],
                                                             ascending=[False, True])
    groups_by_size = df_groups_by_size_sorted['group_name'].values

    #convert to category and reorder
    group_category = group_values.astype("category")
    group_ordered_category = group_category.cat.reorder_categories(groups_by_size,
                                                                   ordered=True)
    return group_ordered_category

def get_coefficients(fit, base_category):

    """
    Extract estimates, significance, and confidence intervals
    for members of the group of interest. The names of the
    predictors are processed to remove the prefix added by Statmodels.
    The name of the base category is added in parenthesis to the Intercept.

    Parameters:
    -----------
    fit: fit
        Linear regression object fitted using statsmodels

    base_category: str
        Name of the group used as reference category when fitting the model

    Returns:
    -------
    df_results: pandas DataFrame
        A dataframe with rows for each category and the following columns:
        - estimate
        - P>[t]
        - Lower and upper boundaries for 95% confidence intervals
    """

    # extract the data we need
    df_results = pd.concat([fit.params, fit.pvalues, fit.conf_int()], axis=1)

    df_results.columns = ['estimate', 'P>[t]', '[0.025', '0.975]']

    # identify the rows we are interested in
    groups = ['Intercept'] + [v for v in df_results.index if 'group' in v]

    df_results = df_results.loc[groups]

    # rename the rows
    df_results.index = [v.split('.')[1].strip(']')
                            if not v == 'Intercept'
                            else 'Intercept ({})'.format(base_category)
                            for v in df_results.index]

    return df_results

def write_fairness_results(fit_dictionary,
                           fairness_container,
                           group,
                           output_dir,
                           experiment_id,
                           file_format):
    """
    Save the results of fairness analysis to disk.

    Parameters
    ----------
    fit_dictionary: dict
        A dictionary of fitted models generated by `get_fairness_analysis()`
    fairness_container: container.DataContainer
        A data container with the results of fairness analysis generated by `get_fairness_analysis()`
    group: str
        The subgroup considered in this analysis
    output_dir: str
        The directory where the results will be saved
    experiment_id: str
        Experiment id
    file_format: str
        File format for data files
    """

    # let's first save model files and summaries
    for model in fit_dictionary:
        fit = fit_dictionary[model]

        ols_file = join(output_dir, '{}_{}_by_{}.ols'.format(experiment_id,
                                                             model,
                                                             group))
        summary_file = join(output_dir, '{}_{}_by_{}_ols_summary.txt'.format(experiment_id,
                                                                             model,
                                                                             group))
        with open(ols_file, 'wb') as olsf, open(summary_file, 'w') as summf:
            pickle.dump(fit, olsf)
            summf.write(str(fit.summary()))

    # Now let's write out the content of the data container
    writer = DataWriter(experiment_id)
    writer.write_experiment_output(output_dir,
                                   fairness_container,
                                   file_format=file_format,
                                   index=True)

def get_fairness_analyses(df,
                          group,
                          system_score_column,
                          human_score_column='sc1'):
    """
    Compute fairness analyses described in Loukina et al. (2019) https://aclweb.org/anthology/papers/W/W19/W19-4401/
    The functions computes how muach variance group memebership explains in
    overall score accuracy (osa), overall score difference (osd) and conditional score difference (csd).
    See paper for further detail.

    Parameters
    ----------
    df: pandas DataFrame
        A dataframe containing columns with numeric human scores, numeric system scores
        and a column showing group membership
    group: str
        Name of the column containing group membership
    system_score_column: str
        Name of the column containing system scores
    human_score_column: str
        Name of the column containing human scores

    Returns
    -------
    model_dict: dictionary
        A dictionary with different proposed metrics as keys
        and fitted models as values
    fairness_container: container.DataContainer
        A datacontainer with the following datasets:
            - `estimates_METRICS_by_GROUP` where `GROUP` corresponds to group
            and `METRICS` can be osa, osd and csd
            Estimates for each group computed by the respective models
            - `fairness_metrics_by_GROUP` - a summary of model fits (R2 and p values)
    """

    # compute error and squared error

    df['error'] = df[system_score_column]-df[human_score_column]
    df['SE'] = df['error']**2

    # convert group values to category and reorder them using
    # the largest category as reference

    df['group'] = convert_to_ordered_category(df[group])
    base_group = df['group'].cat.categories[0]
    df['sc1_cat'] = convert_to_ordered_category(df[human_score_column])

    # Overall score accuracy (OSA)
    # Variance in squared error explained by L1

    # fit the model
    osa_model = smf.ols(formula='SE ~ group', data=df)
    osa_fit = osa_model.fit()

    # collect the results
    osa_dict = {'R2': osa_fit.rsquared_adj,
                'sig': osa_fit.f_pvalue}

    osa_results = pd.Series(osa_dict, name='Overall score accuracy')

    df_coefficients_osa = get_coefficients(osa_fit, base_group)

    # Overall score difference (OSD)
    # variance in signed residuals (raw error) explained by L1

    # fit the model
    osd_model = smf.ols(formula='error ~ group', data=df)
    osd_fit = osd_model.fit()

    # collect the results
    osd_dict = {'R2': osd_fit.rsquared_adj,
                'sig': osd_fit.f_pvalue}

    osd_results = pd.Series(osd_dict, name='Overall score difference')

    df_coefficients_osd = get_coefficients(osd_fit, base_group)

    # conditional score difference CSD
    # Variance in score difference conditioned on Native language

    # fit "null" model with human score only
    csd_null_mod = smf.ols(formula='error ~ sc1_cat', data=df)
    csd_null_fit = csd_null_mod.fit()

    # fit model with both human score and group
    csd_mod = smf.ols(formula='error ~ group + sc1_cat', data=df)
    csd_fit = csd_mod.fit()

    # compare the two models using anova_lm
    # we filter warnings for this function because we get
    # runtime warning due to NaNs in the data.
    # these seem to be by design: https://groups.google.com/forum/#!topic/pystatsmodels/-flY0cNnb3k
    np.warnings.filterwarnings('ignore')
    anova_results = anova_lm(csd_null_fit, csd_fit)
    # we reset warnings
    np.warnings.resetwarnings()

    # collect the results. Note that R2 in this case is a difference
    # in R2 between the two models and significance is obtained from anova
    csd_dict = {'R2': csd_fit.rsquared_adj - csd_null_fit.rsquared_adj,
                'sig': anova_results.values[1][-1]}
    csd_results = pd.Series(csd_dict, name='Conditional score difference')

    df_coefficients_csd = get_coefficients(osd_fit, base_group)

    # create a summary table

    df_r2_all = pd.concat([osa_results, osd_results, csd_results], axis=1, sort=True)
    df_r2_all['base_category'] = base_group

    # assemble all datasets into a DataContainer

    datasets = [{'name': 'estimates_osa_by_{}'.format(group),
                 'frame': df_coefficients_osa},
                {'name': 'estimates_osd_by_{}'.format(group),
                 'frame': df_coefficients_osd},
                {'name': 'estimates_csd_by_{}'.format(group),
                 'frame': df_coefficients_csd},
                {'name': 'fairness_metrics_by_{}'.format(group),
                 'frame': df_r2_all}]

    # assemble all models into a dictionary
    model_dict = {'osa': osa_fit,
                  'osd': osd_fit,
                  'csd': csd_fit}

    return model_dict, DataContainer(datasets=datasets)

<|endoftext|>