<commit_before>"""
This module includes procedures for importing and manipulating emergence data.
"""

import numpy as np
import time

from .yearcalib import uncalib_bloom, c14corr
from .abstractDataClasses import AbsGeoTimeSeries, AbsGeoTimeSeriesContainer

from .. import timestamp

def calcEmergence(sim, emergedata):
    #TODO This function shouldn't have to know what's inside any object, let
    # alone a complicated one, like sim. Consider calculating rsl first and
    # passing it and calculated times in (ts, rsl, emergedata).
    # To reference to present day
    u0 = sim.inputs.harmTrans.spectogrd(sim['upl'][-1])

    uAtLocs = []
    for ut in sim['upl']:
        ut = u0 - sim.inputs.harmTrans.spectogrd(ut)
        interpfunc = sim.inputs.grid.create_interper(ut.T)
        uAtLocs.append(interpfunc.ev(emergedata.lons, emergedata.lats))

    uAtLocs = np.array(uAtLocs).T

    data = {}
    for uAtLoc, loc in zip(uAtLocs, emergedata):
        timeseries = np.array([loc.ts,
                    np.interp(loc.ts, sim.inputs.out_times, uAtLoc)]).T
        data[loc.recnbr] = EmergeDatum(timeseries,
                                        lat=loc.lat,
                                        lon=loc.lon,
                                        desc=loc.desc,
                                        recnbr=loc.recnbr)

    return EmergeData(data)

class EmergeDatum(AbsGeoTimeSeries):
    def __repr__(self):
        return self.desc

class EmergeData(AbsGeoTimeSeriesContainer):
    """
    """
    def __init__(self, data={}):
        self.data = data
        self.form_long_vectors()
        self.W = None
        self.TIMESTAMP = timestamp()

    def transform_locs(self, basemap, inverse=False):
        xs, ys = basemap(self.locs[:,0], self.locs[:,1], inverse=inverse)
        self.locs[:,0], self.locs[:,1] = xs, ys

    def interp(self, simobject, verbose=False):
        """Interpolate uplift surfaces (xyz data at a specific t) to data
        locations (non-grid) and data times (between times calculated).

        Uses progressive linear interpolations: first the uplift at each
        outputted time is interpolated to the data locations in data.locs,
        then they are interpolated to the data times in each location.

        Parameters
        ----------
        simobject
        """
        time_start = time.clock()

        # Extract attributes from simobject
        uplift = simobject.uplift
        out_times = simobject.out_times
        grid = simobject.grid

        # Use map coordinates from sim.grid.basemap for interpolation
        x, y = grid.basemap(self.locs[:,0], self.locs[:,1])

        # interp_data will be an array of size (N_output_times, N_locations)
        # for use in interpolating the calculated emergence to the locations
        # and times at which there are data in data
        interp_data = []
        # Interpolate the calculated uplift at each time on the Lat-Lon grid
        # to the data locations.
        for uplift_at_a_time in uplift:
            interp_func = grid.create_interper(uplift_at_a_time.T)
            interp_data.append(interp_func.ev(x, y))
        interp_data = np.array(interp_data).T

        calc_vector = []
        # Interpolate the calculated uplifted at each time and data location
        # to the times of the data location.
        for interp, loc in zip(interp_data, self.data):
            calc_vector.append(np.interp(loc[:,0],
            out_times[::-1], interp[::-1]))

        # flatten the array
        calc_vector = np.array([item for l in calc_vector for item in l])

        if verbose:
            print 'Interpolation time: {0}s'.format(time.clock()-time_start)

        return calc_vector

    def residual(self, simobject, verbose=False):
        calc_vector = self.interp(simobject, verbose)
        res = calc_vector - self.long_data
        if self.W is None:
            return res/1.0
        else:
            return self.W.dot(res)

    def findLocFromLongdata(self, num):
        """Return location index and point index for index from self.long_data.
        """
        nums = np.array([len(loc) for loc in self])
        cumnums = nums.cumsum()
        index = (np.arange(len(self))[cumnums < num]).max()
        return index+1, num-cumnums[index]-1

    def form_long_vectors(self):
        """Update the long lists: long_data, long_time, and locs with currently
        encapsulated data.

        Certain numbers are stored twice, for convenience in calculating the
        residuals, in the form of long lists. These are every emergence height,
        every emergence time, and every lon,lat pair. The one-time storage
        overhead is worth the time saved from not recreating these lists on
        every iteration of an inversion.
        """

        self.long_data = np.array([None, None])
        self.long_time = np.array([None, None])
        self.locs = np.array([None, None])
        for loc in self:
            self.locs = np.vstack((self.locs, loc.loc))
            self.long_time = np.concatenate((self.long_time, loc.ts))
            self.long_data = np.concatenate((self.long_data, loc.ys))
        self.long_data = self.long_data[1:]
        self.long_time = self.long_time[1:]
        self.locs = self.locs[1:]

def importEmergeDataFromFile(filename):
    """Reads emergence data from a file where it is stored in the form...

    Stores the emergence data in a list self.data of dictionaries with keys:
            lat
            lon
            year      (of publication)
            desc      (of dataset)
            comm      (ent)
            auth      (ors of publication)
            tect      (key for tectonic activity, 1 or 0)
            tectup    (the magnitude of tectonic uplift)
            age       (some sort of age)
            recnbr    (the unique reference number)
            data_dict (a dictionary with keys 'times' and 'emerg' for ordered
                        lists of the data)

    Stores the locations in a list [[lon, lat]]
    Stores all the data in the form

    Parameters
    ----------
    filename - the file containing the data to be imported
    """

    data = {}                         # Initialize the array,

    # A list of bad points
    badpts = [137, 41, 203, 232, 234, 230, 231, 228, 229, 235, 236, 310,
                    200, 183, 318, 319, 320, 295, 296,
                    203, # Sabine high islands - uncertain age, high scatter
                    300 # Malaysia - uncertain calibration
                    ]

    locs = []

    f = open(filename, 'r')                 # open the file,
    line = f.readline()                     # read the first line to initiate
    while line:                             # the while loop.
        # read off the metadata for the location
        lat, lon, year, num = np.array(line.split('\t')[:4]).astype(np.float)
        unit, sign, typ, ageky, tect, tectup, recnbr = np.array(
                                f.readline().split('\t')[:7]).astype(np.float)
        auth = f.readline().split('\t')[0]
        desc = f.readline().split('\t')[0]
        comm = f.readline().split('\t')[0]

        if 100 <= recnbr < 400:         # apply a correction for misentered
            comm, desc = desc, comm     # data.

        locs.append([lon, lat])
        timeseries = np.zeros((num, 2))     # initialize the time series array.
        if recnbr < 400 or recnbr >= 500:
            ncol = 2
            ecol = 1
        else:
            # Newer data were given time bounds to read in as well (we're
            # skipping them for now and reading in the most-likely time).
            ncol = 5
            ecol = 3

        for i in range(int(num)):
            line = np.array(f.readline().split('\t')[:ncol]).astype(np.float)
            if line[0] in timeseries[:,0]:
                pass
            else:
                timeseries[i] = [line[0], line[ecol]]

        # Post processing of data based on metadata
        if unit == 2: timeseries[:,1] = timeseries[:,1]*0.3048      # ft -> m
        if sign >=10:
            timeseries[:,0] = timeseries[:,0]/1000.
            sign = sign/10
        if sign == 1: timeseries[:,0] = -1*timeseries[:,0]

            # Correction for time calibration
        #if ageky == 1:
        #    pass
        #elif ageky == 2:
        #    try:
        #        times = c14corr(times*1000.)/1000.
        #    except:
        #        print desc, recnbr, times
        #        raise
        #elif ageky == 5:
        #    times = uncalib_bloom(times*1000.)/1000.
        #    times = c14corr(times*1000.)/1000.

        # Ignore bad pts
        if int(recnbr) in badpts:
            pass
        # and pts whose locations are already in
        #elif [lon, lat] in locs:
        #    dupnbr = [loc.recnbr for loc in data
        #                if loc.lon==lon and loc.lat==lat]
        #    dupyr = [loc.year for loc in data
        #                if loc.lon==lon and loc.lat==lat]
        #    print ('Recnbr {0} at [{1}, {2}] is a duplicate loc with {3}'\
        #            .format(recnbr, lon, lat, dupnbr))
        #    print ('    pubdates are {0} and {1}').format(year, dupyr)
        else:
            # Form the dictionary entry for this location
            data[int(recnbr)] = EmergeDatum(timeseries, lat=lat, lon=lon,
                                    year=int(year), desc=unicode(desc,
                                    errors='ignore'), comm=comm,
                                    auth=auth, tect=int(tect), tectup=tectup,
                                    ageky=int(ageky), recnbr=recnbr)

        line = f.readline()                 # Step on.
    f.close()                               # Close the file.

    emergeData = EmergeData(data)

    emergeData.form_long_vectors()
    return emergeData
<commit_msg>Fixed error in interpolating calculated fields to data locs/times.<commit_after>"""
This module includes procedures for importing and manipulating emergence data.
"""

import numpy as np
import time

from .yearcalib import uncalib_bloom, c14corr
from .abstractDataClasses import AbsGeoTimeSeries, AbsGeoTimeSeriesContainer

from .. import timestamp

def calcEmergence(sim, emergedata):
    #TODO This function shouldn't have to know what's inside any object, let
    # alone a complicated one, like sim. Consider calculating rsl first and
    # passing it and calculated times in (ts, rsl, emergedata).
    # To reference to present day
    u0 = sim.inputs.harmTrans.spectogrd(sim['upl'][-1])

    uAtLocs = []
    for ut in sim['upl']:
        ut = u0 - sim.inputs.harmTrans.spectogrd(ut)
        interpfunc = sim.inputs.grid.create_interper(ut.T)
        uAtLocs.append(interpfunc.ev(emergedata.lons, emergedata.lats))

    uAtLocs = np.array(uAtLocs).T

    data = {}
    for uAtLoc, loc in zip(uAtLocs, emergedata):
        timeseries = np.array([loc.ts,
                    np.interp(loc.ts,
                                sim.inputs.out_times[::-1], uAtLoc[::-1])]).T
        data[loc.recnbr] = EmergeDatum(timeseries,
                                        lat=loc.lat,
                                        lon=loc.lon,
                                        desc=loc.desc,
                                        recnbr=loc.recnbr)

    return EmergeData(data)

