treetime/treetime.py
def run(self, root=None, infer_gtr=True, relaxed_clock=None, n_iqd = None,
            resolve_polytomies=True, max_iter=0, Tc=None, fixed_clock_rate=None,
            time_marginal=False, sequence_marginal=False, branch_length_mode='auto',
            vary_rate=False, use_covariation=False, **kwargs):

        """
        Run TreeTime reconstruction. Based on the input parameters, it divides
        the analysis into semi-independent jobs and conquers them one-by-one,
        gradually optimizing the tree given the temporal constarints and leaf
        node sequences.

        Parameters
        ----------
        root : str
           Try to find better root position on a given tree. If string is passed,
           the root will be searched according to the specified method. If none,
           use tree as-is.

           See :py:meth:`treetime.TreeTime.reroot` for available rooting methods.

        infer_gtr : bool
           If True, infer GTR model

        relaxed_clock : dic
           If not None, use autocorrelated molecular clock model. Specify the
           clock parameters as :code:`{slack:<slack>, coupling:<coupling>}` dictionary.

        n_iqd : int
           If not None, filter tree nodes which do not obey the molecular clock
           for the particular tree. The nodes, which deviate more than
           :code:`n_iqd` interquantile intervals from the molecular clock
           regression will be marked as 'BAD' and not used in the TreeTime
           analysis

        resolve_polytomies : bool
           If True, attempt to resolve multiple mergers

        max_iter : int
           Maximum number of iterations to optimize the tree

        Tc : float, str
           If not None, use coalescent model to correct the branch lengths by
           introducing merger costs.

           If Tc is float, it is interpreted as the coalescence time scale

           If Tc is str, it should be one of (:code:`opt`, :code:`const`, :code:`skyline`)

        fixed_clock_rate : float
           Fixed clock rate to be used. If None, infer clock rate from the molecular clock.

        time_marginal : bool
           If True, perform a final round of marginal reconstruction of the node's positions.

        sequence_marginal : bool, optional
            use marginal reconstruction for ancestral sequences

        branch_length_mode : str
           Should be one of: :code:`joint`, :code:`marginal`, :code:`input`.

           If 'input', rely on the branch lengths in the input tree and skip directly
           to the maximum-likelihood ancestral sequence reconstruction.
           Otherwise, perform preliminary sequence reconstruction using parsimony
           algorithm and do branch length optimization

        vary_rate : bool or float, optional
            redo the time tree estimation for rates +/- one standard deviation.
            if a float is passed, it is interpreted as standard deviation,
            otherwise this standard deviation is estimated from the root-to-tip regression

        use_covariation : bool, optional
            default False, if False, rate estimates will be performed using simple
            regression ignoring phylogenetic covaration between nodes.

        **kwargs
           Keyword arguments needed by the downstream functions


        Returns
        -------
        TreeTime error/succces code : str
            return value depending on success or error


        """

        # register the specified covaration mode
        self.use_covariation = use_covariation

        if (self.tree is None) or (self.aln is None and self.seq_len is None):
            self.logger("TreeTime.run: ERROR, alignment or tree are missing", 0)
            return ttconf.ERROR
        if (self.aln is None):
            branch_length_mode='input'

        self._set_branch_length_mode(branch_length_mode)

        # determine how to reconstruct and sample sequences
        seq_kwargs = {"marginal_sequences":sequence_marginal or (self.branch_length_mode=='marginal'),
                      "sample_from_profile":"root"}

        tt_kwargs = {'clock_rate':fixed_clock_rate, 'time_marginal':False}
        tt_kwargs.update(kwargs)

        seq_LH = 0
        if "fixed_pi" in kwargs:
            seq_kwargs["fixed_pi"] = kwargs["fixed_pi"]
        if "do_marginal" in kwargs:
            time_marginal=kwargs["do_marginal"]

        # initially, infer ancestral sequences and infer gtr model if desired
        if self.branch_length_mode=='input':
            if self.aln:
                self.infer_ancestral_sequences(infer_gtr=infer_gtr, **seq_kwargs)
                self.prune_short_branches()
        else:
            self.optimize_sequences_and_branch_length(infer_gtr=infer_gtr,
                                                      max_iter=1, prune_short=True, **seq_kwargs)
        avg_root_to_tip = np.mean([x.dist2root for x in self.tree.get_terminals()])

        # optionally reroot the tree either by oldest, best regression or with a specific leaf
        if n_iqd or root=='clock_filter':
            if "plot_rtt" in kwargs and kwargs["plot_rtt"]:
                plot_rtt=True
            else:
                plot_rtt=False
            reroot_mechanism = 'least-squares' if root=='clock_filter' else root
            if self.clock_filter(reroot=reroot_mechanism, n_iqd=n_iqd, plot=plot_rtt)==ttconf.ERROR:
                return ttconf.ERROR
        elif root is not None:
            if self.reroot(root=root)==ttconf.ERROR:
                return ttconf.ERROR

        if self.branch_length_mode=='input':
            if self.aln:
                self.infer_ancestral_sequences(**seq_kwargs)
        else:
            self.optimize_sequences_and_branch_length(max_iter=1, prune_short=False,
                                                      **seq_kwargs)

        # infer time tree and optionally resolve polytomies
        self.logger("###TreeTime.run: INITIAL ROUND",0)
        self.make_time_tree(**tt_kwargs)

        if self.aln:
            seq_LH = self.tree.sequence_marginal_LH if seq_kwargs['marginal_sequences'] else self.tree.sequence_joint_LH
        self.LH =[[seq_LH, self.tree.positional_joint_LH, 0]]

        if root is not None and max_iter:
            if self.reroot(root='least-squares' if root=='clock_filter' else root)==ttconf.ERROR:
                return ttconf.ERROR

        # iteratively reconstruct ancestral sequences and re-infer
        # time tree to ensure convergence.
        niter = 0
        ndiff = 0
        need_new_time_tree=False
        while niter < max_iter:
            self.logger("###TreeTime.run: ITERATION %d out of %d iterations"%(niter+1,max_iter),0)
            # add coalescent prior
            if Tc:
                if Tc=='skyline' and niter<max_iter-1:
                    tmpTc='const'
                else:
                    tmpTc=Tc
                self.add_coalescent_model(tmpTc, **kwargs)
                need_new_time_tree = True

            # estimate a relaxed molecular clock
            if relaxed_clock:
                print("relaxed_clock", relaxed_clock)
                self.relaxed_clock(**relaxed_clock)
                need_new_time_tree = True

            n_resolved=0
            if resolve_polytomies:
                # if polytomies are found, rerun the entire procedure
                n_resolved = self.resolve_polytomies()
                if n_resolved:
                    self.prepare_tree()
                    if self.branch_length_mode!='input': # otherwise reoptimize branch length while preserving branches without mutations
                        self.optimize_sequences_and_branch_length(prune_short=False,
                                                          max_iter=0, **seq_kwargs)

                    need_new_time_tree = True

            if need_new_time_tree:
                self.make_time_tree(**tt_kwargs)
                if self.aln:
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
            else: # no refinements, just iterate
                if self.aln:
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
                self.make_time_tree(**tt_kwargs)

            self.tree.coalescent_joint_LH = self.merger_model.total_LH() if Tc else 0.0

            if self.aln:
                seq_LH = self.tree.sequence_marginal_LH if seq_kwargs['marginal_sequences'] else self.tree.sequence_joint_LH
            self.LH.append([seq_LH, self.tree.positional_joint_LH, self.tree.coalescent_joint_LH])
            niter+=1

            if ndiff==0 and n_resolved==0 and Tc!='skyline':
                self.logger("###TreeTime.run: CONVERGED",0)
                break

        # if the rate is too be varied and the rate estimate has a valid confidence interval
        # rerun the estimation for variations of the rate
        if vary_rate:
            if type(vary_rate)==float:
                res = self.calc_rate_susceptibility(rate_std=vary_rate, params=tt_kwargs)
            elif self.clock_model['valid_confidence']:
                res = self.calc_rate_susceptibility(params=tt_kwargs)
            else:
                res = ttconf.ERROR

            if res==ttconf.ERROR:
                self.logger("TreeTime.run: rate variation failed and can't be used for confidence estimation", 1, warn=True)

        # if marginal reconstruction requested, make one more round with marginal=True
        # this will set marginal_pos_LH, which to be used as error bar estimations
        if time_marginal:
            self.logger("###TreeTime.run: FINAL ROUND - confidence estimation via marginal reconstruction", 0)
            tt_kwargs['time_marginal']=time_marginal
            self.make_time_tree(**tt_kwargs)

        # explicitly print out which branches are bad and whose dates don't correspond to the input dates
        bad_branches =[n for n in self.tree.get_terminals()
                       if n.bad_branch and n.raw_date_constraint]
        if bad_branches:
            self.logger("TreeTime: The following tips don't fit the clock model, "
                        "please remove them from the tree. Their dates have been reset:",0,warn=True)
            for n in bad_branches:
                self.logger("%s, input date: %s, apparent date: %1.2f"%(n.name, str(n.raw_date_constraint), n.numdate),0,warn=True)

        return ttconf.SUCCESS
