<commit_before># Copyright 2021 Google LLC
#
# 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 argparse
import datetime
import hmac
import json
import os
import random
import secrets
import time
import sys

from hashlib import sha1
from urllib.request import Request, urlopen

def make_changes(num_changes, vcs, event_timespan):
    changes = []
    max_time = time.time() - event_timespan
    head_commit = None
    event = None

    for x in range(num_changes):
        change_id = secrets.token_hex(20)
        unix_timestamp = time.time() - random.randrange(0, event_timespan)
        change = {
            "id": change_id,
            "timestamp": datetime.datetime.fromtimestamp(unix_timestamp),
        }

        if unix_timestamp > max_time:
            max_time = unix_timestamp
            head_commit = change

        changes.append(change)

    if vcs == "gitlab":
        event = {
            "object_kind": "push",
            "checkout_sha": head_commit["id"],
            "commits": changes,
        }
    if vcs == "github":
        event = {"head_commit": head_commit, "commits": changes}

    return event

def create_github_deploy_event(change):
    deployment = {
        "deployment_status": {
            "updated_at": change["timestamp"],
            "id": secrets.token_hex(20),
            "state": "success",
        },
        "deployment": {
            "sha": change["id"],
        },
    }
    return deployment

def create_gitlab_pipeline_event(changes):
    pipeline = None
    checkout_sha = changes["checkout_sha"]
    for c in changes["commits"]:
        if c["id"] == checkout_sha:
            pipeline = {
                "object_kind": "pipeline",
                "object_attributes": {
                    "created_at": c["timestamp"],
                    "id": random.randrange(0, 1000),
                    "status": "success",
                },
                "commit": c,
            }
    return pipeline

def create_gitlab_deploy_event(changes, event_num):
    deployment = None
    checkout_sha = changes["checkout_sha"]
    # Create random looking deploy id based on event num, should prevent duplicates until 1000 events
    num_ids = 1000
    coprime_far_away_num = 1787
    deployment_id = (event_num * coprime_far_away_num) % num_ids

    for c in changes["commits"]:
        if c["id"] == checkout_sha:
            deployment = {
                "object_kind": "deployment",
                "status": "success",
                "status_changed_at": c["timestamp"].strftime("%F %T +0200"),
                "deployment_id": deployment_id,
                "commit_url": f"http://example.com/root/test/commit/{checkout_sha}",
            }
    return deployment

def make_github_issue(root_cause):
    event = {
        "issue": {
            "created_at": root_cause["timestamp"],
            "updated_at": datetime.datetime.now(),
            "closed_at": datetime.datetime.now(),
            "number": random.randrange(0, 1000),
            "labels": [{"name": "Incident"}],
            "body": "root cause: %s" % root_cause["id"],
        },
        "repository": {"name": "foobar"},
    }
    return event

def make_gitlab_issue(changes):
    issue = None
    checkout_sha = changes["checkout_sha"]
    for c in changes["commits"]:
        if c["id"] == checkout_sha:
            issue = {
                "object_kind": "issue",
                "object_attributes": {
                    "created_at": c["timestamp"],
                    "updated_at": datetime.datetime.now(),
                    "closed_at": datetime.datetime.now(),
                    "id": random.randrange(0, 1000),
                    "labels": [{"title": "Incident"}],
                    "description": "root cause: %s" % c["id"],
                },
            }
    return issue

def make_webhook_request(vcs, webhook_url, secret, event_type, data, token=None):
    data = json.dumps(data, default=str).encode()
    request = Request(webhook_url, data)

    if vcs == "github":
        signature = hmac.new(secret.encode(), data, sha1)
        request.add_header("X-Github-Event", event_type)
        request.add_header("X-Hub-Signature", "sha1=" + signature.hexdigest())
        request.add_header("User-Agent", "GitHub-Hookshot/mock")

    if vcs == "gitlab":
        request.add_header("X-Gitlab-Event", event_type)
        request.add_header("X-Gitlab-Token", secret)

    request.add_header("Content-Type", "application/json")
    request.add_header("Mock", True)

    if token:
        request.add_header("Authorization", f"Bearer {token}")

    return request

def post_to_webhook(vcs, webhook_url, secret, event_type, data, token=None):

    request = make_webhook_request(vcs, webhook_url, secret, event_type, data, token)

    response = urlopen(request)

    if response.getcode() == 204:
        return 1
    else:
        return 0

if __name__ == "__main__":
    # parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--event_timespan",
        "-t",
        type=int,
        default=604800,
        help="time duration (in seconds) of timestamps of generated events \
                        (from [Now-timespan] to [Now]); default=604800 (1 week)",
    )
    parser.add_argument(
        "--num_events",
        "-e",
        type=int,
        default=20,
        help="number of events to generate; default=20",
    )
    parser.add_argument(
        "--num_issues",
        "-i",
        type=int,
        default=2,
        help="number of issues to generate; default=2",
    )
    parser.add_argument(
        "--vc_system",
        "-v",
        required=True,
        choices=["gitlab", "github"],
        help="version control system (e.g. 'github', 'gitlab')",
    )
    args = parser.parse_args()

    if args.num_issues > args.num_events:
        print("Error: num_issues cannot be greater than num_events")
        sys.exit()

    # get environment vars
    webhook_url = os.environ.get("WEBHOOK")
    secret = os.environ.get("SECRET")
    token = os.environ.get("TOKEN")

    if not webhook_url or not secret:
        print(
            "Error: please ensure the following environment variables are set: WEBHOOK, SECRET"
        )
        sys.exit()

    all_changesets = []
    changes_sent = 0
    changeset_sha = secrets.token_hex(20)
    for x in range(args.num_events):

        # make a change set containing a random number of changes
        changeset = make_changes(
            random.randrange(1, 5),
            args.vc_system,
            args.event_timespan,
        )
        changeset["before"] = changeset_sha  # use the previous one before we set to current one
        changeset_sha = changeset.get("checkout_sha") or changeset.get("head_commit")

        # Send individual changes data
        prev_change_sha = "0000000000000000000000000000000000000000"  # before at the beginning of a branch is just 0s
        for i, c in enumerate(changeset["commits"]):
            curr_change = None
            if args.vc_system == "gitlab":
                curr_change = {
                    "object_kind": "push",
                    "checkout_sha": c["id"],
                    "commits": [c],
                }
            if args.vc_system == "github":
                curr_change = {"head_commit": c, "commits": [c]}

            # We only post individual commits with shas not matching the changeset sha
            if c["id"] != changeset_sha:
                curr_change["before"] = prev_change_sha
                prev_change_sha = c["id"]

                changes_sent += post_to_webhook(
                    args.vc_system, webhook_url, secret, "push", curr_change, token
                )

        # Send fully associated push event
        post_to_webhook(args.vc_system, webhook_url, secret, "push", changeset, token)

        # Make and send a deployment
        if args.vc_system == "gitlab":
            deploy = create_gitlab_deploy_event(changeset, event_num=x)
            post_to_webhook(
                args.vc_system, webhook_url, secret, "deployment", deploy, token
            )

        if args.vc_system == "github":
            deploy = create_github_deploy_event(changeset["head_commit"])
            post_to_webhook(
                args.vc_system, webhook_url, secret, "deployment_status", deploy, token
            )

        all_changesets.append(changeset)

    # randomly create incidents associated to changes
    changesets_with_issues = random.sample(all_changesets, args.num_issues)
    for changeset in changesets_with_issues:
        issue = None
        if args.vc_system == "gitlab":
            issue = make_gitlab_issue(changeset)
        if args.vc_system == "github":
            issue = make_github_issue(changeset["head_commit"])
        post_to_webhook(args.vc_system, webhook_url, secret, "issues", issue, token)

    print(f"{changes_sent} changes successfully sent to event-handler")
<commit_msg>Add it so the generate_data only send deployments for changesets half the time, so that a single deployment can have multiple changesets/branches<commit_after># Copyright 2021 Google LLC
#
# 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 argparse
import datetime
import hmac
import json
import math
import os
import random
import secrets
import time
import sys

from hashlib import sha1
from urllib.request import Request, urlopen

def make_changes(num_changes, vcs, event_timespan):
    changes = []
    max_time = time.time() - event_timespan
    head_commit = None
    event = None

    for x in range(num_changes):
        change_id = secrets.token_hex(20)
        unix_timestamp = time.time() - random.randrange(0, event_timespan)
        change = {
            "id": change_id,
            "timestamp": datetime.datetime.fromtimestamp(unix_timestamp),
        }

        if unix_timestamp > max_time:
            max_time = unix_timestamp
            head_commit = change

        changes.append(change)

    if vcs == "gitlab":
        event = {
            "object_kind": "push",
            "checkout_sha": head_commit["id"],
            "commits": changes,
        }
    if vcs == "github":
        event = {"head_commit": head_commit, "commits": changes}

    return event

def create_github_deploy_event(change):
    deployment = {
        "deployment_status": {
            "updated_at": change["timestamp"],
            "id": secrets.token_hex(20),
            "state": "success",
        },
        "deployment": {
            "sha": change["id"],
        },
    }
    return deployment

def create_gitlab_pipeline_event(changes):
    pipeline = None
    checkout_sha = changes["checkout_sha"]
    for c in changes["commits"]:
        if c["id"] == checkout_sha:
            pipeline = {
                "object_kind": "pipeline",
                "object_attributes": {
                    "created_at": c["timestamp"],
                    "id": random.randrange(0, 1000),
                    "status": "success",
                },
                "commit": c,
            }
    return pipeline

def create_gitlab_deploy_event(changes, deploy_id=None):
    deployment = None
    checkout_sha = changes["checkout_sha"]
    if not deploy_id:
        deploy_id = random.randrange(0, 1000)

    for c in changes["commits"]:
        if c["id"] == checkout_sha:
            deployment = {
                "object_kind": "deployment",
                "status": "success",
                "status_changed_at": c["timestamp"].strftime("%F %T +0200"),
                "deployment_id": deploy_id,
                "commit_url": f"http://example.com/root/test/commit/{checkout_sha}",
            }
    return deployment

def make_github_issue(root_cause):
    event = {
        "issue": {
            "created_at": root_cause["timestamp"],
            "updated_at": datetime.datetime.now(),
            "closed_at": datetime.datetime.now(),
            "number": random.randrange(0, 1000),
            "labels": [{"name": "Incident"}],
            "body": "root cause: %s" % root_cause["id"],
        },
        "repository": {"name": "foobar"},
    }
    return event