class EmergeDatum(AbsGeoTimeSeries):
    def __repr__(self):
        return self.desc

class EmergeData(AbsGeoTimeSeriesContainer):
    """
    """
    def __init__(self, data={}):
        self.data = data
        self.form_long_vectors()
        self.W = None
        self.TIMESTAMP = timestamp()

    def transform_locs(self, basemap, inverse=False):
        xs, ys = basemap(self.locs[:,0], self.locs[:,1], inverse=inverse)
        self.locs[:,0], self.locs[:,1] = xs, ys

    def interp(self, simobject, verbose=False):
        """Interpolate uplift surfaces (xyz data at a specific t) to data
        locations (non-grid) and data times (between times calculated).

        Uses progressive linear interpolations: first the uplift at each
        outputted time is interpolated to the data locations in data.locs,
        then they are interpolated to the data times in each location.

        Parameters
        ----------
        simobject
        """
        time_start = time.clock()

        # Extract attributes from simobject
        uplift = simobject.uplift
        out_times = simobject.out_times
        grid = simobject.grid

        # Use map coordinates from sim.grid.basemap for interpolation
        x, y = grid.basemap(self.locs[:,0], self.locs[:,1])

        # interp_data will be an array of size (N_output_times, N_locations)
        # for use in interpolating the calculated emergence to the locations
        # and times at which there are data in data
        interp_data = []
        # Interpolate the calculated uplift at each time on the Lat-Lon grid
        # to the data locations.
        for uplift_at_a_time in uplift:
            interp_func = grid.create_interper(uplift_at_a_time.T)
            interp_data.append(interp_func.ev(x, y))
        interp_data = np.array(interp_data).T

        calc_vector = []
        # Interpolate the calculated uplifted at each time and data location
        # to the times of the data location.
        for interp, loc in zip(interp_data, self.data):
            calc_vector.append(np.interp(loc[:,0],
            out_times[::-1], interp[::-1]))

        # flatten the array
        calc_vector = np.array([item for l in calc_vector for item in l])

        if verbose:
            print 'Interpolation time: {0}s'.format(time.clock()-time_start)

        return calc_vector

    def residual(self, simobject, verbose=False):
        calc_vector = self.interp(simobject, verbose)
        res = calc_vector - self.long_data
        if self.W is None:
            return res/1.0
        else:
            return self.W.dot(res)

    def findLocFromLongdata(self, num):
        """Return location index and point index for index from self.long_data.
        """
        nums = np.array([len(loc) for loc in self])
        cumnums = nums.cumsum()
        index = (np.arange(len(self))[cumnums < num]).max()
        return index+1, num-cumnums[index]-1

    def form_long_vectors(self):
        """Update the long lists: long_data, long_time, and locs with currently
        encapsulated data.

        Certain numbers are stored twice, for convenience in calculating the
        residuals, in the form of long lists. These are every emergence height,
        every emergence time, and every lon,lat pair. The one-time storage
        overhead is worth the time saved from not recreating these lists on
        every iteration of an inversion.
        """

        self.long_data = np.array([None, None])
        self.long_time = np.array([None, None])
        self.locs = np.array([None, None])
        for loc in self:
            self.locs = np.vstack((self.locs, loc.loc))
            self.long_time = np.concatenate((self.long_time, loc.ts))
            self.long_data = np.concatenate((self.long_data, loc.ys))
        self.long_data = self.long_data[1:]
        self.long_time = self.long_time[1:]
        self.locs = self.locs[1:]

def importEmergeDataFromFile(filename):
    """Reads emergence data from a file where it is stored in the form...

    Stores the emergence data in a list self.data of dictionaries with keys:
            lat
            lon
            year      (of publication)
            desc      (of dataset)
            comm      (ent)
            auth      (ors of publication)
            tect      (key for tectonic activity, 1 or 0)
            tectup    (the magnitude of tectonic uplift)
            age       (some sort of age)
            recnbr    (the unique reference number)
            data_dict (a dictionary with keys 'times' and 'emerg' for ordered
                        lists of the data)

    Stores the locations in a list [[lon, lat]]
    Stores all the data in the form

    Parameters
    ----------
    filename - the file containing the data to be imported
    """

    data = {}                         # Initialize the array,

    # A list of bad points
    badpts = [137, 41, 203, 232, 234, 230, 231, 228, 229, 235, 236, 310,
                    200, 183, 318, 319, 320, 295, 296,
                    203, # Sabine high islands - uncertain age, high scatter
                    300 # Malaysia - uncertain calibration
                    ]

    locs = []

    f = open(filename, 'r')                 # open the file,
    line = f.readline()                     # read the first line to initiate
    while line:                             # the while loop.
        # read off the metadata for the location
        lat, lon, year, num = np.array(line.split('\t')[:4]).astype(np.float)
        unit, sign, typ, ageky, tect, tectup, recnbr = np.array(
                                f.readline().split('\t')[:7]).astype(np.float)
        auth = f.readline().split('\t')[0]
        desc = f.readline().split('\t')[0]
        comm = f.readline().split('\t')[0]

        if 100 <= recnbr < 400:         # apply a correction for misentered
            comm, desc = desc, comm     # data.

        locs.append([lon, lat])
        timeseries = np.zeros((num, 2))     # initialize the time series array.
        if recnbr < 400 or recnbr >= 500:
            ncol = 2
            ecol = 1
        else:
            # Newer data were given time bounds to read in as well (we're
            # skipping them for now and reading in the most-likely time).
            ncol = 5
            ecol = 3

        for i in range(int(num)):
            line = np.array(f.readline().split('\t')[:ncol]).astype(np.float)
            if line[0] in timeseries[:,0]:
                pass
            else:
                timeseries[i] = [line[0], line[ecol]]

        # Post processing of data based on metadata
        if unit == 2: timeseries[:,1] = timeseries[:,1]*0.3048      # ft -> m
        if sign >=10:
            timeseries[:,0] = timeseries[:,0]/1000.
            sign = sign/10
        if sign == 1: timeseries[:,0] = -1*timeseries[:,0]

            # Correction for time calibration
        #if ageky == 1:
        #    pass
        #elif ageky == 2:
        #    try:
        #        times = c14corr(times*1000.)/1000.
        #    except:
        #        print desc, recnbr, times
        #        raise
        #elif ageky == 5:
        #    times = uncalib_bloom(times*1000.)/1000.
        #    times = c14corr(times*1000.)/1000.

        # Ignore bad pts
        if int(recnbr) in badpts:
            pass
        # and pts whose locations are already in
        #elif [lon, lat] in locs:
        #    dupnbr = [loc.recnbr for loc in data
        #                if loc.lon==lon and loc.lat==lat]
        #    dupyr = [loc.year for loc in data
        #                if loc.lon==lon and loc.lat==lat]
        #    print ('Recnbr {0} at [{1}, {2}] is a duplicate loc with {3}'\
        #            .format(recnbr, lon, lat, dupnbr))
        #    print ('    pubdates are {0} and {1}').format(year, dupyr)
        else:
            # Form the dictionary entry for this location
            data[int(recnbr)] = EmergeDatum(timeseries, lat=lat, lon=lon,
                                    year=int(year), desc=unicode(desc,
                                    errors='ignore'), comm=comm,
                                    auth=auth, tect=int(tect), tectup=tectup,
                                    ageky=int(ageky), recnbr=recnbr)

        line = f.readline()                 # Step on.
    f.close()                               # Close the file.

    emergeData = EmergeData(data)

    emergeData.form_long_vectors()
    return emergeData
<|endoftext|>


<commit_before>#!/usr/bin/python
import databag
import click

debug=True

proxy_container="nmdproxytest"

#### Logic for taking servers in and out of rotation ####
#Here's the servers array:
# bag['production']['webcluster01']['servers'] = [
#   "server web01.newmediadenver.com:80;",
#   "server web02.newmediadenver.com:80;",
#   "server web04.newmediadenver.com:80;"
# ]
def remove_server(servers, server_to_remove, cluster):
  """
  Removes a server from the list.

  :param servers: The url entry to add to the databag
  :param server_to_remove: The name of just the server or the full server string

  :returns list of servers
  """
  index_to_remove = None
  if servers == None or type(servers} != type(list):
    print "No servers are in rotation for '{cluster}'".format(cluster=cluster)
    return []
  for index, server in enumerate(servers):
    if server_to_remove in server:
      index_to_remove = index
      break
  if index_to_remove:
    servers.pop(index_to_remove)
  else:
    print "Could not find server {server} in {cluster}[{servers}] to remove.".format(server=server_to_remove,cluster=cluster,servers=",".join(servers))

  print ','.join(servers)
  return servers

def add_server(servers, server_to_add, cluster):
  """
  Add a server to the list.

  :param servers: The url entry to add to the databag
  :param server_to_add: The name of just the server or the full server string

  :returns list of servers
  """
  # Fix the format if necessary
  if "server " not in server_to_add:
    server_to_add="server {server}".format(server=server_to_add)
  if ":80;" not in server_to_add:
    server_to_add="{server}:80;".format(server=server_to_add)

  # Check for the "None" edgecase
  if servers == None or len(servers) == 0:
    print "No servers are in rotation for '{cluster}'".format(cluster=cluster)
    servers=[]
  if any([server_to_add==server for server in servers]):
    print "Server {cluster}['{server}]' already exists in rotation. No action is required".format(cluster=cluster,server=server_to_add)
    return servers
  else:
    print "Adding server '{server}' to {cluster}.".format(server=server_to_add, cluster=cluster)
    servers.append(server_to_add)
  return servers

def get_server_list(environment):
  """
  Get the list of servers for an environment

  :param environment: 'staging' or 'production'

  :returns list of servers
  """
  proxy_databag = databag.get_databag("upstream", container=proxy_container)
  if environment=='production':
    server_list = proxy_databag[environment]['webcluster01']['servers']
  elif environment=="staging":
    server_list = proxy_databag[environment]['web01']['servers']
  else:
    raise Exception("Unrecognized environment of '{environment}. Available options are 'production' and 'development'".format(environment=environment))
    return None
  return server_list

