<commit_before>from django.http import HttpResponse
from django.template.loader import get_template
from django.template import Context
from django.shortcuts import render
import json
import requests

NODE_ADDR = "http://127.0.0.1:9001"
publish_base_url = "http://despina.128.no/publish"

def get_publisher_url():
	r = requests.get(NODE_ADDR + "/" + "publishing")
	response_as_json = json.loads(r.text)
	if response_as_json:
		return response_as_json
	else:
		return None

def homepage(request):
	return render(request, "homepage.html", {});

def editor(request):
	return render(request, "editor_page.html", {});

def upload_article(request):
	dict = request.POST.dict()
	article = {"tags": dict["tags"], "description": dict["description"], "title": dict["title"]}
	article["article"] = dict["article"].replace("src=\"//www.", "src=\"http://www.")

	publisher_url = get_publisher_url()
	if publisher_url:
		r = requests.post("http://"+publisher_url+"/save_article", data = article)
	else:
		# Do some error handling here.
		pass

	return render(request, "editor_page.html", {});

def articles(request):
        r = requests.get(publish_base_url + "/list")

        if r.url != publish_base_url + "/list":
                return HttpResponseRedirect(r.url)

        json_response = r.json()
        if r.status_code == 200 and json_response:
                return render(request, "articles.html", json_response)
        else:
                print(r)

def search(request):
    return render(request, "search.html", {});

def about(request):
    return render(request, "about.html", {});

def article(request):
	id = request.path[-24:]
	r = requests.get(publish_base_url + "/article_json/" + id)
	return render(request, "article.html", r.json())
<commit_msg>pep8-ify the views file.<commit_after>from django.http import HttpResponse
from django.template.loader import get_template
from django.template import Context
from django.shortcuts import render
import json
import requests

NODE_ADDR = "http://127.0.0.1:9001"
publish_base_url = "http://despina.128.no/publish"

def get_publisher_url():
    r = requests.get(NODE_ADDR + "/" + "publishing")
    response_as_json = json.loads(r.text)
    if response_as_json:
        return response_as_json
    else:
        return None

def homepage(request):
    return render(request, "homepage.html", {})

def editor(request):
    return render(request, "editor_page.html", {})

def upload_article(request):
    dict = request.POST.dict()
    article = {"tags": dict["tags"],
               "description": dict["description"],
               "title": dict["title"]}
    article["article"] = dict["article"].replace("src=\"//www.",
                                                 "src=\"http://www.")

    publisher_url = get_publisher_url()
    if publisher_url:
        r = requests.post("http://" + publisher_url + "/save_article",
                          data=article)
    else:
        # Do some error handling here.
        pass

    return render(request, "editor_page.html", {})

def articles(request):
    r = requests.get(publish_base_url + "/list")

    if r.url != publish_base_url + "/list":
        return HttpResponseRedirect(r.url)

    json_response = r.json()
    if r.status_code == 200 and json_response:
        return render(request, "articles.html", json_response)
    else:
        print(r)

def search(request):
    return render(request, "search.html", {})

def about(request):
    return render(request, "about.html", {})

def article(request):
    id = request.path[-24:]
    r = requests.get(publish_base_url + "/article_json/" + id)
    return render(request, "article.html", r.json())
<|endoftext|>


<commit_before>"""
LAMMPS EEX I/O
"""

import pandas as pd
import math
import numpy as np
import eex

from . import lammps_metadata as lmd

import logging
logger = logging.getLogger(__name__)

def write_lammps_file(dl, data_filename, input_filename, unit_style="real", blocksize=110):

    # handle units
    unit_set = lmd.units_style[unit_style]

    term_table = lmd.build_term_table(unit_style)

    nb_term_table = lmd.build_nb_table("real")

    data_file = open(data_filename, 'w')
    input_file = open(input_filename, 'w')

    data_file.write("LAMMPS data file generated by MolSSI EEX\n\n")
    input_file.write("# LAMMPS input file generated by MolSSI EEX\n")

    input_file.write("units\t%s\n" %(unit_style))
    input_file.write("atom_style\tfull\n")

    sizes = {}
    sizes["atoms"] = dl.get_atom_count()
    sizes["bonds"] = dl.get_term_count(2, "total")
    sizes["angles"] = dl.get_term_count(3, "total")
    sizes["dihedrals"] = dl.get_term_count(4, "total") # Not qutie right once we do impropers
    sizes["impropers"] = 0
    sizes["atom types"] = len(dl.list_atom_uids("mass"))

    # All the UID's minus the "total" columns
    sizes["bond types"] = len(dl.get_term_count(2)) - 1
    sizes["angle types"] = len(dl.get_term_count(3)) - 1
    sizes["dihedral types"] = len(dl.get_term_count(4)) - 1
    sizes["improper types"] = 0

    # Write header information
    for k in lmd.size_keys:
        data_file.write(" %d %s\n" % (sizes[k], k))
        # data_file.write(' '.join([str(data["sizes"][k]), k, '\n']))

    # Write box information
    box_size = dl.get_box_size(utype={"a": unit_set["[length]"], "b": unit_set["[length]"], "c": unit_set["[length]"],
                                      "alpha": "degree", "beta": "degree", "gamma": "degree"})

    box_center = dl.get_box_center(utype={"x": "angstrom", "y": "angstrom", "z": "angstrom"})

    lo_hi = {}

    lo_hi["x"] = [box_center["x"] - box_size["a"]/2., box_center["x"] + box_size["a"]/2.]
    lo_hi["y"] = [box_center["y"] - box_size["b"]/2., box_center["y"] + box_size["b"] / 2.]
    lo_hi["z"] = [box_center["z"] - box_size["c"] / 2., box_center["z"] + box_size["c"] / 2.]

    for k, v in lo_hi.items():
        data_file.write("% 8.6f% 8.6f %slo %shi\n" % (v[0],v[1], k, k))
    data_file.write('\n')

    param_fmt = "%10.8f"
    # Loop over Pair Coeffs
    nb_forms = dl.list_stored_nb_types()

    # Handle nonbonds - This needs to be generalized badly - make it work for now.
    if dl.get_mixing_rule() != '':
        data_file.write(("Pair Coeffs\n\n").title())

        stored_nb_parameters = dl.list_nb_parameters(
            nb_name="LJ", nb_model="epsilon/sigma", utype={"epsilon": unit_set["[energy]"], "sigma": unit_set["[length]"]}, itype="single")
    else:
        data_file.write("PairIJ Coeffs\n\n")
        stored_nb_parameters = dl.list_nb_parameters(
            nb_name="LJ", nb_model="epsilon/sigma", utype={"epsilon": unit_set["[energy]"], "sigma": unit_set["[length]"]}, itype="pair")

    for key, value in stored_nb_parameters.items():
        if key[1] == None:
            data_file.write(("%2d %10.8f %10.8f\n" % (key[0], value['epsilon'], value['sigma'])))
        else:
            data_file.write(("%2d %2d %10.8f %10.8f\n" % (key[0], key[1], value['epsilon'], value['sigma'])))

    data_file.write("\n")

    # Loop over all of the parameter data
    for param_order, param_type in zip([2, 3, 4], ["bond", "angle", "dihedral"]):
        param_uids = dl.list_term_uids(param_order)

        if len(param_uids) == 0: continue

        data_file.write(("%s Coeffs\n\n" % param_type).title())
        for uid in param_uids:
            param_coeffs = dl.get_term_parameter(param_order, uid)

            term_data = term_table[param_order][param_coeffs[0]]
            param_coeffs = dl.get_term_parameter(param_order, uid, utype=term_data["utype"])

            # Order the data like lammps wants it
            parameters = [param_coeffs[1][k] for k in term_data["parameters"]]

            data_file.write("%2d " % uid)
            data_file.write(" ".join(param_fmt % f for f in parameters))
            data_file.write("\n")

            input_file.write("%s_style\t%s\n" %(param_type, param_coeffs[0]))

        data_file.write("\n")

    # Write out mass data
    data_file.write(" Masses\n\n")
    for idx in dl.list_atom_uids("mass"):
        mass = dl.get_atom_parameter("mass", idx, utype=lmd.get_context(unit_style, "[mass]"))
        data_file.write("%2d %10.8f\n" % (idx, mass))
    data_file.write('\n')

    # Write out atom data
    data_file.write(" Atoms\n\n")

    atoms = dl.get_atoms(["molecule_index", "atom_type", "charge", "xyz"], by_value=True)
    atoms.index = pd.RangeIndex(start=1, stop=atoms.shape[0] + 1)

    # Build a simple formatter
    def float_fmt(n):
        return "%10.8f" % n
    atoms.to_string(data_file, header=None, float_format=float_fmt)
    data_file.write('\n\n')

    # Write out all of the term data
    for param_order, param_type in zip([2, 3, 4], ["bonds", "angles", "dihedrals"]):
        if sizes[param_type] == 0: continue

        data_file.write((" %s\n\n" % param_type).title())

        # Grab term and reorder
        cols = ["term_index"] + ["atom%s" % d for d in range(1, param_order + 1)]
        term = dl.get_terms(param_type)[cols]
        term.index = pd.RangeIndex(start=1, stop=term.shape[0] + 1)
        # print(term)
        term.to_csv(data_file, header=None, sep=" ")
        data_file.write('\n')

    input_file.write("read_data\t%s\n" %(data_filename))

    data_file.close()
    input_file.close()
    return True