def make_gitlab_issue(changes):
    issue = None
    checkout_sha = changes["checkout_sha"]
    for c in changes["commits"]:
        if c["id"] == checkout_sha:
            issue = {
                "object_kind": "issue",
                "object_attributes": {
                    "created_at": c["timestamp"],
                    "updated_at": datetime.datetime.now(),
                    "closed_at": datetime.datetime.now(),
                    "id": random.randrange(0, 1000),
                    "labels": [{"title": "Incident"}],
                    "description": "root cause: %s" % c["id"],
                },
            }
    return issue

def make_webhook_request(vcs, webhook_url, secret, event_type, data, token=None):
    data = json.dumps(data, default=str).encode()
    request = Request(webhook_url, data)

    if vcs == "github":
        signature = hmac.new(secret.encode(), data, sha1)
        request.add_header("X-Github-Event", event_type)
        request.add_header("X-Hub-Signature", "sha1=" + signature.hexdigest())
        request.add_header("User-Agent", "GitHub-Hookshot/mock")

    if vcs == "gitlab":
        request.add_header("X-Gitlab-Event", event_type)
        request.add_header("X-Gitlab-Token", secret)

    request.add_header("Content-Type", "application/json")
    request.add_header("Mock", True)

    if token:
        request.add_header("Authorization", f"Bearer {token}")

    return request

def post_to_webhook(vcs, webhook_url, secret, event_type, data, token=None):

    request = make_webhook_request(vcs, webhook_url, secret, event_type, data, token)

    response = urlopen(request)

    if response.getcode() == 204:
        return 1
    else:
        return 0

if __name__ == "__main__":
    # parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--event_timespan",
        "-t",
        type=int,
        default=604800,
        help="time duration (in seconds) of timestamps of generated events \
                        (from [Now-timespan] to [Now]); default=604800 (1 week)",
    )
    parser.add_argument(
        "--num_events",
        "-e",
        type=int,
        default=20,
        help="number of events to generate; default=20",
    )
    parser.add_argument(
        "--num_issues",
        "-i",
        type=int,
        default=2,
        help="number of issues to generate; default=2",
    )
    parser.add_argument(
        "--vc_system",
        "-v",
        required=True,
        choices=["gitlab", "github"],
        help="version control system (e.g. 'github', 'gitlab')",
    )
    args = parser.parse_args()

    if args.num_issues > args.num_events:
        print("Error: num_issues cannot be greater than num_events")
        sys.exit()

    # get environment vars
    webhook_url = os.environ.get("WEBHOOK")
    secret = os.environ.get("SECRET")
    token = os.environ.get("TOKEN")

    if not webhook_url or not secret:
        print(
            "Error: please ensure the following environment variables are set: WEBHOOK, SECRET"
        )
        sys.exit()

    all_changesets = []
    changes_sent = 0
    changeset_sha = secrets.token_hex(20)
    gitlab_deployment_id_pool = max(1000, 10**math.ceil(math.log10(args.num_events)))
    gitlab_deploy_ids = random.sample(range(0, gitlab_deployment_id_pool), args.num_events)
    for x, deploy_id in zip(range(args.num_events), gitlab_deploy_ids):

        # make a change set containing a random number of changes
        changeset = make_changes(
            random.randrange(1, 5),
            args.vc_system,
            args.event_timespan,
        )
        changeset["before"] = changeset_sha  # use the previous one before we set to current one
        changeset_sha = changeset.get("checkout_sha") or changeset.get("head_commit")

        # Send individual changes data
        prev_change_sha = "0000000000000000000000000000000000000000"  # before at the beginning of a branch is just 0s
        for i, c in enumerate(changeset["commits"]):
            curr_change = None
            if args.vc_system == "gitlab":
                curr_change = {
                    "object_kind": "push",
                    "checkout_sha": c["id"],
                    "commits": [c],
                }
            if args.vc_system == "github":
                curr_change = {"head_commit": c, "commits": [c]}

            # We only post individual commits with shas not matching the changeset sha
            if c["id"] != changeset_sha:
                curr_change["before"] = prev_change_sha
                prev_change_sha = c["id"]

                changes_sent += post_to_webhook(
                    args.vc_system, webhook_url, secret, "push", curr_change, token
                )

        # Send fully associated push event
        post_to_webhook(args.vc_system, webhook_url, secret, "push", changeset, token)

        # Make and send a deployment half the time, the other half will be change sets without
        # deployments or branches
        if random.choice([True, False]):
            if args.vc_system == "gitlab":
                deploy = create_gitlab_deploy_event(changeset, deploy_id=deploy_id)
                post_to_webhook(
                    args.vc_system, webhook_url, secret, "deployment", deploy, token
                )

            if args.vc_system == "github":
                deploy = create_github_deploy_event(changeset["head_commit"])
                post_to_webhook(
                    args.vc_system, webhook_url, secret, "deployment_status", deploy, token
                )

        all_changesets.append(changeset)

    # randomly create incidents associated to changes
    changesets_with_issues = random.sample(all_changesets, args.num_issues)
    for changeset in changesets_with_issues:
        issue = None
        if args.vc_system == "gitlab":
            issue = make_gitlab_issue(changeset)
        if args.vc_system == "github":
            issue = make_github_issue(changeset["head_commit"])
        post_to_webhook(args.vc_system, webhook_url, secret, "issues", issue, token)

    print(f"{changes_sent} changes successfully sent to event-handler")
<|endoftext|>


<commit_before>#! /usr/bin/env python
#
#  Load a FITS cube , extract the spectrum at a (or reference) pixel
#  and operate and plot some and then more...
#
#
#  22-jun-2017  PJT             summer project - cloned off cubespectrum.py
#
#
#  @todo
#     - have optional RSTFREQ as 3rd argument
#     - resample the gauss finer (not 5 points but may be 10x more?)

import os, sys, math
import numpy as np
import matplotlib.pyplot as plt

from astropy.io import fits
from astropy.units import Quantity
c = 299792.458     # there should be a way to get 'c' from astropy.units ?

if len(sys.argv) == 1:
    print("Usage: %s fitsfile [xpos ypos]" % sys.argv[0])
    sys.exit(1)

if len(sys.argv) > 3:
    fitsfile = sys.argv[1]
    pos = [int(sys.argv[2]),int(sys.argv[3])]
elif len(sys.argv) == 2:
    fitsfile = sys.argv[1]
    pos = None

# open the fits file
hdu = fits.open(fitsfile)
print(len(hdu))

# get a reference to the header and data.  Data should be 3dim numpy array now
h = hdu[0].header
d = hdu[0].data.squeeze()
print(d.shape)

#  grab the restfreq, there are at least two way how this is done
if 'RESTFRQ' in h:
    restfreq=h['RESTFRQ']
elif 'RESTFREQ' in h:
    restfreq=h['RESTFREQ']
else:
    restfreq=None
print("RESTFREQ",restfreq)

if pos == None:
    # the FITS reference pixel is always a good backup
    xpos = int(h['CRPIX1'])
    ypos = int(h['CRPIX2'])
    print("No position given, using reference pixel %g %g" % (xpos,ypos))
else:
    xpos = pos[0]
    ypos = pos[1]

flux     = d[:,ypos,xpos]

nchan    = d.shape[0]
channeln = np.arange(nchan)
zero     = np.zeros(nchan)

cdelt3 = h['CDELT3']
crval3 = h['CRVAL3']
crpix3 = h['CRPIX3']
channelf = (channeln-crpix3+1)*cdelt3 + crval3
channelv = (1.0-channelf/restfreq) * c

# what we plot
channel = channelv

ipeak = flux.argmax()
xpeak = channel[ipeak]
ypeak = flux[ipeak]

# moments around the peak
m = 5
x = channel[ipeak-m:ipeak+m]
y = flux[ipeak-m:ipeak+m]
xmean = (x*y).sum() / y.sum()
xdisp = (x*x*y).sum() / y.sum() - xmean*xmean
if xdisp > 0:
    xdisp = math.sqrt(xdisp)
fwhm = 2.355 * xdisp
print("MEAN/DISP/FWHM:",xmean,xdisp,fwhm)
ymodel = ypeak * np.exp(-0.5*(x-xmean)**2/(xdisp*xdisp))

plt.figure()
plt.plot(channel,flux,'o-',markersize=2,label='data')
plt.plot(channel,zero)
plt.plot(x,ymodel,label='gauss')
plt.xlabel("Velocity (km/s)")
plt.ylabel("Flux")
plt.title("Spectrum at position %g %g" % (xpos,ypos))
plt.legend()
plt.show()

<commit_msg>todo more<commit_after>#! /usr/bin/env python
#
#  Load a FITS cube , extract the spectrum at a (or reference) pixel
#  and operate and plot some and then more...
#
#
#  22-jun-2017  PJT             summer project - cloned off cubespectrum.py
#
#
#  @todo
#     - have optional RSTFREQ as 3rd argument
#     - resample the gauss finer (not 5 points but may be 10x more?)
#     - output the spectrum in a table, much like testCubeSpectrum.tab

import os, sys, math
import numpy as np
import matplotlib.pyplot as plt

from astropy.io import fits
from astropy.units import Quantity
c = 299792.458     # there should be a way to get 'c' from astropy.units ?

if len(sys.argv) == 1:
    print("Usage: %s fitsfile [xpos ypos]" % sys.argv[0])
    sys.exit(1)

if len(sys.argv) > 3:
    fitsfile = sys.argv[1]
    pos = [int(sys.argv[2]),int(sys.argv[3])]
elif len(sys.argv) == 2:
    fitsfile = sys.argv[1]
    pos = None

# open the fits file
hdu = fits.open(fitsfile)
print(len(hdu))

# get a reference to the header and data.  Data should be 3dim numpy array now
h = hdu[0].header
d = hdu[0].data.squeeze()
print(d.shape)

#  grab the restfreq, there are at least two way how this is done
if 'RESTFRQ' in h:
    restfreq=h['RESTFRQ']
elif 'RESTFREQ' in h:
    restfreq=h['RESTFREQ']
else:
    restfreq=None
print("RESTFREQ",restfreq)

if pos == None:
    # the FITS reference pixel is always a good backup
    xpos = int(h['CRPIX1'])
    ypos = int(h['CRPIX2'])
    print("No position given, using reference pixel %g %g" % (xpos,ypos))
else:
    xpos = pos[0]
    ypos = pos[1]

flux     = d[:,ypos,xpos]

nchan    = d.shape[0]
channeln = np.arange(nchan)
zero     = np.zeros(nchan)

cdelt3 = h['CDELT3']
crval3 = h['CRVAL3']
crpix3 = h['CRPIX3']
channelf = (channeln-crpix3+1)*cdelt3 + crval3
channelv = (1.0-channelf/restfreq) * c

# what we plot
channel = channelv

ipeak = flux.argmax()
xpeak = channel[ipeak]
ypeak = flux[ipeak]

