def __setstate__(self, state):
self.bbspec = state[0]
self.bbs = tuple(tuple(_BBlock(*x) for x in bb) for bb in state[1])
self.verts = tuple(_Vertex(*x) for x in state[2])
self.edges = tuple(_Edge(*x) for x in state[3])
_validate_bbs_verts(self.bbs, self.verts)
assert len(self.bbs) == len(self.verts) == len(self.edges) + 1
def _grow_linear_start(verts_pickleable, edges_pickleable, **kwargs):
verts = tuple([_Vertex(*vp) for vp in verts_pickleable])
edges = tuple([_Edge(*ep) for ep in edges_pickleable])
positions = np.empty(shape=(1024, len(verts), 4, 4), dtype=np.float64)
indices = np.empty(shape=(1024, len(verts)), dtype=np.int32)
losses = np.empty(shape=(1024, ), dtype=np.float32)
result = SearchResult(positions=positions, indices=indices, losses=losses)
nresult, result = _grow_linear_recurse(result, verts, edges, **kwargs)
result = SearchResult(*(a[:nresult] for a in result))
return result
def _grow_linear_start(bb_base, verts_pickleable, edges_pickleable, **kwargs):
verts = tuple([_Vertex(*vp) for vp in verts_pickleable])
edges = tuple([_Edge(*ep) for ep in edges_pickleable])
pos = np.empty(shape=(1024, len(verts), 4, 4), dtype=np.float32)
idx = np.empty(shape=(1024, len(verts)), dtype=np.int32)
err = np.empty(shape=(1024, ), dtype=np.float32)
stats = zero_search_stats()
result = ResultJIT(pos=pos, idx=idx, err=err, stats=stats)
bases = np.zeros(len(verts), dtype=np.int64)
nresults, result = _grow_linear_recurse(result=result,
bb_base=bb_base,
verts=verts,
edges=edges,
bases=bases,
**kwargs)
result = ResultJIT(result.pos[:nresults], result.idx[:nresults],
result.err[:nresults], result.stats)
return result
def _grow_linear_mc_start(seconds, verts_pickleable, edges_pickleable,
threadno, pbar, lbl, verbosity, merge_bblock,
pbar_interval, **kwargs):
tstart = time()
verts = tuple([_Vertex(*vp) for vp in verts_pickleable])
edges = tuple([_Edge(*ep) for ep in edges_pickleable])
pos = np.empty(shape=(1024, len(verts), 4, 4), dtype=np.float32)
idx = np.empty(shape=(1024, len(verts)), dtype=np.int32)
err = np.empty(shape=(1024, ), dtype=np.float32)
stats = zero_search_stats()
result = ResultJIT(pos=pos, idx=idx, err=err, stats=stats)
bases = np.zeros(len(verts), dtype=np.int64)
del kwargs['nresults']
if threadno == 0 and pbar:
desc = 'linear search ' + str(lbl)
if merge_bblock is None: merge_bblock = 0
pbar_inst = tqdm(desc=desc,
position=merge_bblock + 1,
total=seconds,
mininterval=pbar_interval)
last = tstart
nbatch = [1000, 330, 100, 33, 10, 3] + [1] * 99
nbatch = nbatch[len(edges)] * 10
nresults = 0
iter = 0
ndups = 0
while time() < tstart + seconds:
if 'pbar_inst' in vars():
pbar_inst.update(time() - last)
last = time()
nresults, result = _grow_linear_mc(nbatch,
result,
verts,
edges,
bases=bases,
nresults=nresults,
**kwargs)
iter += 1
# remove duplicates every 10th iter
if iter % 10 == 0:
nresults_with_dups = nresults
uniq_result = ResultJIT(idx=result.idx[:nresults],
pos=result.pos[:nresults],
err=result.err[:nresults],
stats=result.stats)
uniq_result = remove_duplicate_results(uniq_result)
nresults = len(uniq_result.err)
result.idx[:nresults] = uniq_result.idx
result.pos[:nresults] = uniq_result.pos
result.err[:nresults] = uniq_result.err
ndups += nresults_with_dups - nresults
# print(ndups / nresults)
if nresults >= kwargs['max_linear']: break
if 'pbar_inst' in vars(): pbar_inst.close()
result = ResultJIT(result.pos[:nresults], result.idx[:nresults],
result.err[:nresults], result.stats)
return result