htree 4.1.0
A library for tree reading, embedding, and analysis of phylogenetic trees
License: MIT
Requires Python: >=3.10.16


# Phylogenetic Tree Embedding

This repository provides a Python class for processing and embedding phylogenetic trees in either Euclidean or Hyperbolic geometric spaces. The class utilizes the `treeswift` library for handling tree structures and `torch` for optimizing embeddings. It offers a flexible framework for visualizing, analyzing, and embedding phylogenetic data.

## 📋 Requirements and Installation

To use `htree`, ensure your system meets the following requirements:

### 🖥 System Requirements
- **Python Version**: `>=3.12.2`
- **Operating System**: Compatible with Linux, macOS, and Windows

### 📦 Dependencies
`htree` relies on the following Python libraries:

- `numpy` – Numerical operations  
- `scipy` – Scientific computing  
- `matplotlib` – Visualization  
- `torch` – PyTorch for embedding optimization  
- `treeswift` – Phylogenetic tree processing  
- `tqdm` – Progress bar for operations  
- `imageio` – Image processing  
- `imageio-ffmpeg` – Video processing for animations  

These dependencies **will be installed automatically** when you install `htree` via `pip`:

```bash
pip install htree
```


## Tree Class

The `Tree` class provides an efficient way to load, process, and embed phylogenetic trees into a specified geometric space. It supports both **hyperbolic** and **Euclidean** geometries, allowing for flexible embedding choices depending on the application. The class also includes options for logging, saving intermediate states, and creating visualizations during the embedding process.

## Features

- **Tree Loading**: Load trees from files or directly from `treeswift.Tree` objects.
- **Embedding**: Embed trees in Euclidean or Hyperbolic spaces with customizable dimensions and parameters.
- **Visualization**: Generate figures or movies showing the embedding process.
- **Logging**: Enable detailed logging of operations for better traceability.

---

## Initialization

The `Tree` class can be initialized either by providing a file path to a tree file or by directly passing a `treeswift.Tree` object along with a name.

```python
from htree.tree_collections import Tree
# Initialize from file
tree = Tree("path/to/treefile.tre")
print(tree)
Tree(treefile.tre)
# Initialize from a treeswift Tree object
import treeswift as ts
t = ts.read_tree_newick("path/to/treefile.tre")
tree = Tree("treefile.tre", t)
print(tree)
Tree(treefile.tre)
```

Once initialized, the `Tree` object is ready for operations such as embedding, normalization, or visualization.

## Logging Feature

The class supports logging for debugging and tracing execution. By setting `logger.set_logger(True)`, a log file will be created in the specified directory (default: as configured in your package’s `config`).

```python
import htree.logger as logger
logger.set_logger(True)
# Initialize with logging enabled
tree = Tree("path/to/treefile.tre")
print(tree)
Tree(treefile.tre)
```
The log will store information such as tree initialization details, embedding operations, and error messages, making it easier to trace issues or understand the workflow. Log files are named according to the current timestamp and stored in the configured log directory `conf.LOG_DIRECTORY`.

## Methods

Once a tree is loaded, the `Tree` class provides functions to compute key properties like the distance matrix, diameter, and terminal node names.

* `distance_matrix`: Computes the pairwise distance matrix between terminal nodes.
* `diameter`: Computes the diameter of the tree, defined as the longest path between any two terminal nodes.
* `terminal_names`: Returns a list of terminal node names.
* `embed`: Returns a geomtric embedding for a tree.

```python
# Get top four terminal names
terminals = tree.terminal_names()[:4]
print(terminals)
['Diospyros_malabarica', 'Huperzia_squarrosa', 'Selaginella_moellendorffii_genome', 'Colchicum_autumnale']
# Compute their distance matrix
dist_matrix, names = tree.distance_matrix()
print(dist_matrix[:4,:4])
tensor([[0.0000, 2.0274, 2.1069, 0.9733],
        [2.0274, 0.0000, 1.7451, 1.3855],
        [2.1069, 1.7451, 0.0000, 1.4650],
        [0.9733, 1.3855, 1.4650, 0.0000]])
print(names[:4])
['Diospyros_malabarica', 'Huperzia_squarrosa', 'Selaginella_moellendorffii_genome', 'Colchicum_autumnale']
# Get tree diameter
diameter = tree.diameter()
print(diameter)
tensor(2.5342)
```

## Tree Normalization
The `normalize` method scales the branch lengths of the tree such that the maximum distance (tree diameter) is set to 1. This is particularly useful before performing embeddings, to ensure consistent scaling across different trees.

```python
# Normalize tree to have a diameter of 1
tree.normalize()
diameter = tree.diameter()
print(diameter)
tensor(1.)
```