# moments around the peak
m = 5
x = channel[ipeak-m:ipeak+m]
y = flux[ipeak-m:ipeak+m]
xmean = (x*y).sum() / y.sum()
xdisp = (x*x*y).sum() / y.sum() - xmean*xmean
if xdisp > 0:
    xdisp = math.sqrt(xdisp)
fwhm = 2.355 * xdisp
print("MEAN/DISP/FWHM:",xmean,xdisp,fwhm)
ymodel = ypeak * np.exp(-0.5*(x-xmean)**2/(xdisp*xdisp))

plt.figure()
plt.plot(channel,flux,'o-',markersize=2,label='data')
plt.plot(channel,zero)
plt.plot(x,ymodel,label='gauss')
plt.xlabel("Velocity (km/s)")
plt.ylabel("Flux")
plt.title("Spectrum at position %g %g" % (xpos,ypos))
plt.legend()
plt.show()

<|endoftext|>


<commit_before>"""
User-facing interface to all dimensionality reduction algorithms
"""

def smushplot(data, smusher):
    if isinstance(smusher, str):
        # Need to get appropriate smusher from sklearn given the string
        pass
    else:
        # Assume this is already an initialized sklearn object with the
        # ``fit_transform()`` method
        pass
<commit_msg>Add a bunch of plotting and keyword arguments<commit_after>"""
User-facing interface for plotting all dimensionality reduction algorithms
"""

def smushplot(data, smusher, n_components=2, marker='o', marker_order=None,
              text=False, text_order=None, linewidth=1, linewidth_order=None,
              edgecolor='k', edgecolor_order=None, smusher_kws=None,
              plot_kws=None):
    if isinstance(smusher, str):
        # Need to get appropriate smusher from sklearn given the string
        pass
    else:
        # Assume this is already an initialized sklearn object with the
        # ``fit_transform()`` method
        pass
<|endoftext|>


<commit_before># Copyright (c) 2016 Ultimaker B.V.
# Cura is released under the terms of the AGPLv3 or higher.

import re #To parse container registry names to increment the duplicates-resolving number.

import UM.Application #To link the stack to the global container stack.
import UM.Settings.ContainerRegistry #To search for nozzles, materials, etc.
import UM.Settings.ContainerStack #To create a container stack for this extruder.

class Extruder:
    ##  Creates a new extruder from the specified definition container.
    #
    #   \param definition The definition container defining this extruder.
    def __init__(self, definition):
        self._definition = definition

        container_registry = UM.Settings.ContainerRegistry.getInstance()

        #Find the nozzles that fit on this extruder.
        self._nozzles = container_registry.findInstanceContainers(type = "nozzle", definitions = "*," + self._definition.getId() + ",*") #Extruder needs to be delimited by either a comma or the end of string.
        self._nozzles += container_registry.findInstanceContainers(type = "nozzle", definitions = "*," + self._definition.getId())
        self._nozzles += container_registry.findInstanceContainers(type = "nozzle", definitions = self._definition.getId() + ",*")
        self._nozzles += container_registry.findInstanceContainers(type = "nozzle", definitions = self._definition.getId())

        #Create a container stack for this extruder.
        self._container_stack = UM.Settings.ContainerStack(self._uniqueName(self._definition.getId()))
        self._container_stack.addMetaDataEntry("type", "extruder_train")
        self._container_stack.addContainer(self._definition)

        #Find the nozzle to use for this extruder.
        self._nozzle = container_registry.getEmptyInstanceContainer()
        if self._definition.getMetaDataEntry("has_nozzles", default = "False") == "True":
            if len(self._nozzles) >= 1: #First add any extruder. Later, overwrite with preference if the preference is valid.
                self._nozzle = self._nozzles[0]
            preferred_nozzle_id = self._definition.getMetaDataEntry("preferred_nozzle")
            if preferred_nozzle_id:
                for nozzle in self._nozzles:
                    if nozzle.getId() == preferred_nozzle_id:
                        self._nozzle = nozzle
                        break
            self._container_stack.addContainer(self._nozzle)

        #Find a material to use for this nozzle.
        self._material = container_registry.getEmptyInstanceContainer()
        if self._definition.getMetaDataEntry("has_materials", default = "False") == "True":
            if self._definition.getMetaDataEntry("has_nozzle_materials", default = "False") == "True":
                all_materials = container_registry.findInstanceContainers(type = "material", nozzle = self._nozzle.getId())
            else:
                all_materials = container_registry.findInstanceContainers(type = "material")
            if len(all_materials) >= 1:
                self._material = all_materials[0]
            preferred_material_id = self._definition.getMetaDataEntry("preferred_material")
            if preferred_material_id:
                preferred_material = container_registry.findInstanceContainers(type = "material", id = preferred_material_id.lower())
                if len(preferred_material) >= 1:
                    self._material = preferred_material[0]
            self._container_stack.addContainer(self._material)

        #Find a quality to use for this extruder.
        self._quality = container_registry.getEmptyInstanceContainer()
        if self._definition.getMetaDataEntry("has_machine_quality"):
            all_qualities = container_registry.findInstanceContainers(type = "quality")
            if len(all_qualities) >= 1:
                self._quality = all_qualities[0]
            preferred_quality_id = self._definition.getMetaDataEntry("preferred_quality")
            if preferred_quality_id:
                preferred_quality = container_registry.findInstanceContainers(type = "quality", id = preferred_quality_id.lower())
                if len(preferred_quality) >= 1:
                    self._quality = preferred_quality[0]
            self._container_stack.addContainer(self._quality)

        self._container_stack.setNextStack(UM.Application.getInstance().getGlobalContainerStack())

    ##  Finds a unique name for an extruder stack.
    #
    #   \param extruder An extruder definition to design a name for.
    #   \return A name for an extruder stack that is unique and reasonably
    #   human-readable.
    def _uniqueName(self, extruder):
        container_registry = UM.Settings.ContainerRegistry.getInstance()

        name = extruder.getName().strip()
        num_check = re.compile("(.*?)\s*#\d$").match(name)
        if num_check: #There is a number in the name.
            name = num_check.group(1) #Filter out the number.
        if name == "": #Wait, that deleted everything!
            name = "Extruder"
        unique_name = name

        i = 1
        while container_registry.findContainers(id = unique_name) or container_registry.findContainers(name = unique_name): #A container already has this name.
            i += 1 #Try next numbering.
            unique_name = "%s #%d" % (name, i) #Fill name like this: "Extruder #2".
        return unique_name<commit_msg>Add user profile to extruder stack<commit_after># Copyright (c) 2016 Ultimaker B.V.
# Cura is released under the terms of the AGPLv3 or higher.

import re #To parse container registry names to increment the duplicates-resolving number.

import UM.Application #To link the stack to the global container stack.
import UM.Settings.ContainerRegistry #To search for nozzles, materials, etc.
import UM.Settings.ContainerStack #To create a container stack for this extruder.

class Extruder:
    ##  Creates a new extruder from the specified definition container.
    #
    #   \param definition The definition container defining this extruder.
    def __init__(self, definition):
        self._definition = definition

        container_registry = UM.Settings.ContainerRegistry.getInstance()

        #Find the nozzles that fit on this extruder.
        self._nozzles = container_registry.findInstanceContainers(type = "nozzle", definitions = "*," + self._definition.getId() + ",*") #Extruder needs to be delimited by either a comma or the end of string.
        self._nozzles += container_registry.findInstanceContainers(type = "nozzle", definitions = "*," + self._definition.getId())
        self._nozzles += container_registry.findInstanceContainers(type = "nozzle", definitions = self._definition.getId() + ",*")
        self._nozzles += container_registry.findInstanceContainers(type = "nozzle", definitions = self._definition.getId())

        #Create a container stack for this extruder.
        name = self._uniqueName(self._definition.getId())
        self._container_stack = UM.Settings.ContainerStack(name)
        self._container_stack.addMetaDataEntry("type", "extruder_train")
        self._container_stack.addContainer(self._definition)

        #Find the nozzle to use for this extruder.
        self._nozzle = container_registry.getEmptyInstanceContainer()
        if self._definition.getMetaDataEntry("has_nozzles", default = "False") == "True":
            if len(self._nozzles) >= 1: #First add any extruder. Later, overwrite with preference if the preference is valid.
                self._nozzle = self._nozzles[0]
            preferred_nozzle_id = self._definition.getMetaDataEntry("preferred_nozzle")
            if preferred_nozzle_id:
                for nozzle in self._nozzles:
                    if nozzle.getId() == preferred_nozzle_id:
                        self._nozzle = nozzle
                        break
            self._container_stack.addContainer(self._nozzle)

        #Find a material to use for this nozzle.
        self._material = container_registry.getEmptyInstanceContainer()
        if self._definition.getMetaDataEntry("has_materials", default = "False") == "True":
            if self._definition.getMetaDataEntry("has_nozzle_materials", default = "False") == "True":
                all_materials = container_registry.findInstanceContainers(type = "material", nozzle = self._nozzle.getId())
            else:
                all_materials = container_registry.findInstanceContainers(type = "material")
            if len(all_materials) >= 1:
                self._material = all_materials[0]
            preferred_material_id = self._definition.getMetaDataEntry("preferred_material")
            if preferred_material_id:
                preferred_material = container_registry.findInstanceContainers(type = "material", id = preferred_material_id.lower())
                if len(preferred_material) >= 1:
                    self._material = preferred_material[0]
            self._container_stack.addContainer(self._material)

        #Find a quality to use for this extruder.
        self._quality = container_registry.getEmptyInstanceContainer()
        if self._definition.getMetaDataEntry("has_machine_quality"):
            all_qualities = container_registry.findInstanceContainers(type = "quality")
            if len(all_qualities) >= 1:
                self._quality = all_qualities[0]
            preferred_quality_id = self._definition.getMetaDataEntry("preferred_quality")
            if preferred_quality_id:
                preferred_quality = container_registry.findInstanceContainers(type = "quality", id = preferred_quality_id.lower())
                if len(preferred_quality) >= 1:
                    self._quality = preferred_quality[0]
            self._container_stack.addContainer(self._quality)

        #Add an empty user profile.
        self._user_profile = UM.Settings.InstanceContainer(name + "_current_settings")
        self._user_profile.addMetaDataEntry("type", "user")
        self._container_stack.addContainer(self._user_profile)

        self._container_stack.setNextStack(UM.Application.getInstance().getGlobalContainerStack())

    ##  Finds a unique name for an extruder stack.
    #
    #   \param extruder An extruder definition to design a name for.
    #   \return A name for an extruder stack that is unique and reasonably
    #   human-readable.
    def _uniqueName(self, extruder):
        container_registry = UM.Settings.ContainerRegistry.getInstance()

        name = extruder.getName().strip()
        num_check = re.compile("(.*?)\s*#\d$").match(name)
        if num_check: #There is a number in the name.
            name = num_check.group(1) #Filter out the number.
        if name == "": #Wait, that deleted everything!
            name = "Extruder"
        unique_name = name

        i = 1
        while container_registry.findContainers(id = unique_name) or container_registry.findContainers(name = unique_name): #A container already has this name.
            i += 1 #Try next numbering.
            unique_name = "%s #%d" % (name, i) #Fill name like this: "Extruder #2".
        return unique_name<|endoftext|>


