def search_func(criteria, bbs, monte_carlo, **kw):
stages = [criteria]
if hasattr(criteria, 'stages'):
stages = criteria.stages(bbs=bbs, **kw)
if len(stages) > 1:
assert kw['merge_bblock'] is not None
assert len(monte_carlo) in (1, len(stages))
if len(monte_carlo) != len(stages):
monte_carlo *= len(stages)
results = list()
for i, stage in enumerate(stages):
crit, stage_bbs = stage
if callable(crit): crit = crit(*results[-1][:-1])
lbl = f'stage{i}'
if kw['merge_bblock'] is not None:
lbl = f'stage{i}_mbb{kw["merge_bblock"]:04}'
results.append(
search_single_stage(crit,
monte_carlo=monte_carlo[i],
lbl=lbl,
bbs=stage_bbs,
**kw))
if len(results[-1][2].idx) is 0:
print('mbb', kw['merge_bblock'], 'no results at stage', i)
return None, None, None
if len(results) == 1:
return results[0][1:]
elif len(results) == 2:
# todo: this whole block is very protocol-specific... needs refactoring
mseg = criteria.merge_segment(**kw)
# simple_search_dag getting not-to-simple maybe split?
_____, ssdA, rsltA, logA = results[0]
critB, ssdB, rsltB, logB = results[1]
assert _shared_ssdag
ssdag = simple_search_dag(criteria,
only_seg=mseg,
make_edges=False,
source=_shared_ssdag,
bbs=bbs,
**kw)
ssdag.verts = ssdB.verts[:-1] + (ssdag.verts[mseg], ) + ssdA.verts[1:]
assert len(ssdag.verts) == len(criteria.bbspec)
rslt = merge_results_concat(criteria, ssdag, ssdA, rsltA, critB, ssdB,
rsltB, **kw)
return ssdag, rslt, logA + logB
else:
raise NotImplementedError('dunno more than two stages!')
def search_single_stage(criteria, lbl="", **kw):
if kw["run_cache"]:
if os.path.exists(kw["run_cache"] + lbl + ".pickle"):
with (open(kw["run_cache"] + lbl + ".pickle", "rb")) as inp:
ssdag, result = _pickle.load(inp)
return criteria, ssdag, result, ["from run cache " + lbl]
ssdag = simple_search_dag(criteria, source=_shared_ssdag, lbl=lbl, **kw)
result, tsearch = run_and_time(
grow_linear,
ssdag=ssdag,
loss_function=criteria.jit_lossfunc(),
last_bb_same_as=criteria.from_seg if criteria.is_cyclic else -1,
lbl=lbl,
**kw,
)
Nsparse = result.stats.total_samples[0]
Nsparse_rate = int(Nsparse / tsearch)
log = []
if len(result.idx):
frac_redundant = result.stats.n_redundant_results[0] / len(result.idx)
log = [
f"grow_linear {lbl} done, nresults {len(result.idx):,}, " +
f"total_samples {result.stats.total_samples[0]:,}, " +
f"samp/sec {Nsparse_rate:,}, redundant ratio {frac_redundant}"
]
if log:
print(log[-1])
if kw["run_cache"]:
with (open(kw["run_cache"] + lbl + ".pickle", "wb")) as out:
_pickle.dump((ssdag, result), out)
return criteria, ssdag, result, log
def search_single_stage(criteria, lbl='', **kw):
if kw['run_cache']:
if os.path.exists(kw['run_cache'] + lbl + '.pickle'):
with (open(kw['run_cache'] + lbl + '.pickle', 'rb')) as inp:
ssdag, result = _pickle.load(inp)
return criteria, ssdag, result, ['from run cache ' + lbl]
assert _shared_ssdag
ssdag = simple_search_dag(criteria, source=_shared_ssdag, lbl=lbl, **kw)
# print('number of bblocks:', [len(x) for x in ssdag.bbs])
result, tsearch = run_and_time(
grow_linear,
ssdag=ssdag,
loss_function=criteria.jit_lossfunc(),
last_bb_same_as=criteria.from_seg if criteria.is_cyclic else -1,
lbl=lbl,
**kw)
Nsparse = result.stats.total_samples[0]
Nsparse_rate = int(Nsparse / tsearch)
log = []
if len(result.idx):
frac_redundant = result.stats.n_redundant_results[0] / len(result.idx)
log = [
f'grow_linear {lbl} done, nresults {len(result.idx):,}, ' +
f'samp/sec {Nsparse_rate:,}, redundant ratio {frac_redundant}'
]
if log: print(log[-1])
if kw['run_cache']:
with (open(kw['run_cache'] + lbl + '.pickle', 'wb')) as out:
_pickle.dump((ssdag, result), out)
return criteria, ssdag, result, log
def worms_main2(criteria_list, kw):
print('worms_main,', len(criteria_list), 'criteria, args:')
for k, v in kw.items():
print(' ', k, v)
pyrosetta.init('-mute all -beta -preserve_crystinfo')
blosc.set_releasegil(True)
orig_output_prefix = kw['output_prefix']
for icrit, criteria in enumerate(criteria_list):
if len(criteria_list) > 1:
assert len(criteria_list) is len(kw['config_file'])
name = os.path.basename(kw['config_file'][icrit])
name = name.replace('.config', '')
kw['output_prefix'] = orig_output_prefix + '_' + name
print('================== start job', icrit, '======================')
print('output_prefix:', kw['output_prefix'])
print('criteria:', criteria)
print('bbspec:', criteria.bbspec)
if kw['precache_splices']:
print('precaching splices')
merge_bblock = kw['merge_bblock']
del kw['merge_bblock']
kw['bbs'] = simple_search_dag(criteria,
merge_bblock=None,
precache_only=True,
**kw)
kw['merge_bblock'] = merge_bblock
if kw['precache_splices_and_quit']:
return
global _shared_ssdag
if 'bbs' in kw and (len(kw['bbs']) > 2
or kw['bbs'][0] is not kw['bbs'][1]):
############3
#
# _shared_ssdag = simple_search_dag(
# criteria, print_edge_summary=True, **kw
# )
merge_bblock = kw['merge_bblock']
del kw['merge_bblock']
_shared_ssdag = simple_search_dag(criteria,
merge_bblock=0,
print_edge_summary=True,
**kw)
kw['merge_bblock'] = merge_bblock
print('memuse for global _shared_ssdag:')
_shared_ssdag.report_memory_use()
####
#
if _shared_ssdag:
if not 'bbs' in kw:
kw['bbs'] = _shared_ssdag.bbs
assert len(_shared_ssdag.bbs) == len(kw['bbs'])
for a, b in zip(_shared_ssdag.bbs, kw['bbs']):
for aa, bb in zip(a, b):
assert aa is bb
log = worms_main_each_mergebb(criteria, **kw)
if kw['pbar']:
print('======================== logs ========================')
for msg in log:
print(msg)
print('======================== done ========================')
def worms_main2(criteria_list, kw):
print("worms_main,", len(criteria_list), "criteria, args:")
orig_output_prefix = kw["output_prefix"]
for icrit, criteria in enumerate(criteria_list):
if len(criteria_list) > 1:
assert len(criteria_list) is len(kw["config_file"])
name = os.path.basename(kw["config_file"][icrit])
name = name.replace(".config", "")
kw["output_prefix"] = orig_output_prefix + "_" + name
print("================== start job", icrit, "======================")
print("output_prefix:", kw["output_prefix"])
print("criteria:", criteria)
print("bbspec:", criteria.bbspec)
if kw["precache_splices"]:
print("precaching splices")
merge_bblock = kw["merge_bblock"]
del kw["merge_bblock"]
pbar = kw["pbar"]
del kw["pbar"]
kw["bbs"] = simple_search_dag(criteria,
merge_bblock=None,
precache_only=True,
pbar=True,
**kw)
if kw["only_bblocks"]:
assert len(kw["bbs"]) is len(kw["only_bblocks"])
for i, bb in enumerate(kw["bbs"]):
kw["bbs"][i] = [bb[kw["only_bblocks"][i]]]
print("modified bblock numbers (--only_bblocks)")
print(" ", [len(b) for b in kw["bbs"]])
kw["merge_bblock"] = merge_bblock
kw["pbar"] = pbar
if kw["precache_splices_and_quit"]:
return
global _shared_ssdag
if "bbs" in kw and (len(kw["bbs"]) > 2
or kw["bbs"][0] is not kw["bbs"][1]):
############3
#
# _shared_ssdag = simple_search_dag(
# criteria, print_edge_summary=True, **kw
# )
merge_bblock = kw["merge_bblock"]
del kw["merge_bblock"]
_shared_ssdag = simple_search_dag(criteria,
merge_bblock=0,
print_edge_summary=True,
**kw)
kw["merge_bblock"] = merge_bblock
print("memuse for global _shared_ssdag:")
_shared_ssdag.report_memory_use()
####
#
if _shared_ssdag:
if not "bbs" in kw:
kw["bbs"] = _shared_ssdag.bbs
assert len(_shared_ssdag.bbs) == len(kw["bbs"])
for a, b in zip(_shared_ssdag.bbs, kw["bbs"]):
for aa, bb in zip(a, b):
assert aa is bb
if kw["context_structure"]:
kw["context_structure"] = ClashGrid(kw["context_structure"], **kw)
else:
kw["context_structure"] = None
log = worms_main_each_mergebb(criteria, **kw)
if kw["pbar"]:
print("======================== logs ========================")
for msg in log:
print(msg)
print("======================== done ========================")
def search_func(criteria, bbs, monte_carlo, merge_segment, **kw):
stages = [(criteria, bbs)]
merge = None
if hasattr(criteria, "stages"):
stages, merge = criteria.stages(bbs=bbs, **kw)
if len(stages) > 1:
assert kw["merge_bblock"] is not None
assert len(monte_carlo) in (1, len(stages))
if len(monte_carlo) != len(stages):
monte_carlo *= len(stages)
results = list()
for i, stage in enumerate(stages):
crit, stage_bbs = stage
if callable(crit):
crit = crit(*results[-1][:-1])
lbl = f"stage{i}"
if kw["merge_bblock"] is not None:
lbl = f'stage{i}_mbb{kw["merge_bblock"]:04}'
results.append(
search_single_stage(
crit,
monte_carlo=monte_carlo[i],
lbl=lbl,
bbs=stage_bbs,
merge_segment=merge_segment,
**kw,
))
if not hasattr(crit,
"produces_no_results") and len(results[-1][2].idx) is 0:
print("mbb", kw["merge_bblock"], "no results at stage", i)
return None, None, None
# todo: this whole block is very protocol-specific... needs refactoring
if len(results) is 1:
assert merge is None
return results[0][1:]
elif len(results) is 2:
mseg = merge_segment
if mseg is None:
mseg = criteria.merge_segment(**kw)
# simple_search_dag getting not-to-simple maybe split?
_____, ssdA, rsltA, logA = results[0]
critB, ssdB, rsltB, logB = results[1]
assert _shared_ssdag
ssdag = simple_search_dag(
criteria,
only_seg=mseg,
make_edges=False,
source=_shared_ssdag,
bbs=bbs,
**kw,
)
ssdag.verts = ssdB.verts[:-1] + (ssdag.verts[mseg], ) + ssdA.verts[1:]
assert len(ssdag.verts) == len(criteria.bbspec)
rslt = merge(criteria, ssdag, ssdA, rsltA, critB, ssdB, rsltB, **kw)
return ssdag, rslt, logA + logB
elif len(results) is 3:
# hacky: assume 3stage is brigde protocol
_____, ____, _____, logA = results[0]
_____, ssdB, _____, logB = results[1]
critC, ssdC, rsltC, logC = results[2]
assert _shared_ssdag
mseg = merge_segment
if mseg is None:
mseg = criteria.merge_segment(**kw)
ssdag = simple_search_dag(
criteria,
only_seg=mseg,
make_edges=False,
source=_shared_ssdag,
bbs=bbs,
**kw,
)
ssdag.verts = ssdC.verts[:-1] + (ssdag.verts[mseg], ) + ssdB.verts[1:]
assert len(ssdag.verts) == len(criteria.bbspec)
rslt = merge(criteria, critC, ssdag, ssdB, ssdC, rsltC, **kw)
return ssdag, rslt, logA + logB + logC
else:
assert 0, "unknown 3+ stage protcol"