<commit_before>import asyncio
from graphite_escaping import metrics_name_to_graphite, metrics_name_from_graphite
import pickle
import struct
import socket
import os
import time

DEFAULT_GRAPHITE_PLAINTEXT_PORT = 2003
DEFAULT_GRAPHITE_PICKLE_PORT = 2004
SUPPORT_PROTOCOLS = ["plaintext", "pickle"]

class AioGraphiteSendException(Exception):
    pass

class AsyncioGraphiteSendService(object):

	def __init__(self, graphite_server, graphite_port, protocol = "pickle"):
		"""
			Must call __init__ before use
		"""
		# Graphite Server Address
		self.graphite_server_address = (graphite_server, graphite_port)

		# Connect to Graphite Server
		self.connect_to_graphite()

		# protocol
		self.protocol = protocol

		# Get Event Loop
		self.loop = asyncio.get_event_loop()

	def connect_to_graphite(self):
		"""
			Connect to Graphite Server based on Provided Server Address
		"""
		self.socket = socket.socket()
		self.socket.setblocking(false)
		try:
			self.socket.create_connection(self.graphite_server_address)
		except InterruptedError, e:
			raise e
		except socket.gaierror:
			raise AioGraphiteSendException("Unable to connect to the provided server address %s:%s" % self.graphite_server_address)
		except Exception, e:
			raise AioGraphiteSendException("Unexpected exception while connecting to %s:%s" % self.graphite_server_address)
		return self.socket

	def to_graphite_valid_metric_name(self, metric_dir_list):
		"""
			@purpose:
				Make metric name valid for graphite in case that the metric name includes
				any special character which is not supported by Graphite
			@example:
				Assuming that

					Expected_Metric_Name  =  metaccounts.authentication.password.attempted

				Then input metric_dir_list should be

					metric_dir_list = [metaccounts, authentication, password, attempted]

			@metric_dir_list: List of String
		"""
		return "."join([metrics_name_to_graphite(dir_name) for dir_name in metric_dir_list])

	def send_single_data(self, metric_dir_list, value, timestamp = None):
		"""
			@metric_dir_list: List of string
			@timestamp: the type should be int
			If you're very confident that the metric name is valid, then use <method: send_single_valid_data> instead.
		"""
		valid_metric_name = to_graphite_valid_metric_name(metric_dir_list)
		self.send_single_valid_data(valid_metric_name, value, timestamp)

	def send_dataset_list(self, dataset, timestamp = None):
		"""
			@param: dataset = [(metric1, value1), (metric2, value2), ...]

			If you're very confident that the metric name is valid, then use <method: send_valie_dataset_list> instead.
		"""
		valid_dataset = [(to_graphite_valid_metric_name(metric), value) for metric, value in dataset]
		send_valid_dataset_list(valid_dataset, timestamp)

	def send_single_valid_data(self, metric, value, timestamp = None):
		"""
			@timestamp: the type should be int
			Send a single data(metric value timestamp) to graphite
		"""
        timestamp = int(time.time()) if timestamp is None else int(timestamp)
		message = ""
		if self.protocol == "plaintext":
			message = plaintext_protocol_formatted_data(metric, value, timestamp)
		else:
			listOfMetricTuples = [pickle_protocol_formatted_data(metric, value, timestamp)]
			message = generate_message_for_pickle(listOfMetricTuples)
		loop.run_until_complete(asyncio.ensure_future(self.send_message(message)))

	def send_valid_dataset_list(self, dataset, timestamp = None):
		"""
			@param:
			Support two kinds of dataset
				1)	dataset = [(metric1, value1), (metric2, value2), ...]
				or
				2)	dataset = [(metric1, value1, timestamp1), (metric2, value2, timestamp2), ...]
		"""
		timestamp = int(time.time()) if timestamp is None else int(timestamp)
		message = ""

		# Generate message based on protocol
		# 1. plaintext
		# 2. pickle
		if self.protocol == "plaintext":
			message = generate_message_for_data_list(dataset, timestamp, plaintext_protocol_formatted_data, generate_message_for_plaintext)
		else:
			message = generate_message_for_data_list(dataset, timestamp, pickle_protocol_formatted_data, generate_message_for_pickle)

		# Sending
		loop.run_until_complete(asyncio.ensure_future(self.send_message(message)))

	def send_valid_dataset_dict(self, dataset, timestamp = None):
		"""
			Send data to graphite server when incoming data is in 'dict' format
			@param: dataset = {metric1 : value1, metric2 : value2, ...}
		"""

	def generate_message_for_data_list(self, dataset, timestamp, formate_function, generate_message_function):
		"""
			generate proper formatted message
			@param:
			Support two kinds of dataset
				1)	dataset = [(metric1, value1), (metric2, value2), ...]
				or
				2)	dataset = [(metric1, value1, timestamp1), (metric2, value2, timestamp2), ...]
		"""
		listofData = []
		for data in dataset:
			# unpack metric data
			if len(data) == 2:
				(metric, value) = data
			else:
				(metric, value, data_timestamp) = data
				timestamp = data_timestamp
			listOfData.append(formate_function(metric, value, timestamp))
		message =  generate_message_function(listofData)
		return message

	def generate_message_for_plaintext(self, listOfPlaintext):
		"""
			return the required message formate for protocol 'plaintext'
			@param:
				listOfPlaintext: ["metric1 value1 timestamp1", "metric2 value2 timestamp2", ...]
				type: List of String
		"""
		return "\n".join(listOfPlaintext)

	def generate_message_for_pickle(self, listOfMetricTuples):
		"""
			@param:
				listOfMetricTuples: [(metric1, (timestamp1, value1), (metric2, (timestamp2, value2), ...]
		"""
		payload = pickle.dumps(listOfMetricTuples, protocol=2)
		header = struct.pack("!L", len(payload))
		message = header + payload
		return message

	def plaintext_protocol_formatted_data(self, metric, value, timestamp):
		"""
			@return: required data formate when sending data through 'plaintext' protocol
			@return_type: String
		"""
		return " ".join([metric, value, timestamp])

	def pickle_protocol_formatted_data(self, metric, value, timestamp):
		"""
			@return: required data formate when sending data through 'pickle' protocol
			@return_type: Tuple
		"""
		return (metric, (timestamp, value))

	def send_message(self, message):
		"""
			@message: data ready to sent to graphite server
		"""
		total_sent = 0
		message_size = len(message)
		while total_sent < message_size:
			try:
				sent = self.socket.send(message[total_sent:])
				if sent == 0:
					raise RuntimeError("socket connection broken")
				total_sent = total_sent + sent
			except socket.gaierror as e:
				raise AioGraphiteSendException("Fail to send data to %s, Error: %s" % (self.graphite_server_address, e))
			except Exception, e:
				raise AioGraphiteSendException("Unexpected exception while sending data to %s:%s" % self.graphite_server_address)
		return total_sent

	def disconnect(self):
        """
        	Close the TCP connection
        """
        try:
            self.socket.shutdown(1)
        except AttributeError:
            self.socket = None
        except Exception:
            self.socket = None
        finally:
            self.socket = None

#########################################################
#########################################################
#########################################################

aiographite_send_instance = None

def init(graphite_server, graphite_port = DEFAULT_GRAPHITE_PICKLE_PORT, protocol_type = 'pickle'):

def send_single_data():

def send_valid_single_data():

def send_data_list():

def send_valid_data_list():

def send_valid_data_dic():

def destory():

def reset():

def _dummy_message_plaintext_formate():
	print("Message formate for plaintext protocol")
	print("Metric1 Value TimeStamp1\n Metric2 Value2 Timestamp2")

def main():
	_dummy_message_plaintext_formate()

if __name__ == '__main__':
	main()

<commit_msg>add the feature that we can send by passing dict<commit_after>import asyncio
from graphite_escaping import metrics_name_to_graphite, metrics_name_from_graphite
import pickle
import struct
import socket
import os
import time

DEFAULT_GRAPHITE_PLAINTEXT_PORT = 2003
DEFAULT_GRAPHITE_PICKLE_PORT = 2004
SUPPORT_PROTOCOLS = ["plaintext", "pickle"]

class AioGraphiteSendException(Exception):
    pass