<commit_before>import numpy as np
from multiphenotype_utils import (get_continuous_features_as_matrix, add_id, remove_id_and_get_mat,
    partition_dataframe_into_binary_and_continuous, divide_idxs_into_batches)
import pandas as pd
import tensorflow as tf
from dimreducer import DimReducer
from scipy.special import expit

from general_autoencoder import GeneralAutoencoder
from standard_autoencoder import StandardAutoencoder

class VariationalAutoencoder(StandardAutoencoder):
    """
    Implements a standard variational autoencoder (diagonal Gaussians everywhere).
    """
    def __init__(self,
                 **kwargs):
        super(VariationalAutoencoder, self).__init__(**kwargs)
        self.sigma_scaling = .1 # keeps the sigmas (ie, the std of the normal from which Z is drawn) from getting too large.

    def init_network(self):
        self.weights = {}
        self.biases = {}

        if self.learn_continuous_variance:
            # we exponentiate this because it has to be non-negative.
            self.log_continuous_variance = tf.Variable(self.initialization_function([1]))

        # Encoder layers.
        for encoder_layer_idx, encoder_layer_size in enumerate(self.encoder_layer_sizes):
            if encoder_layer_idx == 0:
                input_dim = len(self.feature_names)
            else:
                input_dim = self.encoder_layer_sizes[encoder_layer_idx - 1]
            output_dim = self.encoder_layer_sizes[encoder_layer_idx]
            print("Added encoder layer with input dimension %i and output dimension %i" % (input_dim, output_dim))
            self.weights['encoder_h%i' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([input_dim, output_dim]))
            self.biases['encoder_b%i' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([output_dim]))
            self.weights['encoder_h%i_sigma' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([input_dim, output_dim]))
            self.biases['encoder_b%i_sigma' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([output_dim]))

        # Decoder layers.
        self.decoder_layer_sizes.append(len(self.feature_names))
        for decoder_layer_idx, decoder_layer_size in enumerate(self.decoder_layer_sizes):
            if decoder_layer_idx == 0:
                input_dim = self.k
            else:
                input_dim = self.decoder_layer_sizes[decoder_layer_idx - 1]
            output_dim = self.decoder_layer_sizes[decoder_layer_idx]
            print("Added decoder layer with input dimension %i and output dimension %i" % (input_dim, output_dim))
            self.weights['decoder_h%i' % decoder_layer_idx] = tf.Variable(
                self.initialization_function([input_dim, output_dim]))
            self.biases['decoder_b%i' % decoder_layer_idx] = tf.Variable(
                self.initialization_function([output_dim]))

    def encode(self, X):
        num_layers = len(self.encoder_layer_sizes)
        # Get mu
        mu = X
        for idx in range(num_layers):
            mu = tf.matmul(mu, self.weights['encoder_h%i' % (idx)]) \
                + self.biases['encoder_b%i' % (idx)]
            # No non-linearity on the last layer
            if idx != num_layers - 1:
                mu = self.non_linearity(mu)
        self.Z_mu = mu

        # Get sigma
        sigma = X
        for idx in range(num_layers):
            sigma = tf.matmul(sigma, self.weights['encoder_h%i_sigma' % (idx)]) \
                + self.biases['encoder_b%i_sigma' % (idx)]
            # No non-linearity on the last layer
            if idx != num_layers - 1:
                sigma = self.non_linearity(sigma)
        sigma = sigma * self.sigma_scaling # scale so sigma doesn't explode when we exponentiate it.
        sigma = tf.exp(sigma)
        self.Z_sigma = sigma

        # Sample from N(mu, sigma)
        self.eps = tf.random_normal(tf.shape(self.Z_mu), dtype=tf.float32, mean=0., stddev=1.0, seed=self.random_seed)
        Z = self.Z_mu + self.Z_sigma * self.eps
        return Z

    def sample_X(self, age, n):
        """
        samples X by first sampling Z from the autoencoder prior, then feeding it through the model.
        Draws n samples for people of a given age.
        Important note: in our age autoencoder formulation, age is zero-centered
        (ie, we train the model with ages whose mean has been subtracted off).
        So you probably want the age you pass in to account for that.
        """
        Z = self.sample_Z(age, n)
        Xr = self.sess.run(self.Xr, feed_dict = {self.Z:Z})
        # for binary features, need to convert logits to 1s and 0s by sampling.
        Xr[:, self.binary_feature_idxs] = np.random.random(Xr[:, self.binary_feature_idxs].shape) < \
                                                           expit(Xr[:, self.binary_feature_idxs])
        # for continuous features, need to add noise.
        if self.learn_continuous_variance:
            std = np.sqrt(tf.exp(self.log_continuous_variance))
        else:
            std = 1

        Xr[:, self.continuous_feature_idxs] = Xr[:, self.continuous_feature_idxs] + \
        np.random.normal(loc=0, scale=std, size=Xr[:, self.continuous_feature_idxs].shape)

        return Xr

    def sample_Z(self, age, n):
        return np.random.multivariate_normal(mean = np.zeros([self.k,]), cov = np.eye(self.k), size = n)

    def get_loss(self):
        """
        Uses self.X, self.Xr, self.Z_sigma, self.Z_mu, self.kl_weighting
        """
        _, binary_loss, continuous_loss, _ = super(VariationalAutoencoder, self).get_loss()

        kl_div_loss = -.5 * (
            1 +
            2 * tf.log(self.Z_sigma) - tf.square(self.Z_mu) - tf.square(self.Z_sigma))
        kl_div_loss = tf.reduce_mean(
            tf.reduce_sum(
                kl_div_loss,
                axis=1),
            axis=0)

        combined_loss = self.combine_loss_components(binary_loss, continuous_loss, kl_div_loss)

        return combined_loss, binary_loss, continuous_loss, kl_div_loss

    def compute_elbo(self, df, continuous_variance=1):
        if self.learn_continuous_variance:
            continuous_variance = np.exp(self.sess.run(self.log_continuous_variance)[0])
            print("Warning: ignoring continuous variance input because we have already learned continuous variance: %2.3f" % continuous_variance)

        data, binary_feature_idxs, continuous_feature_idxs, feature_names = \
            partition_dataframe_into_binary_and_continuous(df)
        ages = None
        if self.need_ages:
            ages = self.get_ages(df)
        assert np.all(binary_feature_idxs == self.binary_feature_idxs)
        assert np.all(continuous_feature_idxs == self.continuous_feature_idxs)
        assert np.all(feature_names == self.feature_names)

        print(("Computing ELBO with %i continuous features, %i binary features, "
            "%i examples, continuous variance = %2.3f") %
              (len(continuous_feature_idxs),
               len(binary_feature_idxs),
               len(data),
               continuous_variance))

        num_iter = 1 # set to more than one if you want to average runs together.
        mean_binary_loss = 0
        mean_continuous_loss = 0
        mean_reg_loss = 0

        for i in range(num_iter):
            _, binary_loss, continuous_loss, reg_loss = self.minibatch_mean_eval(data,
                                                                                ages,
                                                                                regularization_weighting = 1)
            mean_binary_loss += binary_loss / num_iter
            mean_continuous_loss += continuous_loss / num_iter
            mean_reg_loss += reg_loss / num_iter

        # https://www.statlect.com/fundamentals-of-statistics/normal-distribution-maximum-likelihood
        if not self.learn_continuous_variance:
            # in this case, we have to add in the variance term since the mean_continuous_loss is just a squared-error term.
            # if we learn the variance, the continuous loss is actually the full negative Gaussian log likelihood
            # and there is no correction needed.
            constant_offset_per_sample_and_feature = .5 * np.log(2 * np.pi) + .5 * np.log(continuous_variance)
            mean_continuous_loss = constant_offset_per_sample_and_feature * len(continuous_feature_idxs) + mean_continuous_loss / continuous_variance

        # note that we compute the elbo using a weight of 1 for the regularization loss regardless of the regularization weighting.
        mean_combined_loss = mean_binary_loss + mean_continuous_loss + mean_reg_loss
        elbo = -mean_combined_loss
        print("Average ELBO per sample = %2.3f" % elbo)
        return elbo
<commit_msg>tweak<commit_after>import numpy as np
from multiphenotype_utils import (get_continuous_features_as_matrix, add_id, remove_id_and_get_mat,
    partition_dataframe_into_binary_and_continuous, divide_idxs_into_batches)
import pandas as pd
import tensorflow as tf
from dimreducer import DimReducer
from scipy.special import expit

from general_autoencoder import GeneralAutoencoder
from standard_autoencoder import StandardAutoencoder

