treetime/treeanc.py
def optimize_branch_length_global(self, **kwargs):
        """
        EXPERIMENTAL GLOBAL OPTIMIZATION
        """

        self.logger("TreeAnc.optimize_branch_length_global: running branch length optimization...",1)

        def neg_log(s):
            for si, n in zip(s, self.tree.find_clades(order='preorder')):
                n.branch_length = si**2

            self.infer_ancestral_sequences(marginal=True)

            gradient = []
            for si, n in zip(s, self.tree.find_clades(order='preorder')):
                if n.up:
                    pp, pc = self.marginal_branch_profile(n)
                    Qtds = self.gtr.expQsds(si).T
                    Qt = self.gtr.expQs(si).T

                    res = pp.dot(Qt)
                    overlap = np.sum(res*pc, axis=1)

                    res_ds = pp.dot(Qtds)
                    overlap_ds = np.sum(res_ds*pc, axis=1)
                    logP = np.sum(self.multiplicity*overlap_ds/overlap)

                    gradient.append(logP)
                else:
                    gradient.append(2*(si**2-0.001))

            print(-self.tree.sequence_marginal_LH)
            return (-self.tree.sequence_marginal_LH + (s[0]**2-0.001)**2, -1.0*np.array(gradient))

        from scipy.optimize import minimize
        x0 = np.sqrt([n.branch_length for n in self.tree.find_clades(order='preorder')])
        sol = minimize(neg_log, x0, jac=True)

        for new_len, node in zip(sol['x'], self.tree.find_clades()):
            self.logger("Optimization results: old_len=%.4f, new_len=%.4f "
                   " Updating branch length..."%(node.branch_length, new_len), 5)

            node.branch_length = new_len**2
            node.mutation_length=new_len**2

        # as branch lengths changed, the params must be fixed
        self.tree.root.up = None
        self.tree.root.dist2root = 0.0
        self._prepare_nodes()
def optimize_branch_length_global(self, **kwargs):
        """
        EXPERIMENTAL GLOBAL OPTIMIZATION
        """

        self.logger("TreeAnc.optimize_branch_length_global: running branch length optimization...",1)

        def neg_log(s):
            for si, n in zip(s, self.tree.find_clades(order='preorder')):
                n.branch_length = si**2

            self.infer_ancestral_sequences(marginal=True)

            gradient = []
            for si, n in zip(s, self.tree.find_clades(order='preorder')):
                if n.up:
                    pp, pc = self.marginal_branch_profile(n)
                    Qtds = self.gtr.expQsds(si).T
                    Qt = self.gtr.expQs(si).T

                    res = pp.dot(Qt)
                    overlap = np.sum(res*pc, axis=1)

                    res_ds = pp.dot(Qtds)
                    overlap_ds = np.sum(res_ds*pc, axis=1)
                    logP = np.sum(self.multiplicity*overlap_ds/overlap)

                    gradient.append(logP)
                else:
                    gradient.append(2*(si**2-0.001))

            print(-self.tree.sequence_marginal_LH)
            return (-self.tree.sequence_marginal_LH + (s[0]**2-0.001)**2, -1.0*np.array(gradient))

        from scipy.optimize import minimize
        x0 = np.sqrt([n.branch_length for n in self.tree.find_clades(order='preorder')])
        sol = minimize(neg_log, x0, jac=True)

        for new_len, node in zip(sol['x'], self.tree.find_clades()):
            self.logger("Optimization results: old_len=%.4f, new_len=%.4f "
                   " Updating branch length..."%(node.branch_length, new_len), 5)

            node.branch_length = new_len**2
            node.mutation_length=new_len**2

        # as branch lengths changed, the params must be fixed
        self.tree.root.up = None
        self.tree.root.dist2root = 0.0
        self._prepare_nodes()
EXPERIMENTAL GLOBAL OPTIMIZATION