class AsyncioGraphiteSendService(object):

	def __init__(self, graphite_server, graphite_port, protocol = "pickle"):
		"""
			Must call __init__ before use
		"""
		# Graphite Server Address
		self.graphite_server_address = (graphite_server, graphite_port)

		# Connect to Graphite Server
		self.connect_to_graphite()

		# protocol
		self.protocol = protocol

		# Get Event Loop
		self.loop = asyncio.get_event_loop()

	def connect_to_graphite(self):
		"""
			Connect to Graphite Server based on Provided Server Address
		"""
		self.socket = socket.socket()
		self.socket.setblocking(False)
		try:
			self.socket.create_connection(self.graphite_server_address)
		except InterruptedError, e:
			raise e
		except socket.gaierror:
			raise AioGraphiteSendException("Unable to connect to the provided server address %s:%s" % self.graphite_server_address)
		except Exception, e:
			raise AioGraphiteSendException("Unexpected exception while connecting to %s:%s" % self.graphite_server_address)
		return self.socket

	def to_graphite_valid_metric_name(self, metric_dir_list):
		"""
			@purpose:
				Make metric name valid for graphite in case that the metric name includes
				any special character which is not supported by Graphite
			@example:
				Assuming that

					Expected_Metric_Name  =  metaccounts.authentication.password.attempted

				Then input metric_dir_list should be

					metric_dir_list = [metaccounts, authentication, password, attempted]

			@metric_dir_list: List of String
		"""
		return "."join([metrics_name_to_graphite(dir_name) for dir_name in metric_dir_list])

	def send_single_data(self, metric_dir_list, value, timestamp = None):
		"""
			@example:
				Assuming that

					Expected_Metric_Name  =  metaccounts.authentication.password.attempted

				Then input metric_dir_list should be

					metric_dir_list = [metaccounts, authentication, password, attempted]

			@metric_dir_list: List of string
			@timestamp: the type should be int

			If you're very confident that the metric name is valid, then use <method: send_single_valid_data> instead.

		"""
		valid_metric_name = to_graphite_valid_metric_name(metric_dir_list)
		self.send_single_valid_data(valid_metric_name, value, timestamp)

	def send_dataset_list(self, dataset, timestamp = None):
		"""
			@param: dataset = [(metric1, value1), (metric2, value2), ...]

			If you're very confident that the metric name is valid, then use <method: send_valie_dataset_list> instead.

		"""
		valid_dataset = [(to_graphite_valid_metric_name(metric), value) for metric, value in dataset]
		send_valid_dataset_list(valid_dataset, timestamp)

	def send_single_valid_data(self, metric, value, timestamp = None):
		"""
			@timestamp: the type should be int
			Send a single data(metric value timestamp) to graphite
		"""
        timestamp = int(time.time()) if timestamp is None else int(timestamp)
		message = ""
		if self.protocol == "plaintext":
			message = plaintext_protocol_formatted_data(metric, value, timestamp)
		else:
			listOfMetricTuples = [pickle_protocol_formatted_data(metric, value, timestamp)]
			message = generate_message_for_pickle(listOfMetricTuples)
		loop.run_until_complete(asyncio.ensure_future(self.send_message(message)))

	def send_valid_dataset_list(self, dataset, timestamp = None):
		"""
			@param:
			Support two kinds of dataset
				1)	dataset = [(metric1, value1), (metric2, value2), ...]
				or
				2)	dataset = [(metric1, value1, timestamp1), (metric2, value2, timestamp2), ...]
		"""
		timestamp = int(time.time()) if timestamp is None else int(timestamp)
		message = ""

		# Generate message based on protocol
		# 1. plaintext
		# 2. pickle
		if self.protocol == "plaintext":
			message = generate_message_for_data_list(dataset, timestamp, plaintext_protocol_formatted_data, generate_message_for_plaintext)
		else:
			message = generate_message_for_data_list(dataset, timestamp, pickle_protocol_formatted_data, generate_message_for_pickle)

		# Sending
		loop.run_until_complete(asyncio.ensure_future(self.send_message(message)))

	def send_valid_dataset_dict(self, dataset, timestamp = None):
		"""
			Send data to graphite server when incoming data is in 'dict' format
			@param: dataset = {metric1 : value1, metric2 : value2, ...}
		"""

	def generate_message_for_data_list(self, dataset, timestamp, formate_function, generate_message_function):
		"""
			generate proper formatted message
			@param:
			Support two kinds of dataset
				1)	dataset = [(metric1, value1), (metric2, value2), ...]
				or
				2)	dataset = [(metric1, value1, timestamp1), (metric2, value2, timestamp2), ...]
		"""
		listofData = []
		for data in dataset:
			# unpack metric data
			if len(data) == 2:
				(metric, value) = data
			else:
				(metric, value, data_timestamp) = data
				timestamp = data_timestamp
			listOfData.append(formate_function(metric, value, timestamp))
		message =  generate_message_function(listofData)
		return message

	def generate_message_for_plaintext(self, listOfPlaintext):
		"""
			return the required message formate for protocol 'plaintext'
			@param:
				listOfPlaintext: ["metric1 value1 timestamp1", "metric2 value2 timestamp2", ...]
				type: List of String
		"""
		return "\n".join(listOfPlaintext)

	def generate_message_for_pickle(self, listOfMetricTuples):
		"""
			@param:
				listOfMetricTuples: [(metric1, (timestamp1, value1), (metric2, (timestamp2, value2), ...]
		"""
		payload = pickle.dumps(listOfMetricTuples, protocol=2)
		header = struct.pack("!L", len(payload))
		message = header + payload
		return message

	def plaintext_protocol_formatted_data(self, metric, value, timestamp):
		"""
			@return: required data formate when sending data through 'plaintext' protocol
			@return_type: String
		"""
		return " ".join([metric, value, timestamp])

	def pickle_protocol_formatted_data(self, metric, value, timestamp):
		"""
			@return: required data formate when sending data through 'pickle' protocol
			@return_type: Tuple
		"""
		return (metric, (timestamp, value))

	def send_message(self, message):
		"""
			@message: data ready to sent to graphite server
		"""
		total_sent = 0
		message_size = len(message)
		while total_sent < message_size:
			try:
				sent = self.socket.send(message[total_sent:])
				if sent == 0:
					raise RuntimeError("socket connection broken")
				total_sent = total_sent + sent
			except socket.gaierror as e:
				raise AioGraphiteSendException("Fail to send data to %s, Error: %s" % (self.graphite_server_address, e))
			except Exception, e:
				raise AioGraphiteSendException("Unexpected exception while sending data to %s:%s" % self.graphite_server_address)
		return total_sent

	def disconnect(self):
        """
        	Close the TCP connection
        """
        try:
            self.socket.shutdown(1)
        except AttributeError:
            self.socket = None
        except Exception:
            self.socket = None
        finally:
            self.socket = None

#########################################################
#########################################################
#########################################################

# Module Instance Variable
aiographite_send_instance = None

def init(graphite_server, graphite_port = DEFAULT_GRAPHITE_PICKLE_PORT, protocol_type = 'pickle'):
	global aiographite_send_instance
	if aiographite_send_instance is not None:
		destory()

	# Check Init Protocol Type
	if protocol_type not in SUPPORT_PROTOCOLS:
		raise AioGraphiteSendException("%s is not Support Protocol Type!", protocol_type)

	# Construct an aiographite sending service instance
	aiographite_send_instance = AsyncioGraphiteSendService(graphite_server, graphite_port, protocol_type)
	return aiographite_send_instance

def send_single_data(metric_dir_list, value, timestamp = None):
	"""
		@example:
			Assuming that

				Expected_Metric_Name  =  metaccounts.authentication.password.attempted

			Then input metric_dir_list should be

				metric_dir_list = [metaccounts, authentication, password, attempted]

		@metric_dir_list: List of string
		@timestamp: the type should be int

		If you're very confident that the metric name is valid, then use <method: send_single_valid_data> instead.

		"""
	global aiographite_send_instance

	if aiographite_send_instance is None:
		raise AioGraphiteSendException("Must call init before use!")

	# Sending data
	aiographite_send_instance.send_single_data(metric_dir_list, value, timestamp)

def send_single_valid_data(metric, value, timestamp = None):
	"""
		@metric: metric name, type: String
		@timestamp: type should be int
	"""
	global aiographite_send_instance

	if aiographite_send_instance is None:
		raise AioGraphiteSendException("Must call init before use!")

	# Sending data
	aiographite_send_instance.send_single_valid_data(metric, value, timestamp)

def send_data_list(dataset, timestamp = None):

def send_valid_data_list():

def send_valid_data_dic():

def destory():
	"""
		Close TCP connection and destory the instance
	"""
	global aiographite_send_instance
	if aiographite_send_instance is None:
		return False
	aiographite_send_instance.disconnect()
	aiographite_send_instance = None
	return True

def _dummy_message_plaintext_formate():
	print("Message formate for plaintext protocol")
	print("Metric1 Value TimeStamp1\n Metric2 Value2 Timestamp2")

def main():
	_dummy_message_plaintext_formate()

if __name__ == '__main__':
	main()

<|endoftext|>


<commit_before>from astropy.visualization.mpl_normalize import ImageNormalize
from astropy.wcs import WCS
#from rand_bkg import bkg_boxes
from astropy.convolution import Gaussian2DKernel
from astropy.stats import sigma_clipped_stats
from photutils import detect_sources
from photutils.utils import random_cmap
from scipy.ndimage import binary_dilation
import pandas as pd
import numpy as np
import os
import glob
from collections import OrderedDict
from astropy.io import fits
import odi_config as odi
from gaia.tap import cone_search
from astropy.table import Table

def get_sdss_coords_offline(img, ota, inst,output='test.sdss'):
    """
    Pull out and parse the ``CAT.PHOTCALIB`` table from ``img`` header. This
    function will separate the SDSS stars in ``CAT.PHOTCALIB`` based on which
    ``ota`` they fall on.

    Parameters
    ----------
    img : str
        Name of image
    ota : str
        Name of OTA
    int : str
        Version of ODI used, ``podi`` or ``5odi``

    Returns
    -------
    xdim : int
        Size of OTA in the x direction ``NAXIS1``
    ydim : int
        Size of OTA in the y direction ``NAXIS2``

    Note
    ----
    If the images being processed do not fall in the SDSS footprint,
    the QuickReduce pipeline will use PanSTARRS. This function will still
    pull out these stars and treat them as SDSS stars. There will be no
    ``u`` magnitudes available, however.
    """
    hdulist = odi.fits.open(img.f)
    hdu = odi.tan_header_fix(hdulist[ota])

    xdim = hdu.header['NAXIS1']
    ydim = hdu.header['NAXIS2']

    sdss_cat_img = hdulist['CAT.PHOTCALIB']
    cat_img_data = Table.read(sdss_cat_img, format='fits')
    # print cat_img_data.colnames
    # force little-endian byte order to make FITS play nice with pandas
    sdss_cat_img_df = cat_img_data.to_pandas()
    # sdss_cat_img_df = pd.DataFrame.from_dict(cat_img_dict)
    # print sdss_cat_img_df.keys()
    ota = float(ota.strip('OTA.SCI'))
    print 'catalog source:', hdulist[0].header['PHOTMCAT']
    if 'sdss' in hdulist[0].header['PHOTMCAT'] or 'SDSS' in hdulist[0].header['PHOTMCAT']:
        try:
            ota_matches_df = sdss_cat_img_df.iloc[np.where(sdss_cat_img_df[u'ODI_OTA'] == ota)]
            needed_columns = [u'REF_RA',u'REF_DEC',u'REF_U',
                              u'REF_ERR_U', u'REF_G', u'REF_ERR_G', u'REF_R',
                              u'REF_ERR_R', u'REF_I', u'REF_ERR_I', u'REF_Z',
                              u'REF_ERR_Z', u'ODI_OTA']

            output_df = ota_matches_df[needed_columns]
            output_df.to_csv(output,index=False)
        except KeyError:
            oditable = hdulist['CAT.ODI'].data
            oditalbe_df = pd.DataFrame.from_dict(oditable)

            ODI_RA = np.squeeze(np.array(oditalbe_df['RA']))
            ODI_DEC = np.squeeze( np.array(oditalbe_df['DEC']))
            ODI_OTA = np.squeeze( np.array(oditalbe_df['OTA']))

            junkdict = OrderedDict([(u'ODI_RA',ODI_RA),
                                    (u'ODI_DEC',ODI_DEC),
                                    (u'ODI_OTA',ODI_OTA.astype(float))])
            junk_df = pd.DataFrame.from_dict(junkdict)

            matched_df = pd.merge(sdss_cat_img_df,junk_df ,on = [u'ODI_RA',u'ODI_DEC'],how='inner')
            print matched_df.columns
            needed_columns = np.insert(sdss_cat_img_df.columns.values,0,u'ODI_OTA')

            full_df = matched_df[needed_columns]
            ota_matches_df = full_df.iloc[np.where(full_df[u'ODI_OTA'] == ota)]
            needed_columns = [u'SDSS_RA',u'SDSS_DEC',
                              u'SDSS_MAG_U',u'SDSS_ERR_U',
                              u'SDSS_MAG_G', u'SDSS_ERR_G',
                              u'SDSS_MAG_R',u'SDSS_ERR_R',
                              u'SDSS_MAG_I', u'SDSS_ERR_I',
                              u'SDSS_MAG_Z',u'SDSS_ERR_Z',
                              u'ODI_OTA']
            output_df = ota_matches_df[needed_columns]
            output_df.to_csv(output,index=False)
    else:
        ota_matches_df = sdss_cat_img_df.iloc[np.where(sdss_cat_img_df[u'ODI_OTA'] == ota)]
        ota_matches_df = ota_matches_df.reset_index()
        junk_u = np.ones(len(ota_matches_df))
        junk_u_err = np.ones(len(ota_matches_df))
        ota_matches_df[u'IPP_MAG_U'] = junk_u
        ota_matches_df[u'IPP_ERR_U'] = junk_u_err

        needed_columns = [u'IPP_RA', u'IPP_DEC',
                          u'IPP_MAG_U', u'IPP_ERR_U',
                          u'IPP_MAG_G', u'IPP_ERR_G',
                          u'IPP_MAG_R', u'IPP_ERR_R',
                          u'IPP_MAG_I', u'IPP_ERR_I',
                          u'IPP_MAG_Z',u'IPP_ERR_Z',
                          u'ODI_OTA']

        output_df = ota_matches_df[needed_columns]
        output_df.to_csv(output,index=False)

    hdulist.close()
    return xdim, ydim