<commit_msg>lammps: fix writer so that it doesn't fail for non periodic systems<commit_after>"""
LAMMPS EEX I/O
"""

import pandas as pd
import math
import numpy as np
import eex

from . import lammps_metadata as lmd

import logging
logger = logging.getLogger(__name__)

def write_lammps_file(dl, data_filename, input_filename, unit_style="real", blocksize=110):

    # handle units
    unit_set = lmd.units_style[unit_style]

    term_table = lmd.build_term_table(unit_style)

    nb_term_table = lmd.build_nb_table("real")

    data_file = open(data_filename, 'w')
    input_file = open(input_filename, 'w')

    data_file.write("LAMMPS data file generated by MolSSI EEX\n\n")
    input_file.write("# LAMMPS input file generated by MolSSI EEX\n")

    input_file.write("units\t%s\n" %(unit_style))
    input_file.write("atom_style\tfull\n")

    sizes = {}
    sizes["atoms"] = dl.get_atom_count()
    sizes["bonds"] = dl.get_term_count(2, "total")
    sizes["angles"] = dl.get_term_count(3, "total")
    sizes["dihedrals"] = dl.get_term_count(4, "total") # Not qutie right once we do impropers
    sizes["impropers"] = 0
    sizes["atom types"] = len(dl.list_atom_uids("mass"))

    # All the UID's minus the "total" columns
    sizes["bond types"] = len(dl.get_term_count(2)) - 1
    sizes["angle types"] = len(dl.get_term_count(3)) - 1
    sizes["dihedral types"] = len(dl.get_term_count(4)) - 1
    sizes["improper types"] = 0

    # Write header information
    for k in lmd.size_keys:
        data_file.write(" %d %s\n" % (sizes[k], k))
        # data_file.write(' '.join([str(data["sizes"][k]), k, '\n']))

    # Write box information
    box_size = dl.get_box_size(utype={"a": unit_set["[length]"], "b": unit_set["[length]"], "c": unit_set["[length]"],
                                      "alpha": "degree", "beta": "degree", "gamma": "degree"})

    box_center = dl.get_box_center(utype={"x": "angstrom", "y": "angstrom", "z": "angstrom"})

    lo_hi = {}
    if box_center and box_size:
        lo_hi["x"] = [box_center["x"] - box_size["a"]/2., box_center["x"] + box_size["a"]/2.]
        lo_hi["y"] = [box_center["y"] - box_size["b"]/2., box_center["y"] + box_size["b"] / 2.]
        lo_hi["z"] = [box_center["z"] - box_size["c"] / 2., box_center["z"] + box_size["c"] / 2.]

        for k, v in lo_hi.items():
            data_file.write("% 8.6f% 8.6f %slo %shi\n" % (v[0],v[1], k, k))
        data_file.write('\n')

    param_fmt = "%10.8f"
    # Loop over Pair Coeffs
    nb_forms = dl.list_stored_nb_types()

    # Handle nonbonds - This needs to be generalized badly - make it work for now.
    if dl.get_mixing_rule() != '':
        data_file.write(("Pair Coeffs\n\n").title())

        stored_nb_parameters = dl.list_nb_parameters(
            nb_name="LJ", nb_model="epsilon/sigma", utype={"epsilon": unit_set["[energy]"], "sigma": unit_set["[length]"]}, itype="single")
    else:
        data_file.write("PairIJ Coeffs\n\n")
        stored_nb_parameters = dl.list_nb_parameters(
            nb_name="LJ", nb_model="epsilon/sigma", utype={"epsilon": unit_set["[energy]"], "sigma": unit_set["[length]"]}, itype="pair")

    for key, value in stored_nb_parameters.items():
        if key[1] == None:
            data_file.write(("%2d %10.8f %10.8f\n" % (key[0], value['epsilon'], value['sigma'])))
        else:
            data_file.write(("%2d %2d %10.8f %10.8f\n" % (key[0], key[1], value['epsilon'], value['sigma'])))

    data_file.write("\n")

    # Loop over all of the parameter data
    for param_order, param_type in zip([2, 3, 4], ["bond", "angle", "dihedral"]):
        param_uids = dl.list_term_uids(param_order)

        if len(param_uids) == 0: continue

        data_file.write(("%s Coeffs\n\n" % param_type).title())
        for uid in param_uids:
            param_coeffs = dl.get_term_parameter(param_order, uid)

            term_data = term_table[param_order][param_coeffs[0]]
            param_coeffs = dl.get_term_parameter(param_order, uid, utype=term_data["utype"])

            # Order the data like lammps wants it
            parameters = [param_coeffs[1][k] for k in term_data["parameters"]]

            data_file.write("%2d " % uid)
            data_file.write(" ".join(param_fmt % f for f in parameters))
            data_file.write("\n")

            input_file.write("%s_style\t%s\n" %(param_type, param_coeffs[0]))

        data_file.write("\n")

    # Write out mass data
    data_file.write(" Masses\n\n")
    for idx in dl.list_atom_uids("mass"):
        mass = dl.get_atom_parameter("mass", idx, utype=lmd.get_context(unit_style, "[mass]"))
        data_file.write("%2d %10.8f\n" % (idx, mass))
    data_file.write('\n')

    # Write out atom data
    data_file.write(" Atoms\n\n")

    atoms = dl.get_atoms(["molecule_index", "atom_type", "charge", "xyz"], by_value=True)
    atoms.index = pd.RangeIndex(start=1, stop=atoms.shape[0] + 1)

    # Build a simple formatter
    def float_fmt(n):
        return "%10.8f" % n
    atoms.to_string(data_file, header=None, float_format=float_fmt)
    data_file.write('\n\n')

    # Write out all of the term data
    for param_order, param_type in zip([2, 3, 4], ["bonds", "angles", "dihedrals"]):
        if sizes[param_type] == 0: continue

        data_file.write((" %s\n\n" % param_type).title())

        # Grab term and reorder
        cols = ["term_index"] + ["atom%s" % d for d in range(1, param_order + 1)]
        term = dl.get_terms(param_type)[cols]
        term.index = pd.RangeIndex(start=1, stop=term.shape[0] + 1)
        # print(term)
        term.to_csv(data_file, header=None, sep=" ")
        data_file.write('\n')

    input_file.write("read_data\t%s\n" %(data_filename))

    data_file.close()
    input_file.close()
    return True

<|endoftext|>


<commit_before>"""
Wrapper for a loaded scene with properties.
"""
from .shaders import MeshShader, ColorShader, TextureShader
from demosys.opengl import geometry
from demosys.resources import shaders
from pyrr import matrix44

class Scene:
    """Generic scene"""
    def __init__(self, name, loader=None, mesh_shaders=None, **kwargs):
        """
        :param name: Unique name or path for the scene
        :param loader: Loader class for the scene if relevant
        """
        self.name = name
        self.loader = loader
        self.root_nodes = []

        # References resources in the scene
        self.nodes = []
        self.materials = []
        self.meshes = []
        self.cameras = []
        self.mesh_shaders = mesh_shaders or [ColorShader(), TextureShader()]

        self.bbox_min = None
        self.bbox_max = None
        self.bbox_vao = geometry.bbox()
        self.bbox_shader = shaders.get('scene_default/bbox.glsl', create=True)

    def draw(self, m_proj, m_mv):
        for node in self.root_nodes:
            node.draw(m_proj, m_mv)

    def draw_bbox(self, m_proj, m_mv, all=True):
        """Draw scene and mesh bounding boxes"""
        self.bbox_vao.bind(self.bbox_shader)
        self.bbox_shader.uniform_mat4("m_proj", m_proj)
        self.bbox_shader.uniform_mat4("m_mv", m_mv)
        self.bbox_shader.uniform_3fv("bb_min", self.bbox_min)
        self.bbox_shader.uniform_3fv("bb_max", self.bbox_max)
        self.bbox_vao.draw()

        if not all:
            return

        for node in self.root_nodes:
            node.draw_bbox(m_proj, m_mv, self.bbox_shader, self.bbox_vao)

    def apply_mesh_shaders(self, mesh_shaders):
        """Applies mesh shaders to meshes"""
        if mesh_shaders is None:
            return

        for mesh in self.meshes:
            for ms in mesh_shaders:
                instance = ms.apply(mesh)
                if instance is not None:
                    if isinstance(instance, MeshShader):
                        mesh.mesh_shader = ms
                        continue
                    else:
                        raise ValueError("apply() must return a MeshShader instance, not {}".format(type(instance)))
            if not mesh.mesh_shader:
                print("WARING: No mesh shader applied to '{}'".format(mesh.name))

    def calc_scene_bbox(self):
        """Calculate scene bbox"""
        bbox_min, bbox_max = None, None
        for node in self.root_nodes:
            bbox_min, bbox_max = node.calc_global_bbox(
                matrix44.create_identity(),
                bbox_min,
                bbox_max
            )

        self.bbox_min = bbox_min
        self.bbox_max = bbox_max

    def load(self, path):
        """Deferred loading if a loader is specified"""
        if not hasattr(self, 'loader'):
            return

        if self.loader:
            self.loader.load(self, file=path)

        self.apply_mesh_shaders(self.mesh_shaders)

    def __str__(self):
        return "<Scene: {}>".format(self.name)

    def __repr__(self):
        return str(self)
<commit_msg>Use fallback shader by default in scene<commit_after>"""
Wrapper for a loaded scene with properties.
"""
from .shaders import MeshShader, ColorShader, TextureShader, FallbackShader
from demosys.opengl import geometry
from demosys.resources import shaders
from pyrr import matrix44

class Scene:
    """Generic scene"""
    def __init__(self, name, loader=None, mesh_shaders=None, **kwargs):
        """
        :param name: Unique name or path for the scene
        :param loader: Loader class for the scene if relevant
        """
        self.name = name
        self.loader = loader
        self.root_nodes = []

        # References resources in the scene
        self.nodes = []
        self.materials = []
        self.meshes = []
        self.cameras = []
        self.mesh_shaders = mesh_shaders or [ColorShader(), TextureShader(), FallbackShader()]

        self.bbox_min = None
        self.bbox_max = None
        self.bbox_vao = geometry.bbox()
        self.bbox_shader = shaders.get('scene_default/bbox.glsl', create=True)

    def draw(self, m_proj, m_mv):
        for node in self.root_nodes:
            node.draw(m_proj, m_mv)

    def draw_bbox(self, m_proj, m_mv, all=True):
        """Draw scene and mesh bounding boxes"""
        self.bbox_vao.bind(self.bbox_shader)
        self.bbox_shader.uniform_mat4("m_proj", m_proj)
        self.bbox_shader.uniform_mat4("m_mv", m_mv)
        self.bbox_shader.uniform_3fv("bb_min", self.bbox_min)
        self.bbox_shader.uniform_3fv("bb_max", self.bbox_max)
        self.bbox_vao.draw()

        if not all:
            return

        for node in self.root_nodes:
            node.draw_bbox(m_proj, m_mv, self.bbox_shader, self.bbox_vao)

    def apply_mesh_shaders(self):
        """Applies mesh shaders to meshes"""
        if self.mesh_shaders is None:
            return

        for mesh in self.meshes:
            for ms in self.mesh_shaders:
                instance = ms.apply(mesh)
                if instance is not None:
                    if isinstance(instance, MeshShader):
                        mesh.mesh_shader = ms
                        continue
                    else:
                        raise ValueError("apply() must return a MeshShader instance, not {}".format(type(instance)))
            if not mesh.mesh_shader:
                print("WARING: No mesh shader applied to '{}'".format(mesh.name))

    def calc_scene_bbox(self):
        """Calculate scene bbox"""
        bbox_min, bbox_max = None, None
        for node in self.root_nodes:
            bbox_min, bbox_max = node.calc_global_bbox(
                matrix44.create_identity(),
                bbox_min,
                bbox_max
            )

        self.bbox_min = bbox_min
        self.bbox_max = bbox_max

    def load(self, path):
        """Deferred loading if a loader is specified"""
        if not hasattr(self, 'loader'):
            return

        if self.loader:
            self.loader.load(self, file=path)

        self.apply_mesh_shaders()

    def __str__(self):
        return "<Scene: {}>".format(self.name)

    def __repr__(self):
        return str(self)
<|endoftext|>


<commit_before>import os

from fabric.api import *
from fabric.contrib.files import sed

"""
Fabric Script which does the following things:
    # Uploading the SSH keys pair (public/private) on to the remote servers.
    # Creating non-root-users and specific groups for them.
    # Preventing SSH-ing into the remote server as a root-user.