def run(self, root=None, infer_gtr=True, relaxed_clock=None, n_iqd = None,
            resolve_polytomies=True, max_iter=0, Tc=None, fixed_clock_rate=None,
            time_marginal=False, sequence_marginal=False, branch_length_mode='auto',
            vary_rate=False, use_covariation=False, **kwargs):

        """
        Run TreeTime reconstruction. Based on the input parameters, it divides
        the analysis into semi-independent jobs and conquers them one-by-one,
        gradually optimizing the tree given the temporal constarints and leaf
        node sequences.

        Parameters
        ----------
        root : str
           Try to find better root position on a given tree. If string is passed,
           the root will be searched according to the specified method. If none,
           use tree as-is.

           See :py:meth:`treetime.TreeTime.reroot` for available rooting methods.

        infer_gtr : bool
           If True, infer GTR model

        relaxed_clock : dic
           If not None, use autocorrelated molecular clock model. Specify the
           clock parameters as :code:`{slack:<slack>, coupling:<coupling>}` dictionary.

        n_iqd : int
           If not None, filter tree nodes which do not obey the molecular clock
           for the particular tree. The nodes, which deviate more than
           :code:`n_iqd` interquantile intervals from the molecular clock
           regression will be marked as 'BAD' and not used in the TreeTime
           analysis

        resolve_polytomies : bool
           If True, attempt to resolve multiple mergers

        max_iter : int
           Maximum number of iterations to optimize the tree

        Tc : float, str
           If not None, use coalescent model to correct the branch lengths by
           introducing merger costs.

           If Tc is float, it is interpreted as the coalescence time scale

           If Tc is str, it should be one of (:code:`opt`, :code:`const`, :code:`skyline`)

        fixed_clock_rate : float
           Fixed clock rate to be used. If None, infer clock rate from the molecular clock.

        time_marginal : bool
           If True, perform a final round of marginal reconstruction of the node's positions.

        sequence_marginal : bool, optional
            use marginal reconstruction for ancestral sequences

        branch_length_mode : str
           Should be one of: :code:`joint`, :code:`marginal`, :code:`input`.

           If 'input', rely on the branch lengths in the input tree and skip directly
           to the maximum-likelihood ancestral sequence reconstruction.
           Otherwise, perform preliminary sequence reconstruction using parsimony
           algorithm and do branch length optimization

        vary_rate : bool or float, optional
            redo the time tree estimation for rates +/- one standard deviation.
            if a float is passed, it is interpreted as standard deviation,
            otherwise this standard deviation is estimated from the root-to-tip regression

        use_covariation : bool, optional
            default False, if False, rate estimates will be performed using simple
            regression ignoring phylogenetic covaration between nodes.

        **kwargs
           Keyword arguments needed by the downstream functions


        Returns
        -------
        TreeTime error/succces code : str
            return value depending on success or error


        """

        # register the specified covaration mode
        self.use_covariation = use_covariation

        if (self.tree is None) or (self.aln is None and self.seq_len is None):
            self.logger("TreeTime.run: ERROR, alignment or tree are missing", 0)
            return ttconf.ERROR
        if (self.aln is None):
            branch_length_mode='input'

        self._set_branch_length_mode(branch_length_mode)

        # determine how to reconstruct and sample sequences
        seq_kwargs = {"marginal_sequences":sequence_marginal or (self.branch_length_mode=='marginal'),
                      "sample_from_profile":"root"}

        tt_kwargs = {'clock_rate':fixed_clock_rate, 'time_marginal':False}
        tt_kwargs.update(kwargs)

        seq_LH = 0
        if "fixed_pi" in kwargs:
            seq_kwargs["fixed_pi"] = kwargs["fixed_pi"]
        if "do_marginal" in kwargs:
            time_marginal=kwargs["do_marginal"]

        # initially, infer ancestral sequences and infer gtr model if desired
        if self.branch_length_mode=='input':
            if self.aln:
                self.infer_ancestral_sequences(infer_gtr=infer_gtr, **seq_kwargs)
                self.prune_short_branches()
        else:
            self.optimize_sequences_and_branch_length(infer_gtr=infer_gtr,
                                                      max_iter=1, prune_short=True, **seq_kwargs)
        avg_root_to_tip = np.mean([x.dist2root for x in self.tree.get_terminals()])

        # optionally reroot the tree either by oldest, best regression or with a specific leaf
        if n_iqd or root=='clock_filter':
            if "plot_rtt" in kwargs and kwargs["plot_rtt"]:
                plot_rtt=True
            else:
                plot_rtt=False
            reroot_mechanism = 'least-squares' if root=='clock_filter' else root
            if self.clock_filter(reroot=reroot_mechanism, n_iqd=n_iqd, plot=plot_rtt)==ttconf.ERROR:
                return ttconf.ERROR
        elif root is not None:
            if self.reroot(root=root)==ttconf.ERROR:
                return ttconf.ERROR

        if self.branch_length_mode=='input':
            if self.aln:
                self.infer_ancestral_sequences(**seq_kwargs)
        else:
            self.optimize_sequences_and_branch_length(max_iter=1, prune_short=False,
                                                      **seq_kwargs)

        # infer time tree and optionally resolve polytomies
        self.logger("###TreeTime.run: INITIAL ROUND",0)
        self.make_time_tree(**tt_kwargs)

        if self.aln:
            seq_LH = self.tree.sequence_marginal_LH if seq_kwargs['marginal_sequences'] else self.tree.sequence_joint_LH
        self.LH =[[seq_LH, self.tree.positional_joint_LH, 0]]

        if root is not None and max_iter:
            if self.reroot(root='least-squares' if root=='clock_filter' else root)==ttconf.ERROR:
                return ttconf.ERROR

        # iteratively reconstruct ancestral sequences and re-infer
        # time tree to ensure convergence.
        niter = 0
        ndiff = 0
        need_new_time_tree=False
        while niter < max_iter:
            self.logger("###TreeTime.run: ITERATION %d out of %d iterations"%(niter+1,max_iter),0)
            # add coalescent prior
            if Tc:
                if Tc=='skyline' and niter<max_iter-1:
                    tmpTc='const'
                else:
                    tmpTc=Tc
                self.add_coalescent_model(tmpTc, **kwargs)
                need_new_time_tree = True

            # estimate a relaxed molecular clock
            if relaxed_clock:
                print("relaxed_clock", relaxed_clock)
                self.relaxed_clock(**relaxed_clock)
                need_new_time_tree = True

            n_resolved=0
            if resolve_polytomies:
                # if polytomies are found, rerun the entire procedure
                n_resolved = self.resolve_polytomies()
                if n_resolved:
                    self.prepare_tree()
                    if self.branch_length_mode!='input': # otherwise reoptimize branch length while preserving branches without mutations
                        self.optimize_sequences_and_branch_length(prune_short=False,
                                                          max_iter=0, **seq_kwargs)

                    need_new_time_tree = True

            if need_new_time_tree:
                self.make_time_tree(**tt_kwargs)
                if self.aln:
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
            else: # no refinements, just iterate
                if self.aln:
                    ndiff = self.infer_ancestral_sequences('ml',**seq_kwargs)
                self.make_time_tree(**tt_kwargs)

            self.tree.coalescent_joint_LH = self.merger_model.total_LH() if Tc else 0.0

            if self.aln:
                seq_LH = self.tree.sequence_marginal_LH if seq_kwargs['marginal_sequences'] else self.tree.sequence_joint_LH
            self.LH.append([seq_LH, self.tree.positional_joint_LH, self.tree.coalescent_joint_LH])
            niter+=1

            if ndiff==0 and n_resolved==0 and Tc!='skyline':
                self.logger("###TreeTime.run: CONVERGED",0)
                break

        # if the rate is too be varied and the rate estimate has a valid confidence interval
        # rerun the estimation for variations of the rate
        if vary_rate:
            if type(vary_rate)==float:
                res = self.calc_rate_susceptibility(rate_std=vary_rate, params=tt_kwargs)
            elif self.clock_model['valid_confidence']:
                res = self.calc_rate_susceptibility(params=tt_kwargs)
            else:
                res = ttconf.ERROR

            if res==ttconf.ERROR:
                self.logger("TreeTime.run: rate variation failed and can't be used for confidence estimation", 1, warn=True)

        # if marginal reconstruction requested, make one more round with marginal=True
        # this will set marginal_pos_LH, which to be used as error bar estimations
        if time_marginal:
            self.logger("###TreeTime.run: FINAL ROUND - confidence estimation via marginal reconstruction", 0)
            tt_kwargs['time_marginal']=time_marginal
            self.make_time_tree(**tt_kwargs)

        # explicitly print out which branches are bad and whose dates don't correspond to the input dates
        bad_branches =[n for n in self.tree.get_terminals()
                       if n.bad_branch and n.raw_date_constraint]
        if bad_branches:
            self.logger("TreeTime: The following tips don't fit the clock model, "
                        "please remove them from the tree. Their dates have been reset:",0,warn=True)
            for n in bad_branches:
                self.logger("%s, input date: %s, apparent date: %1.2f"%(n.name, str(n.raw_date_constraint), n.numdate),0,warn=True)

        return ttconf.SUCCESS
