def raref_nem(index, tmpdir, beta, sm_degree, free_dispersion, chunk_size, K, krange, seed):
samp = samples[index]
currtmpdir = tmpdir+"/"+str(index)+"/"
if K < 3:
K = ppp.evaluate_nb_partitions(samp, sm_degree, free_dispersion, chunk_size, krange, 0.05, False, 1, tmpdir + "/" + str(index) + "_eval", seed, None)
if len(samp) <= chunk_size:#all good, just write stuff.
edges_weight, nb_fam = ppp.write_nem_input_files(tmpdir=currtmpdir,organisms= set(samp), sm_degree = sm_degree)
cpt_partition = ppp.run_partitioning( currtmpdir, len(samp), beta * (nb_fam/edges_weight), free_dispersion, K = K, seed = seed, init = "param_file")[0]
else:#going to need multiple partitionnings for this sample...
families = set()
cpt_partition = {}
validated = set()
cpt=0
def validate_family(result):
for node, nem_class in result[0].items():
cpt_partition[node][nem_class[0]]+=1
sum_partionning = sum(cpt_partition[node].values())
if (sum_partionning > len(samp)/chunk_size and max(cpt_partition[node].values()) >= sum_partionning*0.5) or (sum_partionning > len(samp)):
if node not in validated:
if max(cpt_partition[node].values()) < sum_partionning*0.5:
cpt_partition[node]["U"] = len(samp)
validated.add(node)
for fam in ppp.pan.geneFamilies:
if not samp.isdisjoint(fam.organisms):#otherwise useless to keep track of
families.add(fam)
cpt_partition[fam.name] = {"P":0,"S":0,"C":0,"U":0}
org_nb_sample = Counter()
for org in samp:
org_nb_sample[org] = 0
condition = len(samp)/chunk_size
while len(validated) < len(families):
org_samples = []
while not all(val >= condition for val in org_nb_sample.values()):#each family must be tested at least len(select_organisms)/chunk_size times.
shuffled_orgs = list(samp)#copy select_organisms
random.shuffle(shuffled_orgs)#shuffle the copied list
while len(shuffled_orgs) > chunk_size:
org_samples.append(set(shuffled_orgs[:chunk_size]))
for org in org_samples[-1]:
org_nb_sample[org] +=1
shuffled_orgs = shuffled_orgs[chunk_size:]
#making arguments for all samples:
for samp in org_samples:
edges_weight, nb_fam = ppp.write_nem_input_files( currtmpdir+"/"+str(cpt)+"/", samp, sm_degree = sm_degree)
validate_family(ppp.run_partitioning( currtmpdir+"/"+str(cpt)+"/", len(samp), beta * (nb_fam/edges_weight), free_dispersion, K = K, seed = seed, init = "param_file"))
cpt+=1
if len(cpt_partition) == 0:
counts = {"persistent":"NA","shell":"NA","cloud":"NA", "undefined":"NA", "K": K}
else:
counts = {"persistent":0,"shell":0,"cloud":0, "undefined":0, "K":K}
for val in cpt_partition.values():
if isinstance(val, str):
part = val
else:
part = max(val, key=val.get)
if part.startswith("P"):
counts["persistent"]+=1
elif part.startswith("C"):
counts["cloud"]+=1
elif part.startswith("S"):
counts["shell"]+=1
else:
counts["undefined"]+=1
return (counts, index)
def makeRarefactionCurve( pangenome, output, tmpdir, beta=2.5, depth = 30, minSampling =1, maxSampling = 100, sm_degree = 10, free_dispersion=False, chunk_size = 500, K=-1, cpu = 1, seed=42, kestimate = False, krange = [3,-1], soft_core = 0.95, show_bar=True):
ppp.pan = pangenome#use the global from partition to store the pangenome, so that it is usable
try:
krange[0] = ppp.pan.parameters["partition"]["K"] if krange[0]<0 else krange[0]
krange[1] = ppp.pan.parameters["partition"]["K"] if krange[1]<0 else krange[1]
except KeyError:
krange=[3,20]
checkPangenomeInfo(pangenome, needAnnotations=True, needFamilies=True, needGraph=True, show_bar=show_bar)
tmpdirObj = tempfile.TemporaryDirectory(dir=tmpdir)
tmpdir = tmpdirObj.name
if float(len(pangenome.organisms)) < maxSampling:
maxSampling = len(pangenome.organisms)
else:
maxSampling = int(maxSampling)
if K < 3 and kestimate is False:#estimate K once and for all.
try:
K = ppp.pan.parameters["partition"]["K"]
logging.getLogger().info(f"Reuse the number of partitions {K}")
except KeyError:
logging.getLogger().info("Estimating the number of partitions...")
K = ppp.evaluate_nb_partitions(pangenome.organisms, sm_degree, free_dispersion, chunk_size, krange, 0.05, False, cpu, tmpdir, seed, None)
logging.getLogger().info(f"The number of partitions has been evaluated at {K}")
logging.getLogger().info("Extracting samples ...")
AllSamples = []
for i in range(minSampling,maxSampling):#each point
for _ in range(depth):#number of samples per points
AllSamples.append(set(random.sample(set(pangenome.organisms), i+1)))
logging.getLogger().info(f"Done sampling organisms in the pangenome, there are {len(AllSamples)} samples")
SampNbPerPart = []
logging.getLogger().info("Computing bitarrays for each family...")
index_org = pangenome.computeFamilyBitarrays()
logging.getLogger().info(f"Done computing bitarrays. Comparing them to get exact and soft core stats for {len(AllSamples)} samples...")
bar = tqdm( range(len(AllSamples) * len(pangenome.geneFamilies)), unit = "gene family", disable=not show_bar)
for samp in AllSamples:
#make the sample's organism bitarray.
sampBitarray = gmpy2.xmpz(0)#pylint: disable=no-member
for org in samp:
sampBitarray[index_org[org]] = 1
part = Counter()
part["soft_core"] = 0
part["exact_core"] = 0
part["exact_accessory"] = 0
part["soft_accessory"] = 0
for fam in pangenome.geneFamilies:
nbCommonOrg = gmpy2.popcount(fam.bitarray & sampBitarray)#pylint: disable=no-member
part["nborgs"] = len(samp)
if nbCommonOrg != 0:#in that case the node 'does not exist'
if nbCommonOrg == len(samp):
part["exact_core"] +=1
else:
part["exact_accessory"] +=1
if float(nbCommonOrg) >= len(samp) * soft_core:
part["soft_core"] +=1
else:
part["soft_accessory"] +=1
bar.update()
SampNbPerPart.append(part)
bar.close()
#done with frequency of each family for each sample.
global samples
samples = AllSamples
args = []
for index, samp in enumerate(samples):
args.append((index, tmpdir, beta, sm_degree, free_dispersion, chunk_size, K, krange, seed))
with Pool(processes = cpu) as p:
#launch partitionnings
logging.getLogger().info("Partitionning all samples...")
bar = tqdm(range(len(args)), unit = "samples partitionned", disable=not show_bar)
random.shuffle(args)#shuffling the processing so that the progress bar is closer to reality.
for result in p.imap_unordered(launch_raref_nem, args):
SampNbPerPart[result[1]] = {**result[0], **SampNbPerPart[result[1]]}
bar.update()
bar.close()
logging.getLogger().info("Done partitionning everything")
warnings.filterwarnings("ignore")
drawCurve(output, maxSampling, SampNbPerPart )
warnings.resetwarnings()
tmpdirObj.cleanup()
logging.getLogger().info("Done making the rarefaction curves")