"""

# Declaring all the env's Globals

# Root User
env.user = 'root'

# List of remote IP's
env.hosts = ['188.166.189.143', '188.166.176.126']

# Password for the remote server
env.password = '<your-server-password>'

# SSH keys path
env.ssh_keys_dir = '~/Projects/python/django-projects/django-deployments-stuff/prod_servers_config'

# Full name of the user
env.full_name_user = 'Saurabh Sharma'

# User group
env.user_group = 'deployers'

# User for the above group
env.user_name = 'deployer'

def start_provision():
    """
    Starts off server provisioning
    """
    env.ssh_keys_name = os.path.join(env.ssh_keys_dir, env.host_string + "_prod_key")
    local('ssh-keygen -t rsa -b 2048 -f {}'.format(env.ssh_keys_name))
    local('cp {} {}/authorized_keys'.format(env.ssh_keys_name + ".pub", env.ssh_keys_dir))

    # Preventing Root SSH-ing into the remote server
    sed('/etc/ssh/sshd_config', '^UsePAM yes', 'UsePAM no')
    sed('/etc/ssh/sshd_config', '^PermitRootLogin yes', 'PermitRootLogin no')
    sed('/etc/ssh/sshd_config', '^#PasswordAuthentication yes', 'PasswordAuthentication no')

    upgrade_server()
    create_deployers_group()
    ceate_deployer_user()
    upload_keys()

def upgrade_server():
    """
    Upgrade the server as a root user
    """
    run('dnf upgrade -y')

def create_deployers_group():
    """
    Creates a group for all the developers on the project
    """
    run('groupadd {}'.format(env.user_group))
    run('mv /etc/sudoers /etc/sudoers-backup')
    run('(cat /etc/sudoers-backup ; echo "%' + env.user_group + 'ALL=(ALL) ALL") > /etc/sudoers')
    run('chmod 440 /etc/sudoers')

def ceate_deployer_user():
    """
    Creates a user for Deployers group
    """
    run('adduser -c "{}" -m -g {} {}'.format(env.full_name_user, env.user_group, env.user_name))
    run('passwd {}'.format(env.user_name))
    run('usermod -a -G {} {}'.format(env.user_group, env.user_name))
    run('mkdir /home/{}/.ssh'.format(env.user_name))
    run('chown -R {} /home/{}/.ssh'.format(env.user_name, env.user_name))
    run('chgrp -R {} /home/{}/.ssh'.format(env.user_group, env.user_name))

def upload_keys():
    """
    Uploads the ssh public/private keys on to the remote server via scp.
    """
    scp_command = "scp {} {}/authorized_keys {}@{}:~/.ssh".format(
                    env.ssh_keys_name + '.pub',
                    env.ssh_keys_dir,
                    env.user_name,
                    env.host_string
                  )
    local(scp_command)
<commit_msg>script for provisioning a single server<commit_after>import os

from fabric.api import *
from fabric.contrib.files import sed

"""
Fabric Script which does the following things:
    # Uploading the SSH keys pair (public/private) on to the remote servers.
    # Creating non-root-users and specific groups for them.
    # Preventing SSH-ing into the remote server as a root-user.
