<commit_before>import requests
from bs4 import BeautifulSoup
from P_Queue import P_Queue

MIN_P = 0
MAX_P = 5000

def parse_source(html, encoding='utf-8'):
    parsed = BeautifulSoup(html, from_encoding=encoding)
    return parsed

def fetch_search_results(keywords="",
                         rh='i:aps',
                         page=None):
    base = 'http://www.amazon.com/s/'
    if len(keywords) == 0:
        raise ValueError("Please enter search keys")
    params = {
        'ref': 'sr_st_relevancerank',
        'keywords': keywords,
        'rh': rh
    }
    if page is not None:
        params['page'] = page
    params['low-price'] = MIN_P
    params['high-price'] = MAX_P
    resp = requests.get(base, params=params, timeout=3)
    resp.raise_for_status()
    return resp.content, resp.encoding

def extract_items(parsed):
    fst = parsed.find_all('div', class_='fstGrid prod celwidget')
    rslt = parsed.find_all('div', class_='rsltGrid prod celwidget')
    for block in fst:
        img = block.find('div', class_='imageBox').find('img')
        link = block.find('h3', class_='newaps').find('a')
        prime_price = block.find('span', class_='bld lrg red')
        new_price = block.find('span', class_='price bld')
        item = item_dictionary(img, link, prime_price, new_price)
        if item is not None:
            yield item
    for block in rslt:
        img = block.find('div', class_='imageBox').find('img')
        link = block.find('h3', class_='newaps').find('a')
        prime_price = block.find('span', class_='bld lrg red')
        new_price = block.find('span', class_='price bld')
        item = item_dictionary(img, link, prime_price, new_price)
        if item is not None:
            yield item

def item_dictionary(img, link, prime_price, new_price):
    item = {}
    item['image'] = img.attrs['src']
    item['link'] = link.attrs['href']
    item['title'] = link.string.strip()
    item['prime_price'] = get_price(prime_price)
    item['new_price'] = get_price(new_price)
    if item['prime_price'] is None:
        item['prime_price'] = u'n/a'
    if item['new_price'] is None:
        item['new_price'] = u'n/a'
    return item if item['new_price'] != 'n/a' \
        or item['prime_price'] != 'n/a' \
        else None

def get_price(price_string):
    if price_string is not None:
        x = price_string.string.strip()
        if '$' in x:
            temp = str(x)[1:-3].replace(',', '')
            if int(temp) >= MIN_P and int(temp) <= MAX_P:
                return x
    return None

def search_results(keywords, category, price, price_range):
    count = 1
    page = 1
    if price is not None:
        price = float(price)
        price_range = float(price_range)
        global MIN_P
        MIN_P = price - (price * price_range / 100)
        global MAX_P
        MAX_P = price + (price * price_range / 100)
    p_queue = P_Queue()
    while count < 10:
        content = fetch_search_results(keywords,
                                       category,
                                       page)
        parsed = parse_source(content[0], content[1])
        print "\nSEARCH RESULTS:\t"+str(len(parsed))+"\n"
        if len(parsed.find_all('div', class_='rsltGrid prod celwidget')) == 0:
            break
        items = extract_items(parsed)
        for i in items:
            pri = 0
            if i['prime_price'] == u'n/a':
                pri = int(i['new_price'][1:-3])
            else:
                pri = int(i['prime_price'][1:-3])
            p_queue.insert(i, pri)
            count += 1
        page += 1
    MIN_P = 0
    MAX_P = 5000
    return p_queue

if __name__ == '__main__':
    keywords = "Apple"
    category = "n:541966"
    price = 300.0
    price_range = 20.0
    a = search_results(keywords, category, price, price_range)
    f = open('extract.txt', 'w')
    while True:
        try:
            i = a.pop()._data
            f.write("\n")
            f.write(str(i['image']))
            f.write("\n")
            f.write(str(i['link']))
            f.write("\n")
            # f.write(str(i['title']))
            # f.write("\n")
            f.write(str(i['prime_price']))
            f.write("\n")
            f.write(str(i['new_price']))
            f.write("\n")
        except IndexError:
            break
    f.close()
<commit_msg>fixed price formating for priority queue insert<commit_after>import requests
from bs4 import BeautifulSoup
from P_Queue import P_Queue

MIN_P = 0
MAX_P = 5000

def parse_source(html, encoding='utf-8'):
    parsed = BeautifulSoup(html, from_encoding=encoding)
    return parsed

def fetch_search_results(keywords="",
                         rh='i:aps',
                         page=None):
    base = 'http://www.amazon.com/s/'
    if len(keywords) == 0:
        raise ValueError("Please enter search keys")
    params = {
        'ref': 'sr_st_relevancerank',
        'keywords': keywords,
        'rh': rh
    }
    if page is not None:
        params['page'] = page
    params['low-price'] = MIN_P
    params['high-price'] = MAX_P
    resp = requests.get(base, params=params, timeout=3)
    resp.raise_for_status()
    return resp.content, resp.encoding

def extract_items(parsed):
    fst = parsed.find_all('div', class_='fstGrid prod celwidget')
    rslt = parsed.find_all('div', class_='rsltGrid prod celwidget')
    for block in fst:
        img = block.find('div', class_='imageBox').find('img')
        link = block.find('h3', class_='newaps').find('a')
        prime_price = block.find('span', class_='bld lrg red')
        new_price = block.find('span', class_='price bld')
        item = item_dictionary(img, link, prime_price, new_price)
        if item is not None:
            yield item
    for block in rslt:
        img = block.find('div', class_='imageBox').find('img')
        link = block.find('h3', class_='newaps').find('a')
        prime_price = block.find('span', class_='bld lrg red')
        new_price = block.find('span', class_='price bld')
        item = item_dictionary(img, link, prime_price, new_price)
        if item is not None:
            yield item

def item_dictionary(img, link, prime_price, new_price):
    item = {}
    item['image'] = img.attrs['src']
    item['link'] = link.attrs['href']
    item['title'] = link.string.strip()
    item['prime_price'] = get_price(prime_price)
    item['new_price'] = get_price(new_price)
    if item['prime_price'] is None:
        item['prime_price'] = u'n/a'
    if item['new_price'] is None:
        item['new_price'] = u'n/a'
    return item if item['new_price'] != 'n/a' \
        or item['prime_price'] != 'n/a' \
        else None

def get_price(price_string):
    if price_string is not None:
        x = price_string.string.strip()
        if '$' in x:
            temp = str(x)[1:-3].replace(',', '')
            if int(temp) >= MIN_P and int(temp) <= MAX_P:
                return x
    return None

def search_results(keywords, category, price, price_range):
    count = 1
    page = 1
    if price is not None:
        price = float(price)
        price_range = float(price_range)
        global MIN_P
        MIN_P = price - (price * price_range / 100)
        global MAX_P
        MAX_P = price + (price * price_range / 100)
    p_queue = P_Queue()
    while count < 10:
        content = fetch_search_results(keywords,
                                       category,
                                       page)
        parsed = parse_source(content[0], content[1])
        print "\nSEARCH RESULTS:\t"+str(len(parsed))+"\n"
        if len(parsed.find_all('div', class_='rsltGrid prod celwidget')) == 0:
            break
        items = extract_items(parsed)
        for i in items:
            pri = 0
            if i['prime_price'] == u'n/a':
                pri = int(i['new_price'][1:-3].replace(',', ''))
            else:
                pri = int(i['prime_price'][1:-3].replace(',', ''))
            p_queue.insert(i, pri)
            count += 1
        page += 1
    MIN_P = 0
    MAX_P = 5000
    return p_queue

if __name__ == '__main__':
    keywords = "Apple"
    category = "n:541966"
    price = 300.0
    price_range = 20.0
    a = search_results(keywords, category, price, price_range)
    f = open('extract.txt', 'w')
    while True:
        try:
            i = a.pop()._data
            f.write("\n")
            f.write(str(i['image']))
            f.write("\n")
            f.write(str(i['link']))
            f.write("\n")
            # f.write(str(i['title']))
            # f.write("\n")
            f.write(str(i['prime_price']))
            f.write("\n")
            f.write(str(i['new_price']))
            f.write("\n")
        except IndexError:
            break
    f.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
#
# Do an amplitude ramp and save outgoing samples
# and samples coming back from the USRP. The generated
# signal is a two-tone signal
#
# Compare the outgoing amplitude with the incoming amplitude
# to measure AM/AM compression.
#
# In a second step, compare phase and measre AM/PM compression.
#
# Copyright (C) 2016
# Matthias P. Braendli, matthias.braendli@mpb.li
# http://www.opendigitalradio.org
# Licence: The MIT License, see LICENCE file

import sys
import os

import traceback
from gnuradio import analog
from gnuradio import filter
from gnuradio import blocks
from gnuradio import gr
from gnuradio import uhd
from grc_gnuradio import blks2 as grc_blks2
import argparse
import time
import socket
import struct
import threading
from Queue import Queue

# TCP ports used to communicate between the flowgraph and the python script
# The flowgraph interleaves 3 float streams :
#  generator magnitude
#  phase difference
#  feedback magnitude
TCP_PORT = 47009

def xrange(start, stop, step):
    x = start
    while x < stop:
        yield x
        x += step