@click.command()
@click.option('--server', prompt="Server name:", help="Name of the server you would like to add/remove")
@click.option('--environment', prompt="Which environment?", help="'staging' or 'production'", type=click.Choice(['staging', 'production']))
@click.option('--add', 'operation', flag_value='add', default=True)
@click.option('--remove', 'operation', flag_value='remove')
@click.option('--debug', 'debug', flag_value=True, default=True)
def modify_server_list(server, environment, operation, debug):
  """
  Get the list of servers for an environment
  """
  if !debug:
    exit(1)
  if debug:
    global proxy_container
    proxy_container="nmdtest"
  proxy_databag = databag.get_databag("upstream", container=proxy_container)
  if environment=="production":
    cluster='webcluster01'
  elif environment=="staging":
    cluster='web01'
  else:
    raise Exception("Unrecognized environment of '{environment}. Available options are 'production' and 'development'".format(environment=environment))
    return None
  server_list = proxy_databag[environment][cluster]['servers']
  if server_list is None:
    print "server_list is None"
  if operation == "add":
    server_list = add_server(server_list, server, cluster)
  elif operation == "remove":
    server_list = remove_server(server_list, server, cluster)
  else:
    print "Unrecognized server operation of '{operation}'".format(operation=operation)
  # Put it all back together
  proxy_databag[environment][cluster]['servers'] = server_list

  databag.save_databag(proxy_databag, bag_name="upstream", container=proxy_container)
  return True

# if __name__ == '__main__':
#   modify_server_list(server="fakeweb07.newmediadenver.com", environment="production", operation="add")
#   modify_server_list(server="fakeweb06.nmdev.us", environment="staging", operation="add")
if __name__ == '__main__':
  modify_server_list()
<commit_msg>Cleaning up debug.:<commit_after>#!/usr/bin/python
import databag
import click

proxy_container="nmdproxytest"

#### Logic for taking servers in and out of rotation ####
#Here's the servers array:
# bag['production']['webcluster01']['servers'] = [
#   "server web01.newmediadenver.com:80;",
#   "server web02.newmediadenver.com:80;",
#   "server web04.newmediadenver.com:80;"
# ]
def remove_server(servers, server_to_remove, cluster):
  """
  Removes a server from the list.

  :param servers: The url entry to add to the databag
  :param server_to_remove: The name of just the server or the full server string

  :returns list of servers
  """
  index_to_remove = None
  if servers == None or not isinstance(servers, list):
    print "No servers are in rotation for '{cluster}'".format(cluster=cluster)
    return []
  for index, server in enumerate(servers):
    if server_to_remove in server:
      index_to_remove = index
      break
  if index_to_remove:
    servers.pop(index_to_remove)
  else:
    print "Could not find server {server} in {cluster}[{servers}] to remove.".format(server=server_to_remove,cluster=cluster,servers=",".join(servers))

  print "Remaining servers: [%s]" % ','.join(servers)
  return servers

def add_server(servers, server_to_add, cluster):
  """
  Add a server to the list.

  :param servers: The url entry to add to the databag
  :param server_to_add: The name of just the server or the full server string

  :returns list of servers
  """
  # Fix the format if necessary
  if "server " not in server_to_add:
    server_to_add="server {server}".format(server=server_to_add)
  if ":80;" not in server_to_add:
    server_to_add="{server}:80;".format(server=server_to_add)

  # Check for the "None" edgecase
  if servers == None or len(servers) == 0:
    print "No servers are in rotation for '{cluster}'".format(cluster=cluster)
    servers=[]
  if any([server_to_add==server for server in servers]):
    print "Server {cluster}['{server}]' already exists in rotation. No action is required".format(cluster=cluster,server=server_to_add)
    return servers
  else:
    print "Adding server '{server}' to {cluster}.".format(server=server_to_add, cluster=cluster)
    servers.append(server_to_add)
  return servers

def get_server_list(environment):
  """
  Get the list of servers for an environment

  :param environment: 'staging' or 'production'

  :returns list of servers
  """
  proxy_databag = databag.get_databag("upstream", container=proxy_container)
  if environment=='production':
    server_list = proxy_databag[environment]['webcluster01']['servers']
  elif environment=="staging":
    server_list = proxy_databag[environment]['web01']['servers']
  else:
    raise Exception("Unrecognized environment of '{environment}. Available options are 'production' and 'development'".format(environment=environment))
    return None
  return server_list

@click.command()
@click.option('--server', prompt="Server name:", help="Name of the server you would like to add/remove")
@click.option('--environment', prompt="Which environment?", help="'staging' or 'production'", type=click.Choice(['staging', 'production']))
@click.option('--add', 'operation', flag_value='add', default=True)
@click.option('--remove', 'operation', flag_value='remove')
@click.option('--debug', 'debug', flag_value=True, default=True)
def modify_server_list(server, environment, operation, debug):
  """
  Get the list of servers for an environment
  """
  if debug:
    global proxy_container
    proxy_container="nmdtest"
  proxy_databag = databag.get_databag("upstream", container=proxy_container)
  if environment=="production":
    cluster='webcluster01'
  elif environment=="staging":
    cluster='web01'
  else:
    raise Exception("Unrecognized environment of '{environment}. Available options are 'production' and 'development'".format(environment=environment))
    return None
  server_list = proxy_databag[environment][cluster]['servers']
  if server_list is None:
    print "server_list is None"
  if operation == "add":
    server_list = add_server(server_list, server, cluster)
  elif operation == "remove":
    server_list = remove_server(server_list, server, cluster)
  else:
    print "Unrecognized server operation of '{operation}'".format(operation=operation)
  # Put it all back together
  proxy_databag[environment][cluster]['servers'] = server_list

  databag.save_databag(proxy_databag, bag_name="upstream", container=proxy_container)
  return True

# if __name__ == '__main__':
#   modify_server_list(server="fakeweb07.newmediadenver.com", environment="production", operation="add")
#   modify_server_list(server="fakeweb06.nmdev.us", environment="staging", operation="add")
if __name__ == '__main__':
  modify_server_list()
<|endoftext|>


<commit_before>import logging
from flask import Flask, Blueprint, g, redirect, request, abort
from flask.ext.babel import Babel
from flask_debugtoolbar import DebugToolbarExtension
from flask.ext.mail import Mail
from sassutils.wsgi import SassMiddleware
from babel.support import LazyProxy
from logic import oauth
from logic.db import MongoHandler

logging.basicConfig(level=logging.DEBUG)

app = Flask(__name__, instance_relative_config=False)

# Load the default configuration
app.config.from_object('config.default')

# Load the configuration from the instance folder
app.config.from_pyfile('../config/settings.py')

# Load the file specified by the APP_CONFIG_FILE environment variable
# Variables defined here will override those in the default configuration
app.config.from_envvar('APP_CONFIG_FILE')

# Setup sass auto-compiling
app.wsgi_app = SassMiddleware(app.wsgi_app, {
	'databasic': ('static/sass', 'static/css', '/static/css')
})

babel = Babel(app)
mongo = MongoHandler(app)
mail = Mail(app)
# uncomment to use toolbar (this slows the app down quite a bit)
# toolbar = DebugToolbarExtension(app)

@app.before_request
def before():
	if request.view_args and 'lang_code' in request.view_args:
		if request.view_args['lang_code'] not in ('es', 'en'):
			return abort(404) # bail on invalid language
		g.current_lang = request.view_args['lang_code']
		request.view_args.pop('lang_code')

@babel.localeselector
def get_locale():
	return g.get('current_lang', 'en')

@app.route('/')
def index():
	return redirect('/' + get_locale())

@app.route('/auth')
def auth():
	if 'code' not in request.args:
		print 'permission was not granted'
	else:
		print request.args['code']
		oauth.authorize(request.args['code'])
	return redirect(oauth.redirect_to())

from databasic.views import home
from databasic.views import samediff
from databasic.views import wordcounter
from databasic.views import wtfcsv
app.register_blueprint(home.mod)
app.register_blueprint(samediff.mod)
app.register_blueprint(wordcounter.mod)#, subdomain='wordcounter')
app.register_blueprint(wtfcsv.mod)

<commit_msg>more basedir stuff<commit_after>import logging, os
from flask import Flask, Blueprint, g, redirect, request, abort
from flask.ext.babel import Babel
from flask_debugtoolbar import DebugToolbarExtension
from flask.ext.mail import Mail
from sassutils.wsgi import SassMiddleware
from babel.support import LazyProxy
from logic import oauth
from logic.db import MongoHandler

logging.basicConfig(level=logging.DEBUG)

app = Flask(__name__, instance_relative_config=False)

# Load the default configuration
app.config.from_object('config.default')

# Load the configuration from the instance folder
app.config.from_pyfile('../config/settings.py')

# Load the file specified by the APP_CONFIG_FILE environment variable
# Variables defined here will override those in the default configuration
app.config.from_envvar('APP_CONFIG_FILE')

# Setup sass auto-compiling
basedir = os.path.dirname(os.path.abspath(__file__))
app.wsgi_app = SassMiddleware(app.wsgi_app, {
	'databasic': ('static/sass', 'static/css', '/static/css')
})

babel = Babel(app)
mongo = MongoHandler(app)
mail = Mail(app)
# uncomment to use toolbar (this slows the app down quite a bit)
# toolbar = DebugToolbarExtension(app)

@app.before_request
def before():
	if request.view_args and 'lang_code' in request.view_args:
		if request.view_args['lang_code'] not in ('es', 'en'):
			return abort(404) # bail on invalid language
		g.current_lang = request.view_args['lang_code']
		request.view_args.pop('lang_code')

@babel.localeselector
def get_locale():
	return g.get('current_lang', 'en')

@app.route('/')
def index():
	return redirect('/' + get_locale())

@app.route('/auth')
def auth():
	if 'code' not in request.args:
		print 'permission was not granted'
	else:
		print request.args['code']
		oauth.authorize(request.args['code'])
	return redirect(oauth.redirect_to())

from databasic.views import home
from databasic.views import samediff
from databasic.views import wordcounter
from databasic.views import wtfcsv
app.register_blueprint(home.mod)
app.register_blueprint(samediff.mod)
app.register_blueprint(wordcounter.mod)#, subdomain='wordcounter')
app.register_blueprint(wtfcsv.mod)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import os
import sys
import shutil
import datetime
import subprocess
import tarfile

from django.conf import settings

def get_base_path():
    if hasattr(settings, 'DATA_DOWNLOADER_PATH'):
        base_path = settings.DATA_DOWNLOADER_PATH
    else:
        base_path = os.path.join(settings.BASE_DIR,
                                 'project/protected_medias/datas')
    return base_path