"""

# Declaring all the env's Globals

# Root User
env.user = 'root'

# List of remote IP's
env.hosts = ['188.166.189.143']

# Password for the remote server
env.password = '<your-server-password>'

# SSH keys path
env.ssh_keys_dir = '~/Projects/python/django-projects/django-deployments-stuff/prod_servers_config'

# Full name of the user
env.full_name_user = 'Saurabh Sharma'

# User group
env.user_group = 'deployers'

# User for the above group
env.user_name = 'deployer'

def start_provision():
    """
    Starts off server provisioning
    """
    env.ssh_keys_name = os.path.join(env.ssh_keys_dir, env.host_string + "_prod_key")
    local('ssh-keygen -t rsa -b 2048 -f {}'.format(env.ssh_keys_name))
    local('cp {} {}/authorized_keys'.format(env.ssh_keys_name + ".pub", env.ssh_keys_dir))

    # Preventing Root SSH-ing into the remote server
    sed('/etc/ssh/sshd_config', '^UsePAM yes', 'UsePAM no')
    sed('/etc/ssh/sshd_config', '^PermitRootLogin yes', 'PermitRootLogin no')
    sed('/etc/ssh/sshd_config', '^#PasswordAuthentication yes', 'PasswordAuthentication no')

    upgrade_server()
    create_deployers_group()
    ceate_deployer_user()
    upload_keys()

def upgrade_server():
    """
    Upgrade the server as a root user
    """
    run('dnf upgrade -y')

def create_deployers_group():
    """
    Creates a group for all the developers on the project
    """
    run('groupadd {}'.format(env.user_group))
    run('mv /etc/sudoers /etc/sudoers-backup')
    run('(cat /etc/sudoers-backup ; echo "%' + env.user_group + 'ALL=(ALL) ALL") > /etc/sudoers')
    run('chmod 440 /etc/sudoers')

def ceate_deployer_user():
    """
    Creates a user for Deployers group
    """
    run('adduser -c "{}" -m -g {} {}'.format(env.full_name_user, env.user_group, env.user_name))
    run('passwd {}'.format(env.user_name))
    run('usermod -a -G {} {}'.format(env.user_group, env.user_name))
    run('mkdir /home/{}/.ssh'.format(env.user_name))
    run('chown -R {} /home/{}/.ssh'.format(env.user_name, env.user_name))
    run('chgrp -R {} /home/{}/.ssh'.format(env.user_group, env.user_name))

def upload_keys():
    """
    Uploads the ssh public/private keys on to the remote server via scp.
    """
    scp_command = "scp {} {}/authorized_keys {}@{}:~/.ssh".format(
                    env.ssh_keys_name + '.pub',
                    env.ssh_keys_dir,
                    env.user_name,
                    env.host_string
                  )
    local(scp_command)
<|endoftext|>


<commit_before>import logging
from distutils.version import LooseVersion

logger = logging.getLogger(__name__)

REQUIREMENTS = {
    "syutil==0.0.2": ["syutil"],
    "matrix_angular_sdk==0.6.0": ["syweb>=0.6.0"],
    "Twisted==14.0.2": ["twisted==14.0.2"],
    "service_identity>=1.0.0": ["service_identity>=1.0.0"],
    "pyopenssl>=0.14": ["OpenSSL>=0.14"],
    "pyyaml": ["yaml"],
    "pyasn1": ["pyasn1"],
    "pynacl": ["nacl"],
    "daemonize": ["daemonize"],
    "py-bcrypt": ["bcrypt"],
    "frozendict>=0.4": ["frozendict"],
    "pillow": ["PIL"],
    "pydenticon": ["pydenticon"],
}

def github_link(project, version, egg):
    return "https://github.com/%s/tarball/%s/#egg=%s" % (project, version, egg)

DEPENDENCY_LINKS=[
    github_link(
        project="matrix-org/syutil",
        version="v0.0.2",
        egg="syutil-0.0.2",
    ),
    github_link(
        project="matrix-org/matrix-angular-sdk",
        version="v0.6.0",
        egg="matrix_angular_sdk-0.6.0",
    ),
    github_link(
        project="pyca/pynacl",
        version="d4d3175589b892f6ea7c22f466e0e223853516fa",
        egg="pynacl-0.3.0",
    )
]

class MissingRequirementError(Exception):
    pass

def check_requirements():
    """Checks that all the modules needed by synapse have been correctly
    installed and are at the correct version"""
    for dependency, module_requirements in REQUIREMENTS.items():
        for module_requirement in module_requirements:
            if ">=" in module_requirement:
                module_name, required_version = module_requirement.split(">=")
                version_test = ">="
            elif "==" in module_requirement:
                module_name, required_version = module_requirement.split("==")
                version_test = "=="
            else:
                module_name = module_requirement
                version_test = None

            try:
                module = __import__(module_name)
            except ImportError:
                logging.exception(
                    "Can't import %r which is part of %r",
                    module_name, dependency
                )
                raise MissingRequirementError(
                    "Can't import %r which is part of %r"
                    % (module_name, dependency)
                )
            version = getattr(module, "__version__", None)
            file_path = getattr(module, "__file__", None)
            logger.info(
                "Using %r version %r from %r to satisfy %r",
                module_name, version, file_path, dependency
            )

            if version_test == ">=":
                if version is None:
                    raise MissingRequirementError(
                        "Version of %r isn't set as __version__ of module %r"
                        % (dependency, module_name)
                    )
                if LooseVersion(version) < LooseVersion(required_version):
                    raise MissingRequirementError(
                        "Version of %r in %r is too old. %r < %r"
                        % (dependency, file_path, version, required_version)
                    )
            elif version_test == "==":
                if version is None:
                    raise MissingRequirementError(
                        "Version of %r isn't set as __version__ of module %r"
                        % (dependency, module_name)
                    )
                if LooseVersion(version) != LooseVersion(required_version):
                    raise MissingRequirementError(
                        "Unexpected version of %r in %r. %r != %r"
                        % (dependency, file_path, version, required_version)
                    )

def list_requirements():
    result = []
    linked = []
    for link in DEPENDENCY_LINKS:
        egg = link.split("#egg=")[1]
        linked.append(egg.split('-')[0])
        result.append(link)
    for requirement in REQUIREMENTS:
        is_linked = False
        for link in linked:
            if requirement.replace('-','_').startswith(link):
                is_linked = True
        if not is_linked:
            result.append(requirement)
    return result

if __name__ == "__main__":
    import sys
    sys.stdout.writelines(req + "\n" for req in list_requirements())
<commit_msg>Revert "Pull in python_dependencies.py from develop"<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2016 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import ujson

from twisted.internet import defer

from ._base import SQLBaseStore

logger = logging.getLogger(__name__)

class DeviceInboxStore(SQLBaseStore):

    @defer.inlineCallbacks
    def add_messages_to_device_inbox(self, local_messages_by_user_then_device,
                                     remote_messages_by_destination):
        """Used to send messages from this server.

        Args:
            sender_user_id(str): The ID of the user sending these messages.
            local_messages_by_user_and_device(dict):
                Dictionary of user_id to device_id to message.
            remote_messages_by_destination(dict):
                Dictionary of destination server_name to the EDU JSON to send.
        Returns:
            A deferred stream_id that resolves when the messages have been
            inserted.
        """

        def add_messages_txn(txn, now_ms, stream_id):
            # Add the local messages directly to the local inbox.
            self._add_messages_to_local_device_inbox_txn(
                txn, stream_id, local_messages_by_user_then_device
            )

            # Add the remote messages to the federation outbox.
            # We'll send them to a remote server when we next send a
            # federation transaction to that destination.
            sql = (
                "INSERT INTO device_federation_outbox"
                " (destination, stream_id, queued_ts, messages_json)"
                " VALUES (?,?,?,?)"
            )
            rows = []
            for destination, edu in remote_messages_by_destination.items():
                edu_json = ujson.dumps(edu)
                rows.append((destination, stream_id, now_ms, edu_json))
            txn.executemany(sql, rows)

        with self._device_inbox_id_gen.get_next() as stream_id:
            now_ms = self.clock.time_msec()
            yield self.runInteraction(
                "add_messages_to_device_inbox",
                add_messages_txn,
                now_ms,
                stream_id,
            )
            for user_id in local_messages_by_user_then_device.keys():
                self._device_inbox_stream_cache.entity_has_changed(
                    user_id, stream_id
                )
            for destination in remote_messages_by_destination.keys():
                self._device_federation_outbox_stream_cache.entity_has_changed(
                    destination, stream_id
                )

        defer.returnValue(self._device_inbox_id_gen.get_current_token())

    @defer.inlineCallbacks
    def add_messages_from_remote_to_device_inbox(
        self, origin, message_id, local_messages_by_user_then_device
    ):
        def add_messages_txn(txn, now_ms, stream_id):
            # Check if we've already inserted a matching message_id for that
            # origin. This can happen if the origin doesn't receive our
            # acknowledgement from the first time we received the message.
            already_inserted = self._simple_select_one_txn(
                txn, table="device_federation_inbox",
                keyvalues={"origin": origin, "message_id": message_id},
                retcols=("message_id",),
                allow_none=True,
            )
            if already_inserted is not None:
                return

            # Add an entry for this message_id so that we know we've processed
            # it.
            self._simple_insert_txn(
                txn, table="device_federation_inbox",
                values={
                    "origin": origin,
                    "message_id": message_id,
                    "received_ts": now_ms,
                },
            )

            # Add the messages to the approriate local device inboxes so that
            # they'll be sent to the devices when they next sync.
            self._add_messages_to_local_device_inbox_txn(
                txn, stream_id, local_messages_by_user_then_device
            )

        with self._device_inbox_id_gen.get_next() as stream_id:
            now_ms = self.clock.time_msec()
            yield self.runInteraction(
                "add_messages_from_remote_to_device_inbox",
                add_messages_txn,
                now_ms,
                stream_id,
            )
            for user_id in local_messages_by_user_then_device.keys():
                self._device_inbox_stream_cache.entity_has_changed(
                    user_id, stream_id
                )

    def _add_messages_to_local_device_inbox_txn(self, txn, stream_id,
                                                messages_by_user_then_device):
        sql = (
            "UPDATE device_max_stream_id"
            " SET stream_id = ?"
            " WHERE stream_id < ?"
        )
        txn.execute(sql, (stream_id, stream_id))

        local_by_user_then_device = {}
        for user_id, messages_by_device in messages_by_user_then_device.items():
            messages_json_for_user = {}
            devices = messages_by_device.keys()
            if len(devices) == 1 and devices[0] == "*":
                # Handle wildcard device_ids.
                sql = (
                    "SELECT device_id FROM devices"
                    " WHERE user_id = ?"
                )
                txn.execute(sql, (user_id,))
                message_json = ujson.dumps(messages_by_device["*"])
                for row in txn.fetchall():
                    # Add the message for all devices for this user on this
                    # server.
                    device = row[0]
                    messages_json_for_user[device] = message_json
            else:
                if not devices:
                    continue
                sql = (
                    "SELECT device_id FROM devices"
                    " WHERE user_id = ? AND device_id IN ("
                    + ",".join("?" * len(devices))
                    + ")"
                )
                # TODO: Maybe this needs to be done in batches if there are
                # too many local devices for a given user.
                txn.execute(sql, [user_id] + devices)
                for row in txn.fetchall():
                    # Only insert into the local inbox if the device exists on
                    # this server
                    device = row[0]
                    message_json = ujson.dumps(messages_by_device[device])
                    messages_json_for_user[device] = message_json

            if messages_json_for_user:
                local_by_user_then_device[user_id] = messages_json_for_user

        if not local_by_user_then_device:
            return

        sql = (
            "INSERT INTO device_inbox"
            " (user_id, device_id, stream_id, message_json)"
            " VALUES (?,?,?,?)"
        )
        rows = []
        for user_id, messages_by_device in local_by_user_then_device.items():
            for device_id, message_json in messages_by_device.items():
                rows.append((user_id, device_id, stream_id, message_json))

        txn.executemany(sql, rows)

    def get_new_messages_for_device(
        self, user_id, device_id, last_stream_id, current_stream_id, limit=100
    ):
        """
        Args:
            user_id(str): The recipient user_id.
            device_id(str): The recipient device_id.
            current_stream_id(int): The current position of the to device
                message stream.
        Returns:
            Deferred ([dict], int): List of messages for the device and where
                in the stream the messages got to.
        """
        has_changed = self._device_inbox_stream_cache.has_entity_changed(
            user_id, last_stream_id
        )
        if not has_changed:
            return defer.succeed(([], current_stream_id))

        def get_new_messages_for_device_txn(txn):
            sql = (
                "SELECT stream_id, message_json FROM device_inbox"
                " WHERE user_id = ? AND device_id = ?"
                " AND ? < stream_id AND stream_id <= ?"
                " ORDER BY stream_id ASC"
                " LIMIT ?"
            )
            txn.execute(sql, (
                user_id, device_id, last_stream_id, current_stream_id, limit
            ))
            messages = []
            for row in txn.fetchall():
                stream_pos = row[0]
                messages.append(ujson.loads(row[1]))
            if len(messages) < limit:
                stream_pos = current_stream_id
            return (messages, stream_pos)

        return self.runInteraction(
            "get_new_messages_for_device", get_new_messages_for_device_txn,
        )

    def delete_messages_for_device(self, user_id, device_id, up_to_stream_id):
        """
        Args:
            user_id(str): The recipient user_id.
            device_id(str): The recipient device_id.
            up_to_stream_id(int): Where to delete messages up to.
        Returns:
            A deferred that resolves when the messages have been deleted.
        """
        def delete_messages_for_device_txn(txn):
            sql = (
                "DELETE FROM device_inbox"
                " WHERE user_id = ? AND device_id = ?"
                " AND stream_id <= ?"
            )
            txn.execute(sql, (user_id, device_id, up_to_stream_id))

        return self.runInteraction(
            "delete_messages_for_device", delete_messages_for_device_txn
        )

    def get_all_new_device_messages(self, last_pos, current_pos, limit):
        """
        Args:
            last_pos(int):
            current_pos(int):
            limit(int):
        Returns:
            A deferred list of rows from the device inbox
        """
        if last_pos == current_pos:
            return defer.succeed([])

        def get_all_new_device_messages_txn(txn):
            sql = (
                "SELECT stream_id FROM device_inbox"
                " WHERE ? < stream_id AND stream_id <= ?"
                " GROUP BY stream_id"
                " ORDER BY stream_id ASC"
                " LIMIT ?"
            )
            txn.execute(sql, (last_pos, current_pos, limit))
            stream_ids = txn.fetchall()
            if not stream_ids:
                return []
            max_stream_id_in_limit = stream_ids[-1]

            sql = (
                "SELECT stream_id, user_id, device_id, message_json"
                " FROM device_inbox"
                " WHERE ? < stream_id AND stream_id <= ?"
                " ORDER BY stream_id ASC"
            )
            txn.execute(sql, (last_pos, max_stream_id_in_limit))
            return txn.fetchall()

        return self.runInteraction(
            "get_all_new_device_messages", get_all_new_device_messages_txn
        )

    def get_to_device_stream_token(self):
        return self._device_inbox_id_gen.get_current_token()

    def get_new_device_msgs_for_remote(
        self, destination, last_stream_id, current_stream_id, limit=100
    ):
        """
        Args:
            destination(str): The name of the remote server.
            last_stream_id(int): The last position of the device message stream
                that the server sent up to.
            current_stream_id(int): The current position of the device
                message stream.
        Returns:
            Deferred ([dict], int): List of messages for the device and where
                in the stream the messages got to.
        """

        has_changed = self._device_federation_outbox_stream_cache.has_entity_changed(
            destination, last_stream_id
        )
        if not has_changed or last_stream_id == current_stream_id:
            return defer.succeed(([], current_stream_id))

        def get_new_messages_for_remote_destination_txn(txn):
            sql = (
                "SELECT stream_id, messages_json FROM device_federation_outbox"
                " WHERE destination = ?"
                " AND ? < stream_id AND stream_id <= ?"
                " ORDER BY stream_id ASC"
                " LIMIT ?"
            )
            txn.execute(sql, (
                destination, last_stream_id, current_stream_id, limit
            ))
            messages = []
            for row in txn.fetchall():
                stream_pos = row[0]
                messages.append(ujson.loads(row[1]))
            if len(messages) < limit:
                stream_pos = current_stream_id
            return (messages, stream_pos)

        return self.runInteraction(
            "get_new_device_msgs_for_remote",
            get_new_messages_for_remote_destination_txn,
        )

    def delete_device_msgs_for_remote(self, destination, up_to_stream_id):
        """Used to delete messages when the remote destination acknowledges
        their receipt.

        Args:
            destination(str): The destination server_name
            up_to_stream_id(int): Where to delete messages up to.
        Returns:
            A deferred that resolves when the messages have been deleted.
        """
        def delete_messages_for_remote_destination_txn(txn):
            sql = (
                "DELETE FROM device_federation_outbox"
                " WHERE destination = ?"
                " AND stream_id <= ?"
            )
            txn.execute(sql, (destination, up_to_stream_id))

        return self.runInteraction(
            "delete_device_msgs_for_remote",
            delete_messages_for_remote_destination_txn
        )
<commit_msg>Fix direct to device messages recieved over federation to notify sync<commit_after># -*- coding: utf-8 -*-
# Copyright 2016 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import ujson

from twisted.internet import defer

from ._base import SQLBaseStore

logger = logging.getLogger(__name__)

class DeviceInboxStore(SQLBaseStore):

    @defer.inlineCallbacks
    def add_messages_to_device_inbox(self, local_messages_by_user_then_device,
                                     remote_messages_by_destination):
        """Used to send messages from this server.

        Args:
            sender_user_id(str): The ID of the user sending these messages.
            local_messages_by_user_and_device(dict):
                Dictionary of user_id to device_id to message.
            remote_messages_by_destination(dict):
                Dictionary of destination server_name to the EDU JSON to send.
        Returns:
            A deferred stream_id that resolves when the messages have been
            inserted.
        """

        def add_messages_txn(txn, now_ms, stream_id):
            # Add the local messages directly to the local inbox.
            self._add_messages_to_local_device_inbox_txn(
                txn, stream_id, local_messages_by_user_then_device
            )

            # Add the remote messages to the federation outbox.
            # We'll send them to a remote server when we next send a
            # federation transaction to that destination.
            sql = (
                "INSERT INTO device_federation_outbox"
                " (destination, stream_id, queued_ts, messages_json)"
                " VALUES (?,?,?,?)"
            )
            rows = []
            for destination, edu in remote_messages_by_destination.items():
                edu_json = ujson.dumps(edu)
                rows.append((destination, stream_id, now_ms, edu_json))
            txn.executemany(sql, rows)

        with self._device_inbox_id_gen.get_next() as stream_id:
            now_ms = self.clock.time_msec()
            yield self.runInteraction(
                "add_messages_to_device_inbox",
                add_messages_txn,
                now_ms,
                stream_id,
            )
            for user_id in local_messages_by_user_then_device.keys():
                self._device_inbox_stream_cache.entity_has_changed(
                    user_id, stream_id
                )
            for destination in remote_messages_by_destination.keys():
                self._device_federation_outbox_stream_cache.entity_has_changed(
                    destination, stream_id
                )

        defer.returnValue(self._device_inbox_id_gen.get_current_token())

    @defer.inlineCallbacks
    def add_messages_from_remote_to_device_inbox(
        self, origin, message_id, local_messages_by_user_then_device
    ):
        def add_messages_txn(txn, now_ms, stream_id):
            # Check if we've already inserted a matching message_id for that
            # origin. This can happen if the origin doesn't receive our
            # acknowledgement from the first time we received the message.
            already_inserted = self._simple_select_one_txn(
                txn, table="device_federation_inbox",
                keyvalues={"origin": origin, "message_id": message_id},
                retcols=("message_id",),
                allow_none=True,
            )
            if already_inserted is not None:
                return

            # Add an entry for this message_id so that we know we've processed
            # it.
            self._simple_insert_txn(
                txn, table="device_federation_inbox",
                values={
                    "origin": origin,
                    "message_id": message_id,
                    "received_ts": now_ms,
                },
            )

            # Add the messages to the approriate local device inboxes so that
            # they'll be sent to the devices when they next sync.
            self._add_messages_to_local_device_inbox_txn(
                txn, stream_id, local_messages_by_user_then_device
            )

        with self._device_inbox_id_gen.get_next() as stream_id:
            now_ms = self.clock.time_msec()
            yield self.runInteraction(
                "add_messages_from_remote_to_device_inbox",
                add_messages_txn,
                now_ms,
                stream_id,
            )
            for user_id in local_messages_by_user_then_device.keys():
                self._device_inbox_stream_cache.entity_has_changed(
                    user_id, stream_id
                )

        defer.returnValue(stream_id)

    def _add_messages_to_local_device_inbox_txn(self, txn, stream_id,
                                                messages_by_user_then_device):
        sql = (
            "UPDATE device_max_stream_id"
            " SET stream_id = ?"
            " WHERE stream_id < ?"
        )
        txn.execute(sql, (stream_id, stream_id))

        local_by_user_then_device = {}
        for user_id, messages_by_device in messages_by_user_then_device.items():
            messages_json_for_user = {}
            devices = messages_by_device.keys()
            if len(devices) == 1 and devices[0] == "*":
                # Handle wildcard device_ids.
                sql = (
                    "SELECT device_id FROM devices"
                    " WHERE user_id = ?"
                )
                txn.execute(sql, (user_id,))
                message_json = ujson.dumps(messages_by_device["*"])
                for row in txn.fetchall():
                    # Add the message for all devices for this user on this
                    # server.
                    device = row[0]
                    messages_json_for_user[device] = message_json
            else:
                if not devices:
                    continue
                sql = (
                    "SELECT device_id FROM devices"
                    " WHERE user_id = ? AND device_id IN ("
                    + ",".join("?" * len(devices))
                    + ")"
                )
                # TODO: Maybe this needs to be done in batches if there are
                # too many local devices for a given user.
                txn.execute(sql, [user_id] + devices)
                for row in txn.fetchall():
                    # Only insert into the local inbox if the device exists on
                    # this server
                    device = row[0]
                    message_json = ujson.dumps(messages_by_device[device])
                    messages_json_for_user[device] = message_json

            if messages_json_for_user:
                local_by_user_then_device[user_id] = messages_json_for_user

        if not local_by_user_then_device:
            return

        sql = (
            "INSERT INTO device_inbox"
            " (user_id, device_id, stream_id, message_json)"
            " VALUES (?,?,?,?)"
        )
        rows = []
        for user_id, messages_by_device in local_by_user_then_device.items():
            for device_id, message_json in messages_by_device.items():
                rows.append((user_id, device_id, stream_id, message_json))

        txn.executemany(sql, rows)

    def get_new_messages_for_device(
        self, user_id, device_id, last_stream_id, current_stream_id, limit=100
    ):
        """
        Args:
            user_id(str): The recipient user_id.
            device_id(str): The recipient device_id.
            current_stream_id(int): The current position of the to device
                message stream.
        Returns:
            Deferred ([dict], int): List of messages for the device and where
                in the stream the messages got to.
        """
        has_changed = self._device_inbox_stream_cache.has_entity_changed(
            user_id, last_stream_id
        )
        if not has_changed:
            return defer.succeed(([], current_stream_id))

        def get_new_messages_for_device_txn(txn):
            sql = (
                "SELECT stream_id, message_json FROM device_inbox"
                " WHERE user_id = ? AND device_id = ?"
                " AND ? < stream_id AND stream_id <= ?"
                " ORDER BY stream_id ASC"
                " LIMIT ?"
            )
            txn.execute(sql, (
                user_id, device_id, last_stream_id, current_stream_id, limit
            ))
            messages = []
            for row in txn.fetchall():
                stream_pos = row[0]
                messages.append(ujson.loads(row[1]))
            if len(messages) < limit:
                stream_pos = current_stream_id
            return (messages, stream_pos)

        return self.runInteraction(
            "get_new_messages_for_device", get_new_messages_for_device_txn,
        )

    def delete_messages_for_device(self, user_id, device_id, up_to_stream_id):
        """
        Args:
            user_id(str): The recipient user_id.
            device_id(str): The recipient device_id.
            up_to_stream_id(int): Where to delete messages up to.
        Returns:
            A deferred that resolves when the messages have been deleted.
        """
        def delete_messages_for_device_txn(txn):
            sql = (
                "DELETE FROM device_inbox"
                " WHERE user_id = ? AND device_id = ?"
                " AND stream_id <= ?"
            )
            txn.execute(sql, (user_id, device_id, up_to_stream_id))

        return self.runInteraction(
            "delete_messages_for_device", delete_messages_for_device_txn
        )

    def get_all_new_device_messages(self, last_pos, current_pos, limit):
        """
        Args:
            last_pos(int):
            current_pos(int):
            limit(int):
        Returns:
            A deferred list of rows from the device inbox
        """
        if last_pos == current_pos:
            return defer.succeed([])

        def get_all_new_device_messages_txn(txn):
            sql = (
                "SELECT stream_id FROM device_inbox"
                " WHERE ? < stream_id AND stream_id <= ?"
                " GROUP BY stream_id"
                " ORDER BY stream_id ASC"
                " LIMIT ?"
            )
            txn.execute(sql, (last_pos, current_pos, limit))
            stream_ids = txn.fetchall()
            if not stream_ids:
                return []
            max_stream_id_in_limit = stream_ids[-1]

            sql = (
                "SELECT stream_id, user_id, device_id, message_json"
                " FROM device_inbox"
                " WHERE ? < stream_id AND stream_id <= ?"
                " ORDER BY stream_id ASC"
            )
            txn.execute(sql, (last_pos, max_stream_id_in_limit))
            return txn.fetchall()

        return self.runInteraction(
            "get_all_new_device_messages", get_all_new_device_messages_txn
        )

    def get_to_device_stream_token(self):
        return self._device_inbox_id_gen.get_current_token()

    def get_new_device_msgs_for_remote(
        self, destination, last_stream_id, current_stream_id, limit=100
    ):
        """
        Args:
            destination(str): The name of the remote server.
            last_stream_id(int): The last position of the device message stream
                that the server sent up to.
            current_stream_id(int): The current position of the device
                message stream.
        Returns:
            Deferred ([dict], int): List of messages for the device and where
                in the stream the messages got to.
        """

        has_changed = self._device_federation_outbox_stream_cache.has_entity_changed(
            destination, last_stream_id
        )
        if not has_changed or last_stream_id == current_stream_id:
            return defer.succeed(([], current_stream_id))

        def get_new_messages_for_remote_destination_txn(txn):
            sql = (
                "SELECT stream_id, messages_json FROM device_federation_outbox"
                " WHERE destination = ?"
                " AND ? < stream_id AND stream_id <= ?"
                " ORDER BY stream_id ASC"
                " LIMIT ?"
            )
            txn.execute(sql, (
                destination, last_stream_id, current_stream_id, limit
            ))
            messages = []
            for row in txn.fetchall():
                stream_pos = row[0]
                messages.append(ujson.loads(row[1]))
            if len(messages) < limit:
                stream_pos = current_stream_id
            return (messages, stream_pos)

        return self.runInteraction(
            "get_new_device_msgs_for_remote",
            get_new_messages_for_remote_destination_txn,
        )

    def delete_device_msgs_for_remote(self, destination, up_to_stream_id):
        """Used to delete messages when the remote destination acknowledges
        their receipt.

        Args:
            destination(str): The destination server_name
            up_to_stream_id(int): Where to delete messages up to.
        Returns:
            A deferred that resolves when the messages have been deleted.
        """
        def delete_messages_for_remote_destination_txn(txn):
            sql = (
                "DELETE FROM device_federation_outbox"
                " WHERE destination = ?"
                " AND stream_id <= ?"
            )
            txn.execute(sql, (destination, up_to_stream_id))

        return self.runInteraction(
            "delete_device_msgs_for_remote",
            delete_messages_for_remote_destination_txn
        )
<|endoftext|>


<commit_before># Copyright 2016 the original author or authors.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

__author__ = "Julian Debatin"
__copyright__ = "The authors"
__license__ = "Apache 2"
__email__ = "juliandebatin@gmail.com"
__status__ = "Production"

from Core.CellType import *
from Core.ExecConfig import ExecConfig
from Core.ModelConfig import ModelConfig
from Steppable.InitializerSteppable import InitializerSteppable
from Steppable.DeathSteppable import DeathSteppable
from Steppable.GrowthMitosisSteppable import GrowthMitosisSteppable
from Steppable.GrowthSteppable import GrowthSteppable
from Steppable.CMTransformationSteppable import CMTransformationSteppable
from Steppable.UrinationSteppable import UrinationSteppable
from Logger.VolumeFitnessSteppable import VolumeFitnessSteppable
from Logger.ArrangementFitnessSteppable import ArrangementFitnessSteppable
from Logger.DummyFitnessSteppable import DummyFitnessSteppable
from Steppable.ColonySteppable import ColonySteppable
from Steppable.MutationSteppable import MutationSteppable

class SpaBpaPcdiInUa(ModelConfig):
    def __init__(self, sim, simthread):
        ModelConfig.__init__(self, sim, simthread)

    def _initModel(self):
        self.name = "SpaBpaPcdiInUa"
        self.cellTypes = self._createCellTypes()
        self.energyMatrix = self._createEnergyMatrix()
        self._run() # Must be the last statement.

    def _createCellTypes(self):
        cellTypes = []
        stem = Stemcell
        stem.setGrowthVolumePerDayRelVolume(0.13)
        self.stemNecrosisProb = stem.necrosisProb = 0

        basal = Basalcell
        basal.setGrowthVolumePerDayRelVolume(0.12)
        basal.apoptosisTimeInDays = 80000000000.0
        self.basalNecrosisProb = basal.necrosisProb = 0.000015

        intermediate = Intermediatecell
        intermediate.setGrowthVolumePerDayRelVolume(0.105)
        intermediate.apoptosisTimeInDays = 200000000.0
        self.intermediateNecrosisProb = intermediate.necrosisProb = 0.000035

        umbrella = Umbrellacell
        umbrella.setGrowthVolumePerDayRelVolume(0.1)
        umbrella.apoptosisTimeInDays = 100000000.0
        self.umbrellaNecrosisProb = umbrella.necrosisProb = 0.00004

        stem.setDescendants(0.90, [stem.id, basal.id])
        stem.setDescendants(0.05, [stem.id, stem.id])
        stem.setDescendants(0.05, [basal.id, basal.id])
        basal.setDescendants(0.90, [basal.id, intermediate.id])
        basal.setDescendants(0.05, [basal.id, basal.id])
        basal.setDescendants(0.05, [intermediate.id, intermediate.id])
        intermediate.setDescendants(1.0, [intermediate.id, intermediate.id])

        cellTypes.extend((Medium, Basalmembrane, stem, basal, intermediate, umbrella))

        return cellTypes

    def _getSteppables(self):
        steppableList = []
        steppableList.append(ColonySteppable(self.sim, self))
        steppableList.append(InitializerSteppable(self.sim, self))
        steppableList.append(GrowthSteppable(self.sim, self))
        steppableList.append(GrowthMitosisSteppable(self.sim, self))
        steppableList.append(CMTransformationSteppable(self.sim, self))
        steppableList.append(UrinationSteppable(self.sim, self, prop=0.02))
        steppableList.append(DeathSteppable(self.sim, self))
        #steppableList.append(OptimumSearchSteppable(self.sim, self))
        steppableList.append(VolumeFitnessSteppable(self.sim, self))
        steppableList.append(ArrangementFitnessSteppable(self.sim, self))
        steppableList.append(DummyFitnessSteppable(self.sim, self))
        steppableList.append(MutationSteppable(self.sim, self, self.stemNecrosisProb, self.basalNecrosisProb,
                                               self.intermediateNecrosisProb, self.umbrellaNecrosisProb))

        return steppableList

    def _createExecConfig(self):
        return ExecConfig(MCSperDay=500,# SEED=10,
                          xLength=500, yLength=150, zLength=0, voxelDensity=.8)
<commit_msg>no message<commit_after># Copyright 2016 the original author or authors.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

__author__ = "Julian Debatin"
__copyright__ = "The authors"
__license__ = "Apache 2"
__email__ = "juliandebatin@gmail.com"
__status__ = "Production"

from Core.CellType import *
from Core.ExecConfig import ExecConfig
from Core.ModelConfig import ModelConfig
from Steppable.InitializerSteppable import InitializerSteppable
from Steppable.DeathSteppable import DeathSteppable
from Steppable.GrowthMitosisSteppable import GrowthMitosisSteppable
from Steppable.GrowthSteppable import GrowthSteppable
from Steppable.CMTransformationSteppable import CMTransformationSteppable
from Steppable.UrinationSteppable import UrinationSteppable
from Logger.VolumeFitnessSteppable import VolumeFitnessSteppable
from Logger.ArrangementFitnessSteppable import ArrangementFitnessSteppable
from Logger.DummyFitnessSteppable import DummyFitnessSteppable
from Steppable.ColonySteppable import ColonySteppable
from Steppable.MutationSteppable import MutationSteppable

class SpaBpaPcdiInUa(ModelConfig):
    def __init__(self, sim, simthread):
        ModelConfig.__init__(self, sim, simthread)

    def _initModel(self):
        self.name = "SpaBpaPcdiInUa"
        self.cellTypes = self._createCellTypes()
        self.energyMatrix = self._createEnergyMatrix()
        self._run() # Must be the last statement.

    def _createCellTypes(self):
        cellTypes = []
        stem = Stemcell
        stem.setGrowthVolumePerDayRelVolume(0.13)
        self.stemNecrosisProb = stem.necrosisProb = 0

        basal = Basalcell
        basal.setGrowthVolumePerDayRelVolume(0.12)
        basal.apoptosisTimeInDays = 80000000000.0
        self.basalNecrosisProb = basal.necrosisProb = 0.000015

        intermediate = Intermediatecell
        intermediate.setGrowthVolumePerDayRelVolume(0.10)
        intermediate.apoptosisTimeInDays = 200000000.0
        self.intermediateNecrosisProb = intermediate.necrosisProb = 0.00004

        umbrella = Umbrellacell
        umbrella.setGrowthVolumePerDayRelVolume(0.1)
        umbrella.apoptosisTimeInDays = 100000000.0
        self.umbrellaNecrosisProb = umbrella.necrosisProb = 0.000045

        stem.setDescendants(0.90, [stem.id, basal.id])
        stem.setDescendants(0.05, [stem.id, stem.id])
        stem.setDescendants(0.05, [basal.id, basal.id])
        basal.setDescendants(0.90, [basal.id, intermediate.id])
        basal.setDescendants(0.05, [basal.id, basal.id])
        basal.setDescendants(0.05, [intermediate.id, intermediate.id])
        intermediate.setDescendants(1.0, [intermediate.id, intermediate.id])

        cellTypes.extend((Medium, Basalmembrane, stem, basal, intermediate, umbrella))

        return cellTypes

    def _getSteppables(self):
        steppableList = []
        steppableList.append(ColonySteppable(self.sim, self))
        steppableList.append(InitializerSteppable(self.sim, self))
        steppableList.append(GrowthSteppable(self.sim, self))
        steppableList.append(GrowthMitosisSteppable(self.sim, self))
        steppableList.append(CMTransformationSteppable(self.sim, self))
        steppableList.append(UrinationSteppable(self.sim, self, prop=0.02))
        steppableList.append(DeathSteppable(self.sim, self))
        #steppableList.append(OptimumSearchSteppable(self.sim, self))
        steppableList.append(VolumeFitnessSteppable(self.sim, self))
        steppableList.append(ArrangementFitnessSteppable(self.sim, self))
        steppableList.append(DummyFitnessSteppable(self.sim, self))
        steppableList.append(MutationSteppable(self.sim, self, self.stemNecrosisProb, self.basalNecrosisProb,
                                               self.intermediateNecrosisProb, self.umbrellaNecrosisProb))

        return steppableList

    def _createExecConfig(self):
        return ExecConfig(MCSperDay=500,# SEED=10,
                          xLength=500, yLength=150, zLength=0, voxelDensity=.8)
<|endoftext|>


<commit_before>"""Class for lightcurves."""

import logging
import itertools

import numpy as np
import matplotlib.pyplot as plt
from astroML import time_series

from k2spin import utils
from k2spin import clean
from k2spin import detrend
from k2spin import plot
from k2spin import prot
from k2spin import evaluate

class LightCurve(object):
    """
    Class to contain lightcurves and run analysis.

    """

    def __init__(self, time, flux, unc_flux, x_pos, y_pos, name,
                 power_threshold=0.5):
        """Clean up the input data and sigma-clip it."""
        # Save the power threshold for later use
        self.power_threshold = power_threshold
        self.name = name

        logging.debug(self.name)
        logging.debug("Threshold %f",self.power_threshold)

        # Clean up the input lightcurve
        cleaned_out = clean.prep_lc(time, flux, unc_flux, clip_at=3.)
        self.time, self.flux, self.unc_flux = cleaned_out[:3]
        self.med, self.stdev, all_kept = cleaned_out[3:]
        self.x_pos, self.y_pos = x_pos[all_kept], y_pos[all_kept]
        logging.debug("len init t %d f %d u %d", len(self.time),
                      len(self.flux),len(self.unc_flux))

        # Detrend the raw flux
        self._bulk_detrend()

    def choose_initial(self):
        """Search raw and detrended LCs for periods, and decide whether there's
        a period there.

        """
        # Run a fit on the raw lc
        r_out = self._run_fit("raw")
        raw_fp, raw_power, raw_prots, raw_pgram, raw_alias, raw_sigma = r_out
        logging.debug("Ran raw fit")

        # Run a fit on the detrended lc
        d_out = self._run_fit("detrended")
        det_fp, det_power, det_prots, det_pgram, det_alias, det_sigma = d_out
        logging.debug("Ran detrended fit")

        # Compare them
        lc_to_use = self._pick_lc(raw_power, det_power)
        if lc_to_use<=1:
            logging.info("Using raw lightcurve")
            self.init_prot , self.init_power = raw_fp, raw_power
            self.init_periods_to_test, self.init_pgram = raw_prots, raw_pgram
            self.use_flux = self.flux / self.med
            self.use_unc = self.unc_flux / self.med
            self.init_sigmas = raw_sigma
            data_labels = ["Raw (Selected)", "Detrended"]
        elif lc_to_use==2:
            logging.info("Using detrended lightcurve")
            self.init_prot , self.init_power = det_fp, det_power
            self.init_periods_to_test, self.init_pgram = det_prots, det_pgram
            self.use_flux = self.det_flux
            self.use_unc = self.unc_flux
            self.init_sigmas = det_sigma
            data_labels = ["Raw", "Detrended (Selected)"]

        logging.info("Initial Prot %f Power %f", self.init_prot,
                     self.init_power)

        # Get aliases for selected period
        eval_out = evaluate.test_pgram(self.init_periods_to_test,
                                       self.init_pgram, self.power_threshold)
        if eval_out[-1]==False:
            logging.warning("Selected lightcurve is not clean")
        else:
            logging.debug("Selected lightcurve is clean")
        plot_aliases = [None, eval_out[2]]

        # Plot them up
        lcs = [[self.time, self.flux/self.med, abs(self.unc_flux/self.med)],
               [self.time, self.det_flux, self.det_unc]]
        pgrams = [[raw_prots, raw_pgram], [det_prots, det_pgram]]
        best_periods = [raw_fp, det_fp]
        sigmas = [raw_sigma, det_sigma]
        logging.debug(sigmas)
        rd_fig, rd_axes = plot.compare_multiple(lcs, pgrams, best_periods,
                                                sigmas,
                                                aliases=plot_aliases,
                                                data_labels=data_labels,
                                                phase_by=self.init_prot)

        rd_fig.suptitle(self.name, fontsize="large", y=0.99)

        rd_fig.delaxes(rd_axes[3])

        plt.savefig("{}_raw_detrend.png".format(self.name))

    def correct_and_fit(self):
        """Position-correct and perform a fit."""
        logging.debug("Fitting corrected lightcurve")

        self._xy_correct()

        fit_out = self._run_fit([self.time, self.corrected_flux,
                                 self.corrected_unc])
        fund_prot, fund_power, periods_to_test, periodogram = fit_out[:4]
        aliases, sigmas = fit_out[4:]

        eval_out =  evaluate.test_pgram(periods_to_test, periodogram,
                                        self.power_threshold)

        self.corrected_prot = fund_prot

        if eval_out[-1]==False:
            logging.warning("Corrected lightcurve is not clean")
        else:
            logging.debug("Corrected lightcurve is clean")
        plot_aliases = [None, eval_out[2]]

        # Plot them up
        lcs = [[self.time, self.use_flux, self.use_unc],
               [self.time, self.corrected_flux, self.corrected_unc]]
        pgrams = [[self.init_periods_to_test, self.init_pgram],
                  [periods_to_test, periodogram]]
        best_periods = [self.init_prot, fund_prot]
        data_labels = ["Initial", "Corrected"]
        sigmas = [self.init_sigmas, sigmas]
        rd_fig, rd_axes = plot.compare_multiple(lcs, pgrams, best_periods,
                                                sigmas,
                                                aliases=plot_aliases,
                                                data_labels=data_labels,
                                                phase_by=fund_prot)

        rd_fig.suptitle(self.name, fontsize="large", y=0.99)

        ptime, fsine = evaluate.fit_sine(self.time, self.corrected_flux,
                                         self.corrected_unc,
                                         self.corrected_prot)

        plotx = np.argsort(ptime)
        rd_axes[2].plot(ptime[plotx], fsine[plotx], color="lightgrey", lw=2)
        rd_axes[2].set_ylim(min(fsine)*0.9, max(fsine)*1.1)

        use_residuals = self.use_flux - fsine
        cor_residuals = self.corrected_flux - fsine

        logging.debug("RESIDUALS")
        logging.debug(use_residuals[:10])
        logging.debug(cor_residuals[:10])

        rd_axes[3].errorbar(self.time % fund_prot, use_residuals,
                            np.zeros_like(self.time), #self.use_unc,
                            fmt=plot.shape1, ms=2, capsize=0,
                            ecolor=plot.color1, color=plot.color1,
                            mec=plot.color1)
        rd_axes[3].errorbar(self.time % fund_prot, cor_residuals,
                            np.zeros_like(self.time), #self.corrected_unc,
                            fmt=plot.shape2, ms=2, capsize=0,
                            ecolor=plot.color2, color=plot.color2,
                            mec=plot.color2)
        rd_axes[3].set_xlim(0, fund_prot)

        plt.savefig("{}_corrected.png".format(self.name))
#        plt.show()
        plt.close("all")

    def _bulk_detrend(self, alpha=8):
        """Smooth the rapid variations in the lightcurve and remove bulk trends.

        inputs
        ------
        alpha: float
            "bass enhancement" for supersmoother.
            Defaults to 8 for now because Kevin was using that.
        """

        logging.debug("Removing bulk trend...")
        det_out = detrend.simple_detrend(self.time, self.flux, self.unc_flux,
                                         kind="supersmoother", phaser=alpha)
        self.det_flux, self.det_unc, self.bulk_trend = det_out

        logging.debug("len detrended t %d f %d u %d", len(self.time),
                      len(self.det_flux),len(self.det_unc))

    def _run_fit(self, use_lc, prot_lims=[0.1,70]):
        """Run a fit on a single lc, either "raw" or "detrended"
        or a array/list of [time, flux, and unc]
        """

        if use_lc=="raw":
            logging.debug("fitting raw lc")
            tt, ff, uu = self.time, self.flux, self.unc_flux
        elif use_lc=="detrended":
            logging.debug("fitting detrended lc")
            tt, ff, uu = self.time, self.det_flux, self.det_unc
        else:
            logging.debug("fitting other lc")
            tt, ff, uu = use_lc

        logging.debug("_run_fit threshold %f", self.power_threshold)

        # Iteratively smooth, clip, and run a periodogram (period_cleaner)
        pc_out = prot.period_cleaner(tt, ff, uu,
                                     pgram_threshold=self.power_threshold,
                                     prot_lims=prot_lims)
        cl_time, cl_flux, cl_unc, sm_flux = pc_out

        logging.debug("Smoothed, now periodogram")
        logging.debug("Cleaned t %d f %d u %d", len(cl_time),
                      len(cl_flux), len(cl_unc))
        # Test the periodogram and pick the best period and power
        ls_out = prot.run_ls(cl_time, cl_flux, cl_unc,
                             threshold=self.power_threshold,
                             prot_lims=prot_lims, run_bootstrap=True)
        #fund_prot, fund_power, periods_to_test, periodogram, aliases, sigmas

        return ls_out

    def _pick_lc(self, fund_power1, fund_power2):
        """Pick the raw or detrended lc to continue with by
        selecting the one with the highest peak in the periodogram
        (no consideration of the *locations* of those peaks)
        """
        # return a integer indicating which lightcurve to use (1 or 2)
        to_use = 0

        if fund_power1 > fund_power2:
            to_use = 1
        elif fund_power2 > fund_power1:
            to_use = 2
        else: # Either something's gone wrong, or they're exactly equal
            to_use = 0

        return to_use

    def _xy_correct(self, n_closest=21):
        """Correct for positional variations in the lightcurve once selected."""

        # Loop through the lightcurve and find the n closest pixels.
        # Then divide by the median flux from those pixels

        num_pts = len(self.use_flux)
        logging.debug("Pixel position correction %d", num_pts)

        self.corrected_flux = np.zeros(num_pts)
        self.corrected_unc = np.zeros(num_pts)

        first_half = self.time<=2102
        x_pos1 = self.x_pos[first_half==True]
        y_pos1 = self.y_pos[first_half==True]
        x_pos2 = self.x_pos[first_half==False]
        y_pos2 = self.y_pos[first_half==False]

        for i, fval, xx, yy in itertools.izip(range(num_pts), self.use_flux,
                                              self.x_pos, self.y_pos):
            if first_half[i]:
                comp_x, comp_y = x_pos1, y_pos1
                comp_f = self.use_flux[first_half==True]
            else:
                comp_x, comp_y = x_pos2, y_pos2
                comp_f = self.use_flux[first_half==False]

#            comp_x, comp_y = self.x_pos, self.y_pos
#            comp_f = self.use_flux

            pix_sep = np.sqrt((xx - comp_x)**2 + (yy - comp_y)**2)
            min_ind = np.argpartition(pix_sep, n_closest)[:n_closest]
            logging.debug(np.median(pix_sep[min_ind]))

            median_nearest = np.median(comp_f[min_ind])
            logging.debug("This flux %f Median Nearest %f",
                          fval, median_nearest)
            self.corrected_flux[i] = fval / median_nearest
            self.corrected_unc[i] = self.use_unc[i] / median_nearest

        logging.debug("Correction completed")

    def _plot_xy(self):
        """Plot some basic informational plots:
        Flux as a function of X-Y position
        Flux as a function of time
        """
        pass

    def _multi_search(self,lc_type):
        """Search a lightcurve for a secondary signal."""
        pass

<commit_msg>trying to implement a more clever selection of the max peak.<commit_after>"""Class for lightcurves."""

import logging
import itertools

import numpy as np
import matplotlib.pyplot as plt
from astroML import time_series

from k2spin import utils
from k2spin import clean
from k2spin import detrend
from k2spin import plot
from k2spin import prot
from k2spin import evaluate

class LightCurve(object):
    """
    Class to contain lightcurves and run analysis.

    """

    def __init__(self, time, flux, unc_flux, x_pos, y_pos, name,
                 power_threshold=0.5):
        """Clean up the input data and sigma-clip it."""
        # Save the power threshold for later use
        self.power_threshold = power_threshold
        self.name = name

        logging.debug(self.name)
        logging.debug("Threshold %f",self.power_threshold)

        # Clean up the input lightcurve
        cleaned_out = clean.prep_lc(time, flux, unc_flux, clip_at=3.)
        self.time, self.flux, self.unc_flux = cleaned_out[:3]
        self.med, self.stdev, all_kept = cleaned_out[3:]
        self.x_pos, self.y_pos = x_pos[all_kept], y_pos[all_kept]
        logging.debug("len init t %d f %d u %d", len(self.time),
                      len(self.flux),len(self.unc_flux))

        # Detrend the raw flux
        self._bulk_detrend()

    def choose_initial(self):
        """Search raw and detrended LCs for periods, and decide whether there's
        a period there.

        """
        # Run a fit on the raw lc
        r_out = self._run_fit("raw")
        raw_fp, raw_power, raw_prots, raw_pgram, raw_alias, raw_sigma = r_out
        logging.debug("Ran raw fit")

        # Run a fit on the detrended lc
        d_out = self._run_fit("detrended")
        det_fp, det_power, det_prots, det_pgram, det_alias, det_sigma = d_out
        logging.debug("Ran detrended fit")

        # Only consider peaks less than ~half the length of the lightcurve
        max_peak_loc = 0.75 * (self.time[-1] - self.time[0])
        logging.info("Max Prot = %f", max_peak_loc)

        raw_loc2 = np.argmax(raw_pgram[raw_prots<max_peak_loc])
        raw_power2 = raw_pgram[raw_prots<max_peak_loc][raw_loc2]
        raw_prot2 = raw_prots[raw_prots<max_peak_loc][raw_loc2]
        logging.info("raw %d FP %f Power %f", raw_loc2, raw_prot2, raw_power2)

        det_loc2 = np.argmax(det_pgram[det_prots<max_peak_loc])
        det_power2 = det_pgram[det_prots<max_peak_loc][det_loc2]
        det_prot2 = det_prots[det_prots<max_peak_loc][det_loc2]
        logging.info("det %d FP %f Power %f", det_loc2, det_prot2, det_power2)

        # Compare them
        lc_to_use = self._pick_lc(raw_power2, det_power2)
        if lc_to_use<=1:
            logging.info("Using raw lightcurve")
            self.init_prot , self.init_power = raw_prot2, raw_power2
            self.init_periods_to_test, self.init_pgram = raw_prots, raw_pgram
            self.use_flux = self.flux / self.med
            self.use_unc = self.unc_flux / self.med
            self.init_sigmas = raw_sigma
            data_labels = ["Raw (Selected)", "Detrended"]
        elif lc_to_use==2:
            logging.info("Using detrended lightcurve")
            self.init_prot , self.init_power = det_prot2, det_power2
            self.init_periods_to_test, self.init_pgram = det_prots, det_pgram
            self.use_flux = self.det_flux
            self.use_unc = self.unc_flux
            self.init_sigmas = det_sigma
            data_labels = ["Raw", "Detrended (Selected)"]

        logging.info("Initial Prot %f Power %f", self.init_prot,
                     self.init_power)

        # Get aliases for selected period
        eval_out = evaluate.test_pgram(self.init_periods_to_test,
                                       self.init_pgram, self.power_threshold)
        if eval_out[-1]==False:
            logging.warning("Selected lightcurve is not clean")
        else:
            logging.debug("Selected lightcurve is clean")
        plot_aliases = [None, eval_out[2]]

        # Plot them up
        lcs = [[self.time, self.flux/self.med, abs(self.unc_flux/self.med)],
               [self.time, self.det_flux, self.det_unc]]
        pgrams = [[raw_prots, raw_pgram], [det_prots, det_pgram]]
        best_periods = [raw_prot2, det_prot2]
        sigmas = [raw_sigma, det_sigma]
        logging.debug(sigmas)
        rd_fig, rd_axes = plot.compare_multiple(lcs, pgrams, best_periods,
                                                sigmas,
                                                aliases=plot_aliases,
                                                data_labels=data_labels,
                                                phase_by=self.init_prot)

        rd_fig.suptitle(self.name, fontsize="large", y=0.99)

        rd_fig.delaxes(rd_axes[3])

        plt.savefig("{}_raw_detrend.png".format(self.name))

    def correct_and_fit(self):
        """Position-correct and perform a fit."""
        logging.debug("Fitting corrected lightcurve")

        self._xy_correct()

        fit_out = self._run_fit([self.time, self.corrected_flux,
                                 self.corrected_unc])
        fund_prot, fund_power, periods_to_test, periodogram = fit_out[:4]
        aliases, sigmas = fit_out[4:]

        eval_out =  evaluate.test_pgram(periods_to_test, periodogram,
                                        self.power_threshold)

        self.corrected_prot = fund_prot

        if eval_out[-1]==False:
            logging.warning("Corrected lightcurve is not clean")
        else:
            logging.debug("Corrected lightcurve is clean")
        plot_aliases = [None, eval_out[2]]

        # Plot them up
        lcs = [[self.time, self.use_flux, self.use_unc],
               [self.time, self.corrected_flux, self.corrected_unc]]
        pgrams = [[self.init_periods_to_test, self.init_pgram],
                  [periods_to_test, periodogram]]
        best_periods = [self.init_prot, fund_prot]
        data_labels = ["Initial", "Corrected"]
        sigmas = [self.init_sigmas, sigmas]
        rd_fig, rd_axes = plot.compare_multiple(lcs, pgrams, best_periods,
                                                sigmas,
                                                aliases=plot_aliases,
                                                data_labels=data_labels,
                                                phase_by=fund_prot)

        rd_fig.suptitle(self.name, fontsize="large", y=0.99)

        ptime, fsine = evaluate.fit_sine(self.time, self.corrected_flux,
                                         self.corrected_unc,
                                         self.corrected_prot)

        plotx = np.argsort(ptime)
        rd_axes[2].plot(ptime[plotx], fsine[plotx], color="lightgrey", lw=2)
        rd_axes[2].set_ylim(min(fsine)*0.9, max(fsine)*1.1)

        use_residuals = self.use_flux - fsine
        cor_residuals = self.corrected_flux - fsine

        logging.debug("RESIDUALS")
        logging.debug(use_residuals[:10])
        logging.debug(cor_residuals[:10])

        rd_axes[3].errorbar(self.time % fund_prot, use_residuals,
                            np.zeros_like(self.time), #self.use_unc,
                            fmt=plot.shape1, ms=2, capsize=0,
                            ecolor=plot.color1, color=plot.color1,
                            mec=plot.color1)
        rd_axes[3].errorbar(self.time % fund_prot, cor_residuals,
                            np.zeros_like(self.time), #self.corrected_unc,
                            fmt=plot.shape2, ms=2, capsize=0,
                            ecolor=plot.color2, color=plot.color2,
                            mec=plot.color2)
        rd_axes[3].set_xlim(0, fund_prot)

        plt.savefig("{}_corrected.png".format(self.name))
        plt.show()
#        plt.close("all")

    def _bulk_detrend(self, alpha=8):
        """Smooth the rapid variations in the lightcurve and remove bulk trends.

        inputs
        ------
        alpha: float
            "bass enhancement" for supersmoother.
            Defaults to 8 for now because Kevin was using that.
        """

        logging.debug("Removing bulk trend...")
        det_out = detrend.simple_detrend(self.time, self.flux, self.unc_flux,
                                         kind="supersmoother", phaser=alpha)
        self.det_flux, self.det_unc, self.bulk_trend = det_out

        logging.debug("len detrended t %d f %d u %d", len(self.time),
                      len(self.det_flux),len(self.det_unc))

    def _run_fit(self, use_lc, prot_lims=[0.1,70]):
        """Run a fit on a single lc, either "raw" or "detrended"
        or a array/list of [time, flux, and unc]
        """

        if use_lc=="raw":
            logging.debug("fitting raw lc")
            tt, ff, uu = self.time, self.flux, self.unc_flux
        elif use_lc=="detrended":
            logging.debug("fitting detrended lc")
            tt, ff, uu = self.time, self.det_flux, self.det_unc
        else:
            logging.debug("fitting other lc")
            tt, ff, uu = use_lc

        logging.debug("_run_fit threshold %f", self.power_threshold)

        # Iteratively smooth, clip, and run a periodogram (period_cleaner)
        pc_out = prot.period_cleaner(tt, ff, uu,
                                     pgram_threshold=self.power_threshold,
                                     prot_lims=prot_lims)
        cl_time, cl_flux, cl_unc, sm_flux = pc_out

        logging.debug("Smoothed, now periodogram")
        logging.debug("Cleaned t %d f %d u %d", len(cl_time),
                      len(cl_flux), len(cl_unc))
        # Test the periodogram and pick the best period and power
        ls_out = prot.run_ls(cl_time, cl_flux, cl_unc,
                             threshold=self.power_threshold,
                             prot_lims=prot_lims, run_bootstrap=True)
        #fund_prot, fund_power, periods_to_test, periodogram, aliases, sigmas

        return ls_out

    def _pick_lc(self, fund_power1, fund_power2):
        """Pick the raw or detrended lc to continue with by
        selecting the one with the highest peak in the periodogram
        (no consideration of the *locations* of those peaks)
        """
        # return a integer indicating which lightcurve to use (1 or 2)
        to_use = 0

        if fund_power1 > fund_power2:
            to_use = 1
        elif fund_power2 > fund_power1:
            to_use = 2
        else: # Either something's gone wrong, or they're exactly equal
            to_use = 0

        return to_use

    def _xy_correct(self, n_closest=21):
        """Correct for positional variations in the lightcurve once selected."""

        # Loop through the lightcurve and find the n closest pixels.
        # Then divide by the median flux from those pixels

        num_pts = len(self.use_flux)
        logging.debug("Pixel position correction %d", num_pts)

        self.corrected_flux = np.zeros(num_pts)
        self.corrected_unc = np.zeros(num_pts)

        first_half = self.time<=2102
        x_pos1 = self.x_pos[first_half==True]
        y_pos1 = self.y_pos[first_half==True]
        x_pos2 = self.x_pos[first_half==False]
        y_pos2 = self.y_pos[first_half==False]

        for i, fval, xx, yy in itertools.izip(range(num_pts), self.use_flux,
                                              self.x_pos, self.y_pos):
            if first_half[i]:
                comp_x, comp_y = x_pos1, y_pos1
                comp_f = self.use_flux[first_half==True]
            else:
                comp_x, comp_y = x_pos2, y_pos2
                comp_f = self.use_flux[first_half==False]

#            comp_x, comp_y = self.x_pos, self.y_pos
#            comp_f = self.use_flux

            pix_sep = np.sqrt((xx - comp_x)**2 + (yy - comp_y)**2)
            min_ind = np.argpartition(pix_sep, n_closest)[:n_closest]
            logging.debug(np.median(pix_sep[min_ind]))

            median_nearest = np.median(comp_f[min_ind])
            logging.debug("This flux %f Median Nearest %f",
                          fval, median_nearest)
            self.corrected_flux[i] = fval / median_nearest
            self.corrected_unc[i] = self.use_unc[i] / median_nearest

        logging.debug("Correction completed")

    def _plot_xy(self):
        """Plot some basic informational plots:
        Flux as a function of X-Y position
        Flux as a function of time
        """
        pass

    def _multi_search(self,lc_type):
        """Search a lightcurve for a secondary signal."""
        pass

<|endoftext|>