class amplitude_ramp(gr.top_block):
    def __init__(self, options):
        gr.top_block.__init__(self, "Amplitude Ramp")

        self.txgain = float(options.txgain)
        self.source_ampl = float(options.ampl_start)

        self.samp_rate = 4e6
        self.rxgain = 0
        self.freq = 222e6
        self.decim = 4000

        # Two-tone signal generator at 1kHz and 2kHz
        self.analog_sig_source_x_0 = analog.sig_source_c(
                self.samp_rate, analog.GR_COS_WAVE, 1000, self.source_ampl, 0)
        self.analog_sig_source_x_1 = analog.sig_source_c(
                self.samp_rate, analog.GR_COS_WAVE, 2000, self.source_ampl, 0)
        self.blocks_add_xx_0 = blocks.add_vcc(1)

        # Connects to both USRP output and mag/phase converter

        self.uhd_usrp_sink_0 = uhd.usrp_sink(
                "",
                uhd.stream_args(
                    cpu_format="fc32",
                    channels=range(1),
                    ),
                )
        self.uhd_usrp_sink_0.set_samp_rate(self.samp_rate)
        self.uhd_usrp_sink_0.set_center_freq(self.freq, 0)
        self.uhd_usrp_sink_0.set_gain(self.txgain, 0)

        self.blocks_complex_to_magphase_0 = blocks.complex_to_magphase(1)

        # mag goes to TCP interleaved, phase goes to subtractor. There is
        # no substraction block, so it is done with multiply by -1 and add

        # Feedback from the USRP, goes to the subtractor after mag/phase
        self.blocks_complex_to_magphase_1 = blocks.complex_to_magphase(1)
        self.uhd_usrp_source_0 = uhd.usrp_source(
                "",
                uhd.stream_args(
                    cpu_format="fc32",
                    channels=range(1),
                    ),
                )
        self.uhd_usrp_source_0.set_samp_rate(self.samp_rate)
        self.uhd_usrp_source_0.set_center_freq(self.freq, 0)
        self.uhd_usrp_source_0.set_gain(self.rxgain, 0)

        self.blocks_invert_signal = blocks.multiply_const_vff((-1, ))
        self.blocks_phase_adder = blocks.add_vff(1)

        # The interleaved takes gen mag, phase diff and feedback mag
        # signals and puts them together. We need to decimate before we interleave
        self.blocks_moving_average_gen = blocks.moving_average_ff(self.decim, 1, 4000)
        self.fir_filter_gen = filter.fir_filter_fff(self.decim, ([1]))
        self.fir_filter_gen.declare_sample_delay(0)

        self.blocks_moving_average_phase = blocks.moving_average_ff(self.decim, 1, 4000)
        self.fir_filter_phase = filter.fir_filter_fff(self.decim, ([1]))
        self.fir_filter_phase.declare_sample_delay(0)

        self.blocks_moving_average_feedback = blocks.moving_average_ff(self.decim, 1, 4000)
        self.fir_filter_feedback = filter.fir_filter_fff(self.decim, ([1]))
        self.fir_filter_feedback.declare_sample_delay(0)

        self.blocks_interleave = blocks.interleave(gr.sizeof_float*1, 1)

        self.blks2_tcp_sink_0 = grc_blks2.tcp_sink(
                itemsize=gr.sizeof_float*1,
                addr="127.0.0.1",
                port=TCP_PORT,
                server=True,
                )

        # Connect outgoing
        self.connect((self.analog_sig_source_x_0, 0), (self.blocks_add_xx_0, 1))
        self.connect((self.analog_sig_source_x_1, 0), (self.blocks_add_xx_0, 0))
        self.connect((self.blocks_add_xx_0, 0), (self.blocks_complex_to_magphase_0, 0))
        self.connect((self.blocks_add_xx_0, 0), (self.uhd_usrp_sink_0, 0))

        self.connect((self.blocks_complex_to_magphase_0, 0),
                (self.blocks_moving_average_gen, 0))
        self.connect((self.blocks_moving_average_gen, 0),
                (self.fir_filter_gen, 0))
        self.connect((self.fir_filter_gen, 0), (self.blocks_interleave, 0))

        self.connect((self.blocks_complex_to_magphase_0, 1),
                (self.blocks_invert_signal, 0))
        self.connect((self.blocks_invert_signal, 0), (self.blocks_phase_adder, 1))

        # Connect feedback
        self.connect((self.uhd_usrp_source_0, 0), (self.blocks_complex_to_magphase_1, 0))
        self.connect((self.blocks_complex_to_magphase_1, 1), (self.blocks_phase_adder, 0))
        self.connect((self.blocks_phase_adder, 0), (self.blocks_moving_average_phase, 0))
        self.connect((self.blocks_moving_average_phase, 0), (self.fir_filter_phase, 0))
        self.connect((self.fir_filter_phase, 0), (self.blocks_interleave, 1))

        self.connect((self.blocks_complex_to_magphase_1, 0),
                (self.blocks_moving_average_feedback, 0))
        self.connect((self.blocks_moving_average_feedback, 0),
                (self.fir_filter_feedback, 0))
        self.connect((self.fir_filter_feedback, 0),
                (self.blocks_interleave, 2))

        # connect interleaver output to TCP socket
        self.connect((self.blocks_interleave, 0), (self.blks2_tcp_sink_0, 0))

    def get_txgain(self):
        return self.txgain

    def set_txgain(self, txgain):
        self.txgain = txgain
        self.uhd_usrp_sink_0.set_gain(self.txgain, 0)

    def get_source_ampl(self):
        return self.source_ampl

    def set_source_ampl(self, source_ampl):
        print("Set amplitude to {}".format(source_ampl))
        self.source_ampl = source_ampl
        self.analog_sig_source_x_0.set_amplitude(self.source_ampl)
        self.analog_sig_source_x_1.set_amplitude(self.source_ampl)

    def get_freq(self):
        return self.freq

    def set_freq(self, freq):
        self.freq = freq
        self.uhd_usrp_sink_0.set_center_freq(self.freq, 0)
        self.uhd_usrp_source_0.set_center_freq(self.freq, 0)

class RampGenerator(threading.Thread):
    def __init__(self, options):
        threading.Thread.__init__(self)
        self.event_queue_ = Queue()
        self.in_queue_ = Queue()

        self.num_meas = int(options.num_meas)
        self.num_meas_to_skip = int(options.num_meas_to_skip)
        self.ampl_start = float(options.ampl_start)
        self.ampl_step = float(options.ampl_step)
        self.ampl_stop = float(options.ampl_stop)

    def set_source_ampl(self, ampl):
        self.event_queue_.put(ampl)
        self.in_queue_.get()

    def wait_on_event(self):
        return self.event_queue_.get()

    def confirm_source_ampl_updated(self):
        self.in_queue_.put(0)

    def run(self):
        try:
            self.run_ex()
        except:
            traceback.print_exc()
        finally:
            self.event_queue_.put("quit")

    def run_ex(self):
        print("Wait before connection")
        time.sleep(3)

        print("Connecting to flowgraph")
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect(("localhost", TCP_PORT))
        print("Connected")

        amplitudes = xrange(self.ampl_start, self.ampl_stop, self.ampl_step)
        measurements = []

        for ampl in amplitudes:
            self.set_source_ampl(ampl)

            mag_gen_sum = 0
            phase_diff_sum = 0
            mag_feedback_sum = 0

            for measurement_ignore in range(self.num_meas_to_skip):
                # Receive and ignore three floats on the socket
                sock.recv(12)

            for measurement_ix in range(self.num_meas):
                # Receive three floats on the socket
                mag_gen, phase_diff, mag_feedback = struct.unpack(
                        "fff",
                        sock.recv(12))

                mag_gen_sum += mag_gen
                phase_diff_sum += phase_diff
                mag_feedback_sum += mag_feedback

                measurements.append((ampl, mag_gen, mag_feedback, phase_diff))

            mag_gen_avg = mag_gen_sum / self.num_meas
            mag_feedback_avg = mag_feedback_sum / self.num_meas
            phase_diff_avg = phase_diff_sum / self.num_meas

            print("Ampl: {} Out: {:10} In: {:10} phase_diff: {:10}".format(
                ampl, mag_gen_avg, mag_feedback_avg, phase_diff_avg))

        self.event_queue_.put("done")
        self.event_queue_.put(measurements)

parser = argparse.ArgumentParser(description='Two-tone amplitude ramp')

parser.add_argument('--ampl-start',
        default='0.1',
        help='Start amplitude',
        required=False)

parser.add_argument('--ampl-stop',
        default='0.8',
        help='Stop amplitude',
        required=False)

parser.add_argument('--ampl-step',
        default='0.02',
        help='Amplitude steps',
        required=False)

parser.add_argument('--txgain',
        default='10',
        help='txgain for USRP sink',
        required=False)

parser.add_argument('--num-meas',
        default='2000',
        help='number of measurements per amplitude',
        required=False)

parser.add_argument('--num-meas-to-skip',
        default='50',
        help='After each amplitude change, ignore num-meas-to-skip measurements',
        required=False)

cli_args = parser.parse_args()

rampgen = RampGenerator(cli_args)
rampgen.start()

# this blocks until the flowgraph is up and running, i.e. all sockets
# got a connection
top = amplitude_ramp(cli_args)
top.set_source_ampl(float(cli_args.ampl_start))
top.start()

while True:
    event = rampgen.wait_on_event()
    if event == "done":
        measurements = rampgen.wait_on_event()
        fd = open("measurements.csv", "w")
        for m in measurements:
            fd.write(",".join("{}".format(x) for x in m) + "\n")
        fd.close()
        break
    elif event == "quit":
        break
    else:
        top.set_source_ampl(event)
        rampgen.confirm_source_ampl_updated()

top.stop()
print("Wait for completion")
top.wait()

print("Done")
<commit_msg>Replace phase measurement by complex multiply<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
#
# Do an amplitude ramp and save outgoing samples
# and samples coming back from the USRP. The generated
# signal is a two-tone signal
#
# Compare the outgoing amplitude with the incoming amplitude
# to measure AM/AM compression.
#
# In a second step, compare phase and measre AM/PM compression.
#
# Copyright (C) 2016
# Matthias P. Braendli, matthias.braendli@mpb.li
# http://www.opendigitalradio.org
# Licence: The MIT License, see LICENCE file

import sys
import os

import traceback
from gnuradio import analog
from gnuradio import filter
from gnuradio import blocks
from gnuradio import gr
from gnuradio import uhd
from grc_gnuradio import blks2 as grc_blks2
import argparse
import time
import socket
import struct
import threading
from Queue import Queue

# TCP ports used to communicate between the flowgraph and the python script
# The flowgraph interleaves 3 float streams :
#  generator magnitude
#  phase difference
#  feedback magnitude
TCP_PORT = 47009

def xrange(start, stop, step):
    x = start
    while x < stop:
        yield x
        x += step