## Embedding in Geometric Spaces
The `embed` method allows embedding the phylogenetic tree into either Euclidean or hyperbolic geometry using different dimensions. This is useful for visualizing or analyzing trees in a geometric space.

`geometry`: The geometry type, either `euclidean` or `hyperbolic`. Defaults to `hyperbolic`.
`dim`: The dimension of the embedding space. 

```python
# Embed tree in 2D hyperbolic space
embedding = tree.embed(dim=2, geometry='hyperbolic')
print(embedding)
HyperbolicEmbedding(curvature=-15.57, model=loid, points_shape=[3, 14])
# Embed tree in 3D Euclidean space
embedding = tree.embed(dim=3, geometry='euclidean')
print(embedding)
EuclideanEmbedding(points_shape=[3, 14])
```

### `embed` Method Parameters

The `embed` has several optional parameters (kwargs) that allow customization of the embedding process. Below is a detailed explanation of each parameter.

#### Parameters:

* **`dim` (int) [Required]:**
  + Defines the number of dimensions for the embedding space. For example, `dim=2` creates a 2D embedding.
  + **Example**: A 2D hyperbolic embedding is common for visualization purposes, while higher dimensions (e.g., `dim=3`) might be used for more complex analyses.
* **`geometry` (str) [Default: `hyperbolic`]:**
  + Specifies the geometric space to use for embedding.
  + Options:
    - `hyperbolic`: Embeds the tree in hyperbolic space, which is often suitable for representing hierarchical or tree-like structures.
    - `euclidean`: Embeds the tree in Euclidean space, more appropriate for flat, linear structures ($$ \ell_2^2 $$ distances).
  + Example:
  ```python
  print(tree.embed(dim=2, geometry='euclidean'))
  EuclideanEmbedding(points_shape=[2, 14])
  ```
* **`precise_opt` (bool) [Default: `False`]:**
  + Determines whether to use a more accurate and optimized embedding method. 
  + Options:
    - `False`: Quick and less computationally expensive based on spectral factorization methods.
    - `True`: Provides a higher quality embedding, but may take longer due to gradient-descent optimization.
    + Example:
    ```python
    print(tree.embed(dim=2, precise_opt=True))
    HyperbolicEmbedding(curvature=-10.58, model=loid, points_shape=[3, 14])
    ```
* **`epochs` (int) [Default: `1000`]:**
  + The number of epochs used for optimization during the embedding process. Relevant when `precise_opt=True` is selected. More epochs typically yield more accurate results at the cost of increased computation time.
  + Example:
  ```python
  print(tree.embed(dim=2, precise_opt=True, epochs=2000))
  HyperbolicEmbedding(curvature=-13.20, model=loid, points_shape=[3, 14])
  ```
* **`lr_init` (float) [Default: `0.01`]:**
  + The initial learning rate for the optimizer during the embedding process. This value controls how quickly the optimization process converges. A higher learning rate can speed up convergence but may overshoot optimal values, while a lower learning rate ensures a smoother, more gradual approach.
  + Example:
  ```python
  print(tree.embed(dim=2, precise_opt=True, lr_init=0.1))
  HyperbolicEmbedding(curvature=-10.58, model=loid, points_shape=[3, 14])
  ```
* **`dist_cutoff` (float) [Default: `10`]:**
  + The distance cutoff used to scale the distance matrix before embedding. This helps to ensure consistent scaling across different trees, which is especially useful for comparing multiple trees embedded in the same space.
+ Example:
  ```python
  print(tree.embed(dim=2, dist_cutoff=5.0))
  HyperbolicEmbedding(curvature=-3.89, model=loid, points_shape=[3, 14])
  ```
* **`save_mode` (bool) [Default: `False`]:**
  + Whether to save intermediate states during the embedding process. This can be useful for tracking progress or debugging, especially when using long-running optimizations with `precise_opt=True`.
  + Example:
  ```python
  print(tree.embed(dim=2, precise_opt=True, save_mode=True))
  HyperbolicEmbedding(curvature=-10.58, model=loid, points_shape=[3, 14])
  ```
  Saved data is stored in the directory `conf.OUTPUT_DIRECTORY/timestamp`, where `timestamp` is the current time when the embedding process starts. Each run generates a unique folder for its saved outputs. The results of the embedding process are saved in the `hyperbolic_embedding_{dim}.pkl` file within the timestamped folder. When `save_mode=True`, the following additional metrics are saved in the same directory during each epoch of the embedding process:

  - **Relative Error**: The normalized difference between the estimated and actual distance matrices at each epoch.
  -  **Weights**: The normalization weights used in the cost during each epoch's optimization.
  -  **Learning Rates**: The learning rate applied during the optimization process.
  -  **Scale Booleans**: Information regarding if the scale learning is enabled in each epoch.
  
  