def get_2mass_coords_offline(img, ota, inst,output='test.mass'):
    """
    Pull out and parse the ``CAT.ODI+2MASS`` table from ``img`` header. This
    function will separate the 2MASS stars in ``CAT.ODI+2MASS`` based on which
    ``ota`` they fall on.

    Parameters
    ----------
    img : str
        Name of image
    ota : str
        Name of OTA
    int : str
        Version of ODI used, ``podi`` or ``5odi``

    Returns
    -------
    xdim : int
        Size of OTA in the x direction ``NAXIS1``
    ydim : int
        Size of OTA in the y direction ``NAXIS2``
    """
    hdulist = odi.fits.open(img.f)
    hdu = odi.tan_header_fix(hdulist[ota])

    xdim = hdu.header['NAXIS1']
    ydim = hdu.header['NAXIS2']

    mass_cat_img = hdulist['CAT.ODI+2MASS']
    mass_cat_img_df = mass_cat_img.data
    hdulist.close()

    ota = float(ota.strip('OTA.SCI'))
    ota_matches_df = mass_cat_img_df[np.where(mass_cat_img_df['OTA'] == ota)]

    junk_df = pd.DataFrame.from_dict(ota_matches_df)

    output_df = junk_df
    output_df.to_csv(output,index=False)
    return xdim, ydim

def get_gaia_coords(img,ota,inst,output='test.gaia',cluster=False,**kwargs):
    """
    Query the online Gaia DR1 based on the central coordinates of the current
    OTA. If the ``cluster`` flag is set to ``True``, the querey will avoid
    a crowded region based on coordinates and a radius set by the user in
    the configuration files.

    Parameters
    ----------
    img : str
        Name of image
    ota : str
        Name of OTA
    int : str
        Version of ODI used, ``podi`` or ``5odi``

    """
    from astropy import units as u
    from astropy.coordinates import SkyCoord
    hdulist = fits.open(img.f)
    hdu_ota = odi.tan_header_fix(hdulist[ota])
    w = WCS(hdu_ota.header)
    ota_center_radec = w.wcs_pix2world([[2018.0,2007.5]],1)

    coord1 = w.wcs_pix2world([[1,1]],1)
    coord2 = w.wcs_pix2world([[4036,1]],1)
    coord3 = w.wcs_pix2world([[4036,4015]],1)
    coord4 = w.wcs_pix2world([[1,4015]],1)
    center_skycoord = SkyCoord(ota_center_radec[0][0]*u.deg,
                               ota_center_radec[0][1]*u.deg,frame='icrs')
    corner_skycoord = SkyCoord(coord2[0][0]*u.deg,
                               coord2[0][1]*u.deg,frame='icrs')
    cone_radius = center_skycoord.separation(corner_skycoord).value

    print 'Retrieving Gaia sources for: ', ota
    ota_gaia_sources = cone_search(ota_center_radec[0][0],
                                   ota_center_radec[0][1],
                                   cone_radius)
    ota_gaia_df = ota_gaia_sources.to_pandas()
    hdulist.close()
    if cluster == True:
        try:
            racenter = kwargs['racenter']
            deccenter = kwargs['deccenter']
            min_radius = kwargs['min_radius']
            G_lim = kwargs['G_lim']
        except KeyError:
            print 'Must provide racenter, deccenter, and min_radius'
        cluster_center = SkyCoord(racenter*u.degree
                                  ,deccenter*u.degree,
                                  frame='icrs')
        gaia_coords = SkyCoord(ota_gaia_df.ra.values*u.deg,
                               ota_gaia_df.dec.values*u.deg,frame='icrs')
        dist_from_center = cluster_center.separation(gaia_coords).arcmin
        ota_gaia_df['dis'] = dist_from_center
        ota_gaia_df = ota_gaia_df[(ota_gaia_df.dis >= min_radius) &
                                  (ota_gaia_df.phot_g_mean_mag <= G_lim)]

    ra_min, ra_max = coord1[0][0],coord2[0][0]
    dec_min, dec_max = coord1[0][1],coord3[0][1]

    ota_gaia_df = ota_gaia_df[(ota_gaia_df.ra > ra_min) &
                              (ota_gaia_df.ra < ra_max) &
                              (ota_gaia_df.dec > dec_min) &
                              (ota_gaia_df.dec < dec_max)]

    gaia_catalog_out = output
    ota_gaia_df.to_csv(gaia_catalog_out,
                       columns=['ra', 'dec','phot_g_mean_mag'],
                       index=False)

def main():
    pass

if __name__ == '__main__':
    main()
<commit_msg>print table columns<commit_after>from astropy.visualization.mpl_normalize import ImageNormalize
from astropy.wcs import WCS
#from rand_bkg import bkg_boxes
from astropy.convolution import Gaussian2DKernel
from astropy.stats import sigma_clipped_stats
from photutils import detect_sources
from photutils.utils import random_cmap
from scipy.ndimage import binary_dilation
import pandas as pd
import numpy as np
import os
import glob
from collections import OrderedDict
from astropy.io import fits
import odi_config as odi
from gaia.tap import cone_search
from astropy.table import Table

def get_sdss_coords_offline(img, ota, inst,output='test.sdss'):
    """
    Pull out and parse the ``CAT.PHOTCALIB`` table from ``img`` header. This
    function will separate the SDSS stars in ``CAT.PHOTCALIB`` based on which
    ``ota`` they fall on.

    Parameters
    ----------
    img : str
        Name of image
    ota : str
        Name of OTA
    int : str
        Version of ODI used, ``podi`` or ``5odi``

    Returns
    -------
    xdim : int
        Size of OTA in the x direction ``NAXIS1``
    ydim : int
        Size of OTA in the y direction ``NAXIS2``

    Note
    ----
    If the images being processed do not fall in the SDSS footprint,
    the QuickReduce pipeline will use PanSTARRS. This function will still
    pull out these stars and treat them as SDSS stars. There will be no
    ``u`` magnitudes available, however.
    """
    hdulist = odi.fits.open(img.f)
    hdu = odi.tan_header_fix(hdulist[ota])

    xdim = hdu.header['NAXIS1']
    ydim = hdu.header['NAXIS2']

    sdss_cat_img = hdulist['CAT.PHOTCALIB']
    cat_img_data = Table.read(sdss_cat_img, format='fits')
    # print cat_img_data.colnames
    # force little-endian byte order to make FITS play nice with pandas
    sdss_cat_img_df = cat_img_data.to_pandas()
    # sdss_cat_img_df = pd.DataFrame.from_dict(cat_img_dict)
    # print sdss_cat_img_df.keys()
    ota = float(ota.strip('OTA.SCI'))
    print 'catalog source:', hdulist[0].header['PHOTMCAT']
    if 'sdss' in hdulist[0].header['PHOTMCAT'] or 'SDSS' in hdulist[0].header['PHOTMCAT']:
        try:
            print sdss_cat_img_df.columns
            ota_matches_df = sdss_cat_img_df.iloc[np.where(sdss_cat_img_df[u'ODI_OTA'] == ota)]
            needed_columns = [u'REF_RA',u'REF_DEC',u'REF_U',
                              u'REF_ERR_U', u'REF_G', u'REF_ERR_G', u'REF_R',
                              u'REF_ERR_R', u'REF_I', u'REF_ERR_I', u'REF_Z',
                              u'REF_ERR_Z', u'ODI_OTA']

            output_df = ota_matches_df[needed_columns]
            output_df.to_csv(output,index=False)
        except KeyError:
            oditable = hdulist['CAT.ODI'].data
            oditalbe_df = pd.DataFrame.from_dict(oditable)

            ODI_RA = np.squeeze(np.array(oditalbe_df['RA']))
            ODI_DEC = np.squeeze( np.array(oditalbe_df['DEC']))
            ODI_OTA = np.squeeze( np.array(oditalbe_df['OTA']))

            junkdict = OrderedDict([(u'ODI_RA',ODI_RA),
                                    (u'ODI_DEC',ODI_DEC),
                                    (u'ODI_OTA',ODI_OTA.astype(float))])
            junk_df = pd.DataFrame.from_dict(junkdict)

            matched_df = pd.merge(sdss_cat_img_df,junk_df ,on = [u'ODI_RA',u'ODI_DEC'],how='inner')
            print matched_df.columns
            needed_columns = np.insert(sdss_cat_img_df.columns.values,0,u'ODI_OTA')

            full_df = matched_df[needed_columns]
            ota_matches_df = full_df.iloc[np.where(full_df[u'ODI_OTA'] == ota)]
            needed_columns = [u'SDSS_RA',u'SDSS_DEC',
                              u'SDSS_MAG_U',u'SDSS_ERR_U',
                              u'SDSS_MAG_G', u'SDSS_ERR_G',
                              u'SDSS_MAG_R',u'SDSS_ERR_R',
                              u'SDSS_MAG_I', u'SDSS_ERR_I',
                              u'SDSS_MAG_Z',u'SDSS_ERR_Z',
                              u'ODI_OTA']
            output_df = ota_matches_df[needed_columns]
            output_df.to_csv(output,index=False)
    else:
        ota_matches_df = sdss_cat_img_df.iloc[np.where(sdss_cat_img_df[u'ODI_OTA'] == ota)]
        ota_matches_df = ota_matches_df.reset_index()
        junk_u = np.ones(len(ota_matches_df))
        junk_u_err = np.ones(len(ota_matches_df))
        ota_matches_df[u'IPP_MAG_U'] = junk_u
        ota_matches_df[u'IPP_ERR_U'] = junk_u_err

        needed_columns = [u'IPP_RA', u'IPP_DEC',
                          u'IPP_MAG_U', u'IPP_ERR_U',
                          u'IPP_MAG_G', u'IPP_ERR_G',
                          u'IPP_MAG_R', u'IPP_ERR_R',
                          u'IPP_MAG_I', u'IPP_ERR_I',
                          u'IPP_MAG_Z',u'IPP_ERR_Z',
                          u'ODI_OTA']

        output_df = ota_matches_df[needed_columns]
        output_df.to_csv(output,index=False)

    hdulist.close()
    return xdim, ydim