class amplitude_ramp(gr.top_block):
    def __init__(self, options):
        gr.top_block.__init__(self, "Amplitude Ramp")

        self.txgain = float(options.txgain)
        self.source_ampl = float(options.ampl_start)

        self.samp_rate = 4e6
        self.rxgain = 0
        self.freq = 222e6
        self.decim = 4000

        # Two-tone signal generator at 1kHz and 2kHz
        self.analog_sig_source_x_0 = analog.sig_source_c(
                self.samp_rate, analog.GR_COS_WAVE, 1000, self.source_ampl, 0)
        self.analog_sig_source_x_1 = analog.sig_source_c(
                self.samp_rate, analog.GR_COS_WAVE, 2000, self.source_ampl, 0)
        self.blocks_add_xx_0 = blocks.add_vcc(1)

        # Connects to both USRP output, mag and phase converter

        self.uhd_usrp_sink_0 = uhd.usrp_sink(
                "",
                uhd.stream_args(
                    cpu_format="fc32",
                    channels=range(1),
                    ),
                )
        self.uhd_usrp_sink_0.set_samp_rate(self.samp_rate)
        self.uhd_usrp_sink_0.set_center_freq(self.freq, 0)
        self.uhd_usrp_sink_0.set_gain(self.txgain, 0)

        self.blocks_complex_to_mag_squared_0 = blocks.complex_to_mag_squared(1)

        self.blocks_multiply_conjugate_cc_0 = blocks.multiply_conjugate_cc(1)
        self.blocks_complex_to_arg_0 = blocks.complex_to_arg(1)

        # mag goes to TCP interleaver, phase goes to complex multiplier
        # and complex-to-arg

        # Feedback from the USRP, goes to the subtractor after mag/phase
        self.blocks_complex_to_mag_squared_1 = blocks.complex_to_mag_squared(1)
        self.uhd_usrp_source_0 = uhd.usrp_source(
                "",
                uhd.stream_args(
                    cpu_format="fc32",
                    channels=range(1),
                    ),
                )
        self.uhd_usrp_source_0.set_samp_rate(self.samp_rate)
        self.uhd_usrp_source_0.set_center_freq(self.freq, 0)
        self.uhd_usrp_source_0.set_gain(self.rxgain, 0)

        # The interleaved takes gen mag, phase diff and feedback mag
        # signals and puts them together. We need to decimate before we interleave
        self.blocks_moving_average_gen = blocks.moving_average_ff(self.decim, 1, 4000)
        self.fir_filter_gen = filter.fir_filter_fff(self.decim, ([1]))
        self.fir_filter_gen.declare_sample_delay(0)

        self.blocks_moving_average_phase = blocks.moving_average_ff(self.decim, 1, 4000)
        self.fir_filter_phase = filter.fir_filter_fff(self.decim, ([1]))
        self.fir_filter_phase.declare_sample_delay(0)

        self.blocks_moving_average_feedback = blocks.moving_average_ff(self.decim, 1, 4000)
        self.fir_filter_feedback = filter.fir_filter_fff(self.decim, ([1]))
        self.fir_filter_feedback.declare_sample_delay(0)

        self.blocks_interleave = blocks.interleave(gr.sizeof_float*1, 1)

        self.blks2_tcp_sink_0 = grc_blks2.tcp_sink(
                itemsize=gr.sizeof_float*1,
                addr="127.0.0.1",
                port=TCP_PORT,
                server=True,
                )

        # Connect outgoing
        self.connect((self.analog_sig_source_x_0, 0), (self.blocks_add_xx_0, 1))
        self.connect((self.analog_sig_source_x_1, 0), (self.blocks_add_xx_0, 0))
        self.connect((self.blocks_add_xx_0, 0), (self.blocks_complex_to_mag_squared_0, 0))
        self.connect((self.blocks_add_xx_0, 0), (self.uhd_usrp_sink_0, 0))
        self.connect((self.blocks_add_xx_0, 0), (self.blocks_multiply_conjugate_cc_0, 0))

        self.connect((self.blocks_complex_to_mag_squared_0, 0),
                (self.blocks_moving_average_gen, 0))
        self.connect((self.blocks_moving_average_gen, 0),
                (self.fir_filter_gen, 0))
        self.connect((self.fir_filter_gen, 0), (self.blocks_interleave, 0))

        # Connect feedback
        self.connect((self.uhd_usrp_source_0, 0),
                (self.blocks_complex_to_mag_squared_1, 0))
        self.connect((self.uhd_usrp_source_0, 0),
                (self.blocks_multiply_conjugate_cc_0, 1))

        # Connect phase comparator
        self.connect((self.blocks_multiply_conjugate_cc_0, 0),
                (self.blocks_complex_to_arg_0, 0))
        self.connect((self.blocks_complex_to_arg_0, 0),
                (self.blocks_moving_average_phase, 0))
        self.connect((self.blocks_moving_average_phase, 0), (self.fir_filter_phase, 0))
        self.connect((self.fir_filter_phase, 0), (self.blocks_interleave, 1))

        self.connect((self.blocks_complex_to_mag_squared_1, 0),
                (self.blocks_moving_average_feedback, 0))
        self.connect((self.blocks_moving_average_feedback, 0),
                (self.fir_filter_feedback, 0))
        self.connect((self.fir_filter_feedback, 0),
                (self.blocks_interleave, 2))

        # connect interleaver output to TCP socket
        self.connect((self.blocks_interleave, 0), (self.blks2_tcp_sink_0, 0))

    def get_txgain(self):
        return self.txgain

    def set_txgain(self, txgain):
        self.txgain = txgain
        self.uhd_usrp_sink_0.set_gain(self.txgain, 0)

    def get_source_ampl(self):
        return self.source_ampl

    def set_source_ampl(self, source_ampl):
        print("Set amplitude to {}".format(source_ampl))
        self.source_ampl = source_ampl
        self.analog_sig_source_x_0.set_amplitude(self.source_ampl)
        self.analog_sig_source_x_1.set_amplitude(self.source_ampl)

    def get_freq(self):
        return self.freq

    def set_freq(self, freq):
        self.freq = freq
        self.uhd_usrp_sink_0.set_center_freq(self.freq, 0)
        self.uhd_usrp_source_0.set_center_freq(self.freq, 0)

class RampGenerator(threading.Thread):
    def __init__(self, options):
        threading.Thread.__init__(self)
        self.event_queue_ = Queue()
        self.in_queue_ = Queue()

        self.num_meas = int(options.num_meas)
        self.num_meas_to_skip = int(options.num_meas_to_skip)
        self.ampl_start = float(options.ampl_start)
        self.ampl_step = float(options.ampl_step)
        self.ampl_stop = float(options.ampl_stop)

    def set_source_ampl(self, ampl):
        self.event_queue_.put(ampl)
        self.in_queue_.get()

    def wait_on_event(self):
        return self.event_queue_.get()

    def confirm_source_ampl_updated(self):
        self.in_queue_.put(0)

    def run(self):
        try:
            self.run_ex()
        except:
            traceback.print_exc()
        finally:
            self.event_queue_.put("quit")

    def run_ex(self):
        print("Wait before connection")
        time.sleep(3)

        print("Connecting to flowgraph")
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.connect(("localhost", TCP_PORT))
        print("Connected")

        amplitudes = xrange(self.ampl_start, self.ampl_stop, self.ampl_step)
        measurements = []

        for ampl in amplitudes:
            self.set_source_ampl(ampl)

            mag_gen_sum = 0
            phase_diff_sum = 0
            mag_feedback_sum = 0

            for measurement_ignore in range(self.num_meas_to_skip):
                # Receive and ignore three floats on the socket
                sock.recv(12)

            for measurement_ix in range(self.num_meas):
                # Receive three floats on the socket
                mag_gen, phase_diff, mag_feedback = struct.unpack(
                        "fff",
                        sock.recv(12))

                mag_gen_sum += mag_gen
                phase_diff_sum += phase_diff
                mag_feedback_sum += mag_feedback

                measurements.append((ampl, mag_gen, mag_feedback, phase_diff))

            mag_gen_avg = mag_gen_sum / self.num_meas
            mag_feedback_avg = mag_feedback_sum / self.num_meas
            phase_diff_avg = phase_diff_sum / self.num_meas

            print("Ampl: {} Out: {:10} In: {:10} phase_diff: {:10}".format(
                ampl, mag_gen_avg, mag_feedback_avg, phase_diff_avg))

        self.event_queue_.put("done")
        self.event_queue_.put(measurements)

parser = argparse.ArgumentParser(description='Two-tone amplitude ramp')

parser.add_argument('--ampl-start',
        default='0.1',
        help='Start amplitude',
        required=False)

parser.add_argument('--ampl-stop',
        default='0.8',
        help='Stop amplitude',
        required=False)

parser.add_argument('--ampl-step',
        default='0.02',
        help='Amplitude steps',
        required=False)

parser.add_argument('--txgain',
        default='10',
        help='txgain for USRP sink',
        required=False)

parser.add_argument('--num-meas',
        default='2000',
        help='number of measurements per amplitude',
        required=False)

parser.add_argument('--num-meas-to-skip',
        default='50',
        help='After each amplitude change, ignore num-meas-to-skip measurements',
        required=False)

cli_args = parser.parse_args()

rampgen = RampGenerator(cli_args)
rampgen.start()

# this blocks until the flowgraph is up and running, i.e. all sockets
# got a connection
top = amplitude_ramp(cli_args)
top.set_source_ampl(float(cli_args.ampl_start))
top.start()

while True:
    event = rampgen.wait_on_event()
    if event == "done":
        measurements = rampgen.wait_on_event()
        fd = open("measurements.csv", "w")
        for m in measurements:
            fd.write(",".join("{}".format(x) for x in m) + "\n")
        fd.close()
        break
    elif event == "quit":
        break
    else:
        top.set_source_ampl(event)
        rampgen.confirm_source_ampl_updated()

top.stop()
print("Wait for completion")
top.wait()

print("Done")
<|endoftext|>


<commit_before>"""
Read/Write AMQP frames over network transports.

2009-01-14 Barry Pederson <bp@barryp.org>

"""
# Copyright (C) 2009 Barry Pederson <bp@barryp.org>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
from __future__ import absolute_import

import errno
import re
import socket

# Jython does not have this attribute
try:
    from socket import SOL_TCP
except ImportError:  # pragma: no cover
    from socket import IPPROTO_TCP as SOL_TCP  # noqa

#
# See if Python 2.6+ SSL support is available
#
try:
    import ssl
    HAVE_PY26_SSL = True
