main_sample_dates.py

# BSD 3-Clause License

# Copyright (c) 2025, Jonathan David Duke

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# 1. Redistributions of source code must retain the above copyright notice, this
#    list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.

# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import sys
import numpy as np

from dated_complete_tree import tree_loading
from dated_complete_tree import tree_labelling
from dated_complete_tree import tree_fixing
from dated_complete_tree import tree_dating
from dated_complete_tree import tree_metrics

import logging
logger = logging.getLogger(__name__)
logging.basicConfig(filename="main.log", filemode="w", force=True, level=logging.DEBUG)

sys.setrecursionlimit(10000)

#####################################################################################################################
# Load and prune tree

# Load metadata for tree from Open Tree and Chronosynth
dates, phylogeny_nodes, taxa = tree_loading.load_metadata()

# Create ete4 tree structure for entire Open Tree of Life, with my annotations
whole_tre_unmodified = tree_loading.build_and_annotate_tree(phylogeny_nodes, taxa)

tree_fixing.strip_birds(whole_tre_unmodified)
tree_fixing.strip_turtles(whole_tre_unmodified)

rng = np.random.default_rng(seed=1)

tree_fixing.remove_subspecies(whole_tre_unmodified, rng)

tree_fixing.impute_species_into_empty_taxa(whole_tre_unmodified)

tree_fixing.fix_taxonomy_ordering(whole_tre_unmodified)

tree_labelling.add_anc_ranks(whole_tre_unmodified)
tree_labelling.add_desc_ranks(whole_tre_unmodified)

tree_fixing.forced_taxa_moves(whole_tre_unmodified)

# Copy tree - we will change the copy, and keep the original unchanged so we can restore it next iteration without
# reloading everything
whole_tre = whole_tre_unmodified.copy()

#####################################################################################################################
# Fix topology

# First, do labelling for steps 1-3:
#  - 1-2 are independent of each other; step 3 collects up nodes not labelled in 1-2.
#  - tree is only labelled at this stage; modifications are made in tree_fixing functions.
genus_dict = {}       # step 1, nodes below genus nodes
nmp_genus_dict = {}   # step 2, non-monophyletic genera
tree_labelling.populate_genus_dict(whole_tre, genus_dict, nmp_genus_dict, None)

tofix_dict = {}       # step 3, all other nodes from taxonomy (not phylogenies) to
                      # be moved to a suitable place in the tree, such that we
                      # generated a plausible hypothetical tree
tree_labelling.populate_tofix_dict(whole_tre, tofix_dict, nmp_genus_dict)

# Second, fix the topology based on the labels.
# Fix steps 1 and 2.
tree_fixing.fix_polyphyly(genus_dict, rng)
tree_fixing.fix_polyphyly(nmp_genus_dict, rng)

tree_fixing.remove_nonspecies_leaves(whole_tre)

# Find and label backbone for step 3, after steps 1 an 2 already fixed.
tree_labelling.populate_tofix_bkb(whole_tre, tofix_dict, [])
fix_dict = tree_labelling.process_tofix_bkb(tofix_dict)

# Finally, fix step 3.
tree_fixing.fix_polyphyly(fix_dict, rng, expand_parent_backbones=True)

tree_fixing.remove_nonspecies_leaves(whole_tre)

# Last of all, polytomy resolution.
tree_fixing.fix_all_polytomies(whole_tre, rng)

# Remove one-child nodes. Gives a fully bifurcating tree.
whole_tre = tree_fixing.delete_one_child_nodes(whole_tre)

#####################################################################################################################
# Assign and interpolate dates

tree_dating.assign_dates(whole_tre, dates)

tree_dating.dq_date_removal(whole_tre)

from dated_complete_tree import tree_checks

tree_checks.check_bifurcating(whole_tre)
tree_checks.count_subspecies(whole_tre)
tree_checks.check_taxonomy_order(whole_tre)
tree_checks.check_zero_dates(whole_tre)
tree_checks.check_inconsistent_dates(whole_tre)

# Date imputation
tree_dating.date_labelling(whole_tre)

tree_checks.check_zero_dates(whole_tre)
tree_checks.check_inconsistent_dates(whole_tre)

tree_dating.impute_missing_dates(whole_tre, l=0.25)

tree_checks.check_zero_dates(whole_tre)
tree_checks.check_inconsistent_dates(whole_tre)

tree_dating.compute_branch_lengths(whole_tre)
tree_metrics.compute_pd(whole_tre)

base_pd = whole_tre.props["pd"]

pd_clades = [cld.strip() for cld in list(open("pd_clades.txt"))]

pd_dict = {}
dates_dict = {}
spp_dict = {}
for clade in pd_clades:
    pd_dict[clade] = []
    dates_dict[clade] = []
    spp_dict[clade] = []

tree_metrics.save_pd_for_clades(whole_tre, pd_clades, pd_dict, dates_dict, spp_dict)


# Assign dates from varying sources
date_source_rng = np.random.default_rng(seed=10)

results = []

import datetime
itr_start = datetime.datetime.now()
num_itrs = 1

for itr in range(num_itrs):
    print("iteration", itr+1, "/ projected end time:", itr_start + (num_itrs)*(datetime.datetime.now() - itr_start)/itr if itr > 0 else "")
    for node in whole_tre.traverse(strategy="preorder"):
        # reset all dates
        node.props["date"] = None
        node.props["imputed_date"] = False
        node.props["imputation_type"] = 0

    date_sources = tree_dating.assign_dates(whole_tre, dates, sample_dates=True, rng=date_source_rng)

    # Date cleaning to ensure time consistency down the tree
    tree_dating.dq_date_removal(whole_tre)

    # Date imputation
    tree_dating.date_labelling(whole_tre)
    dating_rng = np.random.default_rng(seed=100)
    tree_dating.impute_missing_dates(whole_tre, l=0.25, rng=dating_rng)

    tree_dating.compute_branch_lengths(whole_tre)
    tree_metrics.compute_pd(whole_tre)

    results.append(whole_tre.props["pd"])
    tree_metrics.save_pd_for_clades(whole_tre, pd_clades, pd_dict, dates_dict, spp_dict)

tree_metrics.write_pd_dists("figures/figure4/fix_topo", pd_dict, dates_dict, spp_dict)



fout = open("date_source_dists_refreshed.txt", "w")
dates_dist = {}
for n in date_sources:
    if len(date_sources[n]) not in dates_dist:
        dates_dist[len(date_sources[n])] = 1
    else:
        dates_dist[len(date_sources[n])] += 1

    fout.write("%s\t%s\t%s\t%s\t%d" % (n.name, n.props["kingdom"], n.props["phylum"], n.props["clas"], len(date_sources[n])))
    if len(date_sources[n]) == 2 and date_sources[n][0] == date_sources[n][1]:
        fout.write("\tsame")
    fout.write('\n')

fout.close()