class Dump(object):
    def __init__(self, data_type, base_path=None):
        assert data_type in ('data', 'db', 'media')
        self.data_type = data_type
        self.base_path = base_path or get_base_path()

    def get_metadata(self):
        try:
            infos = os.stat(self.path)
            date = datetime.datetime.fromtimestamp(int(infos.st_mtime))
            return {
                'date': date,
                'size': infos.st_size
            }
        except OSError:
            return {
                'date': None,
                'size': None
            }

    @property
    def project_name(self):
        return os.path.basename(settings.BASE_DIR)

    @property
    def archive_name(self):
        return "%s_%s.tar.gz" % (self.project_name, self.data_type)

    @property
    def path(self):
        return os.path.join(self.base_path, self.archive_name)

    def _get_datadump_bin(self):
        if hasattr(settings, 'DATADUMP_BIN_PATH'):
            return settings.DATADUMP_BIN_PATH
        if hasattr(sys, 'real_prefix'):
            # Running in a virtual env
            candidate = os.path.join(sys.prefix, 'bin/datadump')
            if os.path.exists(candidate):
                return candidate

        # Try a globally installed instance
        return 'datadump'

    def _dump_media(self):
        return [
            settings.MEDIA_ROOT,
        ]

    def _dump_db(self):
        self._clean_dumps_path()
        subprocess.check_output(self._get_datadump_bin())
        return [
            os.path.join(settings.BASE_DIR, 'dumps')
        ]

    def _clean_dumps_path(self):
        dumps_path = os.path.join(settings.BASE_DIR, 'dumps')
        if os.path.exists(dumps_path):
            shutil.rmtree(dumps_path)
        os.mkdir(dumps_path)

    def create(self):
        folders = []
        if self.data_type == 'db':
            folders.extend(self._dump_db())
        elif self.data_type == 'media':
            folders.extend(self._dump_media())
        elif self.data_type == 'data':
            folders.extend(self._dump_media())
            folders.extend(self._dump_db())

        with tarfile.open(self.path, "w:gz") as tar:
            for folder in folders:
                tar.add(folder)

    def destroy(self):
        os.remove(self.path)
<commit_msg>use relatif path<commit_after># -*- coding: utf-8 -*-

import os
import sys
import shutil
import datetime
import subprocess
import tarfile

from django.conf import settings

def get_base_path():
    if hasattr(settings, 'DATA_DOWNLOADER_PATH'):
        base_path = settings.DATA_DOWNLOADER_PATH
    else:
        base_path = os.path.join(settings.BASE_DIR,
                                 'project/protected_medias/datas')
    return base_path

class Dump(object):
    def __init__(self, data_type, base_path=None):
        assert data_type in ('data', 'db', 'media')
        self.data_type = data_type
        self.base_path = base_path or get_base_path()

    def get_metadata(self):
        try:
            infos = os.stat(self.path)
            date = datetime.datetime.fromtimestamp(int(infos.st_mtime))
            return {
                'date': date,
                'size': infos.st_size
            }
        except OSError:
            return {
                'date': None,
                'size': None
            }

    @property
    def project_name(self):
        return os.path.basename(settings.BASE_DIR)

    @property
    def archive_name(self):
        return "%s_%s.tar.gz" % (self.project_name, self.data_type)

    @property
    def path(self):
        return os.path.join(self.base_path, self.archive_name)

    def _get_datadump_bin(self):
        if hasattr(settings, 'DATADUMP_BIN_PATH'):
            return settings.DATADUMP_BIN_PATH
        if hasattr(sys, 'real_prefix'):
            # Running in a virtual env
            candidate = os.path.join(sys.prefix, 'bin/datadump')
            if os.path.exists(candidate):
                return candidate

        # Try a globally installed instance
        return 'datadump'

    def _dump_media(self):
        return [settings.MEDIA_ROOT.replace("%s/" % os.getcwd(), ''), ]

    def _dump_db(self):
        self._clean_dumps_path()
        subprocess.check_output(self._get_datadump_bin())
        dump_path = os.path.join(settings.BASE_DIR, 'dumps')
        return [dump_path.replace("%s/" % os.getcwd(), ''), ]

    def _clean_dumps_path(self):
        dumps_path = os.path.join(settings.BASE_DIR, 'dumps')
        if os.path.exists(dumps_path):
            shutil.rmtree(dumps_path)
        os.mkdir(dumps_path)

    def create(self):
        folders = []
        if self.data_type == 'db':
            folders.extend(self._dump_db())
        elif self.data_type == 'media':
            folders.extend(self._dump_media())
        elif self.data_type == 'data':
            folders.extend(self._dump_media())
            folders.extend(self._dump_db())

        with tarfile.open(self.path, "w:gz") as tar:
            for folder in folders:
                tar.add(folder)

    def destroy(self):
        os.remove(self.path)
<|endoftext|>


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

from os.path import isdir, join, basename, dirname
from os import listdir
import numpy as np
import h5py
import yaml
import glob
import re
from copy import copy
from image_toolbox import config
from datafusion import config as df_config
from image_toolbox.image import Image
from image_toolbox.experiment import Experiment
from image_toolbox.util import load_config

'''
Script for fusing measurement data generated by Jterator for use with TissueMAPS.

Jterator generates an HDF5 file for each job. Here we load all files for a
particular experiment (potentially consisting of different sub-experiments,
i.e. cycles) and merge their content into a single HDF5 file.
'''

header_mapper = df_config['CHANNEL_MAPPER']

names = list()
count = 0
for i, cycle in header_mapper.iteritems():
    if isinstance(cycle, dict):
        for old, new in cycle.iteritems():
            count += 1
            names.append(new)

if len(np.unique(names)) < count:
    raise Exception('Names have to be unique.')

def get_jtproject_names(cycle_dir):
    '''
    Get name of a Jterator project, i.e. a sub-folder in the experiment
    directory that contains a .pipe file.
    '''
    return [f for f in listdir(cycle_dir)
            if isdir(join(cycle_dir, f))
            and glob.glob(join(cycle_dir, f, '*.pipe'))]

def image_name_from_joblist(joblist, data_filename):
    job_id = int(re.search(r'_(\d+).data$', data_filename).group(1).lstrip('0'))
    return joblist[job_id].values()[0]

def rename_features(feature_names, mapper):
    '''
    Rename feature, i.e. replace substring in a feature name by another
    as defined in a 'mapper' dictionary.

    Parameters:
        :feature_names:     List of strings.
        :mapper:            Dictionary. keys - old names, values - new names.

    Returns:
        renamed features (list of strings)
    '''
    for i, feature in enumerate(feature_names):
        for j, substring in enumerate(mapper.keys()):
            r = re.compile(substring)
            match = re.search(r, feature)
            if match:
                feature_names[i] = re.sub(mapper.keys()[j],
                                          mapper.values()[j],
                                          feature_names[i])
    return feature_names

def get_data_files(project_dir):
    '''
    List all data HDF5 files for a given Jterator project.
    '''
    return sorted(glob.glob(join(project_dir, 'data', '*.data')))

def merge_data(data_files, names, as_int=False):
    '''
    Merge Jterator data of one experiment cycle stored in several HDF5 files.

    Parameters:
    :data_files:        Paths to Jterator data files : list.
    :names:             Names of a dataset in a HDF5 file : list.
    :as_int:            Convert to integer datatype : bool?

    :returns:           Dataset of shape nxp,
                        where n is the number of objects
                        and p the number of features : ndarray.
    '''
    joblist_filename = glob.glob(join(dirname(data_files[0]), '..',
                                 '*.jobs'))[0]
    joblist = yaml.load(open(joblist_filename).read())
    ids = list()
    feature_names = list()
    for job_ix, filename in enumerate(data_files):

        f = h5py.File(filename, 'r')
        if not f.keys():
            raise Exception('File "%s" is empty' % filename)

        for feat_ix, name in enumerate(names):

            if name not in f.keys():
                raise Exception('Dataset "%s" does not exist in file "%s'
                                % (name, filename))
            ids = re.search(r'.*ObjectIds.*', name, re.IGNORECASE)
            if ids:
                # Get positional information from filename
                im_file = image_name_from_joblist(joblist, filename)
                image = Image(im_file, config)
                (row, column) = image.coordinates  # returns 0-based indices
                site = image.site
                # Translate site specific ids to global ids
                nitems = len(f.values()[0])
                feat = np.vstack((np.repeat(site, nitems),
                                 np.repeat(row, nitems),
                                 np.repeat(column, nitems),
                                 f[ids.group(0)][:nitems]))
                name = ['ID_site', 'ID_row', 'ID_column', 'ID_object']
            else:
                feat = f[name][()]
                name = [name]
            feat = np.matrix(feat)
            if feat.shape[0] < feat.shape[1]:
                feat = feat.H
            if feat_ix == 0:
                feature = feat
            else:
                feature = np.hstack((feature, feat))
            if job_ix == 0:
                feature_names += name

        f.close()

        if job_ix == 0:
            features = feature
        else:
            features = np.vstack((features, feature))

    feature_names = np.array(feature_names, dtype=np.string_)

    if as_int:
        return (features.astype(int), feature_names)
    else:
        return (features, feature_names)