def get_2mass_coords_offline(img, ota, inst,output='test.mass'):
    """
    Pull out and parse the ``CAT.ODI+2MASS`` table from ``img`` header. This
    function will separate the 2MASS stars in ``CAT.ODI+2MASS`` based on which
    ``ota`` they fall on.

    Parameters
    ----------
    img : str
        Name of image
    ota : str
        Name of OTA
    int : str
        Version of ODI used, ``podi`` or ``5odi``

    Returns
    -------
    xdim : int
        Size of OTA in the x direction ``NAXIS1``
    ydim : int
        Size of OTA in the y direction ``NAXIS2``
    """
    hdulist = odi.fits.open(img.f)
    hdu = odi.tan_header_fix(hdulist[ota])

    xdim = hdu.header['NAXIS1']
    ydim = hdu.header['NAXIS2']

    mass_cat_img = hdulist['CAT.ODI+2MASS']
    mass_cat_img_df = mass_cat_img.data
    hdulist.close()

    ota = float(ota.strip('OTA.SCI'))
    ota_matches_df = mass_cat_img_df[np.where(mass_cat_img_df['OTA'] == ota)]

    junk_df = pd.DataFrame.from_dict(ota_matches_df)

    output_df = junk_df
    output_df.to_csv(output,index=False)
    return xdim, ydim

def get_gaia_coords(img,ota,inst,output='test.gaia',cluster=False,**kwargs):
    """
    Query the online Gaia DR1 based on the central coordinates of the current
    OTA. If the ``cluster`` flag is set to ``True``, the querey will avoid
    a crowded region based on coordinates and a radius set by the user in
    the configuration files.

    Parameters
    ----------
    img : str
        Name of image
    ota : str
        Name of OTA
    int : str
        Version of ODI used, ``podi`` or ``5odi``

    """
    from astropy import units as u
    from astropy.coordinates import SkyCoord
    hdulist = fits.open(img.f)
    hdu_ota = odi.tan_header_fix(hdulist[ota])
    w = WCS(hdu_ota.header)
    ota_center_radec = w.wcs_pix2world([[2018.0,2007.5]],1)

    coord1 = w.wcs_pix2world([[1,1]],1)
    coord2 = w.wcs_pix2world([[4036,1]],1)
    coord3 = w.wcs_pix2world([[4036,4015]],1)
    coord4 = w.wcs_pix2world([[1,4015]],1)
    center_skycoord = SkyCoord(ota_center_radec[0][0]*u.deg,
                               ota_center_radec[0][1]*u.deg,frame='icrs')
    corner_skycoord = SkyCoord(coord2[0][0]*u.deg,
                               coord2[0][1]*u.deg,frame='icrs')
    cone_radius = center_skycoord.separation(corner_skycoord).value

    print 'Retrieving Gaia sources for: ', ota
    ota_gaia_sources = cone_search(ota_center_radec[0][0],
                                   ota_center_radec[0][1],
                                   cone_radius)
    ota_gaia_df = ota_gaia_sources.to_pandas()
    hdulist.close()
    if cluster == True:
        try:
            racenter = kwargs['racenter']
            deccenter = kwargs['deccenter']
            min_radius = kwargs['min_radius']
            G_lim = kwargs['G_lim']
        except KeyError:
            print 'Must provide racenter, deccenter, and min_radius'
        cluster_center = SkyCoord(racenter*u.degree
                                  ,deccenter*u.degree,
                                  frame='icrs')
        gaia_coords = SkyCoord(ota_gaia_df.ra.values*u.deg,
                               ota_gaia_df.dec.values*u.deg,frame='icrs')
        dist_from_center = cluster_center.separation(gaia_coords).arcmin
        ota_gaia_df['dis'] = dist_from_center
        ota_gaia_df = ota_gaia_df[(ota_gaia_df.dis >= min_radius) &
                                  (ota_gaia_df.phot_g_mean_mag <= G_lim)]

    ra_min, ra_max = coord1[0][0],coord2[0][0]
    dec_min, dec_max = coord1[0][1],coord3[0][1]

    ota_gaia_df = ota_gaia_df[(ota_gaia_df.ra > ra_min) &
                              (ota_gaia_df.ra < ra_max) &
                              (ota_gaia_df.dec > dec_min) &
                              (ota_gaia_df.dec < dec_max)]

    gaia_catalog_out = output
    ota_gaia_df.to_csv(gaia_catalog_out,
                       columns=['ra', 'dec','phot_g_mean_mag'],
                       index=False)

def main():
    pass

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


<commit_before>from orator.orm import has_many, belongs_to_many

from aioweb.core.model import Model, mutator, accessor
from passlib.handlers.sha2_crypt import sha256_crypt
from ..models import permission
from ..models import group

class AbstractUser(object):
    def is_authenticated(self):
        return False

class User(Model, AbstractUser):
    __guarded__ = ['id']
    __hidden__ = ['password']

    def __init__(self, _attributes=None, **attributes):
        if 'password' in attributes:
            attributes['password'] = User.hash_password(attributes['password'])
        super().__init__(_attributes, **attributes)

    @staticmethod
    def hash_password(value):
        # TODO: configurable hash type
        return sha256_crypt.hash(value)

    @mutator
    def password(self, value):
        if value:
            self.set_raw_attribute('password', User.hash_password(value))
        else:
            self.set_raw_attribute('password', '')

    def can(self, permission):
        # TODO: check it
        return self.permissions.has(permission).get()  # or self.groups.perimissions.has(permission).get()

    @belongs_to_many('user_permissions')
    def permissions(self):
        return permission.Permission

    @belongs_to_many('user_groups')
    def groups(self):
        return group.Group

    @accessor
    def username(self):
        return self.phone if self.phone else self.email

    def __str__(self):
        return self.username
<commit_msg>* agg get_id() to AbstractUser<commit_after>from orator.orm import has_many, belongs_to_many

from aioweb.core.model import Model, mutator, accessor
from passlib.handlers.sha2_crypt import sha256_crypt
from ..models import permission
from ..models import group

class AbstractUser(object):
    def is_authenticated(self):
        return False

    def get_id(self):
        return None

class User(Model, AbstractUser):
    __guarded__ = ['id']
    __hidden__ = ['password']

    def __init__(self, _attributes=None, **attributes):
        if 'password' in attributes:
            attributes['password'] = User.hash_password(attributes['password'])
        super().__init__(_attributes, **attributes)

    def get_id(self):
        return self.id

    @staticmethod
    def hash_password(value):
        # TODO: configurable hash type
        return sha256_crypt.hash(value)

    @mutator
    def password(self, value):
        if value:
            self.set_raw_attribute('password', User.hash_password(value))
        else:
            self.set_raw_attribute('password', '')

    def can(self, permission):
        # TODO: check it
        return self.permissions.has(permission).get()  # or self.groups.perimissions.has(permission).get()

    @belongs_to_many('user_permissions')
    def permissions(self):
        return permission.Permission

    @belongs_to_many('user_groups')
    def groups(self):
        return group.Group

    @accessor
    def username(self):
        return self.phone if self.phone else self.email

    def __str__(self):
        return self.username
<|endoftext|>


<commit_before>from __future__ import print_function
'''
------------------------------------------------------------------------
Last updated 4/7/2015

Miscellaneous functions for SS.py and TPI.py.

This python files calls:
    OUTPUT/Saved_moments/wealth_data_moments.pkl

------------------------------------------------------------------------
'''

# Packages
import os
from io import StringIO
import numpy as np
try:
    import cPickle as pickle
except ImportError:
    import pickle
from pkg_resources import resource_stream, Requirement

EPSILON = 1e-10
PATH_EXISTS_ERRNO = 17

REFORM_DIR = "./OUTPUT_REFORM"
BASELINE_DIR = "./OUTPUT_BASELINE"

# Default year for model runs
DEFAULT_START_YEAR = 2018

# Latest year TaxData extrapolates to
TC_LAST_YEAR = 2027

# Year of data used (e.g. PUF or CPS year)
RECORDS_START_YEAR = 2011

# for f in (REFORM_DIR, BASELINE_DIR):
#     if not os.path.exists(f):
#         os.mkdir(f)

def mkdirs(path):
    '''
    Makes directories to save output.

    Inputs:
        path = string, path name for new directory

    Functions called: None

    Objects in function:

    Returns: N/A
    '''

    try:
        os.makedirs(path)
    except OSError as oe:
        if oe.errno == PATH_EXISTS_ERRNO:
            pass

def pct_diff_func(simul, data):
    '''
    Used to calculate the absolute percent difference between data
    moments and model moments.

    Inputs:
        simul = any shape, model moments
        data  = same shape as simul, data moments

    Functions called: None

    Objects in function:
        frac   = same shape as simul, percent difference between data and model moments
        output = same shape as simul, absolute percent difference between data and model moments

    Returns: output
    '''

    frac = (simul - data) / data
    output = np.abs(frac)
    return output

def convex_combo(var1, var2, nu):
    '''
    Takes the convex combination of two variables, where nu is in [0,1].

    Inputs:
        var1 = any shape, variable 1
        var2 = same shape as var1, variable 2
        nu   = scalar, weight on var1 in convex combination

    Functions called: None

    Objects in function:
        combo = same shape as var1, convex combination of var1 and var2

    Returns: combo
    '''

    combo = nu * var1 + (1 - nu) * var2
    return combo