* **`export_video` (bool) [Default: `False`]:**
  + Whether to generate a visual representation (video) of the embedding process. This can be helpful for understanding the evolution of the embedding as it progresses through optimization steps. 
  + Example:
  ```python
  tree.embed(dim=2, precise_opt=True, export_video=True)
  ```
  The movie is saved in MP4 format in `confg.RESULTS_DIR/Videos/timestamp` where timestamp is the time the embedding process was initiated. In addition to the video, if `export_video=True`, individual frame images captured during the embedding process are saved in `conf.OUTPUT_DIRECTORY/Images/timestamp`. This allows you to view each step of the embedding visually and track its progression at a finer granularity.
  - Video Example: Below is a for a video demonstrating how the embedding evolves over time:
  - [![Watch the video](https://img.youtube.com/vi/bj0z6dPB9Uo/0.jpg)](https://youtu.be/bj0z6dPB9Uo)


* **`lr_fn` (callable, optional):**
  + Specifies a custom learning rate function for the optimization process. The function takes in arguments like the current epoch, total epochs, and a list of losses, and returns a float value representing the learning rate at that epoch. If provided, the actual learning rate is computed as the product of the initial learning rate (`lr_init`) and the value returned by this function. If not provided, an adaptive learning rate is used, adjusting based on the progress of the optimization process.
  + Example:
  ```python
  def custom_learning_rate(epoch: int, total_epochs: int, loss_list: List[float]) -> float:
    """ 
    Calculate a dynamic learning rate based on the current epoch and total number of epochs.
    Parameters:
    - epoch (int): The current epoch in the training process.
    - total_epochs (int): The total number of epochs in the training process.
    - loss_list (list): A list of recorded loss values (can be used for further custom logic).

    Returns:
    - float: The dynamic learning rate for the current epoch.

    Raises:
    - ValueError: If `total_epochs` is less than or equal to 1.
    """

    if total_epochs <= 1:
        raise ValueError("Total epochs must be greater than 1.")

    # Example: Reduce learning rate as training progresses
    decay_factor = 0.5  # Factor by which to decay the learning rate
    loss_threshold = 0.01  # Loss threshold for further reduction
    decay_start_epoch = int(0.7 * total_epochs)  # When to start decaying

    # Reduce learning rate if the epoch is beyond a certain point
    if epoch > decay_start_epoch:
        # Learning rate decays based on the remaining epochs
        decay_rate = 1 - (epoch - decay_start_epoch) / (total_epochs - decay_start_epoch)
    else:
        decay_rate = 1.0  # No decay before the threshold
    # Further adjust learning rate if recent loss has not improved sufficiently
    if len(loss_list) > 1 and loss_list[-1] > loss_threshold:
        decay_rate *= decay_factor
    return  decay_rate
  tree.embed(dim=2, precise_opt=True, lr_fn=custom_learning_rate, lr_init = 0.01)
  ```
* **`scale_fn`(callable, optional):**
  +A custom function that determines whether scale learning should occur during the optimization process. The function takes the current epoch, total epochs, and a list of losses as input, and returns a boolean value. If `True`, scale learning is performed during that epoch; otherwise, it is skipped. If not provided, the default behavior is to enable scale learning during the early epochs of training based on a predefined ratio.
  + Example:
  ```python
  def custom_scale(epoch: int, total_epochs: int, loss_list: List[float]) -> bool:
    """
    Determine whether scale learning should occur based on the current epoch and total number of epochs.

    Parameters:
    - epoch (int): The current epoch in the training process.
    - total_epochs (int): The total number of epochs in the training process.
    - loss_list (list): A list of recorded loss values (can be used for further custom logic).

    Returns:
    - bool: `True` if scale learning should occur, `False` otherwise.
    
    Raises:
    - ValueError: If `total_epochs` is less than or equal to 1.
    """
    if total_epochs <= 1:
        raise ValueError("Total epochs must be greater than 1.")

    # Define the ratio of epochs during which scale learning should be applied
    curv_ratio = 0.3  # For example, learning happens during the first 30% of epochs
    
    return epoch < int(0.6 * total_epochs)
  tree.embed(dim=2, precise_opt=True, scale_fn=custom_scale)
  ```
* **`weight_exp_fn` (callable, optional):**
  + A custom function to dynamically compute the weight exponent during the optimization process. The function takes the current epoch and total epochs as input, and returns a float representing the weight exponent for that epoch. The weight exponent determines how distances are weighted during optimization. If not provided, a default behavior is applied, where weighting changes over the course of the training based on the progress of the optimization process.
  + Example:
  ```python
  def custom_weight_exponent(epoch: int, total_epochs: int,loss_list: List[float]) -> float:
     """
     Calculate the weight exponent based on the current epoch and total number of epochs.
     Parameters:
     - epoch (int): The current epoch in the training process.
     - total_epochs (int): The total number of epochs in the training process.
     - loss_list (list): A list of recorded loss values (can be used for further custom logic).
  
     Returns:
     - float: The calculated weight exponent for the current epoch.

     Raises:
     - ValueError: If `total_epochs` is less than or equal to 1.
     """
     if total_epochs <= 1:
         raise ValueError("Total epochs must be greater than 1.")

     # Define a ratio that determines how long to apply no weights
     no_weight_ratio = 0.3  # Example ratio: first 30% of epochs without weighting
     no_weight_epochs = int(no_weight_ratio * total_epochs)
     # No weighting for the first part of the training
     if epoch < no_weight_epochs:
         return 0.0  # No weighting initially
     # Gradually increase the negative weight exponent after the no-weight phase
     return -(epoch - no_weight_epochs) / (total_epochs - 1 - no_weight_epochs)
  tree.embed(dim=2, precise_opt=True, weight_exp_fn=custom_weight_exponent)
  ```

## Saving and Copying Trees
The class provides utility functions to save and copy trees:

`save()`: Save the tree to a file in Newick format.
`copy()`: Create a deep copy of the tree object, useful for making modifications without affecting the original tree.

```python
# Save tree to a file
tree.save("path/to/treefile2.tre") 
# Copy the tree object
tree_copy = tree.copy()
```

## MultiTree Class
The `MultiTree` class offers a versatile framework for managing multiple phylogenetic trees, allowing them to be loaded from various sources, processed, and embedded into geometric spaces. This class supports both hyperbolic and Euclidean embeddings, providing a flexible approach for different applications. It is subscriptable, enabling efficient selection and manipulation of tree subsets. For example, if multitree contains a collection of phylogenetic trees, an expression like `multitree[:5]` would return a new `MultiTree` instance containing only the first five trees, preserving the original structure and functionality. Additionally, the class includes logging capabilities and utility functions such as distance matrix computation and tree saving.

## Features
- **Flexible Tree Loading**: Load trees from files, lists, or dictionaries of `Tree` objects. 
- **Embedding Options**: **Jointly** embed trees in hyperbolic or Euclidean spaces with customizable dimensions and optimization settings.
- **Distance Matrix**: Compute an aggregate distance matrix for tree comparisons, with optional confidence matrix.
- **Tree Saving**: Save trees in Newick format.
- **Logging**: Enable logging to track the operations performed on the `MultiTree` object for debugging and traceability.


## Initialization
The `MultiTree` class can be initialized from a variety of sources, including file paths, lists of trees of named trees, lists of trees from `treeswift`. 

```python
from htree.tree_collections import MultiTree
# Initialize from a Newick file
multitree = MultiTree("path/to/trees.tre")
print(multitree)
MultiTree(trees.tre, 844 trees)
# Initialize from a list of trees
import treeswift as ts
tree1 = ts.read_tree_newick("path/to/treefile1.tre")
tree2 = ts.read_tree_newick("path/to/treefile2.tre")
tree_list = [tree1, tree2]  # List of trees
multitree = MultiTree('mytrees', tree_list)
print(multitree)
MultiTree(mytrees, 2 trees)
print(multitree.trees)
[Tree(Tree_0), Tree(Tree_1)]
# Initialize from a list of named trees
from htree.tree_collections import Tree
named_trees = [Tree('a', tree1), Tree('b', tree2)]
multitree = MultiTree('mTree', named_trees)
print(multitree)
MultiTree(mTree, 2 trees)
print(multitree.trees)
[Tree(a), Tree(b)]
```
Once initialized, the `MultiTree` object provides methods for tree embedding, saving, and distance matrix computation.

## Logging Feature
Enable logging for detailed tracking of actions and errors during the `MultiTree` operations. When enabled, logs will be written to a file with a timestamp in the filename.
```python
# Initialize with logging enabled
import htree.logger as logger
logger.set_logger(True)
multitree = MultiTree("path/to/trees.tre")
print(multitree)
MultiTree(trees.tre, 844 trees)
```
Logs can be useful for debugging and provide insights into the tree loading, embedding, and saving processes.

## Distance Matrix and Terminal Names
The `terminal_names` method in `MultiTree` returns the union of terminal (leaf) names from all trees, removing duplicates. The `distance_matrix` method computes a matrix that represents the average pairwise distances between nodes corresponding to the terminal names, using a customizable aggregation function `func`. For each tree, a distance matrix is first constructed for the leaf nodes in the union set of terminal names. If some terminal nodes are not present in a particular tree, the corresponding values in the matrix are replaced with `NaN`. The aggregation function `func` is then applied across these matrices to compute the final average distance matrix.

It’s important that the provided `func` automatically removes `NaN` values; otherwise, it will raise an error. Examples of valid functions include `torch.nanmean` (default) and `torch.nanmedian`, which both handle `NaN` values correctly. Optionally, a confidence matrix can be returned, indicating the proportion of non-NaN values for each element in the distance matrix. The `method = 'fp'` computes the average distance based on a weighted scheme where each tree has a weight -- equivalent to the relative similarity of it's distance matrix to the average distance matrix.
```python
multitree = MultiTree("path/to/trees.tre")[:10]
print(multitree)
MultiTree(mTree, 10 trees)
print(multitree.trees)
[Tree(tree_1), Tree(tree_2), Tree(tree_3), Tree(tree_4), Tree(tree_5), Tree(tree_6), Tree(tree_7), Tree(tree_8), Tree(tree_9), Tree(tree_10)]
# Compute the distance matrix with default aggregation (mean)
avg_mat, conf, labels = multitree.distance_matrix()
print(avg_mat[:4,:4])
tensor([[0.0000, 0.7049, 1.2343, 0.5929],
        [0.7049, 0.0000, 1.3234, 0.6870],
        [1.2343, 1.3234, 0.0000, 1.0143],
        [0.5929, 0.6870, 1.0143, 0.0000]])
# Compute the distance matrix with custom aggregation
import torch
avg_mat, conf, labels = multitree.distance_matrix(func=torch.nanmedian)
print(avg_mat[:4,:4])
tensor([[0.0000, 0.5538, 0.9043, 0.5240],
        [0.5538, 0.0000, 1.1598, 0.5902],
        [0.9043, 1.1598, 0.0000, 0.8635],
        [0.5240, 0.5902, 0.8635, 0.0000]])
avg_mat, conf, labels = multitree.distance_matrix(method='fp')
print(avg_mat[:4,:4])
tensor([[0.0000, 0.6760, 1.1487, 0.5696],
        [0.6760, 0.0000, 1.2801, 0.6613],
        [1.1487, 1.2801, 0.0000, 0.9627],
        [0.5696, 0.6613, 0.9627, 0.0000]])
# Compute the union of all terminal names (removes duplicates)
print(multitree.terminal_names()[:4])
['Allamanda_cathartica', 'Alsophila_spinulosa', 'Amborella_trichopoda', 'Aquilegia_formosa']
```


## Embedding Trees

The `embed` method in `MultiTree` allows you to **jointly** embed trees in a hyperbolic or Euclidean space, specifying the geometry, dimension, and other customizable parameters for optimization. It is flexible enough to handle both hyperbolic and Euclidean embeddings, with additional parameters like learning rate, total epochs, and scaling options. The method returns a `MultiEmbedding` object, which can be used for further analysis or visualization.
```python
# Embed trees in a 2D hyperbolic space
multiemb_hyperbolic = multitree.embed(dim=2, geometry='hyperbolic')
print(multiemb_hyperbolic)
MultiEmbedding(10 embeddings)
print(multiemb_hyperbolic.embeddings)
[HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 24]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 76]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 71]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 70]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 63]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 80]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 60]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 76]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 82]), HyperbolicEmbedding(curvature=-0.84, model=loid, points_shape=[3, 51])]
# Embed trees in a 3D Euclidean space
multiemb_euclidean = multitree.embed(dim=3, geometry='euclidean')
print(multiemb_euclidean)
MultiEmbedding(10 embeddings)
print(multiemb_euclidean.embeddings)
[EuclideanEmbedding(points_shape=[3, 24]), EuclideanEmbedding(points_shape=[3, 76]), EuclideanEmbedding(points_shape=[3, 71]), EuclideanEmbedding(points_shape=[3, 70]), EuclideanEmbedding(points_shape=[3, 63]), EuclideanEmbedding(points_shape=[3, 80]), EuclideanEmbedding(points_shape=[3, 60]), EuclideanEmbedding(points_shape=[3, 76]), EuclideanEmbedding(points_shape=[3, 82]), EuclideanEmbedding(points_shape=[3, 51])]
```
The `embed` method takes several key parameters that allow for fine-tuned control of the embedding process:
- `dim`: Specifies the number of dimensions for the embedding space (e.g., 2 for 2D or 3 for 3D).
- `geometry`: Specifies whether the embedding should be in 'hyperbolic' or 'euclidean' space.
- `precise_opt`: Controls the accuracy of the embedding process. This can be passed through kwargs or defaulted to the value in `conf.ENABLE_ACCURATE_OPTIMIZATION`.
- `epochs`: The number of epochs (iterations) to run the optimization for embedding. It defaults to `conf.TOTAL_EPOCHS` but can be overridden.
- `lr_init`: The initial learning rate for the optimization process. It defaults to `conf.LEARNING_RATE_INIT`, but you can specify a custom value.
- `dist_cutoff`: The maximum diameter for scaling the embeddings. This ensures the size of the embedding is bounded and defaults to `conf.MAX_RANGE`.
- `save_mode`: If set to True, the embeddings will be saved automatically during the process. The default is taken from `conf.ENABLE_SAVE_MODE`.
- `lr_fn`: Allows fine-tuning of the learning rate during the process. If not specified, the default learning rate (`lr_init`) will be used.
- `scale_fn`: Determines if scale learning should be enabled, which impacts how distances are scaled in the embedding space.
- `weight_exp_fn`: Controls how weights are adjusted during optimization, potentially affecting the impact of different nodes' distances during learning.

Refer to the `Tree` class for detailed explanations and further usage examples.

## Saving Trees
The `save` method allows you to save the trees in a file (default format is Newick).
```python
# Save trees to a Newick file
multitree.save("path/to/output.tre")
```

# Embedding Class

The `Embedding` class is an abstract base class designed to represent embeddings of points in geometric spaces, either Euclidean or hyperbolic. This class primarily serves as a **foundation**, and specific geometric implementations (e.g., Euclidean and Hyperbolic) are used in practice. The `Embedding` class supports both NumPy arrays and PyTorch tensors for point representation and can log operations if logging is enabled.

## Attributes
- **geometry (str)**: The geometry of the space, either `euclidean` or `hyperbolic`. Default is `hyperbolic`.
- **points (torch.Tensor)**: A PyTorch tensor representing the points in the embedding space.
- **labels (list)**: A list of labels corresponding to the points in the space.

## Initialization 
```python
def __init__(geometry='hyperbolic', points=None, labels=None)
```

## Parameters
- **geometry (`str`)**: The geometry type (`euclidean` or `hyperbolic`). Default is `hyperbolic`.
- **points (`Optional[Union[np.ndarray, torch.Tensor]]`)**: Points to embed, provided as a NumPy array or PyTorch tensor. Default is `None`.
- **labels (`Optional[List[Union[str, int]]]`)**: Labels for the points. Default is `None`.

  
## Methods
- **`save(filename: str)` -> None**: Saves the current instance of the `Embedding` class to a file using pickle, `filename` is the path where the embedding will be saved.
- **`copy() -> Embedding`**: Creates a deep copy of the `Embedding` object.
- **`load(filename: str) -> Embedding`**: Loads an `Embedding` instance from a file using pickle, `filename` is the path from where the embedding will be loaded.

In practice, the `Embedding` class itself is not directly instantiated. Instead, specific subclasses like `EuclideanEmbedding` and `HyperbolicEmbedding` are utilized for the actual geometric embeddings. These subclasses extend the base `Embedding` class, ensuring the correct handling of points and dimensions according to the geometry in use.

## Examples
```python
from htree.embedding import Embedding
import numpy as np
# Create an embedding with hyperbolic geometry (norm requirements are not applied)
n_points = 10
dimension = 2
embedding = Embedding(geometry='hyperbolic', points=np.random.randn(dimension,n_points))
print(embedding)
Embedding(geometry=hyperbolic, points_shape=[2, 10])
# Update points
embedding.points = np.random.randn(2,150)
raise NotImplementedError("update_dimensions must be implemented by a subclass")
# Set labels
labels = [str(i) for i in range(n_points)]
embedding.labels = list(range(n_points))
# Save the embedding
embedding.save('embedding.pkl')
# Load the embedding
loaded_embedding = Embedding.load('embedding.pkl')
print(loaded_embedding)
Embedding(geometry=hyperbolic, points_shape=[2, 10])
import torch
import htree.logger as logger
logger.set_logger(True)
embedding = Embedding(geometry='euclidean', points = np.random.randn(dimension,n_points))
print(embedding)
Embedding(geometry=euclidean, points_shape=[2, 10])
```


# EuclideanEmbedding Class
`EuclideanEmbedding` is a subclass of `Embedding` that represents an embedding in Euclidean space. It provides additional functionality specific to Euclidean geometry, such as translation, rotation, and computation of distance matrices.

## Attributes
- **curvature (`float`)**: The curvature of the Euclidean space, which is always `0`.
- **geometry (`str`)**: The geometry of the space, which is always `euclidean`.
- **points (`torch.Tensor`)**: A PyTorch tensor representing the points in Euclidean space.
- **labels (`List[Union[str, int]]`)**: A list of labels corresponding to the points in the space.

## Methods
- **`translate`**: Translates the points by a given vector in Euclidean space.
- **`rotate`**: Rotates the points by a given rotation matrix in Euclidean space.
- **`center`**: Centers the points by subtracting the centroid from each point.
- **`centroid`**: Computes and returns the centroid of the points. 
- **`distance_matrix`**: Computes and returns the distance matrix for points in Euclidean geometry. 

## `translate` Method
The `translate` method shifts all points in Euclidean space by a specified vector. This method ensures that the points are translated correctly by a vector of matching dimensionality. If the vector’s dimension doesn’t align with the space, it raises an error. The following example demonstrates the translation of points representing the vertices of a regular pentagon.

```python
import matplotlib.pyplot as plt
from htree.embedding import EuclideanEmbedding
import numpy as np
# Function to generate points at the vertices of a regular pentagon
def generate_pentagon_vertices():
     angles = np.linspace(0, 2 * np.pi, 6)[:-1]  # 5 angles for 5 vertices, excluding the last one to avoid duplication
     x = np.cos(angles)
     y = np.sin(angles)
     return np.vstack((x, y))
# Function to plot points with edges, and axes
def plot_pentagon(ax, points, title, color, marker, label_text):
     ax.scatter(points[0, :], points[1, :], c=color, s=100, edgecolors='black', marker=marker, label=label_text)    
     # Draw edges between consecutive points
     for i in range(points.shape[1]):
         start = i
         end = (i + 1) % points.shape[1]
         ax.plot([points[0, start], points[0, end]], [points[1, start], points[1, end]], color=color)
     # Draw the axes
     ax.axhline(0, color='black', linewidth=0.5)
     ax.axvline(0, color='black', linewidth=0.5)
     ax.set_title(title)
     ax.set_xlabel('X-axis')
     ax.set_ylabel('Y-axis')
     ax.set_aspect('equal', adjustable='box')
     ax.legend()
# Define pentagon vertices
points = generate_pentagon_vertices()
dimension = points.shape[0]
n_points = points.shape[1]
# Initialize embedding
embedding = EuclideanEmbedding(points=points)
# Create a figure for the translation
fig1, ax1 = plt.subplots(figsize=(8, 8))
# Plot the original pentagon vertices with edges
plot_pentagon(ax1, embedding.points, "Translation of Pentagon Vertices", color='black', marker='o', label_text='Original Points')
# Translate the points
translation_vector = np.random.randn(dimension)/2  # 2D translation vector
embedding.translate(translation_vector)
# Plot the translated pentagon vertices with edges
plot_pentagon(ax1, embedding.points, "Translation of Pentagon Vertices", color='green', marker='^', label_text='Translated Points')
# Show the plot for translation
plt.show()
```
In this example, a regular pentagon is translated by a randomly generated 2D vector, and the changes in the points are visualized on a plot.


<!-- <img src="https://github.com/puoya/HyperTree/blob/main/images/euclidean_translation.png" alt="GitHub Logo" width="500"/> -->
![GitHub Logo](https://i.imgur.com/5bOyyEA.png)

This demonstrates the effect of the `translate` method on the pentagon's vertices, illustrating how the points move to their new positions.

## `rotate` Method
The `rotate` method allows you to rotate the points in Euclidean space by applying a rotation matrix. The method automatically handles the conversion of NumPy arrays to PyTorch tensors, checks for the validity of the rotation matrix, and ensures that the points are rotated correctly. In this example, we create a set of points (the vertices of a pentagon), rotate them by 30 degrees using a rotation matrix, and plot the points before and after rotation.

```python
embedding = EuclideanEmbedding(points=points, labels=labels)
# Create a figure for the rotation
fig2, ax2 = plt.subplots(figsize=(8, 8))
# Plot the translated pentagon vertices with edges and labels before rotation
plot_pentagon(ax2, embedding.points, "Rotation of Pentagon Vertices", color='black', marker='o', label_text='Before Rotation')
# Convert degrees to radians and define the 2x2 rotation matrix
theta = np.radians(30)
rotation_matrix = np.array([
    [np.cos(theta), -np.sin(theta)],
    [np.sin(theta), np.cos(theta)]
])
# Rotate the points
embedding.rotate(rotation_matrix)
# Plot the pentagon vertices with edges after rotation
plot_pentagon(ax2, embedding.points, "Rotation of Pentagon Vertices", color='red', marker='^', label_text='After Rotation')
# Show the plot for rotation
plt.show()
```


<!-- <img src="https://github.com/puoya/HyperTree/blob/main/images/euclidean_rotation.png" alt="GitHub Logo" width="500"/> -->
![GitHub Logo](https://i.imgur.com/AuqzTjR.png)

## `center`,`centroid`, and `distance_matrix`  Methods
The `center` method centers the points in Euclidean space by subtracting the `centroid` from each point. This operation shifts the points so that their `centroid` becomes the origin. The `centroid` method computes the geometric center of the points in Euclidean space by calculating the mean of each dimension.

```python
# updating points
new_points = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
embedding.points = new_points  # This will automatically update dimensions and other attributes.
print(embedding.points)
tensor([[1., 2., 3.],
        [4., 5., 6.]], dtype=torch.float64)
print(embedding.centroid())
tensor([2., 5.], dtype=torch.float64)
embedding.center()
print(embedding.points)
tensor([[-1.,  0.,  1.],
        [-1.,  0.,  1.]], dtype=torch.float64)
print(embedding.centroid())
tensor([0., 0.], dtype=torch.float64)
```
The `distance_matrix` method computes the pairwise Euclidean distances between all points in the embedding. This results in a symmetric matrix where each element at position `(i, j)` represents the distance between point `i` and point `j`.

```python
distance_matrix,labels = embedding.distance_matrix()
print(distance_matrix)
tensor([[0.0000, 1.4142, 2.8284],
        [1.4142, 0.0000, 1.4142],
        [2.8284, 1.4142, 0.0000]], dtype=torch.float64)
```


# HyperbolicEmbedding Class 

This `HyperbolicEmbedding` class is designed to represent embeddings in hyperbolic space, supporting both the Poincare and Loid models. While the class provides an infrastructure for hyperbolic embeddings, it's primarily intended to be extended through subclasses that specifically handle the Poincare and Loid models.

## Purpose

The key purpose of this class is to allow switching between these two models (`PoincareEmbedding` and `LoidEmbedding`), which are different ways to represent points in hyperbolic space. However, we do not directly use this class in practice. Instead, we rely on its subclasses, which tailor the functionality specifically for Poincare or Loid embeddings.

## Key Features
- **Curvature Scaling**: The curvature of the hyperbolic space can be adjusted, but this does not affect the constraint that hyperbolic points must always have unit norms. The curvature merely scales distances between points.
- **Model Switching**: The `switch_model` method allows converting between Poincare and Loid embeddings seamlessly.
- **Distance Calculations**: Methods like `poincare_distance` compute distances in the Poincare model, while transformation methods (`to_poincare` and `to_loid`) switch between representations.

This structure ensures flexibility in working with hyperbolic embeddings, particularly when needing to switch between models while maintaining consistent point norms.


## Attributes
- **curvature (`torch.Tensor`)**: The curvature of the hyperbolic space, which must always be negative (e.g., -1). It scales distances between points but does not affect the requirement that points in hyperbolic space must have a unit norm..
- **model (`str`)**: Specifies whether the Poincare or Loid model is used to represent the points. Default is `poincare`, but you can switch to `loid`.
- **points (`torch.Tensor or np.ndarray`)**: Contains the points in the hyperbolic space, represented as either NumPy arrays or PyTorch tensors.
- **labels (`List[Union[str, int]]`)**: A list of labels corresponding to the points in the space.

## Methods
- **s`witch_model`**: Switches between the Poincare and Loid models. If the current model is Poincare, it switches to Loid, and vice versa.
- **`poincare_distance(x, y)`**: Computes the Poincare distance between two points `x` and `y` in the Poincare model.
- **`to_poincare(vectors)`**: Transforms vectors from the Loid model to the Poincare model.
- **`to_loid(vectors)`**: Transforms vectors from the Poincare model to the Loid model.

# PoincareEmbedding Class

The `PoincareEmbedding` class represents a hyperbolic embedding space using the Poincare model. It extends from the `HyperbolicEmbedding` class and introduces additional features specific to the Poincare geometry.

## Constructor and Attributes

```python
def __init__(self, curvature: Optional[float] = -1, points: Optional[Union[np.ndarray, torch.Tensor]] = None, labels: Optional[List[Union[str, int]]] = None) -> None
```

- **curvature** (`float`): The curvature of the hyperbolic space (default is `-1`).
- **points** (`torch.Tensor`): Points in the Poincare space.
- **labels** (`List[Union[str, int]]`): Optional labels for the points.

The constructor initializes the embedding in Poincare space and checks the point norms to ensure they adhere to the Poincare constraints.

## Methods
- `_norm2`: Computes and returns the **squared** L2 norm of the points in the embedding. This is useful for validating that points remain within the appropriate norms for the Poincare model.
- `distance_matrix`: Computes the pairwise distance matrix between the points using the Poincare distance formula.
- `centroid`: Computes the centroid of the points in the Poincare embedding. It supports two modes:
  + `default`: Uses the Loid model for initial calculations (based on the definition in [On Procrustes Analysis in Hyperbolic Space](https://ieeexplore.ieee.org/abstract/document/9435083)), then projects the centroid back to Poincare space.
  + `Frechet`: Uses a gradient-based optimization method (Frechet Mean) to find the centroid by minimizing the sum of squared distances.
- `translate`: Translates the points by a given vector, while ensuring that the resulting points still comply with the norm constraints of the Poincare model.
- `rotate`: Rotates the points by a given rotation matrix. The method ensures that the matrix is a valid rotation matrix and adjusts it if necessary.
- `center`: Centers the points by translating them to the centroid calculated by the `centroid` method.

## Example: Initialization
This example illustrates the initialization and manipul