except:
    HAVE_PY26_SSL = False

try:
    bytes
except:
    # Python 2.5 and lower
    bytes = str

from struct import pack, unpack

from .exceptions import AMQPError

AMQP_PORT = 5672

EMPTY_BUFFER = bytes()

# Yes, Advanced Message Queuing Protocol Protocol is redundant
AMQP_PROTOCOL_HEADER = 'AMQP\x01\x01\x00\x09'.encode('latin_1')

# Match things like: [fe80::1]:5432, from RFC 2732
IPV6_LITERAL = re.compile(r'\[([\.0-9a-f:]+)\](?::(\d+))?')

class _AbstractTransport(object):
    """Common superclass for TCP and SSL transports"""

    def __init__(self, host, connect_timeout):
        msg = 'socket.getaddrinfo() for %s returned an empty list' % host
        port = AMQP_PORT

        m = IPV6_LITERAL.match(host)
        if m:
            host = m.group(1)
            if m.group(2):
                port = int(m.group(2))
        else:
            if ':' in host:
                host, port = host.rsplit(':', 1)
                port = int(port)

        self.sock = None
        last_err = None
        for res in socket.getaddrinfo(host, port, 0,
                                      socket.SOCK_STREAM, SOL_TCP):
            af, socktype, proto, canonname, sa = res
            try:
                self.sock = socket.socket(af, socktype, proto)
                self.sock.settimeout(connect_timeout)
                self.sock.connect(sa)
            except socket.error, msg:
                self.sock.close()
                self.sock = None
                last_err = msg
                continue
            break

        if not self.sock:
            # Didn't connect, return the most recent error message
            raise socket.error(last_err)

        self.sock.settimeout(None)
        self.sock.setsockopt(SOL_TCP, socket.TCP_NODELAY, 1)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

        self._setup_transport()

        self._write(AMQP_PROTOCOL_HEADER)

    def __del__(self):
        try:
            self.close()
        except socket.error:
            pass
        finally:
            self.sock = None

    def _read(self, n, initial=False):
        """Read exactly n bytes from the peer"""
        raise NotImplementedError('Must be overriden in subclass')

    def _setup_transport(self):
        """Do any additional initialization of the class (used
        by the subclasses)."""
        pass

    def _shutdown_transport(self):
        """Do any preliminary work in shutting down the connection."""
        pass

    def _write(self, s):
        """Completely write a string to the peer."""
        raise NotImplementedError('Must be overriden in subclass')

    def close(self):
        if self.sock is not None:
            self._shutdown_transport()
            # Call shutdown first to make sure that pending messages
            # reach the AMQP broker if the program exits after
            # calling this method.
            self.sock.shutdown(socket.SHUT_RDWR)
            self.sock.close()
            self.sock = None

    def read_frame(self):
        """Read an AMQP frame."""
        frame_type, channel, size = unpack('>BHI', self._read(7, True))
        payload = self._read(size)
        ch = ord(self._read(1))
        if ch == 206:  # '\xce'
            return frame_type, channel, payload
        else:
            raise AMQPError(
                'Framing Error, received 0x%02x while expecting 0xce' % ch)

    def write_frame(self, frame_type, channel, payload):
        """Write out an AMQP frame."""
        size = len(payload)
        self._write(
            pack('>BHI%dsB' % size, frame_type, channel, size, payload, 0xce),
        )

class SSLTransport(_AbstractTransport):
    """Transport that works over SSL"""

    def __init__(self, host, connect_timeout, ssl):
        if isinstance(ssl, dict):
            self.sslopts = ssl
        self.sslobj = None
        super(SSLTransport, self).__init__(host, connect_timeout)

    def _setup_transport(self):
        """Wrap the socket in an SSL object, either the
        new Python 2.6 version, or the older Python 2.5 and
        lower version."""
        if HAVE_PY26_SSL:
            if hasattr(self, 'sslopts'):
                self.sslobj = ssl.wrap_socket(self.sock, **self.sslopts)
            else:
                self.sslobj = ssl.wrap_socket(self.sock)
            self.sslobj.do_handshake()
        else:
            self.sslobj = socket.ssl(self.sock)

    def _shutdown_transport(self):
        """Unwrap a Python 2.6 SSL socket, so we can call shutdown()"""
        if HAVE_PY26_SSL and (self.sslobj is not None):
            self.sock = self.sslobj.unwrap()
            self.sslobj = None

    def _read(self, n, initial=False):
        """It seems that SSL Objects read() method may not supply as much
        as you're asking for, at least with extremely large messages.
        somewhere > 16K - found this in the test_channel.py test_large
        unittest."""
        result = EMPTY_BUFFER

        while len(result) < n:
            try:
                s = self.sslobj.read(n - len(result))
            except socket.error, exc:
                if not initial and exc.errno in (errno.EAGAIN, errno.EINTR):
                    continue
                raise
            if not s:
                raise IOError('Socket closed')
            result += s

        return result

    def _write(self, s):
        """Write a string out to the SSL socket fully."""
        while s:
            n = self.sslobj.write(s)
            if not n:
                raise IOError('Socket closed')
            s = s[n:]

class TCPTransport(_AbstractTransport):
    """Transport that deals directly with TCP socket."""

    def _setup_transport(self):
        """Setup to _write() directly to the socket, and
        do our own buffered reads."""
        self._write = self.sock.sendall
        self._read_buffer = bytes()

    def _read(self, n, initial=False):
        """Read exactly n bytes from the socket"""
        while len(self._read_buffer) < n:
            try:
                s = self.sock.recv(65536)
            except socket.error, exc:
                if not initial and exc.errno in (errno.EAGAIN, errno.EINTR):
                    continue
                raise
            if not s:
                raise IOError('Socket closed')
            self._read_buffer += s

        result = self._read_buffer[:n]
        self._read_buffer = self._read_buffer[n:]

        return result

def create_transport(host, connect_timeout, ssl=False):
    """Given a few parameters from the Connection constructor,
    select and create a subclass of _AbstractTransport."""
    if ssl:
        return SSLTransport(host, connect_timeout, ssl)
    else:
        return TCPTransport(host, connect_timeout)
<commit_msg>Fix AMQP SSL support<commit_after>"""
Read/Write AMQP frames over network transports.

2009-01-14 Barry Pederson <bp@barryp.org>

"""
# Copyright (C) 2009 Barry Pederson <bp@barryp.org>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
from __future__ import absolute_import

import errno
import re
import socket

# Jython does not have this attribute
try:
    from socket import SOL_TCP
except ImportError:  # pragma: no cover
    from socket import IPPROTO_TCP as SOL_TCP  # noqa

#
# See if Python 2.6+ SSL support is available
#
try:
    import ssl
    HAVE_PY26_SSL = True
except:
    HAVE_PY26_SSL = False

try:
    bytes
except:
    # Python 2.5 and lower
    bytes = str

from struct import pack, unpack

from .exceptions import AMQPError

AMQP_PORT = 5672

EMPTY_BUFFER = bytes()

# Yes, Advanced Message Queuing Protocol Protocol is redundant
AMQP_PROTOCOL_HEADER = 'AMQP\x01\x01\x00\x09'.encode('latin_1')

# Match things like: [fe80::1]:5432, from RFC 2732
IPV6_LITERAL = re.compile(r'\[([\.0-9a-f:]+)\](?::(\d+))?')

class _AbstractTransport(object):
    """Common superclass for TCP and SSL transports"""

    def __init__(self, host, connect_timeout):
        msg = 'socket.getaddrinfo() for %s returned an empty list' % host
        port = AMQP_PORT

        m = IPV6_LITERAL.match(host)
        if m:
            host = m.group(1)
            if m.group(2):
                port = int(m.group(2))
        else:
            if ':' in host:
                host, port = host.rsplit(':', 1)
                port = int(port)

        self._sock = None
        last_err = None
        for res in socket.getaddrinfo(host, port, 0,
                                      socket.SOCK_STREAM, SOL_TCP):
            af, socktype, proto, canonname, sa = res
            try:
                self._sock = socket.socket(af, socktype, proto)
                self._sock.settimeout(connect_timeout)
                self._sock.connect(sa)
            except socket.error, msg:
                self._sock.close()
                self._sock = None
                last_err = msg
                continue
            break

        if not self._sock:
            # Didn't connect, return the most recent error message
            raise socket.error(last_err)

        self._sock.settimeout(None)
        self._sock.setsockopt(SOL_TCP, socket.TCP_NODELAY, 1)
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

        self._setup_transport()

        self._write(AMQP_PROTOCOL_HEADER)

    def __del__(self):
        try:
            self.close()
        except socket.error:
            pass
        finally:
            self._sock = None

    def _read(self, n, initial=False):
        """Read exactly n bytes from the peer"""
        raise NotImplementedError('Must be overriden in subclass')

    def _setup_transport(self):
        """Do any additional initialization of the class (used
        by the subclasses)."""
        pass

    def _shutdown_transport(self):
        """Do any preliminary work in shutting down the connection."""
        pass

    def _write(self, s):
        """Completely write a string to the peer."""
        raise NotImplementedError('Must be overriden in subclass')

    def close(self):
        if self._sock is not None:
            self._shutdown_transport()
            # Call shutdown first to make sure that pending messages
            # reach the AMQP broker if the program exits after
            # calling this method.
            self._sock.shutdown(socket.SHUT_RDWR)
            self._sock.close()
            self._sock = None

    def read_frame(self):
        """Read an AMQP frame."""
        frame_type, channel, size = unpack('>BHI', self._read(7, True))
        payload = self._read(size)
        ch = ord(self._read(1))
        if ch == 206:  # '\xce'
            return frame_type, channel, payload
        else:
            raise AMQPError(
                'Framing Error, received 0x%02x while expecting 0xce' % ch)

    def write_frame(self, frame_type, channel, payload):
        """Write out an AMQP frame."""
        size = len(payload)
        self._write(
            pack('>BHI%dsB' % size, frame_type, channel, size, payload, 0xce),
        )

    @property
    def sock(self):
        return self._sock

    @sock.setter
    def sock(self, v):
        self._sock = v