def read_file(path, fname):
    '''
    Read the contents of 'path'. If it does not exist, assume the file
    is installed in a .egg file, and adjust accordingly.

    Inputs:
        path  = string, path name for new directory
        fname = string, filename

    Functions called: None

    Objects in function:
        path_in_egg
        buf
        _bytes

    Returns: file contents
    '''

    if not os.path.exists(os.path.join(path, fname)):
        path_in_egg = os.path.join("ogusa", fname)
        buf = resource_stream(Requirement.parse("ogusa"), path_in_egg)
        _bytes = buf.read()
        return StringIO(_bytes.decode("utf-8"))
    else:
        return open(os.path.join(path, fname))

def pickle_file_compare(fname1, fname2, tol=1e-3, exceptions={}, relative=False):
    '''
    Read two pickle files and unpickle each. We assume that each resulting
    object is a dictionary. The values of each dict are either numpy arrays
    or else types that are comparable with the == operator.

    Inputs:
        fname1  = string, file name of file 1
        fname2  = string, file name of file 2
        tol     = scalar, tolerance
        exceptions = dictionary, exceptions
        relative = boolean,

    Functions called:
        dict_compare

    Objects in function:
        pkl1 =  dictionary, from first pickle file
        pkl2 = dictionary, from second pickle file

    Returns: difference between dictionaries
    '''
    try:
        pkl1 = pickle.load(open(fname1, 'rb'), encoding='latin1')
    except TypeError:
        pkl1 = pickle.load(open(fname1, 'rb'))
    try:
        pkl2 = pickle.load(open(fname2, 'rb'), encoding='latin1')
    except TypeError:
        pkl2 = pickle.load(open(fname2, 'rb'))

    return dict_compare(fname1, pkl1, fname2, pkl2, tol=tol,
                        exceptions=exceptions, relative=relative)

def comp_array(name, a, b, tol, unequal, exceptions={}, relative=False):
    '''
    Compare two arrays in the L inifinity norm
    Return True if | a - b | < tol, False otherwise
    If not equal, add items to the unequal list
    name: the name of the value being compared

    Inputs:

    Functions called:

    Objects in function:

    Returns: Boolean

    '''

    if name in exceptions:
        tol = exceptions[name]

    if not a.shape == b.shape:
        print("unequal shpaes for {0} comparison ".format(str(name)))
        unequal.append((str(name), a, b))
        return False

    else:

        if np.all(a < EPSILON) and np.all(b < EPSILON):
            return True

        if relative:
            err = abs(a - b)
            mn = np.mean(b)
            err = np.max(err / mn)
        else:
            err = np.max(abs(a - b))

        if not err < tol:
            print("diff for {0} is {1} which is NOT OK".format(str(name), err))
            unequal.append((str(name), a, b))
            return False
        else:
            print("err is {0} which is OK".format(err))
            return True

def comp_scalar(name, a, b, tol, unequal, exceptions={}, relative=False):
    '''
    Compare two scalars in the L inifinity norm
    Return True if abs(a - b) < tol, False otherwise
    If not equal, add items to the unequal list
    '''

    if name in exceptions:
        tol = exceptions[name]

    if (a < EPSILON) and (b < EPSILON):
        return True

    if relative:
        err = float(abs(a - b)) / float(b)
    else:
        err = abs(a - b)

    if not err < tol:
        print("err for {0} is {1} which is NOT OK".format(str(name), err))
        unequal.append((str(name), str(a), str(b)))
        return False
    else:
        print("err is {0} which is OK".format(err))
        return True

def dict_compare(fname1, pkl1, fname2, pkl2, tol, verbose=False,
                 exceptions={}, relative=False):
    '''
    Compare two dictionaries. The values of each dict are either
    numpy arrays
    or else types that are comparable with the == operator.
    For arrays, they are considered the same if |x - y| < tol in
    the L_inf norm.
    For scalars, they are considered the same if x - y < tol
    '''

    keys1 = set(pkl1.keys())
    keys2 = set(pkl2.keys())
    check = True
    if keys1 != keys2:
        if len(keys1) == len(keys2):
            extra1 = keys1 - keys2
            extra2 = keys2 - keys1
            msg1 = "extra items in {0}: {1}"
            print(msg1.format(fname1, extra1))
            print(msg1.format(fname2, extra2))
            return False
        elif len(keys1) > len(keys2):
            bigger = keys1
            bigger_file = fname1
            smaller = keys2
        else:
            bigger = keys2
            bigger_file = fname2
            smaller = keys1
        res = bigger - smaller
        msg = "more items in {0}: {1}"
        print(msg.format(bigger_file, res))
        return False
    else:
        unequal_items = []
        for k, v in pkl1.items():
            if type(v) == np.ndarray:
                check &= comp_array(k, v, pkl2[k], tol, unequal_items,
                                    exceptions=exceptions, relative=relative)
            else:
                try:
                    check &= comp_scalar(k, v, pkl2[k], tol, unequal_items,
                                         exceptions=exceptions, relative=relative)
                except TypeError:
                    check &= comp_array(k, np.array(v), np.array(pkl2[k]), tol,
                                        unequal_items, exceptions=exceptions,
                                        relative=relative)

        if verbose and unequal_items:
            frmt = "Name {0}"
            res = [frmt.format(x[0]) for x in unequal_items]
            print("Different arrays: ", res)
            return False

    return check
<commit_msg>create safe read pickle function<commit_after>from __future__ import print_function
'''
------------------------------------------------------------------------
Last updated 4/7/2015

Miscellaneous functions for SS.py and TPI.py.

This python files calls:
    OUTPUT/Saved_moments/wealth_data_moments.pkl

------------------------------------------------------------------------
'''

# Packages
import os
from io import StringIO
import numpy as np
try:
    import cPickle as pickle
except ImportError:
    import pickle
from pkg_resources import resource_stream, Requirement

EPSILON = 1e-10
PATH_EXISTS_ERRNO = 17

REFORM_DIR = "./OUTPUT_REFORM"
BASELINE_DIR = "./OUTPUT_BASELINE"

# Default year for model runs
DEFAULT_START_YEAR = 2018

# Latest year TaxData extrapolates to
TC_LAST_YEAR = 2027

# Year of data used (e.g. PUF or CPS year)
RECORDS_START_YEAR = 2011

# for f in (REFORM_DIR, BASELINE_DIR):
#     if not os.path.exists(f):
#         os.mkdir(f)

def mkdirs(path):
    '''
    Makes directories to save output.

    Inputs:
        path = string, path name for new directory

    Functions called: None

    Objects in function:

    Returns: N/A
    '''

    try:
        os.makedirs(path)
    except OSError as oe:
        if oe.errno == PATH_EXISTS_ERRNO:
            pass

def pct_diff_func(simul, data):
    '''
    Used to calculate the absolute percent difference between data
    moments and model moments.

    Inputs:
        simul = any shape, model moments
        data  = same shape as simul, data moments

    Functions called: None

    Objects in function:
        frac   = same shape as simul, percent difference between data and model moments
        output = same shape as simul, absolute percent difference between data and model moments

    Returns: output
    '''

    frac = (simul - data) / data
    output = np.abs(frac)
    return output

def convex_combo(var1, var2, nu):
    '''
    Takes the convex combination of two variables, where nu is in [0,1].

    Inputs:
        var1 = any shape, variable 1
        var2 = same shape as var1, variable 2
        nu   = scalar, weight on var1 in convex combination

    Functions called: None

    Objects in function:
        combo = same shape as var1, convex combination of var1 and var2

    Returns: combo
    '''

    combo = nu * var1 + (1 - nu) * var2
    return combo

def read_file(path, fname):
    '''
    Read the contents of 'path'. If it does not exist, assume the file
    is installed in a .egg file, and adjust accordingly.

    Inputs:
        path  = string, path name for new directory
        fname = string, filename

    Functions called: None

    Objects in function:
        path_in_egg
        buf
        _bytes

    Returns: file contents
    '''

    if not os.path.exists(os.path.join(path, fname)):
        path_in_egg = os.path.join("ogusa", fname)
        buf = resource_stream(Requirement.parse("ogusa"), path_in_egg)
        _bytes = buf.read()
        return StringIO(_bytes.decode("utf-8"))
    else:
        return open(os.path.join(path, fname))

def pickle_file_compare(fname1, fname2, tol=1e-3, exceptions={}, relative=False):
    '''
    Read two pickle files and unpickle each. We assume that each resulting
    object is a dictionary. The values of each dict are either numpy arrays
    or else types that are comparable with the == operator.

    Inputs:
        fname1  = string, file name of file 1
        fname2  = string, file name of file 2
        tol     = scalar, tolerance
        exceptions = dictionary, exceptions
        relative = boolean,

    Functions called:
        dict_compare

    Objects in function:
        pkl1 =  dictionary, from first pickle file
        pkl2 = dictionary, from second pickle file

    Returns: difference between dictionaries
    '''
    try:
        pkl1 = pickle.load(open(fname1, 'rb'), encoding='latin1')
    except TypeError:
        pkl1 = pickle.load(open(fname1, 'rb'))
    try:
        pkl2 = pickle.load(open(fname2, 'rb'), encoding='latin1')
    except TypeError:
        pkl2 = pickle.load(open(fname2, 'rb'))

    return dict_compare(fname1, pkl1, fname2, pkl2, tol=tol,
                        exceptions=exceptions, relative=relative)

def comp_array(name, a, b, tol, unequal, exceptions={}, relative=False):
    '''
    Compare two arrays in the L inifinity norm
    Return True if | a - b | < tol, False otherwise
    If not equal, add items to the unequal list
    name: the name of the value being compared

    Inputs:

    Functions called:

    Objects in function:

    Returns: Boolean

    '''

    if name in exceptions:
        tol = exceptions[name]

    if not a.shape == b.shape:
        print("unequal shpaes for {0} comparison ".format(str(name)))
        unequal.append((str(name), a, b))
        return False

    else:

        if np.all(a < EPSILON) and np.all(b < EPSILON):
            return True

        if relative:
            err = abs(a - b)
            mn = np.mean(b)
            err = np.max(err / mn)
        else:
            err = np.max(abs(a - b))

        if not err < tol:
            print("diff for {0} is {1} which is NOT OK".format(str(name), err))
            unequal.append((str(name), a, b))
            return False
        else:
            print("err is {0} which is OK".format(err))
            return True