Run TreeTime reconstruction. Based on the input parameters, it divides
        the analysis into semi-independent jobs and conquers them one-by-one,
        gradually optimizing the tree given the temporal constarints and leaf
        node sequences.

        Parameters
        ----------
        root : str
           Try to find better root position on a given tree. If string is passed,
           the root will be searched according to the specified method. If none,
           use tree as-is.

           See :py:meth:`treetime.TreeTime.reroot` for available rooting methods.

        infer_gtr : bool
           If True, infer GTR model

        relaxed_clock : dic
           If not None, use autocorrelated molecular clock model. Specify the
           clock parameters as :code:`{slack:<slack>, coupling:<coupling>}` dictionary.

        n_iqd : int
           If not None, filter tree nodes which do not obey the molecular clock
           for the particular tree. The nodes, which deviate more than
           :code:`n_iqd` interquantile intervals from the molecular clock
           regression will be marked as 'BAD' and not used in the TreeTime
           analysis

        resolve_polytomies : bool
           If True, attempt to resolve multiple mergers

        max_iter : int
           Maximum number of iterations to optimize the tree

        Tc : float, str
           If not None, use coalescent model to correct the branch lengths by
           introducing merger costs.

           If Tc is float, it is interpreted as the coalescence time scale

           If Tc is str, it should be one of (:code:`opt`, :code:`const`, :code:`skyline`)

        fixed_clock_rate : float
           Fixed clock rate to be used. If None, infer clock rate from the molecular clock.

        time_marginal : bool
           If True, perform a final round of marginal reconstruction of the node's positions.

        sequence_marginal : bool, optional
            use marginal reconstruction for ancestral sequences

        branch_length_mode : str
           Should be one of: :code:`joint`, :code:`marginal`, :code:`input`.

           If 'input', rely on the branch lengths in the input tree and skip directly
           to the maximum-likelihood ancestral sequence reconstruction.
           Otherwise, perform preliminary sequence reconstruction using parsimony
           algorithm and do branch length optimization

        vary_rate : bool or float, optional
            redo the time tree estimation for rates +/- one standard deviation.
            if a float is passed, it is interpreted as standard deviation,
            otherwise this standard deviation is estimated from the root-to-tip regression

        use_covariation : bool, optional
            default False, if False, rate estimates will be performed using simple
            regression ignoring phylogenetic covaration between nodes.

        **kwargs
           Keyword arguments needed by the downstream functions


        Returns
        -------
        TreeTime error/succces code : str
            return value depending on success or error