class SSLTransport(_AbstractTransport):
    """Transport that works over SSL"""

    def __init__(self, host, connect_timeout, ssl):
        if isinstance(ssl, dict):
            self.sslopts = ssl
        self.sslobj = None
        super(SSLTransport, self).__init__(host, connect_timeout)

    def _setup_transport(self):
        """Wrap the socket in an SSL object, either the
        new Python 2.6 version, or the older Python 2.5 and
        lower version."""
        if HAVE_PY26_SSL:
            if hasattr(self, 'sslopts'):
                self.sslobj = ssl.wrap_socket(self._sock, **self.sslopts)
            else:
                self.sslobj = ssl.wrap_socket(self._sock)
            self.sslobj.do_handshake()
        else:
            self.sslobj = socket.ssl(self._sock)

    def _shutdown_transport(self):
        """Unwrap a Python 2.6 SSL socket, so we can call shutdown()"""
        if HAVE_PY26_SSL and (self.sslobj is not None):
            self._sock = self.sslobj.unwrap()
            self.sslobj = None

    def _read(self, n, initial=False):
        """It seems that SSL Objects read() method may not supply as much
        as you're asking for, at least with extremely large messages.
        somewhere > 16K - found this in the test_channel.py test_large
        unittest."""
        result = EMPTY_BUFFER

        while len(result) < n:
            try:
                s = self.sslobj.read(n - len(result))
            except socket.error, exc:
                if not initial and exc.errno in (errno.EAGAIN, errno.EINTR):
                    continue
                raise
            if not s:
                raise IOError('Socket closed')
            result += s

        return result

    def _write(self, s):
        """Write a string out to the SSL socket fully."""
        while s:
            n = self.sslobj.write(s)
            if not n:
                raise IOError('Socket closed')
            s = s[n:]

    @property
    def sock(self):
        return self.sslobj

class TCPTransport(_AbstractTransport):
    """Transport that deals directly with TCP socket."""

    def _setup_transport(self):
        """Setup to _write() directly to the socket, and
        do our own buffered reads."""
        self._write = self._sock.sendall
        self._read_buffer = bytes()

    def _read(self, n, initial=False):
        """Read exactly n bytes from the socket"""
        while len(self._read_buffer) < n:
            try:
                s = self._sock.recv(65536)
            except socket.error, exc:
                if not initial and exc.errno in (errno.EAGAIN, errno.EINTR):
                    continue
                raise
            if not s:
                raise IOError('Socket closed')
            self._read_buffer += s

        result = self._read_buffer[:n]
        self._read_buffer = self._read_buffer[n:]

        return result

def create_transport(host, connect_timeout, ssl=False):
    """Given a few parameters from the Connection constructor,
    select and create a subclass of _AbstractTransport."""
    if ssl:
        return SSLTransport(host, connect_timeout, ssl)
    else:
        return TCPTransport(host, connect_timeout)
<|endoftext|>


<commit_before># Create your tasks here
from __future__ import absolute_import, unicode_literals
from celery import shared_task

from tasks.models import Task
from .models import Analysis
from samples.models import Sample, SampleGroup

@shared_task
def create_analysis_tasks(analysis_id):
    print('analysis_id', analysis_id)
    # samples = []
    # print('hello!')
    #create analysis tasks
    # analysis = Analysis.objects.()
    analysis = Analysis.objects.get(pk=analysis_id)
    print(dir(analysis))
    params = analysis.params

    # if 'sample_groups' in  params:
    #     samples = Sample.objects.filter(samplegroup_members__in=params['sample_groups'])

    #     # sample = Sample.objects.first()
    #     # print(dir(sample))
    # for sample in samples:
    #     print(sample)
    #     for file in sample.files.all():
    #         bam_size = 9223372036854775807
    #         if file.extension == 'bam':
    #             if file.size < bam_size:
    #                 bamfile = file
    #                 bam_size = file.size
    #     print('small bam', bamfile.size)
        #get smallest bam file
    task = Task(user=analysis.user)
    task.manifest = {}
    task.manifest['analyses'] = params['analysis_types']
    task.manifest['files'] = params['files']
    task.status = 'new'
    task.action = 'analysis'
    task.save()
    analysis.tasks.add(task)<commit_msg>checkpoint<commit_after># Create your tasks here
from __future__ import absolute_import, unicode_literals
from celery import shared_task

from tasks.models import Task
from .models import Analysis
from samples.models import Sample, SampleGroup

@shared_task
def create_analysis_tasks(analysis_id):
    print('analysis_id', analysis_id)
    # samples = []
    # print('hello!')
    #create analysis tasks
    # analysis = Analysis.objects.()
    analysis = Analysis.objects.get(pk=analysis_id)
    print(dir(analysis))
    params = analysis.params

    # if 'sample_groups' in  params:
    #     samples = Sample.objects.filter(samplegroup_members__in=params['sample_groups'])

    #     # sample = Sample.objects.first()
    #     # print(dir(sample))
    # for sample in samples:
    #     print(sample)
    #     for file in sample.files.all():
    #         bam_size = 9223372036854775807
    #         if file.extension == 'bam':
    #             if file.size < bam_size:
    #                 bamfile = file
    #                 bam_size = file.size
    #     print('small bam', bamfile.size)
        #get smallest bam file
    task = Task(user=analysis.user)
    task.manifest = {}
    task.manifest = params
    task.status = 'new'
    task.action = 'analysis'
    task.save()
    analysis.tasks.add(task)<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from south.utils import datetime_utils as datetime
from south.db import db
from south.v2 import SchemaMigration
from django.db import models

