Permalink
Fetching contributors…
Cannot retrieve contributors at this time
367 lines (272 sloc) 12.6 KB
from random import sample
from math import pow, sqrt, floor, ceil
from uuid import uuid4
from itertools import chain
from functools import reduce, partial
from operator import sub
from graphviz import Digraph
import logging
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
def batch_build(builders):
return [builder() for builder in builders]
def batch_get_sequence(total, leaf_max):
batches, batch_size = [], ceil(total / max(1, floor(total / (leaf_max * 100))))
for idx in range(0, total, batch_size):
batches.append((idx, min(idx + batch_size, total)))
return batches
def batch_group_keys(sequence, keys, tree_count):
return list(chain.from_iterable([(str(uuid4()), keys[start:end]) for start, end in sequence] for _ in range(tree_count)))
def distance_euclidean(alpha, beta, dimension):
return sqrt(sum(pow(alpha.get(i, 0) - beta.get(i, 0), 2) for i in range(dimension)))
def forest_build(points, pmeta, tree_count, leaf_max=5, n_jobs=1, batch_size=1000):
forest = {}
batches = batch_group_keys(batch_get_sequence(len(points), leaf_max),
list(points.keys()),
tree_count)
meta = {'count': tree_count,
'leaf_max': leaf_max,
'roots': list(zip(*batches))[0]}
builders = [partial(node_build, points, pmeta, keys, root_id, leaf_max)
for root_id, keys
in batches]
if n_jobs == 1:
forest = forest_build_single(builders, len(points) * tree_count)
else:
forest = forest_build_multi(builders, len(points) * tree_count, n_jobs, batch_size)
return meta, forest
def forest_build_multi(builders, forest_total, n_jobs, batch_size):
forest, progress = {}, [0]
with ProcessPoolExecutor(max_workers=n_jobs) as pool:
while builders:
jobs, builders_next = [], []
for idx in range(0, len(builders), batch_size):
jobs.append(pool.submit(batch_build, builders[idx:idx+batch_size]))
for job in jobs:
for node, builders_sub in job.result():
if node['type'] == 'leaf':
progress.append(progress[-1] + node['count'])
progress_log('Forest building progress', progress, forest_total)
forest[node['id']] = node
builders_next.append(builders_sub)
builders = list(chain.from_iterable(builders_next))
return forest
def forest_build_single(builders, forest_total):
forest, progress = {}, [0]
while builders:
builders_next = []
for builder in builders:
node, builders_sub = builder()
if node['type'] == 'leaf':
progress.append(progress[-1] + node['count'])
progress_log('Forest building progress', progress, forest_total)
forest[node['id']] = node
builders_next.append(builders_sub)
builders = list(chain.from_iterable(builders_next))
return forest
def forest_get_dot(forest, fmeta, count=1):
result = Digraph()
for root_id in fmeta['roots']:
result.edge('0', str(root_id))
builders = [partial(node_get_dot, forest, root_id, True, result) for root_id in fmeta['roots']]
while builders:
builders_next = []
for builder in builders:
builders_sub, result = builder()
builders_next = chain(builders_next, builders_sub)
builders = builders_next
return result
def forest_query_neighbourhood(query, forest, fmeta, threshold, n_jobs):
result = []
nodes = [(1., root_id) for root_id in fmeta['roots']]
if n_jobs == 1:
result = query_neighbourhood_single(nodes, forest, threshold)
else:
result = query_neighbourhood_multi(nodes, forest, threshold, n_jobs)
return result
def node_build(points, pmeta, point_ids, node_id, leaf_max=5):
node, children = {}, []
if len(point_ids) <= leaf_max:
node = {
'type': 'leaf',
'count': len(point_ids),
'id': node_id,
'children': point_ids
}
else:
split = split_points([points[idx] for idx in point_ids],
pmeta['dimension'])
distance_calc = plane_point_distance_calculator(
plane_normal(*split, dimension=pmeta['dimension']),
plane_point(*split, dimension=pmeta['dimension']),
pmeta['dimension'])
branches = {False: [], True: []}
child_ids = {False: str(uuid4()), True: str(uuid4())}
for idx in point_ids:
branches[distance_calc(points[idx]) > 0].append(idx)
node = {
'type': 'branch',
'func': distance_calc,
'count': len(point_ids),
'id': node_id,
'children': [idx for _, idx in child_ids.items()]
}
children.append(partial(node_build,
points,
pmeta,
branches[False],
child_ids[False],
leaf_max))
children.append(partial(node_build,
points,
pmeta,
branches[True],
child_ids[True],
leaf_max))
return node, children
def node_get_dot(forest, node_id, is_root, dot):
children = []
if is_root:
dot.node(node_id, '')
if forest[node_id]['type'] == 'branch':
for child_id in forest[node_id]['children']:
if forest[child_id]['type'] == 'branch':
dot.node(child_id, '')
else:
dot.node(child_id, str(forest[child_id]['count']))
dot.edge(node_id, child_id)
children.append(partial(node_get_dot, forest, child_id, False, dot))
return children, dot
def plane_normal(alpha, beta, dimension):
return tuple(alpha.get(i, 0) - beta.get(i, 0) for i in range(dimension))
def plane_point(alpha, beta, dimension):
return tuple((alpha.get(i, 0) + beta.get(i, 0)) / 2. for i in range(dimension))
def plane_point_distance_calculator(normal, point, dimension):
return partial(_ppd_calculator,
normal=normal,
d=reduce(lambda result, incoming: \
result + -(incoming[0] * incoming[1]),
zip(normal, point),
0),
dimension=dimension)
def _ppd_calculator(point, normal, d, dimension):
return ((sum([point.get(i, 0) * normal[i] for i in range(dimension)]) + d) / sqrt(sum([pow(i, 2) for i in normal])))
def point_add(id, vector, dimension, ptype, result={}):
result[id] = point_convert(vector, ptype)
return result
def point_convert(vector, ptype):
ptype_builder = {
'gensim': point_convert_gensim,
'list': point_convert_list
}
return ptype_builder.get(ptype, point_convert_invalid)(vector)
def point_convert_gensim(vector):
return dict(vector)
def point_convert_list(vector):
return dict([(idx, value) for idx, value in enumerate(vector) if value != 0])
def point_convert_invalid(*_):
raise Exception('Not supported')
def points_add(vectors, dimension, ptype, identifiers=None):
meta = {'dimension': dimension}
points = {}
if identifiers:
points = reduce(lambda result, i: point_add(identifiers[i],
vectors[i],
dimension,
ptype,
result),
range(len(identifiers)),
{})
else:
points = reduce(lambda result, vector: point_add(uuid4(),
vector,
dimension,
ptype,
result),
vectors,
{})
return meta, points
def progress_log(caption, progress, total):
to_print = False
if progress[-1] == total:
to_print = True
elif len(progress) > 1 and floor(progress[-2] / total * 10) < floor(progress[-1] / total * 10):
to_print = True
to_print and logging.info('{} {: >6.2f}%'.format(caption, progress[-1] / total * 100))
def query_get_children(query, forest, threshold, multiplier, node_id):
result, delta = [], forest[node_id]['func'](query)
if threshold > 0 and delta > 0 and delta <= threshold:
result.append((multiplier * 0.9,
forest[node_id]['children'][False]))
result.append((multiplier * 1.,
forest[node_id]['children'][True]))
elif threshold > 0 and delta <= 0 and -threshold <= delta:
result.append((multiplier * 1.,
forest[node_id]['children'][False]))
result.append((multiplier * 0.9,
forest[node_id]['children'][True]))
else:
result.append((multiplier * 1.,
forest[node_id]['children'][delta > 0]))
return result
def query_neighbourhood_multi(nodes, forest, threshold, n_jobs):
result = []
with ThreadPoolExecutor(max_workers=n_jobs) as pool:
while nodes:
nodes_next, jobs = [], []
for idx in range(0, len(nodes), 1000):
_job = []
for multiplier, node_id in nodes[idx:1000]:
if forest[node_id]['type'] == 'leaf':
result.append((1 - multiplier, forest[node_id]))
else:
_job.append(partial(query_get_children,
query,
forest,
threshold,
multiplier,
node_id))
jobs.append(pool.submit(batch_build, _job))
for job in jobs:
nodes_next.append(chain.from_iterable(job.result()))
nodes = list(chain.from_iterable(nodes_next))
return result
def query_neighbourhood_single(nodes, forest, threshold):
result = []
while nodes:
nodes_next = []
for multiplier, node_id in nodes:
if forest[node_id]['type'] == 'leaf':
result.append((1 - multiplier, forest[node_id]))
else:
delta = forest[node_id]['func'](query)
if threshold > 0 and delta > 0 and delta <= threshold:
nodes_next.append((multiplier * 0.9,
forest[node_id]['children'][False]))
nodes_next.append((multiplier * 1.,
forest[node_id]['children'][True]))
elif threshold > 0 and delta <= 0 and -threshold <= delta:
nodes_next.append((multiplier * 1.,
forest[node_id]['children'][False]))
nodes_next.append((multiplier * 0.9,
forest[node_id]['children'][True]))
else:
nodes_next.append((multiplier * 1.,
forest[node_id]['children'][delta > 0]))
nodes = nodes_next
return result
def search(query, points, pmeta, forest, fmeta, n=1, threshold=None, n_jobs=1):
neighbourhood = forest_query_neighbourhood(query,
forest,
fmeta,
threshold if threshold else (pmeta['dimension'] * 2e-2),
n_jobs)
result, candidate_max = {}, max(n, fmeta['leaf_max']) * fmeta['count']
for idx in chain.from_iterable(leaf['children'] for _, leaf in sorted(neighbourhood, key=lambda _: _[0])):
if idx not in result:
result[idx] = distance_euclidean(query, points[idx], pmeta['dimension'])
if len(result) >= candidate_max:
break
return sorted(result.items(), key=lambda _: _[-1])[:n]
def split_points(points, dimension):
result = sample(points, 2)
return result if reduce(lambda _result, incoming: _result or sub(*[result[i].get(i, 0) for i in range(2)]) != 0, range(dimension), False) else split_points(points)