class VariationalAutoencoder(StandardAutoencoder):
    """
    Implements a standard variational autoencoder (diagonal Gaussians everywhere).
    """
    def __init__(self,
                 **kwargs):
        super(VariationalAutoencoder, self).__init__(**kwargs)
        self.sigma_scaling = .1 # keeps the sigmas (ie, the std of the normal from which Z is drawn) from getting too large.

    def init_network(self):
        self.weights = {}
        self.biases = {}

        if self.learn_continuous_variance:
            # we exponentiate this because it has to be non-negative.
            self.log_continuous_variance = tf.Variable(self.initialization_function([1]))

        # Encoder layers.
        for encoder_layer_idx, encoder_layer_size in enumerate(self.encoder_layer_sizes):
            if encoder_layer_idx == 0:
                input_dim = len(self.feature_names)
            else:
                input_dim = self.encoder_layer_sizes[encoder_layer_idx - 1]
            output_dim = self.encoder_layer_sizes[encoder_layer_idx]
            print("Added encoder layer with input dimension %i and output dimension %i" % (input_dim, output_dim))
            self.weights['encoder_h%i' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([input_dim, output_dim]))
            self.biases['encoder_b%i' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([output_dim]))
            self.weights['encoder_h%i_sigma' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([input_dim, output_dim]))
            self.biases['encoder_b%i_sigma' % encoder_layer_idx] = tf.Variable(
                self.initialization_function([output_dim]))

        # Decoder layers.
        self.decoder_layer_sizes.append(len(self.feature_names))
        for decoder_layer_idx, decoder_layer_size in enumerate(self.decoder_layer_sizes):
            if decoder_layer_idx == 0:
                input_dim = self.k
            else:
                input_dim = self.decoder_layer_sizes[decoder_layer_idx - 1]
            output_dim = self.decoder_layer_sizes[decoder_layer_idx]
            print("Added decoder layer with input dimension %i and output dimension %i" % (input_dim, output_dim))
            self.weights['decoder_h%i' % decoder_layer_idx] = tf.Variable(
                self.initialization_function([input_dim, output_dim]))
            self.biases['decoder_b%i' % decoder_layer_idx] = tf.Variable(
                self.initialization_function([output_dim]))

    def encode(self, X):
        num_layers = len(self.encoder_layer_sizes)
        # Get mu
        mu = X
        for idx in range(num_layers):
            mu = tf.matmul(mu, self.weights['encoder_h%i' % (idx)]) \
                + self.biases['encoder_b%i' % (idx)]
            # No non-linearity on the last layer
            if idx != num_layers - 1:
                mu = self.non_linearity(mu)
        self.Z_mu = mu

        # Get sigma
        sigma = X
        for idx in range(num_layers):
            sigma = tf.matmul(sigma, self.weights['encoder_h%i_sigma' % (idx)]) \
                + self.biases['encoder_b%i_sigma' % (idx)]
            # No non-linearity on the last layer
            if idx != num_layers - 1:
                sigma = self.non_linearity(sigma)
        sigma = sigma * self.sigma_scaling # scale so sigma doesn't explode when we exponentiate it.
        sigma = tf.exp(sigma)
        self.Z_sigma = sigma

        # Sample from N(mu, sigma)
        self.eps = tf.random_normal(tf.shape(self.Z_mu), dtype=tf.float32, mean=0., stddev=1.0, seed=self.random_seed)
        Z = self.Z_mu + self.Z_sigma * self.eps
        return Z

    def sample_X(self, age, n):
        """
        samples X by first sampling Z from the autoencoder prior, then feeding it through the model.
        Draws n samples for people of a given age.
        Important note: in our age autoencoder formulation, age is zero-centered
        (ie, we train the model with ages whose mean has been subtracted off).
        So you probably want the age you pass in to account for that.
        """
        Z = self.sample_Z(age, n)
        Xr = self.sess.run(self.Xr, feed_dict = {self.Z:Z})
        # for binary features, need to convert logits to 1s and 0s by sampling.
        Xr[:, self.binary_feature_idxs] = np.random.random(Xr[:, self.binary_feature_idxs].shape) < \
                                                           expit(Xr[:, self.binary_feature_idxs])
        # for continuous features, need to add noise.
        if self.learn_continuous_variance:
            continuous_variance = np.exp(self.sess.run(self.log_continuous_variance)[0])
            std = np.sqrt(continuous_variance)
        else:
            std = 1

        Xr[:, self.continuous_feature_idxs] = Xr[:, self.continuous_feature_idxs] + \
        np.random.normal(loc=0, scale=std, size=Xr[:, self.continuous_feature_idxs].shape)

        return Xr

    def sample_Z(self, age, n):
        return np.random.multivariate_normal(mean = np.zeros([self.k,]), cov = np.eye(self.k), size = n)

    def get_loss(self):
        """
        Uses self.X, self.Xr, self.Z_sigma, self.Z_mu, self.kl_weighting
        """
        _, binary_loss, continuous_loss, _ = super(VariationalAutoencoder, self).get_loss()

        kl_div_loss = -.5 * (
            1 +
            2 * tf.log(self.Z_sigma) - tf.square(self.Z_mu) - tf.square(self.Z_sigma))
        kl_div_loss = tf.reduce_mean(
            tf.reduce_sum(
                kl_div_loss,
                axis=1),
            axis=0)

        combined_loss = self.combine_loss_components(binary_loss, continuous_loss, kl_div_loss)

        return combined_loss, binary_loss, continuous_loss, kl_div_loss

    def compute_elbo(self, df, continuous_variance=1):
        if self.learn_continuous_variance:
            continuous_variance = np.exp(self.sess.run(self.log_continuous_variance)[0])
            print("Warning: ignoring continuous variance input because we have already learned continuous variance: %2.3f" % continuous_variance)

        data, binary_feature_idxs, continuous_feature_idxs, feature_names = \
            partition_dataframe_into_binary_and_continuous(df)
        ages = None
        if self.need_ages:
            ages = self.get_ages(df)
        assert np.all(binary_feature_idxs == self.binary_feature_idxs)
        assert np.all(continuous_feature_idxs == self.continuous_feature_idxs)
        assert np.all(feature_names == self.feature_names)

        print(("Computing ELBO with %i continuous features, %i binary features, "
            "%i examples, continuous variance = %2.3f") %
              (len(continuous_feature_idxs),
               len(binary_feature_idxs),
               len(data),
               continuous_variance))

        num_iter = 1 # set to more than one if you want to average runs together.
        mean_binary_loss = 0
        mean_continuous_loss = 0
        mean_reg_loss = 0

        for i in range(num_iter):
            _, binary_loss, continuous_loss, reg_loss = self.minibatch_mean_eval(data,
                                                                                ages,
                                                                                regularization_weighting = 1)
            mean_binary_loss += binary_loss / num_iter
            mean_continuous_loss += continuous_loss / num_iter
            mean_reg_loss += reg_loss / num_iter

        # https://www.statlect.com/fundamentals-of-statistics/normal-distribution-maximum-likelihood
        if not self.learn_continuous_variance:
            # in this case, we have to add in the variance term since the mean_continuous_loss is just a squared-error term.
            # if we learn the variance, the continuous loss is actually the full negative Gaussian log likelihood
            # and there is no correction needed.
            constant_offset_per_sample_and_feature = .5 * np.log(2 * np.pi) + .5 * np.log(continuous_variance)
            mean_continuous_loss = constant_offset_per_sample_and_feature * len(continuous_feature_idxs) + mean_continuous_loss / continuous_variance

        # note that we compute the elbo using a weight of 1 for the regularization loss regardless of the regularization weighting.
        mean_combined_loss = mean_binary_loss + mean_continuous_loss + mean_reg_loss
        elbo = -mean_combined_loss
        print("Average ELBO per sample = %2.3f" % elbo)
        return elbo
<|endoftext|>


<commit_before>import os

from twod_materials.utils import is_converged, write_runjob, get_status

from pymatgen.io.vasp.outputs import Vasprun
from pymatgen.io.vasp.inputs import Kpoints, Incar
from pymatgen.symmetry.bandstructure import HighSymmKpath
from pymatgen.core.structure import Structure

HSE_INCAR_DICT = {}

def remove_z_kpoints():
    """
    Strips all paths from a linemode KPOINTS that include a z-component,
    since these are not relevant for 2D materials.
    """

    kpoint_lines = open('KPOINTS').readlines()
    with open('KPOINTS', 'w') as kpts:
        for line in kpoint_lines[:4]:
            kpts.write(line)
        i = 4
        while i < len(kpoint_lines):
            if (
                not float(kpoint_lines[i].split()[2]) and
                not float(kpoint_lines[i+1].split()[2])
                    ):
                kpts.write(kpoint_lines[i])
                kpts.write(kpoint_lines[i+1])
                kpts.write('\n')
            i += 3

def run_linemode_calculation(submit=True):
    """
    Setup and submit a normal PBE calculation for band structure along
    high symmetry k-paths.
    """

    directory = os.getcwd().split('/')[-1]
    if is_converged(os.getcwd()) and not get_status(os.getcwd()):
        if not os.path.isdir('bandstructure'):
            os.mkdir('bandstructure')
        if not is_converged('bandstructure'):
            os.chdir('bandstructure')
            os.system('cp ../CONTCAR ./POSCAR')
            os.system('cp ../POTCAR ./')
            os.system('cp ../vdw_kernel.bindat ./')
            incar_dict = Incar.from_file('../INCAR').as_dict()
            incar_dict.update(
                {'NSW': 0, 'LVTOT': True, 'LVHAR': True, 'LORBIT': 11})
            Incar.from_dict(incar_dict).write_file('INCAR')
            structure = Structure.from_file('POSCAR')
            kpath = HighSymmKpath(structure)
            Kpoints.automatic_linemode(40, kpath).write_file('KPOINTS')
            remove_z_kpoints()
            write_runjob('{}_pbebands'.format(
                directory), 1, 16, '600mb', '6:00:00', 'vasp')

            if submit:
                os.system('qsub runjob')

            os.chdir('../')

def run_hse_calculation(submit=True):
    """
    Setup/submit an HSE06 calculation to get an accurate band structure.
    Requires a previous WAVECAR and IBZKPT from a PBE run. See
    http://cms.mpi.univie.ac.at/wiki/index.php/Si_bandstructure for more
    details.
    """

    directory = os.getcwd().split('/')[-1]
    vasprun = Vasprun('vasprun.xml')
    band_gap = vasprun.get_band_structure().get_band_gap()
    kpath = []
    if band_gap['energy']:
        transition = band_gap['transition'].split('-')

        # Amount to increment along the kpath to result in 10
        # points total.
        increment = ((transition[1][0] - transition[0][0]) / 9,
                     (transition[1][1] - transition[0][1]) / 9,
                     (transition[1][2] - transition[0][2] / 9))

        for i in range(10):
            kpath.append((transition[0] + increment[0] * i),
                         (transition[1] + increment[1] * i),
                         (transition[2] + increment[2] * i))

        if not os.path.isdir('HSE'):
            os.mkdir('HSE')
        os.chdir('HSE')
        os.system('cp ../CONTCAR ./POSCAR')
        os.system('cp ../POTCAR ./POTCAR')
        os.system('cp ../vdw_kernel.bindat ./')
        Incar.from_dict(HSE_INCAR_DICT).write_file('INCAR')
        write_runjob('{}_hsebands'.format(
            directory), 1, 32, '1200mb', '150:00:00', 'vasp')

        # Re-use the irreducible brillouin zone KPOINTS from a
        # previous GGA run.
        os.system('cp ../IBZKPT ./KPOINTS')
        kpoints_lines = open('KPOINTS').readlines()
        n_kpts = int(kpoints_lines[1].split()[0])
        with open('KPOINTS', 'w') as kpts:
            kpts.write(kpoints_lines[0])
            kpts.write('{}\n'.format(n_kpts + 10))
            for line in kpoints_lines[2:]:
                kpts.write(line)
            for kpoint in kpath:
                kpts.write('{}\n'.format(' '.join(kpoint)))

        if submit:
            os.system('qsub runjob')