def comp_scalar(name, a, b, tol, unequal, exceptions={}, relative=False):
    '''
    Compare two scalars in the L inifinity norm
    Return True if abs(a - b) < tol, False otherwise
    If not equal, add items to the unequal list
    '''

    if name in exceptions:
        tol = exceptions[name]

    if (a < EPSILON) and (b < EPSILON):
        return True

    if relative:
        err = float(abs(a - b)) / float(b)
    else:
        err = abs(a - b)

    if not err < tol:
        print("err for {0} is {1} which is NOT OK".format(str(name), err))
        unequal.append((str(name), str(a), str(b)))
        return False
    else:
        print("err is {0} which is OK".format(err))
        return True

def dict_compare(fname1, pkl1, fname2, pkl2, tol, verbose=False,
                 exceptions={}, relative=False):
    '''
    Compare two dictionaries. The values of each dict are either
    numpy arrays
    or else types that are comparable with the == operator.
    For arrays, they are considered the same if |x - y| < tol in
    the L_inf norm.
    For scalars, they are considered the same if x - y < tol
    '''

    keys1 = set(pkl1.keys())
    keys2 = set(pkl2.keys())
    check = True
    if keys1 != keys2:
        if len(keys1) == len(keys2):
            extra1 = keys1 - keys2
            extra2 = keys2 - keys1
            msg1 = "extra items in {0}: {1}"
            print(msg1.format(fname1, extra1))
            print(msg1.format(fname2, extra2))
            return False
        elif len(keys1) > len(keys2):
            bigger = keys1
            bigger_file = fname1
            smaller = keys2
        else:
            bigger = keys2
            bigger_file = fname2
            smaller = keys1
        res = bigger - smaller
        msg = "more items in {0}: {1}"
        print(msg.format(bigger_file, res))
        return False
    else:
        unequal_items = []
        for k, v in pkl1.items():
            if type(v) == np.ndarray:
                check &= comp_array(k, v, pkl2[k], tol, unequal_items,
                                    exceptions=exceptions, relative=relative)
            else:
                try:
                    check &= comp_scalar(k, v, pkl2[k], tol, unequal_items,
                                         exceptions=exceptions, relative=relative)
                except TypeError:
                    check &= comp_array(k, np.array(v), np.array(pkl2[k]), tol,
                                        unequal_items, exceptions=exceptions,
                                        relative=relative)

        if verbose and unequal_items:
            frmt = "Name {0}"
            res = [frmt.format(x[0]) for x in unequal_items]
            print("Different arrays: ", res)
            return False

    return check

def safe_read_pickle(file_path):
    '''
    This function reads a pickle from Python 2 into Python 2 or Python 3
    '''
    with open(file_path, 'rb') as f:
        try:
            obj = pickle.load(f, encoding='latin1')
        except TypeError:
            obj = pickle.load(f)
    return obj
<|endoftext|>


<commit_before># Create your views here.

from django.shortcuts import render
from django.contrib.auth.decorators import login_required
from django.core.urlresolvers import reverse, reverse_lazy
from django.shortcuts import render_to_response
from django.template import RequestContext

from django.views.generic import ListView, DetailView, CreateView, UpdateView, DeleteView, TemplateView
from braces import views
from django.views import View

from . import response_msg

from django.http import HttpResponse
from django.http import JsonResponse
from django.template import loader
from django.views.decorators.csrf import csrf_exempt
from django.core.exceptions import ObjectDoesNotExist
import json
import ast

from django.contrib.gis.geos import Point

from .models import Ambulances, TrackableDevice, Status
from .forms import AmbulanceUpdateForm, AmbulanceCreateForm, StatusCreateForm

class AmbulanceView(CreateView):
    model = Ambulances
    context_object_name = "ambulance_form"
    form_class = AmbulanceCreateForm
    success_url = reverse_lazy('ambulance_create')

    def get_context_data(self, **kwargs):
        context = super(AmbulanceView, self).get_context_data(**kwargs)
        context['ambulances'] = Ambulances.objects.all().order_by('license_plate')
        return context

class AmbulanceInfoView(views.JSONResponseMixin, View):

    def build_json(self, request):

        idFilter = request.GET.get('id')
        statusFilter = request.GET.get('status')

        ambulances = Ambulances.objects.all()
        if (idFilter):
            ambulances = ambulances.filter(license_plate=idFilter)
        if (statusFilter):
            ambulances = ambulances.filter(status=statusFilter)

        json = []

        for ambulance in ambulances:
            json.append({
                "id": ambulance.license_plate,
                "status": ambulance.status,
                "lat": repr(ambulance.location.y),
                "long": repr(ambulance.location.x)
            })
        return json

    def get_ajax(self, request):
        json = self.build_json(request)
        return self.render_json_response(json)

    def get(self, request):
        json = self.build_json(request)
        return self.render_json_response(json)

class StatusCreateView(CreateView):
    model = Status
    context_object_name = "status_form"
    form_class = StatusCreateForm
    success_url = reverse_lazy("status_create")

    def get_context_data(self, **kwargs):
        context = super(StatusCreateView, self).get_context_data(**kwargs)
        context['statuses'] = Status.objects.all().order_by('status_string')
        return context

class AmbulanceUpdateView(views.JSONResponseMixin, View):

    def update_ambulance(self, pk):
        try:
            record = Ambulances.objects.get(pk=pk)
        except ObjectDoesNotExist:
            return {
                "message": response_msg.NO_AMBULANCE_MSG["message"],
                "result": response_msg.NO_AMBULANCE_MSG["result"]
            }

        # lookup status
        status = self.request.GET.get('status')
        if status:
            # update record
            record.status = status

        # lookup location
        longitude = float(self.request.GET.get('long'))
        latitude = float(self.request.GET.get('lat'))
        if longitude and latitude:
            # update record
            location = Point(longitude, latitude, srid=4326)
            record.location = location

        # save updated record
        record.save()
        return {
            "message": response_msg.AMBULANCE_UPDATE_SUCCESS["message"],
            "result": response_msg.AMBULANCE_UPDATE_SUCCESS["result"]
        }

    def get_ajax(self, request, pk):
        record = self.update_ambulance(pk)

        json = {"status": record.status,
                "long": record.location.x,
                "lat": record.location.y
                }
        return self.render_json_response(json)

    # Through the browser, can render HTML for human-friendly viewing
    def get(self, request, pk):
        response = self.update_ambulance(pk)
        return self.render_json_response(response)

class CreateRoute(views.JSONResponseMixin, View):
    def post(self, request):
        # json_data = json.loads(request.body)
        points = ast.literal_eval(request.body)
        text = ""
        for p in points:
            text = text + p["alex"] + "\n"
        return HttpResponse(text)

class AmbulanceMap(views.JSONResponseMixin, views.AjaxResponseMixin, ListView):
    template_name = 'ambulances/ambulance_map.html'

    def get_queryset(self):
        return Ambulances.objects.all()

    #def get(self, request):
        #return render(request, 'ambulances/ambulance_map.html')

<commit_msg>long/lat fix<commit_after># Create your views here.

from django.shortcuts import render
from django.contrib.auth.decorators import login_required
from django.core.urlresolvers import reverse, reverse_lazy
from django.shortcuts import render_to_response
from django.template import RequestContext

from django.views.generic import ListView, DetailView, CreateView, UpdateView, DeleteView, TemplateView
from braces import views
from django.views import View

from . import response_msg

from django.http import HttpResponse
from django.http import JsonResponse
from django.template import loader
from django.views.decorators.csrf import csrf_exempt
from django.core.exceptions import ObjectDoesNotExist
import json
import ast

from django.contrib.gis.geos import Point

from .models import Ambulances, TrackableDevice, Status
from .forms import AmbulanceUpdateForm, AmbulanceCreateForm, StatusCreateForm

class AmbulanceView(CreateView):
    model = Ambulances
    context_object_name = "ambulance_form"
    form_class = AmbulanceCreateForm
    success_url = reverse_lazy('ambulance_create')

    def get_context_data(self, **kwargs):
        context = super(AmbulanceView, self).get_context_data(**kwargs)
        context['ambulances'] = Ambulances.objects.all().order_by('license_plate')
        return context

class AmbulanceInfoView(views.JSONResponseMixin, View):

    def build_json(self, request):

        idFilter = request.GET.get('id')
        statusFilter = request.GET.get('status')

        ambulances = Ambulances.objects.all()
        if (idFilter):
            ambulances = ambulances.filter(license_plate=idFilter)
        if (statusFilter):
            ambulances = ambulances.filter(status=statusFilter)

        json = []

        for ambulance in ambulances:
            json.append({
                "id": ambulance.license_plate,
                "status": ambulance.status,
                "lat": repr(ambulance.location.y),
                "long": repr(ambulance.location.x)
            })
        return json

    def get_ajax(self, request):
        json = self.build_json(request)
        return self.render_json_response(json)

    def get(self, request):
        json = self.build_json(request)
        return self.render_json_response(json)

class StatusCreateView(CreateView):
    model = Status
    context_object_name = "status_form"
    form_class = StatusCreateForm
    success_url = reverse_lazy("status_create")

    def get_context_data(self, **kwargs):
        context = super(StatusCreateView, self).get_context_data(**kwargs)
        context['statuses'] = Status.objects.all().order_by('status_string')
        return context

class AmbulanceUpdateView(views.JSONResponseMixin, View):

    def update_ambulance(self, pk):
        try:
            record = Ambulances.objects.get(pk=pk)
        except ObjectDoesNotExist:
            return {
                "message": response_msg.NO_AMBULANCE_MSG["message"],
                "result": response_msg.NO_AMBULANCE_MSG["result"]
            }

        # lookup status
        status = self.request.GET.get('status')
        if status:
            # update record
            record.status = status

        # lookup location
        longitude = self.request.GET.get('long')
        latitude = self.request.GET.get('lat')
        if longitude and latitude:
            # update record
            location = Point(float(longitude), float(latitude), srid=4326)
            record.location = location

        # save updated record
        record.save()
        return {
            "message": response_msg.AMBULANCE_UPDATE_SUCCESS["message"],
            "result": response_msg.AMBULANCE_UPDATE_SUCCESS["result"]
        }

    def get_ajax(self, request, pk):
        record = self.update_ambulance(pk)

        json = {"status": record.status,
                "long": record.location.x,
                "lat": record.location.y
                }
        return self.render_json_response(json)

    # Through the browser, can render HTML for human-friendly viewing
    def get(self, request, pk):
        response = self.update_ambulance(pk)
        return self.render_json_response(response)

