phylo2vec 1.2.0__tar.gz → 1.3.0__tar.gz

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/Cargo.lock

@@ -544,7 +544,7 @@ dependencies = [
 
 [[package]]
 name = "py-phylo2vec"
-version = "1.2.0"
+version = "1.3.0"
 dependencies = [
  "phylo2vec",
  "pyo3",

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: phylo2vec
-Version: 1.2.0
+Version: 1.3.0
 Classifier: Development Status :: 3 - Alpha
 Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
 Classifier: Programming Language :: Python
@@ -11,12 +11,17 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)
 Classifier: Operating System :: Unix
 Classifier: Operating System :: Microsoft :: Windows
-Requires-Dist: numba>=0.56.4
 Requires-Dist: numpy>=1.22,<2.1
 Requires-Dist: biopython>=1.84.0
 Requires-Dist: joblib>=1.2.0
 Requires-Dist: ete3==3.1.3
 Requires-Dist: six>=1.16.0
+Requires-Dist: jax ; extra == 'opt'
+Requires-Dist: joblib ; extra == 'opt'
+Requires-Dist: optax ; extra == 'opt'
+Requires-Dist: rpy2 ; extra == 'opt'
+Requires-Dist: tqdm ; extra == 'opt'
+Provides-Extra: opt
 License-File: LICENSE
 Summary: Phylo2Vec: integer vector representation of binary (phylogenetic) trees
 Author: Neil Scheidwasser <neil.clow@sund.ku.dk>
@@ -85,8 +90,10 @@ in Rust.
 
 ### Installing R package
 
-To install the R package, first you need to retrieve one of the compiled file
-within the package [releases](https://github.com/sbhattlab/phylo2vec/releases).
+#### Option 1: from a release (Windows, Mac, Ubuntu >= 22.04)
+
+Retrieve one of the compiled binaries from the
+[releases](https://github.com/sbhattlab/phylo2vec/releases) that fits your OS.
 Once the file is downloaded, simply run `install.packages` in your R command
 line.
 
@@ -94,6 +101,41 @@ line.
 install.packages("/path/to/package_file", repos = NULL, type = 'source')
 ```
 
+#### Option 2: using `devtools`
+
+⚠️ This requires installing [Rust](https://www.rust-lang.org/tools/install) to
+build the core package.
+
+```R
+devtools::install_github("sbhattlab/phylo2vec", subdir="./r-phylo2vec", build = FALSE)
+```
+
+Note: to download a specific version, use:
+
+```R
+devtools::install_github("sbhattlab/phylo2vec@vX.Y.Z", subdir="./r-phylo2vec", build = FALSE)
+```
+
+#### Option 3: manual installation
+
+⚠️ This requires installing [Rust](https://www.rust-lang.org/tools/install) to
+build the core package.
+
+Clone the repository and run the following `install.packages` in your R command
+line.
+
+Note: to download a specific version, you can use `git checkout` to a desired
+tag.
+
+```bash
+git clone https://github.com/sbhattlab/phylo2vec
+cd phylo2vec
+```
+
+```R
+install.packages("./r-phylo2vec", repos = NULL, type = 'source')
+```
+
 ## Basic Usage
 
 ### Python

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/README.md

@@ -58,8 +58,10 @@ in Rust.
 
 ### Installing R package
 
-To install the R package, first you need to retrieve one of the compiled file
-within the package [releases](https://github.com/sbhattlab/phylo2vec/releases).
+#### Option 1: from a release (Windows, Mac, Ubuntu >= 22.04)
+
+Retrieve one of the compiled binaries from the
+[releases](https://github.com/sbhattlab/phylo2vec/releases) that fits your OS.
 Once the file is downloaded, simply run `install.packages` in your R command
 line.
 
@@ -67,6 +69,41 @@ line.
 install.packages("/path/to/package_file", repos = NULL, type = 'source')
 ```
 
+#### Option 2: using `devtools`
+
+⚠️ This requires installing [Rust](https://www.rust-lang.org/tools/install) to
+build the core package.
+
+```R
+devtools::install_github("sbhattlab/phylo2vec", subdir="./r-phylo2vec", build = FALSE)
+```
+
+Note: to download a specific version, use:
+
+```R
+devtools::install_github("sbhattlab/phylo2vec@vX.Y.Z", subdir="./r-phylo2vec", build = FALSE)
+```
+
+#### Option 3: manual installation
+
+⚠️ This requires installing [Rust](https://www.rust-lang.org/tools/install) to
+build the core package.
+
+Clone the repository and run the following `install.packages` in your R command
+line.
+
+Note: to download a specific version, you can use `git checkout` to a desired
+tag.
+
+```bash
+git clone https://github.com/sbhattlab/phylo2vec
+cd phylo2vec
+```
+
+```R
+install.packages("./r-phylo2vec", repos = NULL, type = 'source')
+```
+
 ## Basic Usage
 
 ### Python

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/phylo2vec/base/ancestry.py

@@ -20,7 +20,7 @@ def from_ancestry(ancestry: np.ndarray) -> np.ndarray:
 
     Returns
     -------
-    numpy.array
+    numpy.ndarray
         Phylo2Vec vector
     """
     v_list = core.from_ancestry(ancestry)
@@ -50,12 +50,12 @@ def to_ancestry(v: np.ndarray) -> np.ndarray:
 
     Parameters
     ----------
-    v : numpy.array
+    v : numpy.ndarray
         Phylo2Vec vector
 
     Returns
     -------
-    ancestry : numpy.array
+    ancestry : numpy.ndarray
         Ancestry matrix
         1st column: child 1
         2nd column: child 2

{phylo2vec-1.2.0/py-phylo2vec → phylo2vec-1.3.0}/phylo2vec/base/edges.py

@@ -21,7 +21,7 @@ def from_edges(edges: List[Tuple[int, int]]) -> np.ndarray:
 
     Returns
     -------
-    v : numpy.array
+    v : numpy.ndarray
         Phylo2Vec vector
     """
     v = core.from_edges(edges)
@@ -35,7 +35,7 @@ def to_edges(v: np.ndarray) -> List[Tuple[int, int]]:
 
     Parameters
     ----------
-    v : numpy.array
+    v : numpy.ndarray
         Phylo2Vec vector
 
     Returns

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/phylo2vec/base/newick.py

@@ -31,7 +31,7 @@ def to_newick(vector_or_matrix: np.ndarray) -> str:
 
     Parameters
     ----------
-    vector_or_matrix : numpy.array
+    vector_or_matrix : numpy.ndarray
         Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
 
     Returns

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/phylo2vec/base/pairs.py

@@ -11,10 +11,9 @@ from typing import List, Tuple
 def from_pairs(pairs: List[Tuple[int, int]]) -> np.ndarray:
     """Convert a list of pairs to a Phylo2Vec vector
 
-    Each pair is represented as a list of pairs (B, L)
+    Each pair is represented as a tuple (B, L)
     indicating that leaf L descends from branch B.
 
-
     Parameters
     ----------
     pairs : List[Tuple[int, int]]
@@ -22,7 +21,7 @@ def from_pairs(pairs: List[Tuple[int, int]]) -> np.ndarray:
 
     Returns
     -------
-    v : numpy.array
+    v : numpy.ndarray
         Phylo2Vec vector
     """
     v = core.from_pairs(pairs)
@@ -32,12 +31,12 @@ def from_pairs(pairs: List[Tuple[int, int]]) -> np.ndarray:
 def to_pairs(v: np.ndarray) -> List[Tuple[int, int]]:
     """Convert a Phylo2Vec vector to a list of pairs
 
-    Each pair is represented as a list of pairs (B, L)
+    Each pair is represented as a tuple (B, L)
     indicating that leaf L descends from branch B.
 
     Parameters
     ----------
-    v : numpy.array
+    v : numpy.ndarray
         Phylo2Vec vector
 
     Returns

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/phylo2vec/io/writer.py

@@ -44,7 +44,7 @@ def save_newick(
 
     Parameters
     ----------
-    vector_or_matrix : numpy.array
+    vector_or_matrix : numpy.ndarray
         Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
     filepath : str
         Path to the output file

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/phylo2vec/metrics/pairwise.py

@@ -4,6 +4,7 @@ import warnings
 import numpy as np
 
 from phylo2vec import _phylo2vec_core as core
+from phylo2vec.utils.matrix import check_matrix
 from phylo2vec.utils.vector import check_vector
 
 
@@ -14,7 +15,7 @@ def cophenetic_distances(vector_or_matrix, unrooted=False):
 
     Parameters
     ----------
-    vector_or_matrix : numpy.array
+    vector_or_matrix : numpy.ndarray
         Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
 
     Returns
@@ -46,7 +47,7 @@ def cophenetic_distances(vector_or_matrix, unrooted=False):
 PAIRWISE_DISTANCES = {"cophenetic": cophenetic_distances}
 
 
-def pairwise_distances(v, metric="cophenetic"):
+def pairwise_distances(vector_or_matrix, metric="cophenetic"):
     """
     Compute a pairwise distance matrix
     for tree nodes from a Phylo2Vec vector.
@@ -55,8 +56,8 @@ def pairwise_distances(v, metric="cophenetic"):
 
     Parameters
     ----------
-    v : numpy.ndarray
-        Phylo2Vec vector
+    vector_or_matrix : numpy.ndarray
+        Phylo2Vec vector (ndim == 1)/matrix (ndim == 2)
     metric : str, optional
         Pairwise distance metric, by default "cophenetic"
 
@@ -65,8 +66,15 @@ def pairwise_distances(v, metric="cophenetic"):
     numpy.ndarray
         Distance matrix
     """
-    check_vector(v)
+    if vector_or_matrix.ndim == 2:
+        check_matrix(vector_or_matrix)
+    elif vector_or_matrix.ndim == 1:
+        check_vector(vector_or_matrix)
+    else:
+        raise ValueError(
+            "vector_or_matrix should either be a vector (ndim == 1) or matrix (ndim == 2)"
+        )
 
     func = PAIRWISE_DISTANCES[metric]
 
-    return func(v)
+    return func(vector_or_matrix)

{phylo2vec-1.2.0/py-phylo2vec → phylo2vec-1.3.0}/phylo2vec/opt/__init__.py

@@ -1,6 +1,7 @@
 """Phylo2Vec-based optimisation methods."""
 
 from ._base import BaseOptimizer
+from ._gradme import GradMEOptimizer
 from ._hc import HillClimbingOptimizer
 
-__all__ = ["BaseOptimizer", "HillClimbingOptimizer"]
+__all__ = ["BaseOptimizer", "GradMEOptimizer", "HillClimbingOptimizer"]

{phylo2vec-1.2.0 → phylo2vec-1.3.0}/phylo2vec/opt/_base.py

@@ -1,16 +1,46 @@
 """Base class for all optimisation methods in Phylo2Vec."""
 
+import multiprocessing
 import random
+import time
+
+from dataclasses import dataclass
+from typing import Dict, List, Final
 
-import numba as nb
 import numpy as np
 
 from phylo2vec.datasets import read_fasta
 from phylo2vec.utils.vector import sample_vector
 
+# Multiprocessing
+DEFAULT_N_JOBS: Final = multiprocessing.cpu_count() // 4
+MIN_N_JOBS: Final = 4
+# Seeding
 MAX_SEED = 42
 
 
+@dataclass
+class BaseResult:
+    """Result of the optimization process.
+
+    Attributes
+    ----------
+    v_opt : numpy.ndarray
+        Optimized phylo2vec vector.
+    label_mapping : Dict[int, str]
+        Mapping of leaf labels (integer) to taxa.
+    best_score : float
+        The best score achieved during optimization.
+    scores : List[float]
+        List of scores obtained during optimization.
+    """
+
+    v: np.ndarray
+    label_mapping: Dict[int, str]
+    best_score: float
+    scores: List[float]
+
+
 class BaseOptimizer:
     """
     Base class for all phylo2vec-based optimizers
@@ -21,25 +51,28 @@ class BaseOptimizer:
         Random seed, by default None
     """
 
-    def __init__(self, random_seed=None):
+    def __init__(self, random_seed=None, verbose=False, n_jobs=None):
         self.random_seed = (
             random.randint(0, MAX_SEED) if random_seed is None else random_seed
         )
         random.seed(self.random_seed)
         np.random.seed(self.random_seed)
 
+        self.verbose = verbose
+
+        self.n_jobs = self._infer_n_jobs(n_jobs)
+
     @staticmethod
-    def _make_label_mapping(records):
-        label_mapping = nb.typed.Dict.empty(
-            key_type=nb.types.int64, value_type=nb.types.unicode_type
-        )
+    def _infer_n_jobs(n_jobs=None):
+        return n_jobs or max(MIN_N_JOBS, DEFAULT_N_JOBS)
 
-        for i, r in enumerate(records):
-            label_mapping[i] = r.id.replace(" ", ".")
+    @staticmethod
+    def _make_label_mapping(records):
+        label_mapping = dict(enumerate(r.id.replace(" ", ".") for r in records))
 
         return label_mapping
 
-    def fit(self, fasta_path):
+    def fit(self, fasta_path) -> BaseResult:
         """Fit an optimizer to a fasta file
 
         Parameters
@@ -51,7 +84,7 @@ class BaseOptimizer:
         -------
         v_opt : numpy.ndarray
             Optimized phylo2vec vector
-        label_mapping : dict[int, str]
+        label_mapping : List[str]
             Mapping of leaf labels (integer) to taxa
         losses : array-like
             List/Array of collected losses
@@ -67,9 +100,19 @@ class BaseOptimizer:
 
         v_init = sample_vector(n_leaves)
 
-        v_opt, label_mapping, losses = self._optimise(fasta_path, v_init, label_mapping)
+        start_time = time.time()
+
+        result = self._optimise(fasta_path, v_init, label_mapping)
+
+        end_time = time.time()
+
+        if self.verbose:
+            print(
+                f"Optimisation finished in {end_time - start_time:.2f} seconds "
+                f"with {len(result.scores)} loss evaluations."
+            )
 
-        return v_opt, label_mapping, losses
+        return result
 
     def _optimise(self, fasta_path, v, label_mapping):
         raise NotImplementedError

phylo2vec-1.3.0/phylo2vec/opt/_gradme.py

@@ -0,0 +1,200 @@
+"""Methods for GradME optimisation."""
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+import optax
+import rpy2
+import rpy2.robjects as ro
+
+from jax import jit, value_and_grad
+from rpy2.robjects import pandas2ri
+from rpy2.robjects.conversion import localconverter
+from rpy2.robjects.packages import importr
+from tqdm import tqdm
+
+from phylo2vec.opt._base import BaseOptimizer, BaseResult
+from phylo2vec.opt._gradme_losses import gradme_loss
+from phylo2vec.utils.vector import queue_shuffle, reroot_at_random
+
+
+# Disable rpy2 warning
+rpy2.rinterface_lib.callbacks.consolewrite_warnerror = lambda *args: None
+
+
+@dataclass
+class GradMEResult(BaseResult):
+    """Result of the GradME optimization.
+
+    See BaseResult for more details.
+
+    Attributes
+    ----------
+    W : jax.numpy.ndarray
+        The optimized weight matrix representing the phylogenetic tree.
+    """
+
+    W: jnp.ndarray
+
+
+class GradMEOptimizer(BaseOptimizer):
+    """GradME Optimizer for phylogenetic trees.
+
+    This optimizer uses the GradME algorithm to optimize phylogenetic trees.
+    It computes the loss using the GradME loss function and updates the tree
+    representation accordingly.
+
+    Parameters
+    ----------
+    random_seed : int, optional
+        Random seed for reproducibility, by default None
+    n_jobs : int, optional
+        Number of parallel jobs, by default None
+    verbose : bool, optional
+        Verbosity level, by default False
+    """
+
+    def __init__(
+        self,
+        model,
+        solver="adafactor",
+        learning_rate=1.5,
+        rooted=False,
+        n_shuffles=100,
+        n_iter_per_step=5000,
+        tol=1e-8,
+        random_seed=None,
+        n_jobs=None,
+        verbose=False,
+    ):
+        super().__init__(random_seed=random_seed, n_jobs=n_jobs, verbose=verbose)
+
+        self.model = model
+
+        self.optimizer = getattr(optax, solver)(learning_rate=learning_rate)
+        self.learning_rate = learning_rate
+        self.rooted = rooted
+        self.n_shuffles = n_shuffles
+        self.n_iter_per_step = n_iter_per_step
+        self.tol = tol
+
+    def _optimise(
+        self,
+        fasta_path,
+        v,
+        label_mapping,
+    ):
+        data = self.pdist(fasta_path, self.model)
+        dm = jnp.asarray(data)
+        k = dm.shape[0] - 1
+
+        # Forward and backward pass function
+        value_and_grad_fun = jit(value_and_grad(gradme_loss))
+
+        # Initial "best" score, set as an arbitrarily high value
+        best_score = 1e8
+
+        # List of scores obtained during optimization
+        scores = []
+
+        iterator = range(self.n_shuffles)
+
+        if self.verbose:
+            iterator = tqdm(iterator)
+
+        for _ in iterator:
+            w_in = self._init_W(k)
+
+            w_out = self._step(
+                w_in,
+                dm,
+                value_and_grad_fun,
+            )
+
+            v = w_out.argmax(1)
+
+            w_discrete = jnp.eye(w_out.shape[0])[v]
+
+            score = gradme_loss(w_discrete, dm, rooted=True)
+
+            best_score = min(best_score, score)
+
+            scores.append(best_score)
+
+            if not self.rooted:
+                v = reroot_at_random(v)
+
+            # Queue shuffle
+            _, vec_mapping = queue_shuffle(v, shuffle_cherries=True)
+
+            # Re-arrange the label mapping and the distance matrix
+            col_order = []
+            for i, idx in enumerate(vec_mapping):
+                label_mapping[i] = label_mapping[idx]
+                col_order.append(label_mapping[i])
+
+            dm = jnp.asarray(data.loc[col_order, col_order])
+
+            if self.verbose:
+                iterator.set_postfix({"\033[95m Best score ": best_score})
+
+        v = jnp.eye(w_out.shape[0])[w_out.argmax(1)]
+
+        best_params = GradMEResult(
+            v=v,
+            best_score=best_score,
+            scores=scores,
+            W=w_out,
+            label_mapping=label_mapping,
+        )
+
+        return best_params
+
+    @staticmethod
+    def _init_W(k, eps=1e-8):
+        x = jnp.tril(jnp.ones((k, k)))
+
+        w_init = x / (x.sum(1)[:, jnp.newaxis] + eps)
+
+        return w_init
+
+    def _step(self, w, dm, value_and_grad_fun):
+        state = self.optimizer.init(w)
+
+        prev_loss = 1e8
+
+        for _ in range(self.n_iter_per_step):
+            loss, gradients = value_and_grad_fun(w, dm, self.rooted)
+
+            if jnp.abs(loss - prev_loss) < self.tol:
+                break
+
+            prev_loss = loss
+
+            updates, state = self.optimizer.update(gradients, state, w)
+
+            w = optax.apply_updates(w, updates)
+
+        return w
+
+    @staticmethod
+    def pdist(fasta_path, model):
+        with localconverter(ro.default_converter + pandas2ri.converter):
+            importr("ape")
+
+            ro.globalenv["fasta_path"] = fasta_path
+            ro.globalenv["model"] = model
+
+            # DNA Evolution model: F81 + Gamma
+            dm = ro.r(
+                """
+                aln <- read.FASTA(fasta_path, type = "DNA")
+
+                dm <- dist.dna(aln, model = model)
+
+                D <- as.data.frame(as.matrix(dm))
+                D
+                """
+            )
+
+        return dm