<commit_msg>Update run_hse_calculation.<commit_after>import os

from twod_materials.utils import is_converged, write_runjob, get_status

from pymatgen.io.vasp.inputs import Kpoints, Incar
from pymatgen.symmetry.bandstructure import HighSymmKpath
from pymatgen.core.structure import Structure

HSE_INCAR_DICT = {}

def remove_z_kpoints():
    """
    Strips all paths from a linemode KPOINTS that include a z-component,
    since these are not relevant for 2D materials.
    """

    kpoint_lines = open('KPOINTS').readlines()
    with open('KPOINTS', 'w') as kpts:
        for line in kpoint_lines[:4]:
            kpts.write(line)
        i = 4
        while i < len(kpoint_lines):
            if (
                not float(kpoint_lines[i].split()[2]) and
                not float(kpoint_lines[i+1].split()[2])
                    ):
                kpts.write(kpoint_lines[i])
                kpts.write(kpoint_lines[i+1])
                kpts.write('\n')
            i += 3

def run_linemode_calculation(submit=True):
    """
    Setup and submit a normal PBE calculation for band structure along
    high symmetry k-paths.
    """

    directory = os.getcwd().split('/')[-1]
    if is_converged(os.getcwd()) and not get_status(os.getcwd()):
        if not os.path.isdir('bandstructure'):
            os.mkdir('bandstructure')
        if not is_converged('bandstructure'):
            os.chdir('bandstructure')
            os.system('cp ../CONTCAR ./POSCAR')
            os.system('cp ../POTCAR ./')
            os.system('cp ../vdw_kernel.bindat ./')
            incar_dict = Incar.from_file('../INCAR').as_dict()
            incar_dict.update(
                {'NSW': 0, 'LVTOT': True, 'LVHAR': True, 'LORBIT': 11})
            Incar.from_dict(incar_dict).write_file('INCAR')
            structure = Structure.from_file('POSCAR')
            kpath = HighSymmKpath(structure)
            Kpoints.automatic_linemode(40, kpath).write_file('KPOINTS')
            remove_z_kpoints()
            write_runjob('{}_pbebands'.format(
                directory), 1, 16, '600mb', '6:00:00', 'vasp')

            if submit:
                os.system('qsub runjob')

            os.chdir('../')

def run_hse_calculation(submit=True):
    """
    Setup/submit an HSE06 calculation to get an accurate band structure.
    Requires a previous WAVECAR and IBZKPT from a PBE run. See
    http://cms.mpi.univie.ac.at/wiki/index.php/Si_bandstructure for more
    details.
    """

    HSE_INCAR_DICT = {'LHFCALC': True, 'HFSCREEN': 0.2, 'AEXX': 0.25,
                      'ALGO': 'D', 'TIME': 0.4, 'LDIAG': True}

    if not os.path.isdir('HSE'):
        os.mkdir('HSE')
    os.chdir('HSE')
    os.system('cp ../CONTCAR ./POSCAR')
    os.system('cp ../POTCAR ./POTCAR')
    os.system('cp ../vdw_kernel.bindat ./')
    os.system('cp ../INCAR ./')
    incar_dict = Incar.from_file('INCAR').as_dict()
    incar_dict.update(HSE_INCAR_DICT)
    Incar.from_dict(incar_dict).write_file('INCAR')
    write_runjob('{}_hsebands'.format(
        os.getcwd().split('/')[-3]), 1, 32, '1200mb', '150:00:00', 'vasp')

    # Re-use the irreducible brillouin zone KPOINTS from a
    # previous GGA run.
    os.system('cp ../../IBZKPT ./KPOINTS')
    kpoints_lines = open('KPOINTS').readlines()
    n_ibz_kpts = int(kpoints_lines[1].split()[0])
    kpath = HighSymmKpath(Structure.from_file('POSCAR'))
    twod_kpts = [kpt for kpt in kpath.get_kpoints(line_density=10)[0]
                 if not kpt[2]]
    with open('KPOINTS', 'w') as kpts:
        kpts.write(kpoints_lines[0])
        kpts.write('{}\n'.format(n_ibz_kpts + len(twod_kpts)))
        for line in kpoints_lines[2:]:
            kpts.write(line)
        for kpoint in twod_kpts:
            kpts.write('{} 0\n'.format(' '.join([str(c) for c in kpoint])))

    if submit:
        os.system('qsub runjob')
<|endoftext|>


<commit_before># Copyright 2020 Google LLC
#
# 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
#
#    https://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.

"""Training Script for STFTWaveGAN on MAESTRO.

Follows the same setup as SpecPhaseGAN, but
generates STFTs instead of Magnitude and Instantaneous
Frequency. Additionally, a Waveform critic is added
"""

import os
import tensorflow as tf
from audio_synthesis.structures import spec_gan, wave_gan
from audio_synthesis.models import wgan
from audio_synthesis.datasets import maestro_dataset
from audio_synthesis.utils import maestro_save_helper as save_helper
from audio_synthesis.utils import spectral

# Setup Paramaters
D_UPDATES_PER_G = 5
Z_DIM = 64
BATCH_SIZE = 64
EPOCHS = 300
SAMPLING_RATE = 16000
FFT_FRAME_LENGTH = 512
FFT_FRAME_STEP = 128
Z_IN_SHAPE = [4, 8, 1024]
SPECTOGRAM_IMAGE_SHAPE = [-1, 128, 256, 2]
SIGNAL_LENGTH = 2**14
WAVEFORM_SHAPE = [-1, SIGNAL_LENGTH, 1]
CRITIC_WEIGHTINGS = [1.0, 1.0/1000.0]
CHECKPOINT_DIR = '_results/representation_study/STFTWaveGAN_HR/training_checkpoints/'
RESULT_DIR = '_results/representation_study/STFTWaveGAN_HR/audio/'
MAESTRO_PATH = 'data/MAESTRO_6h.npz'

def _get_discriminator_input_representations(stft_in):
    """Computes the input representations for the STFTWaveGAN discriminator models,
    returning the input waveform and coresponding spectogram representations

    Args:
        x_in: A batch of waveforms with shape (-1, SIGNAL_LENGTH).

    Returns:
        A tuple containing the waveform and spectogram representaions of
        x_in.
    """

    stft_in = tf.squeeze(stft_in)

    waveform = spectral.stft_2_waveform(
        stft_in, frame_length=FFT_FRAME_LENGTH, frame_step=FFT_FRAME_STEP
    )
    waveform = tf.squeeze(waveform)
    waveform = waveform[:, 0:SIGNAL_LENGTH]

    return (stft_in, waveform)

def main():
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    print('Num GPUs Available: ', len(tf.config.experimental.list_physical_devices('GPU')))

    raw_maestro = maestro_dataset.get_maestro_stft_dataset(
        MAESTRO_PATH, frame_length=FFT_FRAME_LENGTH, frame_step=FFT_FRAME_STEP
    )
    print(raw_maestro.shape)

    generator = spec_gan.Generator(channels=2, in_shape=Z_IN_SHAPE)
    discriminator = spec_gan.Discriminator(input_shape=SPECTOGRAM_IMAGE_SHAPE, weighting=CRITIC_WEIGHTINGS[1])
    waveform_discriminator = wave_gan.Discriminator(input_shape=WAVEFORM_SHAPE, weighting=CRITIC_WEIGHTINGS[0])

    generator_optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.5, beta_2=0.9)
    discriminator_optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.5, beta_2=0.9)

    get_waveform = lambda stft:\
        spectral.stft_2_waveform(
            stft, FFT_FRAME_LENGTH, FFT_FRAME_STEP
        )[0]

    save_examples = lambda epoch, real, generated:\
        save_helper.save_wav_data(
            epoch, real, generated, SAMPLING_RATE, RESULT_DIR, get_waveform
        )

    spec_phase_gan_model = wgan.WGAN(
        raw_maestro, generator, [discriminator, waveform_discriminator], Z_DIM,
        generator_optimizer, discriminator_optimizer, discriminator_training_ratio=D_UPDATES_PER_G,
        batch_size=BATCH_SIZE, epochs=EPOCHS, checkpoint_dir=CHECKPOINT_DIR, fn_save_examples=save_examples,
        fn_get_discriminator_input_representations=_get_discriminator_input_representations
    )

    spec_phase_gan_model.train()

if __name__ == '__main__':
    main()
<commit_msg>Refactor Variable Name<commit_after># Copyright 2020 Google LLC
#
# 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
#
#    https://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.

"""Training Script for STFTWaveGAN on MAESTRO.

Follows the same setup as SpecPhaseGAN, but
generates STFTs instead of Magnitude and Instantaneous
Frequency. Additionally, a Waveform critic is added
"""

import os
import tensorflow as tf
from audio_synthesis.structures import spec_gan, wave_gan
from audio_synthesis.models import wgan
from audio_synthesis.datasets import maestro_dataset
from audio_synthesis.utils import maestro_save_helper as save_helper
from audio_synthesis.utils import spectral

# Setup Paramaters
D_UPDATES_PER_G = 5
Z_DIM = 64
BATCH_SIZE = 64
EPOCHS = 300
SAMPLING_RATE = 16000
FFT_FRAME_LENGTH = 512
FFT_FRAME_STEP = 128
Z_IN_SHAPE = [4, 8, 1024]
SPECTOGRAM_IMAGE_SHAPE = [-1, 128, 256, 2]
SIGNAL_LENGTH = 2**14
WAVEFORM_SHAPE = [-1, SIGNAL_LENGTH, 1]
CRITIC_WEIGHTINGS = [1.0, 1.0/1000.0]
CHECKPOINT_DIR = '_results/representation_study/STFTWaveGAN_HR/training_checkpoints/'
RESULT_DIR = '_results/representation_study/STFTWaveGAN_HR/audio/'
MAESTRO_PATH = 'data/MAESTRO_6h.npz'

def _get_discriminator_input_representations(stft_in):
    """Computes the input representations for the STFTWaveGAN discriminator models,
    returning the input waveform and coresponding spectogram representations

    Args:
        x_in: A batch of waveforms with shape (-1, SIGNAL_LENGTH).

    Returns:
        A tuple containing the waveform and spectogram representaions of
        x_in.
    """

    stft_in = tf.squeeze(stft_in)

    waveform = spectral.stft_2_waveform(
        stft_in, frame_length=FFT_FRAME_LENGTH, frame_step=FFT_FRAME_STEP
    )
    waveform = tf.squeeze(waveform)
    waveform = waveform[:, 0:SIGNAL_LENGTH]

    return (stft_in, waveform)