def build_global_ids(ids, id_names):
    '''
    Build global, continuous ids from local (image specific) ids in combination
    with row and column ids (position of the image in the acquisition "snake").

    Parameters:
    :ids:               numpy matrix with dimensions nx3, where n is the number
                        of objects
    :id_names:          numpy array of strings specifying the 3 ids:
                        "ID_row", "ID_column", "ID_object"

    Returns:
        numpy array with dimensions 1xn
    '''
    global_ids = list()
    ix_row = np.where(id_names == 'ID_row')[0]
    ix_col = np.where(id_names == 'ID_column')[0]
    ix_local = np.where(id_names == 'ID_object')[0]
    # avoid 'matrix' !!!
    rows = np.squeeze(np.asarray(ids[:, ix_row]))
    columns = np.squeeze(np.asarray(ids[:, ix_col]))
    n_row = np.max(ids[:, ix_row])
    n_col = np.max(ids[:, ix_col])
    global_ids = np.squeeze(np.asarray(ids[:, ix_local])).copy()

    offset = 0
    for r in range(1, n_row+1):
        for c in range(1, n_col+1):
            # Get the indices of the rows that belong to the site with row r
            # and column c
            ix = np.where(np.logical_and(rows == r, columns == c))[0]
            local_ids = np.squeeze(np.asarray(ids[ix, ix_local]))
            # Add a constant to the local ids and save them in the vector
            # global_ids
            global_ids[ix] = local_ids + offset
            offset = np.max(local_ids + offset)

    return global_ids

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        description='Merge data generated by Jterator for use in tissueMAPS.')

    parser.add_argument('experiment_dir', nargs='*',
                        help='absolute path to experiment directory')

    parser.add_argument('-o', '--output', dest='output_dir', required=True,
                        help='directory where the HDF5 file should be saved')

    parser.add_argument('-c', '--config', dest='config',
                        help='path to custom yaml configuration file \
                        (defaults to "datafusion" configuration)')

    args = parser.parse_args()

    if args.config:
        # Overwrite default "datafusion" configuration
        df_config = load_config(args.config)

    exp_dir = args.experiment_dir[0]
    output_dir = args.output_dir

    segmentation_project_dir = join(exp_dir,
                                    df_config['SEGMENTATION_PROJECT'].format(
                                            experiment=basename(exp_dir)))

    if not exp_dir:
        raise Exception('Project directory "%s" does not exist.'
                        % exp_dir)

    if not output_dir:
        raise Exception('Output directory "%s" does not exist.' % output_dir)

    output_filename = join(output_dir, 'data.h5')

    cycles = Experiment(exp_dir, config).subexperiments

    data_header = np.array(list(), dtype=np.string_)
    for i, c in enumerate(cycles):

        print('. Extracting features of cycle #%d: %s' %
              (c.cycle, c.name))

        cycle_dir = join(exp_dir, c.name)
        project_names = get_jtproject_names(cycle_dir)

        for project_name in project_names:

            project_dir = join(cycle_dir, project_name)

            if project_dir == segmentation_project_dir:
                # data of segmentation project is handled separately
                continue

            print '.. Merging data of Jterator project "%s"' % project_name

            data_files = get_data_files(project_dir)

            f = h5py.File(data_files[0], 'r')
            feature_names = f.keys()
            f.close()

            if 'OriginalObjectIds' not in feature_names:
                raise Exception('Files must contain a dataset called "%s"' %
                                'OriginalObjectIds')

            (features, feature_names) = merge_data(data_files,
                                                   names=feature_names)

            if features.shape[1] != len(feature_names):
                raise Exception('Number of features in dataset and length of '
                                'list with feature names must be identical.')

            if i == 0:
                data = features
            else:
                data = np.hstack((data, features))

            feature_names = rename_features(feature_names,
                                            header_mapper[c.cycle])
            # # convert to safe string format for HDF5
            # feature_names = np.array(map(np.string_, feature_names))
            data_header = np.hstack((data_header, feature_names))

    print '. Combining data from different cycles'
    (data_header, unique_ix) = np.unique(data_header, return_index=True)
    # TODO: sanity checks
    data = data[:, unique_ix]

    # separate object ids from dataset
    print '. Building global object ids'
    ids_ix = np.where(np.array([re.search('ID', i) is not None
                                for i in data_header]))[0]
    ids = data[:, ids_ix].astype(int)
    ids_header = data_header[ids_ix]
    data = np.delete(data, ids_ix, axis=1)
    data_header = np.delete(data_header, ids_ix)
    global_ids = build_global_ids(ids, ids_header)

    print '. Separate features belonging to different objects'
    objects = [re.match(r'^([^_]+)', name).group(1) for name in data_header]
    (objects, object_ix) = np.unique(objects, return_inverse=True)

    print '. Writing fused data into HDF5 file "%s"' % output_filename
    f = h5py.File(output_filename, 'w')  # truncate file if exists
    f.create_dataset('parent', data=np.string_('cells'))  # hard-coded for now
    for i, obj in enumerate(objects):

        print '. Writing data for object "%s"' % obj

        obj_ix = object_ix == i
        obj_name = obj.lower()  # use lower case consistently

        location = 'objects/%s/original-ids' % obj_name
        f.create_dataset(location, data=ids)
        f[location].attrs.__setitem__('names', ids_header)

        f.create_dataset('objects/%s/ids' % obj_name, data=global_ids)

        # data_files = get_data_files(segmentation_project_dir)

        # centroids = merge_data(data_files, names=['%s_Centroids' % obj])
        # # TODO: centroids have to be global!!!
        # # illuminati.segment.compute_cell_centroids()
        # location = 'objects/%s/centroids' % obj_name
        # f.create_dataset(location, data=centroids)
        # f[location].attrs.__setitem__('names', np.array(['y', 'x']))

        # border = merge_data(data_files, names=['%s_BorderIx' % obj],
        #                     as_int=True)
        # f.create_dataset('objects/%s/border' % obj_name, data=border)

        location = 'objects/%s/features' % obj_name
        f.create_dataset(location, data=data[:, obj_ix])
        # Add the 'data_header' as an attribute
        f[location].attrs.__setitem__('names', data_header[obj_ix])

    f.close()
<commit_msg>Fix bugs and add documentation in datafusion<commit_after>#! /usr/bin/env python
# encoding: utf-8

from os.path import isdir, join, basename, dirname
from os import listdir
import numpy as np
import h5py
import yaml
import glob
import re
from copy import copy
from image_toolbox import config
from datafusion import config as df_config
from image_toolbox.image import Image
from image_toolbox.experiment import Experiment
from image_toolbox.util import load_config

'''
Script for fusing measurement data generated by Jterator for use with TissueMAPS.

Jterator generates an HDF5 file for each job. Here we load all files for a
particular experiment (potentially consisting of different sub-experiments,
i.e. cycles) and merge their content into a single HDF5 file.
'''

header_mapper = df_config['CHANNEL_MAPPER']

names = list()
count = 0
for i, cycle in header_mapper.iteritems():
    if isinstance(cycle, dict):
        for old, new in cycle.iteritems():
            count += 1
            names.append(new)

if len(np.unique(names)) < count:
    raise Exception('Names have to be unique.')

def get_jtproject_names(cycle_dir):
    '''
    Get name of a Jterator project, i.e. a sub-folder in the experiment
    directory that contains a .pipe file.
    '''
    return [f for f in listdir(cycle_dir)
            if isdir(join(cycle_dir, f))
            and glob.glob(join(cycle_dir, f, '*.pipe'))]

def image_name_from_joblist(joblist, data_filename):
    job_id = int(re.search(r'_(\d+).data$', data_filename).group(1).lstrip('0'))
    return joblist[job_id].values()[0]

def rename_features(feature_names, mapper):
    '''
    Rename feature, i.e. replace substring in a feature name by another
    as defined in a 'mapper' dictionary.

    Parameters:
        :feature_names:     List of strings.
        :mapper:            Dictionary. keys - old names, values - new names.

    Returns:
        renamed features (list of strings)
    '''
    for i, feature in enumerate(feature_names):
        for j, substring in enumerate(mapper.keys()):
            r = re.compile(substring)
            match = re.search(r, feature)
            if match:
                feature_names[i] = re.sub(mapper.keys()[j],
                                          mapper.values()[j],
                                          feature_names[i])
    return feature_names

def get_data_files(project_dir):
    '''
    List all data HDF5 files for a given Jterator project.
    '''
    return sorted(glob.glob(join(project_dir, 'data', '*.data')))

def fuse_data(data_files, names, as_int=False):
    '''
    Fuse Jterator data of one experiment cycle stored in several HDF5 files.

    Parameters:
    :data_files:        Paths to Jterator data files : list.
    :names:             Names of a dataset in a HDF5 file : list.
    :as_int:            Convert to integer datatype : bool?

    :returns:           Dataset of shape nxp,
                        where n is the number of objects
                        and p the number of features : ndarray.
    '''
    joblist_files = glob.glob(join(dirname(data_files[0]), '..', '*.jobs'))

    joblist_filename = joblist_files[0]  # there should only be one
    joblist = yaml.load(open(joblist_filename).read())
    ids = list()
    feature_names = list()
    for job_ix, filename in enumerate(data_files):

        f = h5py.File(filename, 'r')
        if not f.keys():
            raise Exception('File "%s" is empty' % filename)

        for feat_ix, name in enumerate(names):
            # TODO: use pandas instead of numpy
            if name not in f.keys():
                raise ValueError('Dataset "%s" does not exist in file "%s'
                                 % (name, filename))
            ids = re.search(r'.*ObjectIds.*', name, re.IGNORECASE)
            if ids:
                # Get positional information from filename
                im_file = image_name_from_joblist(joblist, filename)
                image = Image(im_file, config)
                (row, column) = image.coordinates  # returns 0-based indices
                site = image.site
                # Translate site specific ids to global ids
                nitems = len(f.values()[0])
                feat = np.vstack((np.repeat(site, nitems),
                                 np.repeat(row, nitems),
                                 np.repeat(column, nitems),
                                 f[ids.group(0)][:nitems]))
                name = ['ID_site', 'ID_row', 'ID_column', 'ID_object']
            else:
                feat = f[name][()]
                name = [name]
            feat = np.matrix(feat)
            if feat.shape[0] < feat.shape[1]:
                feat = feat.H
            if feat_ix == 0:
                feature = feat
            else:
                feature = np.hstack((feature, feat))
            if job_ix == 0:
                feature_names += name

        f.close()

        if job_ix == 0:
            features = feature
        else:
            features = np.vstack((features, feature))

    feature_names = np.array(feature_names, dtype=np.string_)

    # TODO: convert back to array (avoid matrix)

    if as_int:
        return (features.astype(int), feature_names)
    else:
        return (features, feature_names)

def build_global_ids(ids, id_names):
    '''
    Build global, continuous ids from local (image specific) ids in combination
    with row and column ids (position of the image in the acquisition "snake").

    Parameters:
    :ids:               numpy matrix with dimensions nx3, where n is the number
                        of objects
    :id_names:          numpy array of strings specifying the 4 ids:
                        "ID_site", "ID_row", "ID_column", "ID_object"

    Returns:
        numpy array with dimensions 1xn
    '''
    global_ids = list()
    ix_row = np.where(id_names == 'ID_row')[0]
    ix_col = np.where(id_names == 'ID_column')[0]
    ix_local = np.where(id_names == 'ID_object')[0]
    # avoid 'matrix' !!!
    rows = np.squeeze(np.asarray(ids[:, ix_row]))
    columns = np.squeeze(np.asarray(ids[:, ix_col]))
    n_row = np.max(ids[:, ix_row])
    n_col = np.max(ids[:, ix_col])
    global_ids = np.squeeze(np.asarray(ids[:, ix_local])).copy()

    offset = 0
    for r in range(1, n_row+1):
        for c in range(1, n_col+1):
            # Get the indices that belong to the site with row r and column c
            ix = np.where(np.logical_and(rows == r, columns == c))[0]
            local_ids = np.squeeze(np.asarray(ids[ix, ix_local]))
            # Add a constant to the local ids and save them in the vector
            # global_ids
            global_ids[ix] = local_ids + offset
            offset = np.max(local_ids + offset)

    return global_ids

