From 59466a69bc55582d2ea651160e91c53fbd6a47aa Mon Sep 17 00:00:00 2001 From: Andy Eschbacher Date: Mon, 24 Apr 2017 14:40:21 -0400 Subject: [PATCH] moves functions to rely on queries instead of tqblenames --- .../crankshaft/analysis_data_provider.py | 77 ++++++++++++------- .../crankshaft/optimization/optim.py | 12 +-- 2 files changed, 56 insertions(+), 33 deletions(-) diff --git a/src/py/crankshaft/crankshaft/analysis_data_provider.py b/src/py/crankshaft/crankshaft/analysis_data_provider.py index e2da80e..d733757 100644 --- a/src/py/crankshaft/crankshaft/analysis_data_provider.py +++ b/src/py/crankshaft/crankshaft/analysis_data_provider.py @@ -14,7 +14,7 @@ class AnalysisDataProvider(object): return pu.empty_zipped_array(4) else: return result - except plpy.SPIError, err: + except plpy.SPIError as err: plpy.error('Analysis failed: %s' % err) def get_markov(self, w_type, params): @@ -27,7 +27,7 @@ class AnalysisDataProvider(object): return pu.empty_zipped_array(4) return data - except plpy.SPIError, err: + except plpy.SPIError as err: plpy.error('Analysis failed: %s' % err) def get_moran(self, w_type, params): @@ -40,8 +40,8 @@ class AnalysisDataProvider(object): if len(data) == 0: return pu.empty_zipped_array(2) return data - except plpy.SPIError, err: - plpy.error('Analysis failed: %s' % e) + except plpy.SPIError as err: + plpy.error('Analysis failed: %s' % err) return pu.empty_zipped_array(2) def get_nonspatial_kmeans(self, query): @@ -49,7 +49,7 @@ class AnalysisDataProvider(object): try: data = plpy.execute(query) return data - except plpy.SPIError, err: + except plpy.SPIError as err: plpy.error('Analysis failed: %s' % err) def get_spatial_kmeans(self, params): @@ -63,43 +63,66 @@ class AnalysisDataProvider(object): try: data = plpy.execute(query) return data - except plpy.SPIError, err: + except plpy.SPIError as err: plpy.error('Analysis failed: %s' % err) - def get_column(self, table, column, dtype=float): + def get_column(self, subquery, column, dtype=float, id_col='cartodb_id'): """ - Retrieve the column from the specified table + Retrieve the column from the specified table from a connected + PostgreSQL database. - :param table: table to retrieve column from - :type table: text - :param column: column to retrieve - :type column: text - :param dtype: data type in column (e.g, float, int, str) - :type dtype: type - :returns: column from table as a NumPy array - :rtype: NumPy array + Args: + subquery (str): subquery to retrieve column from + column (str): column to retrieve + dtype (type): data type in column (e.g, float, int, str) + id_col (str, optional): Column name for index. Defaults to + `cartodb_id`. + + Returns: + numpy.array: column from table as a NumPy array """ query = ''' - SELECT array_agg("{column}" ORDER BY "cartodb_id" ASC) as col - FROM "{table}" - '''.format(table=table, column=column) + SELECT array_agg("{column}" ORDER BY "{id_col}" ASC) as col + FROM ({subquery}) As _wrap + '''.format(subquery=subquery, + column=column, + id_col=id_col) + resp = plpy.execute(query) return np.array(resp[0]['col'], dtype=dtype) - def get_pairwise_distances(self, drain, source): - """retuns the pairwise distances between row i and j for all i in table1 and j in table1""" + def get_pairwise_distances(self, drain_query, source_query, + id_col='cartodb_id'): + """Retuns the pairwise distances between row i and j for all i in + drain_query and j in source_query + + Args: + drain_query (str): Query that exposes the `the_geom` and + `cartodb_id` (or what is specified in `id_col`) of the dataset + for 'drain' locations + source_query (str): Query that exposes the `the_geom` and + `cartodb_id` (or what is specified in `id_col`) of the dataset + for 'source' locations + id_col (str, optional): Column name for table index. Defaults to + `cartodb_id`. + + Returns: + numpy.array: A len(s) by len(d) array of distances from source i to + drain j + """ query = ''' SELECT array_agg(ST_Distance(d."the_geom"::geography, s."the_geom"::geography) / 1000.0 - ORDER BY d."cartodb_id" ASC) as dist - FROM "{drain}" as d, "{source}" as s - GROUP BY s."cartodb_id" - ORDER BY s."cartodb_id" ASC - '''.format(drain=drain, source=source) + ORDER BY d."{id_col}" ASC) as dist + FROM ({drain_query}) AS d, ({source_query}) AS s + GROUP BY s."{id_col}" + ORDER BY s."{id_col}" ASC + '''.format(drain_query=drain_query, + source_query=source_query, + id_col=id_col) resp = plpy.execute(query) # len(s) x len(d) matrix return np.array([np.array(row['dist'], dtype=float) for row in resp], dtype=float) - diff --git a/src/py/crankshaft/crankshaft/optimization/optim.py b/src/py/crankshaft/crankshaft/optimization/optim.py index 14f0386..ba91b69 100644 --- a/src/py/crankshaft/crankshaft/optimization/optim.py +++ b/src/py/crankshaft/crankshaft/optimization/optim.py @@ -14,10 +14,10 @@ class Optim(object): That is, `cost ~ population * distance` """ - def __init__(self, drain_table, source_table, capacity_column, # pylint: disable=too-many-arguments + def __init__(self, drain_query, source_table, capacity_column, # pylint: disable=too-many-arguments production_column, marginal_column, **kwargs): - # set data provider (defaults to SQL database access + # set data provider - defaults to SQL database access self.data_provider = kwargs.get('data_provider', AnalysisDataProvider()) # model parameters @@ -30,21 +30,21 @@ class Optim(object): # model data self.model_data = { - 'drain_capacity': self.data_provider.get_column(drain_table, + 'drain_capacity': self.data_provider.get_column(drain_query, capacity_column), 'source_amount': (self.model_params['amount_per_unit'] * (1. - self.model_params['recycle_rate']) * self.data_provider.get_column(source_table, production_column)), - 'marginal_cost': self.data_provider.get_column(drain_table, + 'marginal_cost': self.data_provider.get_column(drain_query, marginal_column), 'distance': self.data_provider.get_pairwise_distances(source_table, - drain_table) + drain_query) } self.model_data['cost'] = self.calc_cost() # database ids self.ids = { - 'drain': self.data_provider.get_column(drain_table, + 'drain': self.data_provider.get_column(drain_query, 'cartodb_id', dtype=int), 'source': self.data_provider.get_column(source_table,