def main():
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    print('Num GPUs Available: ', len(tf.config.experimental.list_physical_devices('GPU')))

    raw_maestro = maestro_dataset.get_maestro_stft_dataset(
        MAESTRO_PATH, frame_length=FFT_FRAME_LENGTH, frame_step=FFT_FRAME_STEP
    )
    print(raw_maestro.shape)

    generator = spec_gan.Generator(channels=2, in_shape=Z_IN_SHAPE)
    discriminator = spec_gan.Discriminator(input_shape=SPECTOGRAM_IMAGE_SHAPE, weighting=CRITIC_WEIGHTINGS[1])
    waveform_discriminator = wave_gan.Discriminator(input_shape=WAVEFORM_SHAPE, weighting=CRITIC_WEIGHTINGS[0])

    generator_optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.5, beta_2=0.9)
    discriminator_optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.5, beta_2=0.9)

    get_waveform = lambda stft:\
        spectral.stft_2_waveform(
            stft, FFT_FRAME_LENGTH, FFT_FRAME_STEP
        )[0]

    save_examples = lambda epoch, real, generated:\
        save_helper.save_wav_data(
            epoch, real, generated, SAMPLING_RATE, RESULT_DIR, get_waveform
        )

    stft_wave_gan_model = wgan.WGAN(
        raw_maestro, generator, [discriminator, waveform_discriminator], Z_DIM,
        generator_optimizer, discriminator_optimizer, discriminator_training_ratio=D_UPDATES_PER_G,
        batch_size=BATCH_SIZE, epochs=EPOCHS, checkpoint_dir=CHECKPOINT_DIR, fn_save_examples=save_examples,
        fn_get_discriminator_input_representations=_get_discriminator_input_representations
    )

    stft_wave_gan_model.train()

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


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

import os
import shutil
import datetime
import sys
import argparse
import bitarray
import itertools

from ete3 import Tree

import logging

DEFAULT_FORMAT = 1

class Read:
	def __init__(self, tree, simulate_lca=False, annotate=False):
		self.tree=tree
		self.k=tree.k
		self.simulate_lca=simulate_lca
		self.annotate=annotate

	def process_krakline(self,krakline,form):
		self.load_krakline(krakline)
		self.find_assignments()
		#print(self.asgs)
		self.filter_assignments()
		self.print_assignments(form)

	def load_krakline(self,krakline):
		_,self.qname,_,qlen,self.krakmers=krakline.strip().split("\t")
		self.qlen=int(qlen)
		self.seq=None
		self.qual=None

		#self.hitmasks=None
		self.asgs={}

		# list of (count,list of nodes)
		self.kmer_blocks=[]

		b_sum=0
		for b in self.krakmers.split(" "):
			(ids,count)=b.split(":")
			count=int(count)
			b_sum+=count
			self.kmer_blocks.append((ids.split(","),count))
		assert self.qlen==b_sum+self.k-1, krakline

	def find_assignments(self):
		# hits before top-down propagation
		hitmasks,covmasks=self.tree.masks_from_kmer_blocks(self.kmer_blocks,simulate_lca=self.simulate_lca)

		rnames=hitmasks.keys()

		# hits after top-down propagation
		for rname in rnames:
			if rname=="0" or rname=="A":
				continue

			#(hm,cm)=self.hitmasks[rname].copy()

			self.asgs[rname]={
					'hitmask' : hitmasks[rname].copy(),
					'covmask' : covmasks[rname].copy(),
				}

			node=self.tree.name_dict[rname]
			for p_rname in self.tree.upnodes_dict[rname] & rnames:
				self.asgs[rname]['hitmask']|=hitmasks[p_rname]
				self.asgs[rname]['covmask']|=covmasks[p_rname]

			#while node.up:
			#	node=node.up
			#	#print("node up",node.name,file=sys.stderr)
			#	try:
			#		self.asgs[rname]['hitmask']|=hitmasks[node.name]
			#		self.asgs[rname]['covmask']|=covmasks[node.name]
			#	except KeyError:
			#		pass

	def filter_assignments(self):
		"""
		Annotate & filter assignment to a node.

		rname=None => unassigned
		"""

		self.max_hit=0
		self.max_cov=0

		self.max_hit_rnames=[]
		self.max_cov_rnames=[]

		for rname in self.asgs:
			"""
			1. hit count
			"""
			hit=self.asgs[rname]['hitmask'].count()
			self.asgs[rname]['h1']=hit

			if hit>self.max_hit:
				self.max_hit=hit
				self.max_hit_rnames=[rname]
			elif hit==self.max_hit:
				self.max_hit_rnames.append(rname)

			"""
			2. coverage
			"""
			cov=self.asgs[rname]['covmask'].count()
			self.asgs[rname]['c1']=cov

			if cov>self.max_cov:
				self.max_cov=cov
				self.max_cov_rnames=[rname]
			elif cov==self.max_cov:
				self.max_cov_rnames.append(rname)

			#x=self.asgs[rname]['covmask'].replace("01","0\t1").replace("10","1\t0")
			#y=[]
			#for b in x.split():
			#	y.extend([str(len(b)),"=" if b[0]=="1" else "X"])
			#self.asgs[rname]['cigar']="".join(y)
			#self.asgs[rname]['cigar']="X"

	def annotate_assignment(self, rname, full_annotation=False):
		asg=self.asgs[rname]

		if full_annotation:
			try:
				asg['gi']=self.tree.name_dict[rname].gi
			except AttributeError:
				pass

			try:
				asg['ti']=self.tree.name_dict[rname].taxid
			except AttributeError:
				pass

			try:
				asg['sn']=self.tree.name_dict[rname].sci_name
			except AttributeError:
				pass

			try:
				asg['ra']=self.tree.name_dict[rname].rank
			except AttributeError:
				pass

		c=[]
		runs=itertools.groupby(asg['covmask'])
		for run in runs:
			c.append(str(len(list(run[1]))))
			c.append('=' if run[0] else 'X')
		self.asgs[rname]['cigar']="".join(c)

	def print_assignments(self, form):
		if len(self.max_hit_rnames)>0:
			for rname in self.max_hit_rnames:
				if form=="sam":
					self.annotate_assignment(rname,full_annotation=self.annotate)
					self.print_sam_line(rname)
				elif form=="kraken":
					self.print_kraken_line(rname)

		else:
			if form=="sam":
				self.print_sam_line(None)
			elif form=="kraken":
				self.print_kraken_line(None)

	def print_sam_line(self,rname,file=sys.stdout):
		tags=[]
		qname=self.qname
		if rname is None:
			flag=4
			rname="*"
			pos="0"
			mapq="0"
			cigar="*"
		else:
			flag=0
			mapq="60"
			pos="1"
			cigar=self.asgs[rname]['cigar']

		columns=[
				qname,str(flag),rname,
				pos,mapq,cigar,
				"*","0", "0","*","*",
			]
		columns.extend(self.sam_tags(rname))
		print("\t".join(columns),file=file)

	def sam_tags(self,rname):
		tags=[]

		if rname!="*":
			asg=self.asgs[rname]
			if self.annotate:
				for tag in ['gi','ti','sn','ra']:
					try:
						tags.append("".join( [tag,":Z:",asg[tag]] ))
					except KeyError:
						pass

			for tag in ['h1','c1']:
				tags.append("".join( [tag,":i:",str(asg[tag])] ))

		return tags

	def print_sam_header(self,file=sys.stdout):
		print("@HD\tVN:1.5\tSO:unsorted",file=file)
		for node in self.tree.tree.traverse("postorder"):
			self.tree.name_dict[node.name]=node

			try:
				ur="\tUR:{}".format(node.fastapath)
			except:
				ur=""

			try:
				sp="\tSP:{}".format(node.sci_name)
			except:
				sp=""

			try:
				as_="\tAS:{}".format(node.gi)
			except:
				as_=""

			if node.name!='':
				print("@SQ\tSN:{rname}\tLN:{rlen}{as_}{ur}{sp}".format(
						rname=node.name,
						rlen=4242,
						as_=as_,
						ur=ur,
						sp=sp,
					),file=file)

	def print_kraken_line(self,rname,file=sys.stdout):
		if rname is None:
			stat="U"
			rname="0"
		else:
			stat="C"
		columns=[stat,self.qname,rname,str(self.qlen),self.krakmers]
		print("\t".join(columns),file=file)

class TreeIndex:

	def __init__(self,tree_newick_fn,k,format=DEFAULT_FORMAT):
		self.tree_newick_fn=tree_newick_fn
		self.tree=Tree(tree_newick_fn,format=1)

		self.k=k

		self.name_dict={}
		self.upnodes_dict={}

		for node in self.tree.traverse("postorder"):
			rname=node.name
			self.name_dict[rname]=node
			self.upnodes_dict[rname]=set()
			while node.up:
				node=node.up
				self.upnodes_dict[rname].add(node.name)

		#print (self.name_dict)

	"""
		kmers_assigned_l:
			list of (list_of_nodes, count)

		dict:
			[
				node => hit_vector,
				node => cov_vector
			]
	"""
	def masks_from_kmer_blocks(self,kmers_assigned_l,simulate_lca=False):
		d_h={}
		d_c={}

		npos=sum([x[1] for x in kmers_assigned_l])

		h_len=npos
		c_len=npos+k-1

		pos=0
		for (rname_l, count) in kmers_assigned_l:
			if rname_l!=["0"] and rname_l!=["A"]:
				if lca:
					rname_l=[self.lca(rname_l)]

				v_h=bitarray.bitarray(pos*[False] + count*[True] + (npos-pos-count)*[False])
				v_c=bitarray.bitarray(pos*[False] + (count+k-1)*[True] + (npos-pos-count)*[False])

				#assert len(v_h)==h_len, v_h
				#assert len(v_c)==c_len, v_c

				for rname in rname_l:
					try:
						#assert len(d_h[noden])==h_len
						#assert len(d_c[noden])==c_len
						d_h[rname]|=v_h
						d_c[rname]|=v_c
					except KeyError:
						d_h[rname]=v_h.copy()
						d_c[rname]=v_c.copy()

			pos+=count
		return (d_h, d_c)

	def lca(self,noden_l):
		"""Return LCA for a given list of nodes.
		"""
		assert len(noden_l)>0, noden_l
		if len(noden_l)==1:
			return noden_l[0]
		nodes_l=list(map(lambda x:self.name_dict[x],noden_l))
		lca=nodes_l[0].get_common_ancestor(nodes_l)

		if lca.is_root() and len(lca.children)==1:
			lca=lca.children[0]

		return lca.name