def calc_global_centroids(local_centroids, ids, id_names, image_size):
    '''
    Calculate global centroid coordinates from site specific coordinates.

    Parameters:
    :local_centroids:   numpy array of of size nx2, where n is the number of
                        objects, specifying the row, column coordinates
                        of each object at the local site (i.e. image)
    :ids:               numpy matrix with dimensions nx3, where n is the number
                        of objects
    :id_names:          numpy array of strings specifying the 4 ids:
                        "ID_site", "ID_row", "ID_column", "ID_object"
    :image_size:        dimensions of an individual site : (int, int)

    Returns:
        numpy array with dimensions nx2
    '''
    ix_row = np.where(id_names == 'ID_row')[0]
    ix_col = np.where(id_names == 'ID_column')[0]

    rows = np.squeeze(np.asarray(ids[:, ix_row]))
    columns = np.squeeze(np.asarray(ids[:, ix_col]))
    n_row = np.max(ids[:, ix_row])
    n_col = np.max(ids[:, ix_col])
    global_centroids = local_centroids.copy()

    y_offset = 0
    x_offset = 0
    for r in range(1, n_row+1):
        for c in range(1, n_col+1):
            # Get the centroids that belong to the site with row r and column c
            ix = np.where(np.logical_and(rows == r, columns == c))[0]
            global_centroids[ix, 0] = (local_centroids[ix, 0] + y_offset).reshape(-1)
            global_centroids[ix, 1] = (local_centroids[ix, 1] + x_offset).reshape(-1)
            y_offset = np.max(image_size[0] + y_offset)
            x_offset = np.max(image_size[1] + x_offset)

    return global_centroids

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        description='Merge data generated by Jterator for use in tissueMAPS.')

    parser.add_argument('experiment_dir', nargs='*',
                        help='absolute path to experiment directory')

    parser.add_argument('-o', '--output', dest='output_dir', required=True,
                        help='directory where the HDF5 file should be saved')

    parser.add_argument('-c', '--config', dest='config',
                        help='path to custom yaml configuration file \
                        (defaults to "datafusion" configuration)')

    args = parser.parse_args()

    if args.config:
        # Overwrite default "datafusion" configuration
        df_config = load_config(args.config)

    exp_dir = args.experiment_dir[0]
    output_dir = args.output_dir

    if not exp_dir:
        raise ValueError('Project directory "%s" does not exist.'
                        % exp_dir)

    if not output_dir:
        raise ValueError('Output directory "%s" does not exist.' % output_dir)

    output_filename = join(output_dir, 'data.h5')

    experiment = Experiment(exp_dir, config)
    cycles = experiment.subexperiments

    output_filename = experiment.data_filename

    # Determine the Jterator project that was used to segment objects.
    # It will be treated separately, because it outputs different data.
    segmentation_project_dir = join(exp_dir,
                                    df_config['SEGMENTATION_PROJECT'].format(
                                            experiment=basename(exp_dir)))

    ###########################################################################
    #                 load individual datasets and fuse them                  #
    ###########################################################################

    data_header = np.array(list(), dtype=np.string_)
    for i, c in enumerate(cycles):

        print('. Extracting features of cycle #%d: %s' %
              (c.cycle, c.name))

        cycle_dir = join(exp_dir, c.name)
        project_names = get_jtproject_names(cycle_dir)

        for project_name in project_names:

            project_dir = join(cycle_dir, project_name)

            if project_dir == segmentation_project_dir:
                # data of segmentation project is handled separately
                segm_ix = i
                continue

            print '.. Merging data of Jterator project "%s"' % project_name

            data_files = get_data_files(project_dir)

            f = h5py.File(data_files[0], 'r')
            feature_names = f.keys()
            f.close()

            if 'OriginalObjectIds' not in feature_names:
                raise ValueError('Files must contain a dataset called "%s"' %
                                'OriginalObjectIds')

            (features, feature_names) = fuse_data(data_files,
                                                  names=feature_names)

            if features.shape[1] != len(feature_names):
                raise ValueError('Number of features in dataset and length of '
                                'list with feature names must be identical.')

            if i == 0:
                data = features
            else:
                data = np.hstack((data, features))

            feature_names = rename_features(feature_names,
                                            header_mapper[c.cycle])
            # # convert to safe string format for HDF5
            # feature_names = np.array(map(np.string_, feature_names))
            data_header = np.hstack((data_header, feature_names))

    print '. Combining data from different cycles'
    (data_header, unique_ix) = np.unique(data_header, return_index=True)
    # TODO: sanity checks
    data = data[:, unique_ix]

    # separate object ids from dataset
    print '. Building global object ids'
    ids_ix = np.where(np.array([re.search('ID', i) is not None
                                for i in data_header]))[0]
    ids = data[:, ids_ix].astype(int)
    ids_header = data_header[ids_ix]
    data = np.delete(data, ids_ix, axis=1)
    data_header = np.delete(data_header, ids_ix)
    global_ids = build_global_ids(ids, ids_header)

    print '. Separate features belonging to different objects'
    objects = [re.match(r'^([^_]+)', name).group(1) for name in data_header]
    (objects, object_ix) = np.unique(objects, return_inverse=True)

    ###########################################################################
    #                          write fused dataset                            #
    ###########################################################################

    print '. Writing fused data into HDF5 file "%s"' % output_filename
    f = h5py.File(output_filename, 'w')  # truncate file if exists
    for i, obj in enumerate(objects):

        print '. Writing data for object "%s"' % obj

        obj_ix = object_ix == i
        obj_name = obj.lower()  # use lower case consistently

        # Measurement data
        print '.. Writing measurement data'

        # Store global object ids
        f.create_dataset('objects/%s/ids' % obj_name, data=global_ids)

        # Also store local "original" object ids for use outside of TissueMAPS
        location = 'objects/%s/original-ids' % obj_name
        f.create_dataset(location, data=ids)
        f[location].attrs.__setitem__('names', ids_header)

        # Store the actual features
        location = 'objects/%s/features' % obj_name
        f.create_dataset(location, data=data[:, obj_ix])
        f[location].attrs.__setitem__('names', data_header[obj_ix])

        # Segmentation data
        print '.. Writing segmentation data'

        data_files = get_data_files(segmentation_project_dir)

        # Determine parent-child relationship between objects
        try:
            parent_ids, names = fuse_data(data_files,
                                          names=['%s_ParentIds' % obj],
                                          as_int=True)
            location = 'objects/%s/parent-ids' % obj_name
            f.create_dataset(location, data=global_ids)
        except ValueError as e:
            # If there are no parent ids, this object is the parent
            f.create_dataset('parent', data=np.string_(obj_name))

        # Calculate global centroids
        # For this we also need the dimensions of the images.
        # It's sufficient to obtain them from a single image.
        centroids, names = fuse_data(data_files,
                                     names=['%s_Centroids' % obj],
                                     as_int=True)
        image_size = cycles[segm_ix].project.segmentation_files[0].dimensions
        global_centroids = calc_global_centroids(centroids, ids, ids_header,
                                                 image_size)

        # Store global centroids
        location = 'objects/%s/centroids' % obj_name
        f.create_dataset(location, data=global_centroids)
        f[location].attrs.__setitem__('names', np.array(['y', 'x']))

        # Store index of objects at the border of images
        border, names = fuse_data(data_files,
                                  names=['%s_BorderIx' % obj],
                                  as_int=True)
        f.create_dataset('objects/%s/border' % obj_name, data=border)

    f.close()
<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Astronomical and physics constants in SI units.  See :mod:`astropy.constants`
for a complete listing of constants defined in Astropy.
"""

import numpy as np

from .constant import Constant, EMConstant

# PHYSICAL CONSTANTS

class CODATA2014(Constant):
    default_reference = 'CODATA 2014'
    _registry = {}
    _has_incompatible_units = set()

class EMCODATA2014(CODATA2014, EMConstant):
    _registry = CODATA2014._registry

h = CODATA2014('h', "Planck constant", 6.626070040e-34,
               'J s', 0.000000081e-34, system='si')

hbar = CODATA2014('hbar', "Reduced Planck constant", 1.054571800e-34,
                  'J s', 0.000000013e-34, system='si')

k_B = CODATA2014('k_B', "Boltzmann constant", 1.38064852e-23,
                 'J / (K)', 0.00000079e-23, system='si')

c = CODATA2014('c', "Speed of light in vacuum", 299792458.,
               'm / (s)', 0.0, system='si')

G = CODATA2014('G', "Gravitational constant", 6.67408e-11,
               'm3 / (kg s2)', 0.00031e-11, system='si')

g0 = CODATA2014('g0', "Standard acceleration of gravity", 9.80665,
                'm / s2', 0.0, system='si')

m_p = CODATA2014('m_p', "Proton mass", 1.672621898e-27,
                 'kg', 0.000000021e-27, system='si')

m_n = CODATA2014('m_n', "Neutron mass", 1.674927471e-27,
                 'kg', 0.000000021e-27, system='si')

m_e = CODATA2014('m_e', "Electron mass", 9.10938356e-31,
                 'kg', 0.00000011e-31, system='si')

u = CODATA2014('u', "Atomic mass", 1.660539040e-27,
               'kg', 0.000000020e-27, system='si')

sigma_sb = CODATA2014('sigma_sb', "Stefan-Boltzmann constant", 5.670367e-8,
                      'W / (K4 m2)', 0.000013e-8, system='si')

e = EMCODATA2014('e', 'Electron charge', 1.6021766208e-19,
                 'C', 0.0000000098e-19, system='si')

eps0 = EMCODATA2014('eps0', 'Electric constant', 8.854187817e-12,
                    'F/m', 0.0, system='si')

N_A = CODATA2014('N_A', "Avogadro's number", 6.022140857e23,
                 '1 / (mol)', 0.000000074e23, system='si')

R = CODATA2014('R', "Gas constant", 8.3144598,
               'J / (K mol)', 0.0000048, system='si')

Ryd = CODATA2014('Ryd', 'Rydberg constant', 10973731.568508,
                 '1 / (m)', 0.000065, system='si')

a0 = CODATA2014('a0', "Bohr radius", 0.52917721067e-10,
                'm', 0.00000000012e-10, system='si')

muB = CODATA2014('muB', "Bohr magneton", 927.4009994e-26,
                 'J/T', 0.00002e-26, system='si')

alpha = CODATA2014('alpha', "Fine-structure constant", 7.2973525664e-3,
                 '', 0.0000000017e-3, system='si')

atm = CODATA2014('atm', "Standard atmosphere", 101325,
                 'Pa', 0.0, system='si')

mu0 = CODATA2014('mu0', "Magnetic constant", 4.0e-7 * np.pi, 'N/A2', 0.0,
                 system='si')

sigma_T = CODATA2014('sigma_T', "Thomson scattering cross-section",
                     0.66524587158e-28, 'm2', 0.00000000091e-28,
                     system='si')

b_wien = CODATA2014('b_wien', 'Wien wavelength displacement law constant',
                    2.8977729e-3, 'm K', 00.0000017e-3, system='si')

# cgs constants
# Only constants that cannot be converted directly from S.I. are defined here.

e_esu = EMCODATA2014(e.abbrev, e.name, e.value * c.value * 10.0,
                     'statC', e.uncertainty * c.value * 10.0, system='esu')

e_emu = EMCODATA2014(e.abbrev, e.name, e.value / 10, 'abC',
                     e.uncertainty / 10, system='emu')

e_gauss = EMCODATA2014(e.abbrev, e.name, e.value * c.value * 10.0,
                     'Fr', e.uncertainty * c.value * 10.0, system='gauss')
<commit_msg>Fix codata2014 formatting<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Astronomical and physics constants in SI units.  See :mod:`astropy.constants`
for a complete listing of constants defined in Astropy.
"""