class CreateRoute(views.JSONResponseMixin, View):
    def post(self, request):
        # json_data = json.loads(request.body)
        points = ast.literal_eval(request.body)
        text = ""
        for p in points:
            text = text + p["alex"] + "\n"
        return HttpResponse(text)

class AmbulanceMap(views.JSONResponseMixin, views.AjaxResponseMixin, ListView):
    template_name = 'ambulances/ambulance_map.html'

    def get_queryset(self):
        return Ambulances.objects.all()

    #def get(self, request):
        #return render(request, 'ambulances/ambulance_map.html')
<|endoftext|>


<commit_before>from mozdef_util.plugin_set import PluginSet
from mozdef_util.utilities.logger import logger

class AlertPluginSet(PluginSet):

    def send_message_to_plugin(self, plugin_class, message, metadata=None):
        if 'utctimestamp' in message and 'summary' in message:
            message_log_str = '{0} received message: ({1}) {2}'.format(plugin_class.__module__, message['utctimestamp'], message['summary'])
            logger.info(message_log_str)

        return plugin_class.onMessage(message), metadata
<commit_msg>Remove logger entry for alert plugins receiving alerts<commit_after>from mozdef_util.plugin_set import PluginSet

class AlertPluginSet(PluginSet):

    def send_message_to_plugin(self, plugin_class, message, metadata=None):
        return plugin_class.onMessage(message), metadata
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# dials.algorithms.indexing.indexer.py
#
#  Copyright (C) 2013 Diamond Light Source
#
#  Author: Graeme Winter, Nick Sauter, Richard Gildea
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.

from __future__ import division
import math

# we need these things

def discover_better_experimental_model(spot_positions, detector, beam,
                                       goniometer, scan):
    '''Given an attempt at indexing derive a more likely model for the
    experimental geometry.'''
    # the spot_position data are altered.  New mm centroid positions are calculated
    # first map pixel to mm
    Indexer._map_spots_pixel_to_mm_rad(
      spots=spot_positions,detector=detector,scan=scan)
    # derive a max_cell from mm spots
    # derive a grid sampling from spots

    from rstbx.indexing_api.lattice import DPS_primitive_lattice
    DPS = DPS_primitive_lattice(max_cell = 300, # max possible cell in Angstroms; guess for now
          recommended_grid_sampling_rad = 0.029, # grid sampling in radians; guess for now
          horizon_phil = phil_set)
    DPS.S0_vector = beam.get_S0_vector()
    DPS.inv_wave = 1./beam.get_wavelength()
    DPS.axis = goniometer.get_rotation_axis()
    DPS.set_detector(detector)

    # transform input into what Nick needs
    # i.e., construct a flex.vec3 double consisting of mm spots, phi in degrees
    from rstbx.array_family import flex

    data = flex.vec3_double()
    for spot in spot_positions:
      data.append((spot.centroid_position[0],
                  spot.centroid_position[1],
                  spot.centroid_position[2]*180./math.pi))
    DPS.index(raw_spot_input = data)

    # perform calculation
    if params.indexing.improve_local_scope=="origin_offset":
      new_detector = DPS.optimize_origin_offset_local_scope()
    elif params.indexing.improve_local_scope=="S0_vector":
      new_S0_vector = DPS.optimize_S0_local_scope()
      import copy
      new_beam = copy.copy(beam)
      new_beam.set_s0(new_S0_vector)

    # return these: either detector or beam

def candidate_basis_vectors_fft1d():
    pass

def candidate_basis_vectors_fft3d(rs_positions_xyz, params):
    # FIXME implement this
    candidate_basis_vectors_lattice_1 = []
    candidate_basis_vectors = [candidate_basis_vectors_lattice_1]
    return candidate_basis_vectors

def determine_basis_set(rs_positions_xyz,
                        candidate_basis_vectors_one_lattice):
    # given a list of 3 or more candidate basis vectors, decide on a good
    # basis set for indexing [triclinic A matrix]
    triclinic_a_matrix = []
    return triclinic_a_matrix

# this is some kind of bodge to define a container class type

def bravais_lattice_factory(bravais_lattice_type, unit_cell,
                            a_axis, b_axis, c_axis, penalty):
    from collections import namedtuple
    BravaisLattice = namedtuple(
        'bravais_lattice_type unit_cell a_axis b_axis c_axis penalty')
    return BravaisLattice(bravais_lattice_type, unit_cell,
                          a_axis, b_axis, c_axis, penalty)

def possible_lattices_given_basis_set(rs_positions_xyz,
                                      basis_set):
    # given rs_positions_xyz decide list of possible Bravais lattice, unit
    # cell constant combinations from an input triclinic basis

    # FIXME generate list of Bravais lattice types from the
    # bravais_lattice_factory

    return

def refine_lattice_from_raw_spot_positions(indexer, spot_positions, lattice):
    # FIXME interface for doing the refinement
    positions_mm_rad = indexer._map_spots_pixel_to_mm_strategy(spot_positions)
    # do some refinement
    return

# up to here...

class toy_validate_spots_detector(object):
    def __init__(self):
        pass

    def __call__(self, spots, detector):
        '''Make sure that the input is sensible.'''

        # FIXME make this work as follows:
        # - work through spot list; for each spot verify that panel is defined
        #   and that the position is within the bounds of the panel

        return

class Indexer(object):
    ''' The indexer base class. '''

    def __init__(self, strategies, parameters):
        ''' Initialise the indexer base class.

        Params:
            strategies: TBD; strategies for e.g. basis discovery
            parameters: TBD; the set of phil parameters

        '''
        # FIXME make these work properly by naming the strategies etc.
        self.strategies = strategies
        self.parameters = parameters
        # will work like this:
        # self.do_this = do_this
        # however N.B. that this will probably mean you cannot pickle an
        # instance of this

    def __call__(self, spots, detector, beam, goniometer = None, scan = None):
        ''' Call to index.

        Params:
            spots: The spot list inc. panel number, space for lattice ID
            detectors: dxtbx detector
            beam: beam information
            goniometer: the goniometer; optional (for e.g. still images)
            scan: the scan information; optional (for e.g. still images)

        Returns:
            TBD

        '''

        # structured loops within loops to employ input strategies to -
        #
        # - validate input (setup) *** this is not a strategy optional ***
        # - discover beam centre (setup)
        # - map spots pixels to mm (setup) *** this will depend on geometry ***
        # - map spots to RS (setup)
        # - determine candidate basis vectors (index)
        # - determine basis sets ([P1_matrix], spots_and_lattice_id) (analyse)
        # - score possible lattice for each solution (analyse)
        # - refine lattice for each solution (refine) *1
        # - reject outliers for each solution (refine)
        #
        # *1 this may also be performed with or without the lattice constraints
        #    so will need as input a BravaisLattice type...

        while self._refined:
            while not self._analysed:
                while not self._indexed:
                    while not self._setuped:
                        self._setup()
                    self._index()
                self._analyse()
            self._refine()

        return

    def _setup(self):
        # FIXME one day we should implement this
        # self._discover_beam_centre_strategy()
        # FIXME this should probably be delegated to the detector object
        self._map_spots_pixel_to_mm_rad()
        self._map_centroids_to_reciprocal_space()
        return

    def _index(self):
        self._index_strategy(self)
        # if lots of unindexed reflections:
        #     consider indexing unindexed spots
        #     store results in self._lattices
        return

    def _analyse(self):
        for lattice in self._lattices:
            self._analyse_strategy(lattice)
        return

    def _refine(self):
        for lattice in self._lattices:
            for bravais_lattice in lattice.bravais_lattices:
                self._refine_strategy(bravais_lattice, spots)

        if not self._refined:
            # perhaps need to wind back to the mapping to reciprocal space and
            # try reindexing
            return

        # FIXME may well have a think about which Bravais lattice is correct
        # here

        for lattice in self._lattices:
            # this will perform outlier searching on each lattice and on every
            # spot - we may start going back to indexing at this stage to
            # mop up unindexed reflections
            self._outlier_strategy(lattice, spots)

        return

    def set_target_cell_lattice(self, cell, lattice):
        self._indexer_cell = cell
        self._indexer_lattice = lattice
        return

    def set_max_primitive_cell(self, max_primitive_cell):
        self._indexer_max_primitive_cell = max_primitive_cell
        return

    @staticmethod
    def _map_spots_pixel_to_mm_rad(spots, detector, scan):
        from dials.algorithms.centroid import centroid_px_to_mm_panel
        # ideally don't copy, but have separate spot attributes for mm and pixel
        spots_mm = spots.deep_copy()
        for i_spot, spot in enumerate(spots_mm):
            # just a quick check for now that the reflections haven't come from
            # somewhere else
            assert spot.image_coord_mm == (0,0)

            # set reflection properties that might be needed by the dials refinement
            # engine, and convert values from pixels and image number to mm/rads
            spot.frame_number = spot.centroid_position[2]
            # XXX nasty hack - why are the centroid variances ever zero in the
            # first place?
            centroid_variance = list(spot.centroid_variance)
            for i in range(3):
                if centroid_variance[i] == 0:
                    centroid_variance[i] = 0.25
            spot.centroid_variance = centroid_variance
            centroid_position, centroid_variance, _ = centroid_px_to_mm_panel(
                detector[spot.panel_number], scan,
                spot.centroid_position,
                spot.centroid_variance,
                (1,1,1))
            spot.centroid_position = centroid_position
            spot.centroid_variance = centroid_variance
            spot.rotation_angle = centroid_position[2]
        return spots_mm

    @staticmethod
    def _map_centroids_to_reciprocal_space(spots_mm, detector, beam, goniometer):
        assert(len(detector) == 1)
        x, y, _ = spots_mm.centroid_position().parts()
        s1 = detector[0].get_lab_coord(flex.vec2_double(x,y))
        beam_vectors = s1/s1.norms() * (1/beam.get_wavelength())
        spots_mm.set_beam_vector(beam_vectors) # needed by refinement
        S = s1 - beam.get_s0()
        reciprocal_space_points = S.rotate_around_origin(
          goniometer.get_rotation_axis(),
          -reflections.rotation_angle())
        return reciprocal_space_points

    # etc.