if __name__ == "__main__":

	parser = argparse.ArgumentParser(description='Implementation of assignment algorithm')

	parser.add_argument('-i', '--input',
			type=argparse.FileType('r'),
			required=True,
			dest='input_file',
			help='input file',
		)

	parser.add_argument('-k', '--kmer-size',
			type=int,
			required=True,
			dest='k',
			help='k-mer size',
		)

	parser.add_argument('-n', '--newick-tree',
			type=str,
			metavar='str',
			required=True,
			dest='newick_fn',
			help='newick tree',
		)

	parser.add_argument('-f', '--oformat',
			choices=['kraken','sam'],
			default='kraken',
			dest='format',
			help='format of output',
		)

	parser.add_argument('-l', '--sim-lca',
			action='store_true',
			dest='lca',
			help='simulate LCA',
		)

	parser.add_argument('-a', '--annotate',
			action='store_true',
			dest='annotate',
			help='annotate assignments',
		)

	args = parser.parse_args()

	newick_fn=args.newick_fn
	inp_fo=args.input_file
	lca=args.lca
	form=args.format
	k=args.k
	annotate=args.annotate

	ti=TreeIndex(
			tree_newick_fn=newick_fn,
			k=k,
		)

	read=Read(
			tree=ti,
			simulate_lca=lca,
			annotate=annotate
		)
	if form=="sam":
		read.print_sam_header()
	for x in inp_fo:
		read.process_krakline(x,form=form)

<commit_msg>Precompute annotations<commit_after>#! /usr/bin/env python3

import os
import shutil
import datetime
import sys
import argparse
import bitarray
import itertools

from ete3 import Tree

import logging

DEFAULT_FORMAT = 1

class Read:
	def __init__(self, tree, simulate_lca=False, annotate=False):
		self.tree=tree
		self.k=tree.k
		self.simulate_lca=simulate_lca
		self.annotate=annotate

	def process_krakline(self,krakline,form):
		self.load_krakline(krakline)
		self.find_assignments()
		#print(self.asgs)
		self.filter_assignments()
		self.print_assignments(form)

	def load_krakline(self,krakline):
		_,self.qname,_,qlen,self.krakmers=krakline.strip().split("\t")
		self.qlen=int(qlen)
		self.seq=None
		self.qual=None

		#self.hitmasks=None
		self.asgs={}

		# list of (count,list of nodes)
		self.kmer_blocks=[]

		b_sum=0
		for b in self.krakmers.split(" "):
			(ids,count)=b.split(":")
			count=int(count)
			b_sum+=count
			self.kmer_blocks.append((ids.split(","),count))
		assert self.qlen==b_sum+self.k-1, krakline

	def find_assignments(self):
		# hits before top-down propagation
		hitmasks,covmasks=self.tree.masks_from_kmer_blocks(self.kmer_blocks,simulate_lca=self.simulate_lca)

		rnames=hitmasks.keys()

		# hits after top-down propagation
		for rname in rnames:
			if rname=="0" or rname=="A":
				continue

			#(hm,cm)=self.hitmasks[rname].copy()

			self.asgs[rname]={
					'hitmask' : hitmasks[rname].copy(),
					'covmask' : covmasks[rname].copy(),
				}

			node=self.tree.name_dict[rname]
			for p_rname in self.tree.upnodes_dict[rname] & rnames:
				self.asgs[rname]['hitmask']|=hitmasks[p_rname]
				self.asgs[rname]['covmask']|=covmasks[p_rname]

	def filter_assignments(self):
		"""
		Annotate & filter assignment to a node.

		rname=None => unassigned
		"""

		self.max_hit=0
		self.max_cov=0

		self.max_hit_rnames=[]
		self.max_cov_rnames=[]

		for rname in self.asgs:
			asg=self.asgs[rname]

			"""
			1. hit count
			"""
			hit=asg['hitmask'].count()
			asg['h1']=hit

			if hit>self.max_hit:
				self.max_hit=hit
				self.max_hit_rnames=[rname]
			elif hit==self.max_hit:
				self.max_hit_rnames.append(rname)

			"""
			2. coverage
			"""
			cov=asg['covmask'].count()
			asg['c1']=cov

			if cov>self.max_cov:
				self.max_cov=cov
				self.max_cov_rnames=[rname]
			elif cov==self.max_cov:
				self.max_cov_rnames.append(rname)

	def print_assignments(self, form):
		if len(self.max_hit_rnames)>0:
			for rname in self.max_hit_rnames:
				asg=self.asgs[rname]
				if form=="sam":
					# compute cigar
					c=[]
					runs=itertools.groupby(asg['covmask'])
					for run in runs:
						c.append(str(len(list(run[1]))))
						c.append('=' if run[0] else 'X')
					asg['cigar']="".join(c)

					self.print_sam_line(rname,self.tree.sam_annotations_dict[rname] if self.annotate else "")
				elif form=="kraken":
					self.print_kraken_line(rname)

		else:
			if form=="sam":
				self.print_sam_line(None)
			elif form=="kraken":
				self.print_kraken_line(None)

	def print_sam_line(self,rname,suffix="",file=sys.stdout):
		tags=[]
		qname=self.qname
		if rname is None:
			flag=4
			rname="*"
			pos="0"
			mapq="0"
			cigar="*"
		else:
			flag=0
			mapq="60"
			pos="1"
			cigar=self.asgs[rname]['cigar']

		columns=[
				qname,str(flag),rname,
				pos,mapq,cigar,
				"*","0", "0","*","*",
			]

		if rname!="*":
			asg=self.asgs[rname]
			for tag in ['h1','c1']:
				columns.append("".join( [tag,":i:",str(asg[tag])] ))

		print("\t".join(columns),suffix,file=file,sep="")

	def print_sam_header(self,file=sys.stdout):
		print("@HD\tVN:1.5\tSO:unsorted",file=file)
		for node in self.tree.tree.traverse("postorder"):
			self.tree.name_dict[node.name]=node

			try:
				ur="\tUR:{}".format(node.fastapath)
			except:
				ur=""

			try:
				sp="\tSP:{}".format(node.sci_name)
			except:
				sp=""

			try:
				as_="\tAS:{}".format(node.gi)
			except:
				as_=""

			if node.name!='':
				print("@SQ\tSN:{rname}\tLN:{rlen}{as_}{ur}{sp}".format(
						rname=node.name,
						rlen=4242,
						as_=as_,
						ur=ur,
						sp=sp,
					),file=file)

	def print_kraken_line(self,rname,file=sys.stdout):
		if rname is None:
			stat="U"
			rname="0"
		else:
			stat="C"
		columns=[stat,self.qname,rname,str(self.qlen),self.krakmers]
		print("\t".join(columns),file=file)

class TreeIndex:

	def __init__(self,tree_newick_fn,k,format=DEFAULT_FORMAT):
		self.tree_newick_fn=tree_newick_fn
		self.tree=Tree(tree_newick_fn,format=1)

		self.k=k

		self.name_dict={}
		self.sam_annotations_dict={}
		self.upnodes_dict={}

		for node in self.tree.traverse("postorder"):
			rname=node.name
			self.name_dict[rname]=node

			# annotations
			tags_parts=[]
			try:
				tags_parts.extend(["\tgi:Z:",node.gi])
			except AttributeError:
				pass

			try:
				tags_parts.extend(["\tti:Z:",node.taxid])
			except AttributeError:
				pass

			try:
				tags_parts.extend(["\tsn:Z:",node.sci_name])
			except AttributeError:
				pass

			try:
				tags_parts.extend(["\tra:Z:",node.rank])
			except AttributeError:
				pass

			self.sam_annotations_dict[rname]="".join(tags_parts)

			# set of upper nodes
			self.upnodes_dict[rname]=set()
			while node.up:
				node=node.up
				self.upnodes_dict[rname].add(node.name)

	"""
		kmers_assigned_l:
			list of (list_of_nodes, count)

		dict:
			[
				node => hit_vector,
				node => cov_vector
			]
	"""
	def masks_from_kmer_blocks(self,kmers_assigned_l,simulate_lca=False):
		d_h={}
		d_c={}

		npos=sum([x[1] for x in kmers_assigned_l])

		h_len=npos
		c_len=npos+k-1

		pos=0
		for (rname_l, count) in kmers_assigned_l:
			if rname_l!=["0"] and rname_l!=["A"]:
				if lca:
					rname_l=[self.lca(rname_l)]

				v_h=bitarray.bitarray(pos*[False] + count*[True] + (npos-pos-count)*[False])
				v_c=bitarray.bitarray(pos*[False] + (count+k-1)*[True] + (npos-pos-count)*[False])

				#assert len(v_h)==h_len, v_h
				#assert len(v_c)==c_len, v_c

				for rname in rname_l:
					try:
						#assert len(d_h[noden])==h_len
						#assert len(d_c[noden])==c_len
						d_h[rname]|=v_h
						d_c[rname]|=v_c
					except KeyError:
						d_h[rname]=v_h.copy()
						d_c[rname]=v_c.copy()

			pos+=count
		return (d_h, d_c)

	def lca(self,noden_l):
		"""Return LCA for a given list of nodes.
		"""
		assert len(noden_l)>0, noden_l
		if len(noden_l)==1:
			return noden_l[0]
		nodes_l=list(map(lambda x:self.name_dict[x],noden_l))
		lca=nodes_l[0].get_common_ancestor(nodes_l)

		if lca.is_root() and len(lca.children)==1:
			lca=lca.children[0]

		return lca.name

if __name__ == "__main__":

	parser = argparse.ArgumentParser(description='Implementation of assignment algorithm')

	parser.add_argument('-i', '--input',
			type=argparse.FileType('r'),
			required=True,
			dest='input_file',
			help='input file',
		)

	parser.add_argument('-k', '--kmer-size',
			type=int,
			required=True,
			dest='k',
			help='k-mer size',
		)

	parser.add_argument('-n', '--newick-tree',
			type=str,
			metavar='str',
			required=True,
			dest='newick_fn',
			help='newick tree',
		)

	parser.add_argument('-f', '--oformat',
			choices=['kraken','sam'],
			default='kraken',
			dest='format',
			help='format of output',
		)

	parser.add_argument('-l', '--sim-lca',
			action='store_true',
			dest='lca',
			help='simulate LCA',
		)

	parser.add_argument('-a', '--annotate',
			action='store_true',
			dest='annotate',
			help='annotate assignments',
		)

	args = parser.parse_args()

	newick_fn=args.newick_fn
	inp_fo=args.input_file
	lca=args.lca
	form=args.format
	k=args.k
	annotate=args.annotate

	ti=TreeIndex(
			tree_newick_fn=newick_fn,
			k=k,
		)

	read=Read(
			tree=ti,
			simulate_lca=lca,
			annotate=annotate
		)
	if form=="sam":
		read.print_sam_header()
	for x in inp_fo:
		read.process_krakline(x,form=form)

<|endoftext|>