class Migration(SchemaMigration):

    def forwards(self, orm):
        # Adding model 'RegistrationProfile'
        db.create_table(u'registration_registrationprofile', (
            (u'id', self.gf('django.db.models.fields.AutoField')(primary_key=True)),
            ('user', self.gf('django.db.models.fields.related.OneToOneField')(to=orm['user.User'], unique=True)),
            ('activation_key', self.gf('django.db.models.fields.CharField')(max_length=40)),
        ))
        db.send_create_signal(u'registration', ['RegistrationProfile'])

    def backwards(self, orm):
        # Deleting model 'RegistrationProfile'
        db.delete_table(u'registration_registrationprofile')

    models = {
        u'auth.group': {
            'Meta': {'object_name': 'Group'},
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'unique': 'True', 'max_length': '80'}),
            'permissions': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['auth.Permission']", 'symmetrical': 'False', 'blank': 'True'})
        },
        u'auth.permission': {
            'Meta': {'ordering': "(u'content_type__app_label', u'content_type__model', u'codename')", 'unique_together': "((u'content_type', u'codename'),)", 'object_name': 'Permission'},
            'codename': ('django.db.models.fields.CharField', [], {'max_length': '100'}),
            'content_type': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['contenttypes.ContentType']"}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '50'})
        },
        u'contenttypes.contenttype': {
            'Meta': {'ordering': "('name',)", 'unique_together': "(('app_label', 'model'),)", 'object_name': 'ContentType', 'db_table': "'django_content_type'"},
            'app_label': ('django.db.models.fields.CharField', [], {'max_length': '100'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'model': ('django.db.models.fields.CharField', [], {'max_length': '100'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '100'})
        },
        u'core.accommodation': {
            'Meta': {'object_name': 'Accommodation'},
            'address': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'address_complement': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'address_neighborhood': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'address_number': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'city': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.City']"}),
            'email': ('django.db.models.fields.EmailField', [], {'max_length': '75', 'blank': 'True'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'image_url': ('django.db.models.fields.URLField', [], {'max_length': '200', 'blank': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'telephone': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'text': ('django.db.models.fields.TextField', [], {}),
            'the_geom': ('django.contrib.gis.db.models.fields.PointField', [], {}),
            'type': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.AccommodationType']"}),
            'website': ('django.db.models.fields.URLField', [], {'max_length': '200', 'blank': 'True'})
        },
        u'core.accommodationtype': {
            'Meta': {'ordering': "(u'name',)", 'object_name': 'AccommodationType'},
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'})
        },
        u'core.attraction': {
            'Meta': {'ordering': "(u'name',)", 'object_name': 'Attraction'},
            'city': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.City']"}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'image_url': ('django.db.models.fields.URLField', [], {'max_length': '200', 'blank': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'text': ('django.db.models.fields.TextField', [], {}),
            'the_geom': ('django.contrib.gis.db.models.fields.PointField', [], {}),
            'type': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.AttractionType']"})
        },
        u'core.attractiontype': {
            'Meta': {'ordering': "(u'name',)", 'object_name': 'AttractionType'},
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'})
        },
        u'core.city': {
            'Meta': {'ordering': "(u'name',)", 'object_name': 'City'},
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'image_url': ('django.db.models.fields.URLField', [], {'max_length': '200', 'blank': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'og_description': ('django.db.models.fields.TextField', [], {'null': 'True', 'blank': 'True'}),
            'og_image': ('django.db.models.fields.files.ImageField', [], {'max_length': '100', 'null': 'True', 'blank': 'True'}),
            'og_title': ('django.db.models.fields.CharField', [], {'max_length': '100', 'null': 'True', 'blank': 'True'}),
            'slug': ('django.db.models.fields.SlugField', [], {'unique': 'True', 'max_length': '255'}),
            'text': ('django.db.models.fields.TextField', [], {}),
            'the_geom': ('django.contrib.gis.db.models.fields.PointField', [], {})
        },
        u'core.feeding': {
            'Meta': {'object_name': 'Feeding'},
            'address': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'address_complement': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'address_neighborhood': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'address_number': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'city': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.City']"}),
            'email': ('django.db.models.fields.EmailField', [], {'max_length': '75', 'blank': 'True'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'image_url': ('django.db.models.fields.URLField', [], {'max_length': '200', 'blank': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'telephone': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'text': ('django.db.models.fields.TextField', [], {}),
            'the_geom': ('django.contrib.gis.db.models.fields.PointField', [], {}),
            'type': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.FeedingType']"}),
            'website': ('django.db.models.fields.URLField', [], {'max_length': '200', 'blank': 'True'})
        },
        u'core.feedingtype': {
            'Meta': {'ordering': "(u'name',)", 'object_name': 'FeedingType'},
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'})
        },
        u'core.roadmap': {
            'Meta': {'ordering': "(u'-modified',)", 'object_name': 'Roadmap'},
            'cities': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['core.City']", 'symmetrical': 'False', 'through': u"orm['core.RoadmapCity']", 'blank': 'True'}),
            'created': ('django.db.models.fields.DateTimeField', [], {'auto_now_add': 'True', 'blank': 'True'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'modified': ('django.db.models.fields.DateTimeField', [], {'auto_now': 'True', 'blank': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'og_description': ('django.db.models.fields.TextField', [], {'null': 'True', 'blank': 'True'}),
            'og_image': ('django.db.models.fields.files.ImageField', [], {'max_length': '100', 'null': 'True', 'blank': 'True'}),
            'og_title': ('django.db.models.fields.CharField', [], {'max_length': '100', 'null': 'True', 'blank': 'True'}),
            'public': ('django.db.models.fields.BooleanField', [], {'default': 'False'}),
            'user': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['user.User']"})
        },
        u'core.roadmapcity': {
            'Meta': {'unique_together': "((u'roadmap', u'city'),)", 'object_name': 'RoadmapCity'},
            'accommodations': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['core.Accommodation']", 'symmetrical': 'False', 'through': u"orm['core.RoadmapCityAccommodation']", 'blank': 'True'}),
            'attractions': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['core.Attraction']", 'symmetrical': 'False', 'through': u"orm['core.RoadmapCityAttraction']", 'blank': 'True'}),
            'city': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.City']"}),
            'feedings': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['core.Feeding']", 'symmetrical': 'False', 'through': u"orm['core.RoadmapCityFeeding']", 'blank': 'True'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'roadmap': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.Roadmap']"})
        },
        u'core.roadmapcityaccommodation': {
            'Meta': {'unique_together': "((u'roadmapcity', u'accommodation'),)", 'object_name': 'RoadmapCityAccommodation'},
            'accommodation': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.Accommodation']"}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'roadmapcity': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.RoadmapCity']"})
        },
        u'core.roadmapcityattraction': {
            'Meta': {'unique_together': "((u'roadmapcity', u'attraction'),)", 'object_name': 'RoadmapCityAttraction'},
            'attraction': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.Attraction']"}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'roadmapcity': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.RoadmapCity']"})
        },
        u'core.roadmapcityfeeding': {
            'Meta': {'unique_together': "((u'roadmapcity', u'feeding'),)", 'object_name': 'RoadmapCityFeeding', 'db_table': "u'roadmap_city_feeding'"},
            'feeding': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.Feeding']"}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'roadmapcity': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['core.RoadmapCity']"})
        },
        u'registration.registrationprofile': {
            'Meta': {'object_name': 'RegistrationProfile'},
            'activation_key': ('django.db.models.fields.CharField', [], {'max_length': '40'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'user': ('django.db.models.fields.related.OneToOneField', [], {'to': u"orm['user.User']", 'unique': 'True'})
        },
        u'user.user': {
            'Meta': {'ordering': "(u'name', u'email')", 'unique_together': "((u'email',),)", 'object_name': 'User'},
            'access_count': ('django.db.models.fields.PositiveIntegerField', [], {'default': '0'}),
            'avatar': ('django.db.models.fields.files.ImageField', [], {'max_length': '100', 'null': 'True', 'blank': 'True'}),
            'birthday': ('django.db.models.fields.CharField', [], {'max_length': '10'}),
            'body': ('django.db.models.fields.TextField', [], {}),
            'city': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'country': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'date_joined': ('django.db.models.fields.DateTimeField', [], {'default': 'datetime.datetime.now'}),
            'email': ('django.db.models.fields.EmailField', [], {'unique': 'True', 'max_length': '255', 'db_index': 'True'}),
            'favorites_roadmaps': ('django.db.models.fields.related.ManyToManyField', [], {'related_name': "u'sgusers_favorited'", 'symmetrical': 'False', 'to': u"orm['core.Roadmap']"}),
            'gender': ('django.db.models.fields.CharField', [], {'max_length': '10'}),
            'groups': ('django.db.models.fields.related.ManyToManyField', [], {'symmetrical': 'False', 'related_name': "u'sguser_set'", 'blank': 'True', 'to': u"orm['auth.Group']"}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'is_active': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
            'is_staff': ('django.db.models.fields.BooleanField', [], {'default': 'False'}),
            'is_superuser': ('django.db.models.fields.BooleanField', [], {'default': 'False'}),
            'last_access': ('django.db.models.fields.DateTimeField', [], {'auto_now_add': 'True', 'blank': 'True'}),
            'last_login': ('django.db.models.fields.DateTimeField', [], {'default': 'datetime.datetime.now'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '255', 'null': 'True'}),
            'password': ('django.db.models.fields.CharField', [], {'max_length': '128'}),
            'photo': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'photo_dir': ('django.db.models.fields.CharField', [], {'max_length': '255', 'blank': 'True'}),
            'region': ('django.db.models.fields.CharField', [], {'max_length': '255'}),
            'user_permissions': ('django.db.models.fields.related.ManyToManyField', [], {'symmetrical': 'False', 'related_name': "u'sguser_set'", 'blank': 'True', 'to': u"orm['auth.Permission']"})
        }
    }

    complete_apps = ['registration']<commit_msg>Fix south migrations, previous checkin had models that don't exist in the django tree nor in django-registration-redux.<commit_after># -*- coding: utf-8 -*-
import datetime
from south.db import db
from south.v2 import SchemaMigration
from django.db import models

class Migration(SchemaMigration):

    def forwards(self, orm):
        # Adding model 'RegistrationProfile'
        db.create_table(u'registration_registrationprofile', (
            (u'id', self.gf('django.db.models.fields.AutoField')(primary_key=True)),
            ('user', self.gf('django.db.models.fields.related.OneToOneField')(to=orm['auth.User'], unique=True)),
            ('activation_key', self.gf('django.db.models.fields.CharField')(max_length=40)),
        ))
        db.send_create_signal(u'registration', ['RegistrationProfile'])

    def backwards(self, orm):
        # Deleting model 'RegistrationProfile'
        db.delete_table(u'registration_registrationprofile')

    models = {
        u'auth.group': {
            'Meta': {'object_name': 'Group'},
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'unique': 'True', 'max_length': '80'}),
            'permissions': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['auth.Permission']", 'symmetrical': 'False', 'blank': 'True'})
        },
        u'auth.permission': {
            'Meta': {'ordering': "(u'content_type__app_label', u'content_type__model', u'codename')", 'unique_together': "((u'content_type', u'codename'),)", 'object_name': 'Permission'},
            'codename': ('django.db.models.fields.CharField', [], {'max_length': '100'}),
            'content_type': ('django.db.models.fields.related.ForeignKey', [], {'to': u"orm['contenttypes.ContentType']"}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '50'})
        },
        u'auth.user': {
            'Meta': {'object_name': 'User'},
            'date_joined': ('django.db.models.fields.DateTimeField', [], {'default': 'datetime.datetime.now'}),
            'email': ('django.db.models.fields.EmailField', [], {'max_length': '75', 'blank': 'True'}),
            'first_name': ('django.db.models.fields.CharField', [], {'max_length': '30', 'blank': 'True'}),
            'groups': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['auth.Group']", 'symmetrical': 'False', 'blank': 'True'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'is_active': ('django.db.models.fields.BooleanField', [], {'default': 'True'}),
            'is_staff': ('django.db.models.fields.BooleanField', [], {'default': 'False'}),
            'is_superuser': ('django.db.models.fields.BooleanField', [], {'default': 'False'}),
            'last_login': ('django.db.models.fields.DateTimeField', [], {'default': 'datetime.datetime.now'}),
            'last_name': ('django.db.models.fields.CharField', [], {'max_length': '30', 'blank': 'True'}),
            'password': ('django.db.models.fields.CharField', [], {'max_length': '128'}),
            'user_permissions': ('django.db.models.fields.related.ManyToManyField', [], {'to': u"orm['auth.Permission']", 'symmetrical': 'False', 'blank': 'True'}),
            'username': ('django.db.models.fields.CharField', [], {'unique': 'True', 'max_length': '30'})
        },
        u'contenttypes.contenttype': {
            'Meta': {'ordering': "('name',)", 'unique_together': "(('app_label', 'model'),)", 'object_name': 'ContentType', 'db_table': "'django_content_type'"},
            'app_label': ('django.db.models.fields.CharField', [], {'max_length': '100'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'model': ('django.db.models.fields.CharField', [], {'max_length': '100'}),
            'name': ('django.db.models.fields.CharField', [], {'max_length': '100'})
        },
        u'registration.registrationprofile': {
            'Meta': {'object_name': 'RegistrationProfile'},
            'activation_key': ('django.db.models.fields.CharField', [], {'max_length': '40'}),
            u'id': ('django.db.models.fields.AutoField', [], {'primary_key': 'True'}),
            'user': ('django.db.models.fields.related.OneToOneField', [], {'to': u"orm['auth.User']", 'unique': 'True'})
        }
    }

    complete_apps = ['registration']<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division

""" Convert 2dfdr reduced images to extracted spectra in GALAH FITS format. """

__author__ = "Andy Casey <arc@ast.cam.ac.uk>"

__all__ = ["from_2dfdr"]

from subprocess import check_output

import numpy as np

from astropy.io import fits
from astropy.constants import c as speed_of_light
import motions

# [TODO] Read the GAP version from elsewhere.
WG6_VER = 6.4
WG6_GIT_HASH = check_output("git rev-parse --short HEAD".split()).strip()

def get_ccd_number(image):
    ccd_name = image[0].header.comments["SPLYTEMP"].split(" ")[0]
    ccd_number = {
        "BLUE": 1,
        "GREEN": 2,
        "RED": 3,
        "IR": 4
    }[ccd_name]

    return (ccd_number, ccd_name)

def verify_hermes_origin(image):
    assert image[0].header["ORIGIN"].strip() == "AAO"
    assert image[0].header["INSTRUME"].strip() == "HERMES-2dF"

def read_2dfdr_extensions(image):
    """
    Parse the image extensions from a 2dfdr-reduced image.

    :param image:
        The 2dfdr-reduced image.

    :type image:
        :class:`astropy.io.fits.hdu.hdulist.HDUList`
    """

    extensions = {
        "data": 0,
    }
    extnames = map(str.strip, [hdu.header.get("EXTNAME", "") for hdu in image])
    extensions["variance"] = extnames.index("VARIANCE")
    extensions["fibres"] = extnames.index("FIBRES")
    return extensions

def dummy_normalisation_hdus(hdulist=None):
    """
    Return dummy HDUs for the normalised flux and associated uncertainty.
    """

    # Return a dummy normalisation HDU
    hdu_normed_flux = fits.ImageHDU(data=None, header=None,
        do_not_scale_image_data=True)

    # Add header information
    if hdulist is not None:
        for key in ("CRVAL1", "CDELT1", "CRPIX1", "CTYPE1", "CUNIT1"):
            hdu_normed_flux.header[key] = hdulist[0].header[key]
            hdu_normed_flux.header.comments[key] = hdulist[0].header.comments[key]

    hdu_normed_flux.header["EXTNAME"] = "normalised_spectrum"
    hdu_normed_flux.header.comments["EXTNAME"] = "Normalised spectrum flux"
    hdu_normed_flux.add_checksum()

    hdu_normed_sigma = fits.ImageHDU(data=None, header=None,
        do_not_scale_image_data=True)

    # Add header information
    if hdulist is not None:
        for key in ("CRVAL1", "CDELT1", "CRPIX1", "CTYPE1", "CUNIT1"):
            hdu_normed_sigma.header[key] = hdulist[0].header[key]
            hdu_normed_sigma.header.comments[key] = hdulist[0].header.comments[key]

    hdu_normed_sigma.header["EXTNAME"] = "normalised_sigma"
    hdu_normed_sigma.header.comments["EXTNAME"] = "Normalised flux sigma"
    hdu_normed_sigma.add_checksum()

    return [hdu_normed_flux, hdu_normed_sigma]

def dummy_ccf_hdu(hdulist=None):
    """
    Return a dummy HDU for the best-fitting cross-correlation function.
    """

    hdu_ccf = fits.ImageHDU(data=None, header=None,
        do_not_scale_image_data=True)

    hdu_ccf.header["EXTNAME"] = "CCF"
    hdu_ccf.header.comments["EXTNAME"] = "CCF from best-fitting template"

    # Add empty default values.
    default_headers = [
        ("VRAD", "Radial velocity (km/s)"),
        ("U_VRAD", "Uncertainty on radial velocity (km/s)"),
        ("TEFF", "Effective temperature"),
        ("LOGG", "Surface gravity"),
        ("FE_H", "Metallicity ([Fe/H])"),
        ("ALPHA_FE", "Alpha-enhancement ([alpha/Fe])")
    ]
    for key, comment in default_headers:
        hdu_ccf.header[key] = "NaN"
        hdu_ccf.header.comments[key] = comment

    hdu_ccf.add_checksum()

    return hdu_ccf

def from_2dfdr(reduced_filename, dummy_hdus=True):
    """
    Returns a list of `astropy.io.fits.hdu.hdulist.HDUList` objects
    (1 HDU List per 2dfdr program object).

    :param reduced_filename:
        The path of the 2dfdr-reduced combined filename.

    :type reduced_filename:
        str
    """

    image = fits.open(reduced_filename)
    ext = read_2dfdr_extensions(image)

    # Verifications
    verify_hermes_origin(image)

    # Header modifications.
    keywords = ["NAME", "RA", "DEC", "PMRA", "PMDEC", ("MAGNITUDE", "MAG"),
        ("COMMENT", "DESCR")]
    keyword_comments = {
        "NAME": "Input source name",
        "DESCR": "Input source comment",
        "MAG": "Input source magnitude",
        "RA": "Right Ascension (degrees)",
        "DEC": "Declination (degrees)",
        "PMRA": "Proper motion in RA (mas/yr)",
        "PMDEC": "Proper motion in DEC (mas/yr)",
        "FIBRE": "Fibre number"
    }
    header_template = image[ext["data"]].header.copy()
    for column in ("CDELT2", "CRPIX2", "CRVAL2", "CTYPE2", "CUNIT2"):
        del header_template[column]

    ccd_number, ccd_name = get_ccd_number(image)
    header_template["CCD"] = ccd_number
    header_template.comments["CCD"] = "{0} camera".format(ccd_name)
    header_template["WG6_VER"] = WG6_VER
    header_template["WG6_HASH"] = WG6_GIT_HASH
    header_template.comments["WG6_VER"] = "WG6 standardisation code version"
    header_template.comments["WG6_HASH"] = "WG6 standardisation commit hash"

    extracted_sources = []
    for program_index in np.where(image[ext["fibres"]].data["TYPE"] == "P")[0]:

        flux = image[ext["data"]].data[program_index, :]
        variance = image[ext["variance"]].data[program_index, :]
        header = header_template.copy()

        # Add header information from the fibre table.
        header["FIBRE"] = program_index + 1
        fibre_header = image[ext["fibres"]].data[program_index]
        for keyword in keywords:
            if isinstance(keyword, (str, unicode)):
                header[keyword] = fibre_header[keyword]
            else:
                from_keyword, to_keyword = keyword
                header[to_keyword] = fibre_header[from_keyword]

        # The 2dfdr fibre table has the RA and DEC in radians. Who does that?!
        if "RA" in keywords and "DEC" in keywords:
            header["RA"] *= 180./np.pi
            header["DEC"] *= 180./np.pi

        # Add associated comments for the fibre table information
        for keyword, comment in keyword_comments.items():
            header.comments[keyword] = comment

        # Create the HDUList.
        hdu_flux = fits.PrimaryHDU(data=flux, header=header,
            do_not_scale_image_data=True)
        hdu_flux.header["EXTNAME"] = "input_spectrum"
        hdu_flux.header.comments["EXTNAME"] = "Spectrum flux"

        hdu_sigma = fits.ImageHDU(data=variance**0.5, header=None,
            do_not_scale_image_data=True)

        for key in ("CRVAL1", "CDELT1", "CRPIX1", "CTYPE1", "CUNIT1"):
            hdu_sigma.header[key] = hdu_flux.header[key]
            hdu_sigma.header.comments[key] = hdu_flux.header.comments[key]
        hdu_sigma.header["EXTNAME"] = "input_sigma"
        hdu_sigma.header.comments["EXTNAME"] = "Flux sigma"

        # Calculate the barycentric motion.
        v_bary, v_helio = motions.from_header(header)

        hdu_flux.header["V_BARY"] = v_bary.to("km/s").value
        hdu_flux.header["V_HELIO"] = v_helio.to("km/s").value
        hdu_flux.header.comments["V_BARY"] = "Barycentric motion (km/s)"
        hdu_flux.header.comments["V_HELIO"] = "Heliocentric motion (km/s)"
        hdu_flux.header["HISTORY"] = "Corrected for barycentric motion (V_BARY)"

        # Add motion correction information.
        for hdu in (hdu_flux, hdu_sigma):
            hdu.header["CRVAL1"] *= 1. + (v_bary/speed_of_light).value
            hdu.header["CDELT1"] *= 1. + (v_bary/speed_of_light).value

        # Create HDUList and add checksums.
        hdulist = fits.HDUList([hdu_flux, hdu_sigma])
        [hdu.add_checksum() for hdu in hdulist]

        # Add dummy extensions
        if dummy_hdus:
            hdulist.extend(dummy_normalisation_hdus(hdulist))
            hdulist.append(dummy_ccf_hdu(hdulist))

        extracted_sources.append(hdulist)

    image.close()

    return extracted_sources

<commit_msg>Add median sky HDU<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division

""" Convert 2dfdr reduced images to extracted spectra in GALAH FITS format. """

__author__ = "Andy Casey <arc@ast.cam.ac.uk>"

__all__ = ["from_2dfdr"]

from subprocess import check_output

import numpy as np

from astropy.io import fits
from astropy.constants import c as speed_of_light
import motions

WG6_GIT_HASH = check_output("git rev-parse --short HEAD".split()).strip()

def get_ccd_number(image):
    crval1 = image[0].header["CRVAL1"]
    #4801, 5760, 6610, 7741
    if crval1 > 7000:
        ccd_number, ccd_name = 4, "IR"
    elif crval1 > 6000:
        ccd_number, ccd_name = 3, "RED",
    elif crval1 > 5000:
        ccd_number, ccd_name = 2, "GREEN",
    elif crval1 > 4000:
        ccd_number, ccd_name = 1, "BLUE"
    else:
        raise ValueError("CRVAL1 value is weird: {0}".format(crval1))
    return (ccd_number, ccd_name)

def verify_hermes_origin(image):
    assert image[0].header["ORIGIN"].strip() == "AAO"
    assert image[0].header["INSTRUME"].strip() == "HERMES-2dF"

def read_2dfdr_extensions(image):
    """
    Parse the image extensions from a 2dfdr-reduced image.

    :param image:
        The 2dfdr-reduced image.

    :type image:
        :class:`astropy.io.fits.hdu.hdulist.HDUList`
    """

    extensions = {
        "data": 0,
    }
    extnames = map(str.strip, [hdu.header.get("EXTNAME", "") for hdu in image])
    extensions["variance"] = extnames.index("VARIANCE")
    extensions["fibres"] = extnames.index("FIBRES")
    return extensions

def dummy_normalisation_hdus(hdulist=None):
    """
    Return dummy HDUs for the normalised flux and associated uncertainty.
    """

    # Return a dummy normalisation HDU
    hdu_normed_flux = fits.ImageHDU(data=None, header=None,
        do_not_scale_image_data=True)

    # Add header information
    if hdulist is not None:
        for k in ("CRVAL1", "CDELT1", "CRPIX1", "CTYPE1", "CUNIT1"):
            hdu_normed_flux.header[k] = hdulist[0].header[k]
            hdu_normed_flux.header.comments[k] = hdulist[0].header.comments[k]

    hdu_normed_flux.header["EXTNAME"] = "normalised_spectrum"
    hdu_normed_flux.header.comments["EXTNAME"] = "Normalised spectrum flux"
    hdu_normed_flux.add_checksum()

    hdu_normed_sigma = fits.ImageHDU(data=None, header=None,
        do_not_scale_image_data=True)

    # Add header information
    if hdulist is not None:
        for k in ("CRVAL1", "CDELT1", "CRPIX1", "CTYPE1", "CUNIT1"):
            hdu_normed_sigma.header[k] = hdulist[0].header[k]
            hdu_normed_sigma.header.comments[k] = hdulist[0].header.comments[k]

    hdu_normed_sigma.header["EXTNAME"] = "normalised_sigma"
    hdu_normed_sigma.header.comments["EXTNAME"] = "Normalised flux sigma"
    hdu_normed_sigma.add_checksum()

    return [hdu_normed_flux, hdu_normed_sigma]

def dummy_ccf_hdu(hdulist=None):
    """
    Return a dummy HDU for the best-fitting cross-correlation function.
    """

    hdu_ccf = fits.ImageHDU(data=None, header=None,
        do_not_scale_image_data=True)

    hdu_ccf.header["EXTNAME"] = "CCF"
    hdu_ccf.header.comments["EXTNAME"] = "CCF from best-fitting template"

    # Add empty default values.
    default_headers = [
        ("VRAD", "Radial velocity (km/s)"),
        ("U_VRAD", "Uncertainty on radial velocity (km/s)"),
        ("TEFF", "Effective temperature"),
        ("LOGG", "Surface gravity"),
        ("FE_H", "Metallicity ([Fe/H])"),
        ("ALPHA_FE", "Alpha-enhancement ([alpha/Fe])")
    ]
    for key, comment in default_headers:
        hdu_ccf.header[key] = "NaN"
        hdu_ccf.header.comments[key] = comment

    hdu_ccf.add_checksum()

    return hdu_ccf

def from_2dfdr(reduced_filename, dummy_hdus=True):
    """
    Returns a list of `astropy.io.fits.hdu.hdulist.HDUList` objects
    (1 HDU List per 2dfdr program object).

    :param reduced_filename:
        The path of the 2dfdr-reduced combined filename.

    :type reduced_filename:
        str
    """

    image = fits.open(reduced_filename)
    ext = read_2dfdr_extensions(image)

    # Verifications
    verify_hermes_origin(image)

    # Header modifications.
    keywords = ["NAME", "RA", "DEC", "PMRA", "PMDEC", ("MAGNITUDE", "MAG"),
        ("COMMENT", "DESCR")]
    keyword_comments = {
        "NAME": "Input source name",
        "DESCR": "Input source comment",
        "MAG": "Input source magnitude",
        "RA": "Right Ascension (degrees)",
        "DEC": "Declination (degrees)",
        "PMRA": "Proper motion in RA (mas/yr)",
        "PMDEC": "Proper motion in DEC (mas/yr)",
        "FIBRE": "Fibre number"
    }
    header_template = image[ext["data"]].header.copy()
    for column in ("CDELT2", "CRPIX2", "CRVAL2", "CTYPE2", "CUNIT2"):
        del header_template[column]

    ccd_number, ccd_name = get_ccd_number(image)
    header_template["CCD"] = ccd_number
    header_template.comments["CCD"] = "{0} camera".format(ccd_name)
    header_template["WG6_HASH"] = WG6_GIT_HASH
    header_template.comments["WG6_HASH"] = "WG6 standardisation commit hash"

    # Get the median sky flux
    sky_fibres = np.where(image[ext["fibres"]].data["TYPE"] == "S")[0]
    if not sky_fibres:
        median_sky_flux = np.nan * np.ones(image[ext["data"]].data.shape[1])
    else:
        median_sky_flux = np.nanmedian(image[ext["data"]].data[sky_fibres, :],
            axis=0)

    extracted_sources = []
    for program_index in np.where(image[ext["fibres"]].data["TYPE"] == "P")[0]:

        flux = image[ext["data"]].data[program_index, :]
        variance = image[ext["variance"]].data[program_index, :]
        header = header_template.copy()

        # Add header information from the fibre table.
        header["FIBRE"] = program_index + 1
        fibre_header = image[ext["fibres"]].data[program_index]
        for keyword in keywords:
            if isinstance(keyword, (str, unicode)):
                header[keyword] = fibre_header[keyword]
            else:
                from_keyword, to_keyword = keyword
                header[to_keyword] = fibre_header[from_keyword]

        # The 2dfdr fibre table has the RA and DEC in radians. Who does that?!
        if "RA" in keywords and "DEC" in keywords:
            header["RA"] *= 180./np.pi
            header["DEC"] *= 180./np.pi

        # Add associated comments for the fibre table information
        for keyword, comment in keyword_comments.items():
            header.comments[keyword] = comment

        # Create the HDUList.
        hdu_flux = fits.PrimaryHDU(data=flux, header=header,
            do_not_scale_image_data=True)
        hdu_flux.header["EXTNAME"] = "input_spectrum"
        hdu_flux.header.comments["EXTNAME"] = "Spectrum flux"

        hdu_sigma = fits.ImageHDU(data=variance**0.5, header=None,
            do_not_scale_image_data=True)

        for key in ("CRVAL1", "CDELT1", "CRPIX1", "CTYPE1", "CUNIT1"):
            hdu_sigma.header[key] = hdu_flux.header[key]
            hdu_sigma.header.comments[key] = hdu_flux.header.comments[key]
        hdu_sigma.header["EXTNAME"] = "input_sigma"
        hdu_sigma.header.comments["EXTNAME"] = "Flux sigma"

        # Calculate the barycentric motion.
        v_bary, v_helio = motions.from_header(header)

        hdu_flux.header["V_BARY"] = v_bary.to("km/s").value
        hdu_flux.header["V_HELIO"] = v_helio.to("km/s").value
        hdu_flux.header.comments["V_BARY"] = "Barycentric motion (km/s)"
        hdu_flux.header.comments["V_HELIO"] = "Heliocentric motion (km/s)"
        hdu_flux.header["HISTORY"] \
            = "Corrected for heliocentric motion (V_HELIO)"

        # Add motion correction information.
        for hdu in (hdu_flux, hdu_sigma):
            hdu.header["CRVAL1"] *= 1. + (v_helio/speed_of_light).value
            hdu.header["CDELT1"] *= 1. + (v_helio/speed_of_light).value

        # Create HDUList and add checksums.
        hdulist = fits.HDUList([hdu_flux, hdu_sigma])
        [hdu.add_checksum() for hdu in hdulist]

        # Add dummy extensions
        if dummy_hdus:
            hdulist.extend(dummy_normalisation_hdus(hdulist))
            hdulist.append(dummy_ccf_hdu(hdulist))

        # Add median sky hdu
        hdu_sky = fits.ImageHDU(data=median_sky_flux, header=None,
            do_not_scale_image_data=True)
        hdu_sky.header["EXTNAME"] = "median_sky"
        hdu_sky.header.comments["EXTNAME"] \
            = "Median sky flux from {} fibres".format(sky_fibres.size)
        hdulist.append(hdu_sky)

        extracted_sources.append(hdulist)

    image.close()

    return extracted_sources

<|endoftext|>


<commit_before>#pragma repy

"""
Author: Armon Dadgar
Description:
  Tries to fire up a lot of threads and checks that all the thread names returned are unique.
  This test launches threads via incoming UDP messages, which will check that the EventDeliverer
  is working properly.
"""

# Store the names that we've already seen
THREAD_NAMES = set([])

# Allowed Events
ALLOWED_EVENTS = 9

# How many times do we re-launch all the threads?
CHECK_TIMES = 3

# Entry point for new threads
def new_mess(ip,port,mess,ch1):
  # Get our name
  name = get_thread_name()

  # Check if this name has been used before
  if name in THREAD_NAMES:
    raise Exception, "Re-used thread name '"+str(name)+"'"

  # Record this name
  else:
    THREAD_NAMES.add(name)

# Try re-using the same 9 threads 3 times
if callfunc == "initialize":
  # IP and port
  ip = getmyip()
  port = 12345

  # Setup recvmess
  recvh = recvmess(ip,port,new_mess)

  # Add our own name
  THREAD_NAMES.add(get_thread_name())

  # Run for each interval
  for round in xrange(CHECK_TIMES):
    # Launch all the threads (via messages)
    for event in xrange(ALLOWED_EVENTS):
      sendmess(ip,port,"ping")

    # Wait for the threads to terminate
    sleep(1)

  # Expected thread count is EVENTS * CHECK_TIMES + main thread
  expected = ALLOWED_EVENTS * CHECK_TIMES + 1

  if len(THREAD_NAMES) != expected:
    raise Exception,"Unexpected number of thread names! Thread names: "+str(THREAD_NAMES)

  # Stop listening and exit
  stopcomm(recvh)
  exitall()

<commit_msg>Test should take port as a parameter...<commit_after>#pragma repy

"""
Author: Armon Dadgar
Description:
  Tries to fire up a lot of threads and checks that all the thread names returned are unique.
  This test launches threads via incoming UDP messages, which will check that the EventDeliverer
  is working properly.
"""

# Store the names that we've already seen
THREAD_NAMES = set([])

# Allowed Events
ALLOWED_EVENTS = 9

# How many times do we re-launch all the threads?
CHECK_TIMES = 3

# Entry point for new threads
def new_mess(ip,port,mess,ch1):
  # Get our name
  name = get_thread_name()

  # Check if this name has been used before
  if name in THREAD_NAMES:
    raise Exception, "Re-used thread name '"+str(name)+"'"

  # Record this name
  else:
    THREAD_NAMES.add(name)

# Try re-using the same 9 threads 3 times
if callfunc == "initialize":
  # IP and port
  ip = getmyip()
  port = <messport>

  # Setup recvmess
  recvh = recvmess(ip,port,new_mess)

  # Add our own name
  THREAD_NAMES.add(get_thread_name())

  # Run for each interval
  for round in xrange(CHECK_TIMES):
    # Launch all the threads (via messages)
    for event in xrange(ALLOWED_EVENTS):
      sendmess(ip,port,"ping")

    # Wait for the threads to terminate
    sleep(1)

  # Expected thread count is EVENTS * CHECK_TIMES + main thread
  expected = ALLOWED_EVENTS * CHECK_TIMES + 1

  if len(THREAD_NAMES) != expected:
    raise Exception,"Unexpected number of thread names! Thread names: "+str(THREAD_NAMES)

  # Stop listening and exit
  stopcomm(recvh)
  exitall()

<|endoftext|>