class IndexerFactory(object):
    ''' Factory class to create indexers '''

    @staticmethod
    def get_from_somewhere_else(params):
        '''Get a different indexer implementation, which for example may
        overload __call__ internally, or something'''

        # FIXME in here check with the registry for one of these based on the
        # input PHIL parameters

        return False

    @staticmethod
    def from_parameters(params):
        ''' Given a set of parameters, construct the indexer

        Params:
            params The input parameters

        Returns:
            The indexer instance

        '''

        one_from_somewhere_else = IndexerFactory.get_from_somewhere_else(params)
        if one_from_somewhere_else:
            return one_from_somewhere_else

        # else configure a standard one with strategies

        # this is deliberately not implemented
        strategies = IndexerFactory.get_strategies_from_somewhere(params)

        # Return the indexer with the given strategies
        return Indexer(strategies, params)

    @staticmethod
    def get_strategies_from_somewhere(params):
        '''Get the strategies from somewhere, for example a registry.'''

        indexing = params.indexing

        return { } # or whatever; TBD
<commit_msg>better indexer experimental model<commit_after>#!/usr/bin/env python
#
# dials.algorithms.indexing.indexer.py
#
#  Copyright (C) 2013 Diamond Light Source
#
#  Author: Graeme Winter, Nick Sauter, Richard Gildea
#
#  This code is distributed under the BSD license, a copy of which is
#  included in the root directory of this package.

from __future__ import division
import math

# we need these things

def discover_better_experimental_model(spot_positions, detector, beam,
                                       goniometer, scan, params):
    '''Given an attempt at indexing derive a more likely model for the
    experimental geometry.'''
    # the spot_position data are altered.  New mm centroid positions are calculated
    # first map pixel to mm
    Indexer._map_spots_pixel_to_mm_rad(
      spots=spot_positions,detector=detector,scan=scan)
    # derive a max_cell from mm spots
    # derive a grid sampling from spots

    from rstbx.indexing_api.lattice import DPS_primitive_lattice
    DPS = DPS_primitive_lattice(max_cell = 300,  # max possible cell in Angstroms; guess for now
          recommended_grid_sampling_rad = 0.029, # grid sampling in radians; guess for now
          horizon_phil = params)
    DPS.S0_vector = beam.get_S0_vector()
    DPS.inv_wave = 1./beam.get_wavelength()
    DPS.axis = goniometer.get_rotation_axis()
    DPS.set_detector(detector)

    # transform input into what Nick needs
    # i.e., construct a flex.vec3 double consisting of mm spots, phi in degrees
    from rstbx.array_family import flex

    data = flex.vec3_double()
    for spot in spot_positions:
      data.append((spot.centroid_position[0],
                  spot.centroid_position[1],
                  spot.centroid_position[2]*180./math.pi))
    DPS.index(raw_spot_input = data)

    # perform calculation
    if params.indexing.improve_local_scope=="origin_offset":
      new_detector = DPS.optimize_origin_offset_local_scope()
    elif params.indexing.improve_local_scope=="S0_vector":
      new_S0_vector = DPS.optimize_S0_local_scope()
      import copy
      new_beam = copy.copy(beam)
      new_beam.set_s0(new_S0_vector)

    # return these: either detector or beam

def candidate_basis_vectors_fft1d():
    pass

def candidate_basis_vectors_fft3d(rs_positions_xyz, params):
    # FIXME implement this
    candidate_basis_vectors_lattice_1 = []
    candidate_basis_vectors = [candidate_basis_vectors_lattice_1]
    return candidate_basis_vectors

def determine_basis_set(rs_positions_xyz,
                        candidate_basis_vectors_one_lattice):
    # given a list of 3 or more candidate basis vectors, decide on a good
    # basis set for indexing [triclinic A matrix]
    triclinic_a_matrix = []
    return triclinic_a_matrix

# this is some kind of bodge to define a container class type

def bravais_lattice_factory(bravais_lattice_type, unit_cell,
                            a_axis, b_axis, c_axis, penalty):
    from collections import namedtuple
    BravaisLattice = namedtuple(
        'bravais_lattice_type unit_cell a_axis b_axis c_axis penalty')
    return BravaisLattice(bravais_lattice_type, unit_cell,
                          a_axis, b_axis, c_axis, penalty)

def possible_lattices_given_basis_set(rs_positions_xyz,
                                      basis_set):
    # given rs_positions_xyz decide list of possible Bravais lattice, unit
    # cell constant combinations from an input triclinic basis

    # FIXME generate list of Bravais lattice types from the
    # bravais_lattice_factory

    return

def refine_lattice_from_raw_spot_positions(indexer, spot_positions, lattice):
    # FIXME interface for doing the refinement
    positions_mm_rad = indexer._map_spots_pixel_to_mm_strategy(spot_positions)
    # do some refinement
    return

# up to here...

class toy_validate_spots_detector(object):
    def __init__(self):
        pass

    def __call__(self, spots, detector):
        '''Make sure that the input is sensible.'''

        # FIXME make this work as follows:
        # - work through spot list; for each spot verify that panel is defined
        #   and that the position is within the bounds of the panel

        return

class Indexer(object):
    ''' The indexer base class. '''

    def __init__(self, strategies, parameters):
        ''' Initialise the indexer base class.

        Params:
            strategies: TBD; strategies for e.g. basis discovery
            parameters: TBD; the set of phil parameters

        '''
        # FIXME make these work properly by naming the strategies etc.
        self.strategies = strategies
        self.parameters = parameters
        # will work like this:
        # self.do_this = do_this
        # however N.B. that this will probably mean you cannot pickle an
        # instance of this

    def __call__(self, spots, detector, beam, goniometer = None, scan = None):
        ''' Call to index.

        Params:
            spots: The spot list inc. panel number, space for lattice ID
            detectors: dxtbx detector
            beam: beam information
            goniometer: the goniometer; optional (for e.g. still images)
            scan: the scan information; optional (for e.g. still images)

        Returns:
            TBD

        '''

        # structured loops within loops to employ input strategies to -
        #
        # - validate input (setup) *** this is not a strategy optional ***
        # - discover beam centre (setup)
        # - map spots pixels to mm (setup) *** this will depend on geometry ***
        # - map spots to RS (setup)
        # - determine candidate basis vectors (index)
        # - determine basis sets ([P1_matrix], spots_and_lattice_id) (analyse)
        # - score possible lattice for each solution (analyse)
        # - refine lattice for each solution (refine) *1
        # - reject outliers for each solution (refine)
        #
        # *1 this may also be performed with or without the lattice constraints
        #    so will need as input a BravaisLattice type...

        while self._refined:
            while not self._analysed:
                while not self._indexed:
                    while not self._setuped:
                        self._setup()
                    self._index()
                self._analyse()
            self._refine()

        return

    def _setup(self):
        # FIXME one day we should implement this
        # self._discover_beam_centre_strategy()
        # FIXME this should probably be delegated to the detector object
        self._map_spots_pixel_to_mm_rad()
        self._map_centroids_to_reciprocal_space()
        return

    def _index(self):
        self._index_strategy(self)
        # if lots of unindexed reflections:
        #     consider indexing unindexed spots
        #     store results in self._lattices
        return

    def _analyse(self):
        for lattice in self._lattices:
            self._analyse_strategy(lattice)
        return

    def _refine(self):
        for lattice in self._lattices:
            for bravais_lattice in lattice.bravais_lattices:
                self._refine_strategy(bravais_lattice, spots)

        if not self._refined:
            # perhaps need to wind back to the mapping to reciprocal space and
            # try reindexing
            return

        # FIXME may well have a think about which Bravais lattice is correct
        # here

        for lattice in self._lattices:
            # this will perform outlier searching on each lattice and on every
            # spot - we may start going back to indexing at this stage to
            # mop up unindexed reflections
            self._outlier_strategy(lattice, spots)

        return

    def set_target_cell_lattice(self, cell, lattice):
        self._indexer_cell = cell
        self._indexer_lattice = lattice
        return

    def set_max_primitive_cell(self, max_primitive_cell):
        self._indexer_max_primitive_cell = max_primitive_cell
        return

    @staticmethod
    def _map_spots_pixel_to_mm_rad(spots, detector, scan):
        from dials.algorithms.centroid import centroid_px_to_mm_panel
        # ideally don't copy, but have separate spot attributes for mm and pixel
        spots_mm = spots.deep_copy()
        for i_spot, spot in enumerate(spots_mm):
            # just a quick check for now that the reflections haven't come from
            # somewhere else
            assert spot.image_coord_mm == (0,0)

            # set reflection properties that might be needed by the dials refinement
            # engine, and convert values from pixels and image number to mm/rads
            spot.frame_number = spot.centroid_position[2]
            # XXX nasty hack - why are the centroid variances ever zero in the
            # first place?
            centroid_variance = list(spot.centroid_variance)
            for i in range(3):
                if centroid_variance[i] == 0:
                    centroid_variance[i] = 0.25
            spot.centroid_variance = centroid_variance
            centroid_position, centroid_variance, _ = centroid_px_to_mm_panel(
                detector[spot.panel_number], scan,
                spot.centroid_position,
                spot.centroid_variance,
                (1,1,1))
            spot.centroid_position = centroid_position
            spot.centroid_variance = centroid_variance
            spot.rotation_angle = centroid_position[2]
        return spots_mm

    @staticmethod
    def _map_centroids_to_reciprocal_space(spots_mm, detector, beam, goniometer):
        assert(len(detector) == 1)
        x, y, _ = spots_mm.centroid_position().parts()
        s1 = detector[0].get_lab_coord(flex.vec2_double(x,y))
        beam_vectors = s1/s1.norms() * (1/beam.get_wavelength())
        spots_mm.set_beam_vector(beam_vectors) # needed by refinement
        S = s1 - beam.get_s0()
        reciprocal_space_points = S.rotate_around_origin(
          goniometer.get_rotation_axis(),
          -reflections.rotation_angle())
        return reciprocal_space_points

    # etc.

class IndexerFactory(object):
    ''' Factory class to create indexers '''

    @staticmethod
    def get_from_somewhere_else(params):
        '''Get a different indexer implementation, which for example may
        overload __call__ internally, or something'''

        # FIXME in here check with the registry for one of these based on the
        # input PHIL parameters

        return False

    @staticmethod
    def from_parameters(params):
        ''' Given a set of parameters, construct the indexer

        Params:
            params The input parameters

        Returns:
            The indexer instance

        '''

        one_from_somewhere_else = IndexerFactory.get_from_somewhere_else(params)
        if one_from_somewhere_else:
            return one_from_somewhere_else

        # else configure a standard one with strategies

        # this is deliberately not implemented
        strategies = IndexerFactory.get_strategies_from_somewhere(params)

        # Return the indexer with the given strategies
        return Indexer(strategies, params)

    @staticmethod
    def get_strategies_from_somewhere(params):
        '''Get the strategies from somewhere, for example a registry.'''

        indexing = params.indexing

        return { } # or whatever; TBD
<|endoftext|>