import numpy as np

from .constant import Constant, EMConstant

# PHYSICAL CONSTANTS

class CODATA2014(Constant):
    default_reference = 'CODATA 2014'
    _registry = {}
    _has_incompatible_units = set()

class EMCODATA2014(CODATA2014, EMConstant):
    _registry = CODATA2014._registry

h = CODATA2014('h', "Planck constant", 6.626070040e-34,
               'J s', 0.000000081e-34, system='si')

hbar = CODATA2014('hbar', "Reduced Planck constant", 1.054571800e-34,
                  'J s', 0.000000013e-34, system='si')

k_B = CODATA2014('k_B', "Boltzmann constant", 1.38064852e-23,
                 'J / (K)', 0.00000079e-23, system='si')

c = CODATA2014('c', "Speed of light in vacuum", 299792458.,
               'm / (s)', 0.0, system='si')

G = CODATA2014('G', "Gravitational constant", 6.67408e-11,
               'm3 / (kg s2)', 0.00031e-11, system='si')

g0 = CODATA2014('g0', "Standard acceleration of gravity", 9.80665,
                'm / s2', 0.0, system='si')

m_p = CODATA2014('m_p', "Proton mass", 1.672621898e-27,
                 'kg', 0.000000021e-27, system='si')

m_n = CODATA2014('m_n', "Neutron mass", 1.674927471e-27,
                 'kg', 0.000000021e-27, system='si')

m_e = CODATA2014('m_e', "Electron mass", 9.10938356e-31,
                 'kg', 0.00000011e-31, system='si')

u = CODATA2014('u', "Atomic mass", 1.660539040e-27,
               'kg', 0.000000020e-27, system='si')

sigma_sb = CODATA2014('sigma_sb', "Stefan-Boltzmann constant", 5.670367e-8,
                      'W / (K4 m2)', 0.000013e-8, system='si')

e = EMCODATA2014('e', 'Electron charge', 1.6021766208e-19,
                 'C', 0.0000000098e-19, system='si')

eps0 = EMCODATA2014('eps0', 'Electric constant', 8.854187817e-12,
                    'F/m', 0.0, system='si')

N_A = CODATA2014('N_A', "Avogadro's number", 6.022140857e23,
                 '1 / (mol)', 0.000000074e23, system='si')

R = CODATA2014('R', "Gas constant", 8.3144598,
               'J / (K mol)', 0.0000048, system='si')

Ryd = CODATA2014('Ryd', 'Rydberg constant', 10973731.568508,
                 '1 / (m)', 0.000065, system='si')

a0 = CODATA2014('a0', "Bohr radius", 0.52917721067e-10,
                'm', 0.00000000012e-10, system='si')

muB = CODATA2014('muB', "Bohr magneton", 927.4009994e-26,
                 'J/T', 0.00002e-26, system='si')

alpha = CODATA2014('alpha', "Fine-structure constant", 7.2973525664e-3,
                   '', 0.0000000017e-3, system='si')

atm = CODATA2014('atm', "Standard atmosphere", 101325,
                 'Pa', 0.0, system='si')

mu0 = CODATA2014('mu0', "Magnetic constant", 4.0e-7 * np.pi, 'N/A2', 0.0,
                 system='si')

sigma_T = CODATA2014('sigma_T', "Thomson scattering cross-section",
                     0.66524587158e-28, 'm2', 0.00000000091e-28,
                     system='si')

b_wien = CODATA2014('b_wien', 'Wien wavelength displacement law constant',
                    2.8977729e-3, 'm K', 0.0000017e-3, system='si')

# cgs constants
# Only constants that cannot be converted directly from S.I. are defined here.

e_esu = EMCODATA2014(e.abbrev, e.name, e.value * c.value * 10.0,
                     'statC', e.uncertainty * c.value * 10.0, system='esu')

e_emu = EMCODATA2014(e.abbrev, e.name, e.value / 10, 'abC',
                     e.uncertainty / 10, system='emu')

e_gauss = EMCODATA2014(e.abbrev, e.name, e.value * c.value * 10.0,
                       'Fr', e.uncertainty * c.value * 10.0, system='gauss')
<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
This module contains functions for matching coordinate catalogs.
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np

from ..extern import six
from .representation import UnitSphericalRepresentation

__all__ = ['match_coordinates_3d', 'match_coordinates_sky']

def match_coordinates_3d(matchcoord, catalogcoord, nthneighbor=1, storekdtree='_kdtree_3d'):
    """
    Finds the nearest 3-dimensional matches of a coordinate or coordinates in
    a set of catalog coordinates.

    This finds the 3-dimensional closest neighbor, which is only different
    from the on-sky distance if ``distance`` is set in either ``matchcoord``
    or ``catalogcoord``.

    Parameters
    ----------
    matchcoord : `~astropy.coordinates.BaseCoordinateFrame`
        The coordinate(s) to match to the catalog.
    catalogcoord : `~astropy.coordinates.BaseCoordinateFrame`
        The base catalog in which to search for matches. Typically this will
        be a coordinate object that is an array (i.e.,
        ``catalogcoord.isscalar == False``)
    nthneighbor : int, optional
        Which closest neighbor to search for.  Typically ``1`` is desired here,
        as that is correct for matching one set of coordinates to another.
        The next likely use case is ``2``, for matching a coordinate catalog
        against *itself* (``1`` is inappropriate because each point will find
        itself as the closest match).
    storekdtree : bool or str, optional
        If a string, will store the KD-Tree used for the computation
        in the ``catalogcoord``, as an attrbute in ``catalogcoord`` with the
        provided name.  This dramatically speeds up subsequent calls with the
        same catalog. If False, the KD-Tree is discarded after use.

    Returns
    -------
    idx : integer array
        Indecies into ``catalogcoord`` to get the matched points for each
        ``matchcoord``. Shape matches ``matchcoord``.
    sep2d : `~astropy.coordinates.Angle`
        The on-sky separation between the closest match for each ``matchcoord``
        and the ``matchcoord``. Shape matches ``matchcoord``.
    dist3d : `~astropy.units.Quantity`
        The 3D distance between the closest match for each ``matchcoord`` and
        the ``matchcoord``. Shape matches ``matchcoord``.

    Notes
    -----
    This function requires `SciPy <http://www.scipy.org>`_ to be installed
    or it will fail.
    """
    from warnings import warn

    #without scipy this will immediately fail
    from scipy import spatial
    try:
        KDTree = spatial.cKDTree
    except:
        warn('C-base KD tree not found, falling back on (much slower) '
             'python implementation')
        KDTree = spatial.KDTree

    if storekdtree is True:  # backwards compatibility for pre v0.4
        storekdtree = '_kdtree'

    # figure out where any cached KDTree might be
    if isinstance(storekdtree, six.string_types):
        kdt = getattr(catalogcoord, storekdtree, None)
        if kdt is not None and not isinstance(kdt, KDTree):
            raise ValueError('Invalid `storekdtree` string:' + storekdtree)
    elif isinstance(storekdtree, KDTree):
        kdt = storekdtree
        storekdtree = None
    elif not storekdtree:
        kdt = None
    else:
        raise ValueError('Invalid `storekdtree` argument:' +
                          str(storekdtree))

    if kdt is None:
        #need to build the cartesian KD-tree for the catalog
        cartxyz = catalogcoord.cartesian.xyz
        flatxyz = cartxyz.reshape((3, np.prod(cartxyz.shape) // 3))
        kdt = KDTree(flatxyz.value.T)

    #make sure coordinate systems match
    matchcoord = matchcoord.transform_to(catalogcoord)

    #make sure units match
    catunit = catalogcoord.cartesian.x.unit
    matchxyz = matchcoord.cartesian.xyz.to(catunit)

    matchflatxyz = matchxyz.reshape((3, np.prod(matchxyz.shape) // 3))
    dist, idx = kdt.query(matchflatxyz.T, nthneighbor)

    if nthneighbor > 1:  # query gives 1D arrays if k=1, 2D arrays otherwise
        dist = dist[:, -1]
        idx = idx[:, -1]

    if storekdtree:
        #cache the kdtree in `catalogcoord`
        setattr(catalogcoord, storekdtree, kdt)

    sep2d = catalogcoord[idx].separation(matchcoord)
    return idx.reshape(matchxyz.shape[1:]), sep2d, dist.reshape(matchxyz.shape[1:]) * catunit

def match_coordinates_sky(matchcoord, catalogcoord, nthneighbor=1, storekdtree='_kdtree_sky'):
    """
    Finds the nearest on-sky matches of a coordinate or coordinates in
    a set of catalog coordinates.

    This finds the on-sky closest neighbor, which is only different from the
    3-dimensional match if ``distance`` is set in either ``matchcoord``
    or ``catalogcoord``.

    Parameters
    ----------
    matchcoord : `~astropy.coordinates.BaseCoordinateFrame`
        The coordinate(s) to match to the catalog.
    catalogcoord : `~astropy.coordinates.BaseCoordinateFrame`
        The base catalog in which to search for matches. Typically this will
        be a coordinate object that is an array (i.e.,
        ``catalogcoord.isscalar == False``)
    nthneighbor : int, optional
        Which closest neighbor to search for.  Typically ``1`` is desired here,
        as that is correct for matching one set of coordinates to another.
        The next likely use case is ``2``, for matching a coordinate catalog
        against *itself* (``1`` is inappropriate because each point will find
        itself as the closest match).
    storekdtree : bool or str, optional
        If a string, will store the KD-Tree used for the computation
        in the ``catalogcoord``, as an attrbute in ``catalogcoord`` with the
        provided name.  This dramatically speeds up subsequent calls with the
        same catalog. If False, the KD-Tree is discarded after use.

    Returns
    -------
    idx : integer array
        Indecies into ``catalogcoord`` to get the matched points for each
        ``matchcoord``. Shape matches ``matchcoord``.
    sep2d : `~astropy.coordinates.Angle`
        The on-sky separation between the closest match for each
        ``matchcoord`` and the ``matchcoord``. Shape matches ``matchcoord``.
    dist3d : `~astropy.units.Quantity`
        The 3D distance between the closest match for each ``matchcoord`` and
        the ``matchcoord``. Shape matches ``matchcoord``.  If either
        `matchcoord` or `catalogcoord` don't have a distance, this is the 3D
        distance on the unit sphere, rather than a true distance.

    Notes
    -----
    This function requires `SciPy <http://www.scipy.org>`_ to be installed
    or it will fail.
    """

    # send to catalog frame
    newmatch = matchcoord.transform_to(catalogcoord)

    #strip out distance info
    match_urepr = newmatch.data.represent_as(UnitSphericalRepresentation)
    newmatch_u = newmatch.realize_frame(match_urepr)

    cat_urepr = catalogcoord.data.represent_as(UnitSphericalRepresentation)
    newcat_u = catalogcoord.realize_frame(cat_urepr)

    idx, sep2d, sep3d = match_coordinates_3d(newmatch_u, newcat_u, nthneighbor, storekdtree)
    # sep3d is *wrong* above, because the distance information was removed,
    # unless one of the catalogs doesn't have a real distance
    if not (isinstance(catalogcoord.data, UnitSphericalRepresentation) or
            isinstance(newmatch.data, UnitSphericalRepresentation)):
        sep3d = catalogcoord[idx].separation_3d(newmatch)

    #update the kdtree on the actual passed-in coordinate
    if storekdtree:
        setattr(catalogcoord, storekdtree, getattr(newcat_u, storekdtree))

    return idx, sep2d, sep3d
<commit_msg>added note info to docstrings that the matching functions support SkyCoord<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
This module contains functions for matching coordinate catalogs.
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np

from ..extern import six
from .representation import UnitSphericalRepresentation

__all__ = ['match_coordinates_3d', 'match_coordinates_sky']

def match_coordinates_3d(matchcoord, catalogcoord, nthneighbor=1, storekdtree='_kdtree_3d'):
    """
    Finds the nearest 3-dimensional matches of a coordinate or coordinates in
    a set of catalog coordinates.

    This finds the 3-dimensional closest neighbor, which is only different
    from the on-sky distance if ``distance`` is set in either ``matchcoord``
    or ``catalogcoord``.

    Parameters
    ----------
    matchcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
        The coordinate(s) to match to the catalog.
    catalogcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
        The base catalog in which to search for matches. Typically this will
        be a coordinate object that is an array (i.e.,
        ``catalogcoord.isscalar == False``)
    nthneighbor : int, optional
        Which closest neighbor to search for.  Typically ``1`` is desired here,
        as that is correct for matching one set of coordinates to another.
        The next likely use case is ``2``, for matching a coordinate catalog
        against *itself* (``1`` is inappropriate because each point will find
        itself as the closest match).
    storekdtree : bool or str, optional
        If a string, will store the KD-Tree used for the computation
        in the ``catalogcoord``, as an attribute in ``catalogcoord`` with the
        provided name.  This dramatically speeds up subsequent calls with the
        same catalog. If False, the KD-Tree is discarded after use.

    Returns
    -------
    idx : integer array
        Indecies into ``catalogcoord`` to get the matched points for each
        ``matchcoord``. Shape matches ``matchcoord``.
    sep2d : `~astropy.coordinates.Angle`
        The on-sky separation between the closest match for each ``matchcoord``
        and the ``matchcoord``. Shape matches ``matchcoord``.
    dist3d : `~astropy.units.Quantity`
        The 3D distance between the closest match for each ``matchcoord`` and
        the ``matchcoord``. Shape matches ``matchcoord``.

    Notes
    -----
    This function requires `SciPy <http://www.scipy.org>`_ to be installed
    or it will fail.
    """
    from warnings import warn

    #without scipy this will immediately fail
    from scipy import spatial
    try:
        KDTree = spatial.cKDTree
    except:
        warn('C-base KD tree not found, falling back on (much slower) '
             'python implementation')
        KDTree = spatial.KDTree

    if storekdtree is True:  # backwards compatibility for pre v0.4
        storekdtree = '_kdtree'

    # figure out where any cached KDTree might be
    if isinstance(storekdtree, six.string_types):
        kdt = getattr(catalogcoord, storekdtree, None)
        if kdt is not None and not isinstance(kdt, KDTree):
            raise ValueError('Invalid `storekdtree` string:' + storekdtree)
    elif isinstance(storekdtree, KDTree):
        kdt = storekdtree
        storekdtree = None
    elif not storekdtree:
        kdt = None
    else:
        raise ValueError('Invalid `storekdtree` argument:' +
                          str(storekdtree))

    if kdt is None:
        #need to build the cartesian KD-tree for the catalog
        cartxyz = catalogcoord.cartesian.xyz
        flatxyz = cartxyz.reshape((3, np.prod(cartxyz.shape) // 3))
        kdt = KDTree(flatxyz.value.T)

    #make sure coordinate systems match
    matchcoord = matchcoord.transform_to(catalogcoord)

    #make sure units match
    catunit = catalogcoord.cartesian.x.unit
    matchxyz = matchcoord.cartesian.xyz.to(catunit)

    matchflatxyz = matchxyz.reshape((3, np.prod(matchxyz.shape) // 3))
    dist, idx = kdt.query(matchflatxyz.T, nthneighbor)

    if nthneighbor > 1:  # query gives 1D arrays if k=1, 2D arrays otherwise
        dist = dist[:, -1]
        idx = idx[:, -1]

    if storekdtree:
        #cache the kdtree in `catalogcoord`
        setattr(catalogcoord, storekdtree, kdt)

    sep2d = catalogcoord[idx].separation(matchcoord)
    return idx.reshape(matchxyz.shape[1:]), sep2d, dist.reshape(matchxyz.shape[1:]) * catunit

def match_coordinates_sky(matchcoord, catalogcoord, nthneighbor=1, storekdtree='_kdtree_sky'):
    """
    Finds the nearest on-sky matches of a coordinate or coordinates in
    a set of catalog coordinates.

    This finds the on-sky closest neighbor, which is only different from the
    3-dimensional match if ``distance`` is set in either ``matchcoord``
    or ``catalogcoord``.

    Parameters
    ----------
    matchcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
        The coordinate(s) to match to the catalog.
    catalogcoord : `~astropy.coordinates.BaseCoordinateFrame` or `~astropy.coordinates.SkyCoord`
        The base catalog in which to search for matches. Typically this will
        be a coordinate object that is an array (i.e.,
        ``catalogcoord.isscalar == False``)
    nthneighbor : int, optional
        Which closest neighbor to search for.  Typically ``1`` is desired here,
        as that is correct for matching one set of coordinates to another.
        The next likely use case is ``2``, for matching a coordinate catalog
        against *itself* (``1`` is inappropriate because each point will find
        itself as the closest match).
    storekdtree : bool or str, optional
        If a string, will store the KD-Tree used for the computation
        in the ``catalogcoord``, as an attrbute in ``catalogcoord`` with the
        provided name.  This dramatically speeds up subsequent calls with the
        same catalog. If False, the KD-Tree is discarded after use.

    Returns
    -------
    idx : integer array
        Indecies into ``catalogcoord`` to get the matched points for each
        ``matchcoord``. Shape matches ``matchcoord``.
    sep2d : `~astropy.coordinates.Angle`
        The on-sky separation between the closest match for each
        ``matchcoord`` and the ``matchcoord``. Shape matches ``matchcoord``.
    dist3d : `~astropy.units.Quantity`
        The 3D distance between the closest match for each ``matchcoord`` and
        the ``matchcoord``. Shape matches ``matchcoord``.  If either
        `matchcoord` or `catalogcoord` don't have a distance, this is the 3D
        distance on the unit sphere, rather than a true distance.

    Notes
    -----
    This function requires `SciPy <http://www.scipy.org>`_ to be installed
    or it will fail.
    """

    # send to catalog frame
    newmatch = matchcoord.transform_to(catalogcoord)

    #strip out distance info
    match_urepr = newmatch.data.represent_as(UnitSphericalRepresentation)
    newmatch_u = newmatch.realize_frame(match_urepr)

    cat_urepr = catalogcoord.data.represent_as(UnitSphericalRepresentation)
    newcat_u = catalogcoord.realize_frame(cat_urepr)

    idx, sep2d, sep3d = match_coordinates_3d(newmatch_u, newcat_u, nthneighbor, storekdtree)
    # sep3d is *wrong* above, because the distance information was removed,
    # unless one of the catalogs doesn't have a real distance
    if not (isinstance(catalogcoord.data, UnitSphericalRepresentation) or
            isinstance(newmatch.data, UnitSphericalRepresentation)):
        sep3d = catalogcoord[idx].separation_3d(newmatch)

    #update the kdtree on the actual passed-in coordinate
    if storekdtree:
        setattr(catalogcoord, storekdtree, getattr(newcat_u, storekdtree))

    return idx, sep2d, sep3d
<|endoftext|>